2
$\begingroup$

I have a model of a physical system, $f(\pmb{x})$, where $f$ is the output of a mathematical model and $\pmb{x}$ are inputs to the model, which are available as observations. My goal is to find the expectation of $f(\pmb{x})$ with respect to the observed distribution of $\pmb{x}$ values via Bayesian quadrature.

Since my post is quite long, I have broken it into three sections:

In the introduction, I give an overview of how I determine the probability density of $\pmb{x}$, given observations, and provide a reference problem/solution using simple rectangular quadrature to compute the expectation of $f$.

In the Main Problem section, I give an overview of Bayesian quadrature and explain how I have applied it to my reference problem. In particular, I use a change of variables to make the Bayesian optimization tractable. I show that my example Bayesian quadrature solution does not match the reference solution.

In the Question section, I state my question. I am concerned that I have made an incorrect assumption regarding the Gaussian Process used by the Bayesian quadrature and that my approach is invalid. If my assumptions are valid, I would like to track down why my implementation is not working.

Introduction

In my problem, I observe different values of $\pmb{x}$. Collecting many observations, I can estimate the density of $\pmb{x}$, $P(\pmb{x})$, using a kernel density estimator:

$P(\pmb{x})\sim \sum_{i=1}^{N_{\mathrm{obs}}} \mathcal{N}(\pmb{x}, \pmb{X}_i, l_{\mathrm{KDE}})$,

where $P(\pmb{x})$ represents the probability density of $\pmb{x}$, $N_{\mathrm{obs}}$ is the number of observed inputs, $\pmb{X}_i$ is observed input $i$, $\mathcal{N}(\pmb{x}, \pmb{X}_i, l_{\mathrm{KDE}})$ is a multivariate Gaussian probability density function evaluated at $\pmb{x}$ with mean $\pmb{X}_i$ and standard deviation $l_\mathrm{KDE}$:

$\mathcal{N}(\pmb{x}, \pmb{X}_i, l_{\mathrm{KDE}}) = \prod_{j=1}^{\mathrm{dim}(\pmb{x})} \frac{1}{l_{\mathrm{KDE}}\sqrt{2\pi}}\exp\left[-\frac{1}{2}\left(\frac{x_j - X_{ij}}{l_{\mathrm{KDE}}}\right)^2\right]$

Given this kernel density estimator, it is simple to find the expectation of $f(\pmb{x})$ using standard quadrature methods. The following snippet of code finds the expectation of an example $f$ and set of observed $\pmb{x}$ using rectangular quadrature as a reference before I attempt the corresponding Bayesian Quadrature problem. The correct integral expectation is about 31.

import numpy as np
import pandas as pd
from scipy.stats import norm
import matplotlib.pyplot as plt


# define function to integrate
def f(x):
   return (np.sum(x ** 2))

# generate input samples to create reference ground truth
XSAMPS = 20000
train_x = np.array([np.random.normal((2, 5), (1, 1)) for __ in range(XSAMPS // 2)])
train_x = np.append(train_x, [np.random.normal((-2, -5), (1, 1)) for __ in range(XSAMPS // 2)], 0)

# compute ground truth
evals = np.array([f(x) for x in train_x])
print('ground truth is ', np.mean(evals))

# generate input samples to create KDE
XSAMPS = 200
train_x = np.array([np.random.normal((2, 5), (1, 1)) for __ in range(XSAMPS // 2)])
train_x = np.append(train_x, [np.random.normal((-2, -5), (1, 1)) for __ in range(XSAMPS // 2)], 0)

# define length scale of KDE estimate
LSCALE = 0.4

# create KDE of data
def kde(x, lscale=1):
   density = np.sum(norm.pdf(x[0], loc=train_x[:, 0], scale=lscale) * norm.pdf(x[1], loc=train_x[:, 1], scale=lscale))
   density /= train_x.shape[0]
   return density

# perform quadrature
print('Number of samples, mean')
for NX in range(6, 400, 2):
   x = np.linspace(-20, 20, NX)
   X, Y = np.meshgrid(x, x)
   x_plot = np.array([X, Y]).reshape(2, -1)
   evals = np.array([f(x) for x in x_plot.T])
   ps = np.array([kde(x, LSCALE) for x in x_plot.T])
   print(evals.size,  np.sum(evals * ps) / np.sum(ps))

Main Problem

My goal is to perform Bayesian quadrature of this function to find the expectation of $f$ with respect to the observed probability density.

In Bayesian quadrature, a Gaussian Process is used to represent the observed outputs as a function of the observed inputs. Then, an acquisition function is defined as the difference in the integral variance between the current Gaussian Process and a Gaussian Process where we have added one more point. The integral variance is the integral of the Gaussian Process variance across the range of allowable inputs.

The emukit package offers a fast implementation of this acquisition function--the integral variance is computed directly from the Gaussian Process, instead of performing numerical integration of the Gaussian Process. Without this clever step, numerical integration would greatly slow down the optimization of the acquisition function.

However, this fast implementation of the acquisition function is only available when estimating the integral of a function, without a corresponding probability density, $\int f(\pmb{u}) d\pmb{u}$.

Given this limitation, my strategy has been to map from a uniform probability distribution $P(\pmb{u}) \sim U[0, 1]^{\mathrm{dim}(u}$ to the kernel destiny estimator probability distribution $P(\pmb{x})$, $\pmb{x}=\Psi(\pmb{u})$. In order to accomplish this, I have defined $\pmb{u}$ to have one more dimension than $\pmb{x}$, $\mathrm{dim}(\pmb{u})=\mathrm{dim}(\pmb{x}) + 1$. In the $\Psi$ function, the first component of $\pmb{u}$ is used to determine the percentile of the observed $\pmb{x}$ values to be used in the sample. The remaining components offset that observed sample according to the inverse normal distribution of each component.

I have verified that my method is able to fairly map from uniform samples to samples from the kernel density estimator. This image, taken from the output of the code below, verifies this, by comparing observed $\pmb{x}$ samples to $\pmb{x}$ samples generated from uniform samples.

enter image description here

In my attempt, I perform the Bayesian optimization of $f$ with respect to $\pmb{u}$, using $\Psi$ to map from $\pmb{u}$ to $\pmb{x}$ whenever evaluating $f$. The expectations I estimate are not correct (my code computes the expectation as around 28). I worry there is a fundamental flaw in my formulation: the Gaussian process assumes the inputs are correlated according to their distance from other observed points. However, the distance between $\pmb{u}$ points may be confusing the Gaussian Process used in the Bayesian quadrature, as it is the distance between $\pmb{x}$ points that truly matters.

Here is the full code of the Bayesian quadrature attempt:

import numpy as np
import pandas as pd
from sklearn.datasets  import load_diabetes
from scipy.stats import norm
import matplotlib.pyplot as plt
from emukit.quadrature.methods import VanillaBayesianQuadrature
from emukit.quadrature.acquisitions import IntegralVarianceReduction
from emukit.core.optimization import GradientAcquisitionOptimizer
from emukit.core.parameter_space import ParameterSpace
import GPy
from emukit.model_wrappers.gpy_quadrature_wrappers import BaseGaussianProcessGPy, RBFGPy
from emukit.quadrature.kernels import QuadratureRBFLebesgueMeasure

# define function to integrate
def f(x):
   return (np.sum(x ** 2))

# generate input samples to create KDE
XSAMPS = 500
train_x = np.array([np.random.normal((2, 5), (1, 1)) for __ in range(XSAMPS // 2)])
train_x = np.append(train_x, [np.random.normal((-2, -5), (1, 1)) for __ in range(XSAMPS // 2)], 0)

# normalize data
mu = train_x.mean(0)
std = train_x.std(0)
train_x -= mu
train_x /= std
dimension = train_x.shape[1]
data_size = train_x.shape[0]

# define length scale of KDE estimate
LSCALE = 0.1

# create KDE of data
def kde(x, lscale=1):
   density = np.sum(norm.pdf(x[0], loc=train_x[:, 0], scale=lscale) * norm.pdf(x[1], loc=train_x[:, 1], scale=lscale))
   density /= train_x.shape[0]
   return density

# the next section of code defines the inverse PDF, mapping from uniform distributions to the KDE distribution
# the idea is to use the one more input than there are dimensions to the input,
#   using the first input the determine the cumulative density associated with the CDF
#   and the remaining components to determine offset from that percentile according to the C matrix

# Compute Cholesky factors
# these will be used to map from uniform samples to KDE samples
# they are computed using the local correlation matrix for each sample
C = np.zeros((data_size, dimension, dimension))
for imat in range(data_size):
  x = train_x[imat, 0] * mu[0] + std[0]
  y = train_x[imat, 1]* mu[0] + std[0]
  subtrain = train_x[np.logical_and(train_x[:, 0] < x + 1, train_x[:, 0] > x - 1), :]
  subtrain = subtrain[np.logical_and(subtrain[:, 1] < y + 1, subtrain[:, 1] > (y) - 1), :]
  if subtrain.shape[0] == 1:
     C[imat, :, :] = np.eye(2)
  else: C[imat, :, :] = np.linalg.cholesky(np.cov(subtrain.T)) 
  if np.any(np.isnan(C)): hey

# it seems like the covariance should depend on the length scale of the GPs. 
# So, I multiplied it by the length scale
C *= LSCALE

# Function Psi doing the transformation
def psi(u):

  # cumulative probabilities, here all equal 1/nsample for simplicity
  qprob = np.arange(data_size) / data_size

  # determine component according to first coordinate
  comp = sum(qprob < u[0]) - 1

  # determine normal according to the remaining coordinates
  Z = norm.ppf(u[1:])

  return((train_x[comp, :] + C[comp,:,:]@Z) * std + mu)



# draw samples from U]0,1[ to sample space using phi
np.random.seed(10)
u_samps = np.random.uniform(0, 1, (dimension + 1, 1000)).T
generated_samps = np.array([psi(u) for u in u_samps])
evals = np.array([f(x) for x in generated_samps])

# compare samples generated from U]0, 1[ --> KDE versus training samples
if True:
   plt.scatter(generated_samps[:, 0], generated_samps[:, 1], c='r', s=10, label='Generated')
   plt.scatter(train_x[:, 0] * std[0] + mu[0], train_x[:, 1] * std[1] + mu[1], c='k', s=10, label='Observed')
   plt.legend()
   plt.savefig('density')
   plt.clf()

# sample KDE for BQ
u_samps = np.random.uniform(0, 1, (dimension + 1, 5)).T
generated_samps = np.array([psi(u) for u in u_samps])

# evaluate function using KDE samples
evals = np.atleast_2d([f(x) for x in generated_samps]).T
DIM = u_samps.shape[1] - 1

# set parameters for BQ
delta = 1e-6 # how close to 0 and 1 to integrate
lb, ub = (delta, 1 - delta)
LENGTH = 0.1 # initial guess of length scale of BQ GP

# Begin BQ
for ii in range(200):

   # create GP model relating u to f[psi(u)]
   gpy_model = GPy.models.GPRegression(X=u_samps, Y=evals, kernel=GPy.kern.RBF(input_dim=DIM + 1, lengthscale=LENGTH, variance=1.0))

   # optimizae length scale
   gpy_model.constrain_bounded(0.001, 0.2)
   gpy_model.optimize()

   # minimize BQ acquisition function
   emukit_rbf = RBFGPy(gpy_model.kern)
   emukit_qrbf = QuadratureRBFLebesgueMeasure(emukit_rbf, integral_bounds=np.array([[lb, ub] for rr in range(DIM + 1)]))
   emukit_model = BaseGaussianProcessGPy(kern=emukit_qrbf, gpy_model=gpy_model)
   emukit_method = VanillaBayesianQuadrature(base_gp=emukit_model, X=u_samps, Y=evals)

   # compute \int f(u) du
   initial_integral_mean, initial_integral_variance = emukit_method.integrate()
   print('iteration ', ii, initial_integral_mean, initial_integral_variance)

   # find next u/x point to sample by minimizing variance reduction acquisition function
   # unfortunatly, this fast computation is not available for \int P(x) f(x) dx
   # computing the acquisition function for \int P(x) f(x) dx is much more intensive than \int f(u) du
   # that's why we needed to do this u --> x transform
   ivr_acquisition = IntegralVarianceReduction(emukit_method)
   space = ParameterSpace(emukit_method.reasonable_box_bounds.convert_to_list_of_continuous_parameters())
   optimizer = GradientAcquisitionOptimizer(space)
   u_new,_ = optimizer.optimize(ivr_acquisition)

   # add optimal u and f[psi(u)] to observed inputs and outputs
   u_samps = np.append(u_samps, np.atleast_2d(u_new), 0)
   generated_samps = np.array([psi(u) for u in u_samps])
   evals = np.append(evals, np.atleast_2d(f(psi(u_new[0]))), 0)

Question

Is there a fundamental flaw in my Bayesian Quadrature formulation? If there is not a fundamental flaw, why am I not getting the correct expectation of this integral? If there is a fundamental flaw, is it possible to correct my formulation?

$\endgroup$
2
  • $\begingroup$ Acquisition functions are used in order to select the next point in the input space ($x$) on which to evaluate the target function. But according to your introduction, the $x$'s are generated randomly from some unknown distribution - so what do you need the acquisition function for ? $\endgroup$
    – J. Delaney
    Commented Jun 1, 2022 at 13:35
  • $\begingroup$ @J.Delaney The idea is to initially sample the space before starting the algorithm, then to use the acquisition function to choose the next points to add $\endgroup$
    – kilojoules
    Commented Jun 3, 2022 at 18:54

1 Answer 1

2
+150
$\begingroup$

The problem with your approach is that the transformation you use from $\pmb u$ to $\pmb x$ is non continuous: basically it chooses one of the samples $\pmb x_i$ based on the first component of $\pmb u$ and shifts it a little bit. But the samples have no particular order, so a small change in $\pmb u$ will lead to a completely different $\pmb x$. This means that the function $f(\pmb u)$ is also highly discontinuous.

The idea behind using a Gaussian process to model an unknown function is to exploit some assumptions on the smoothness of the function, which are captured by the length scale of the covariance function. But if the function is highly discontinuous, then this doesn't make a lot of sense.

There are more straightforward and robust methods you can use to calculate the expectation of a function. Obviously the sample mean

$$ \frac{1}{N} \sum_{i=1}^N f(\pmb x_i) $$

where $\pmb x_i$ are selected randomly from the underlying distribution, is an unbiased estimator of the expectation of $f(\pmb x) $. If you want to make use of the additional data you have on the distribution of $\pmb x$, you can apply some clustering algorithm and approximate the density with a piecewise constant function, where the probability of each Voronoi cell is estimated by the fraction of samples falling in that cell. Then the estimate of the expectation would be

$$ \frac{n_k}{N} \sum_{k=1}^M f(\pmb x_k) $$

where $M$ is the number of clusters and $x_k$ are their centroids. You can extend the same approach to using a Gaussian process estimate of the function $f(\pmb x) $ (if you can assume it is smooth with some typical length scale), namely use the Gaussian process interpolations $\tilde f(\pmb x_k)$, generated from a smaller set of samples.

$\endgroup$
3
  • $\begingroup$ Thanks. I agree the discontinuity in psi goes against the assumption of the Gaussian Process. It made me think: perhaps I can interpolate between the observed samples when performing the percentile function, making that part of the process smooth w.r.t. $\pmb{x}$. Perhaps there is some way to also interpolate the $C$ matrix. With these changes, hopefully the method will converge. I will report back after trying this. $\endgroup$
    – kilojoules
    Commented Jun 6, 2022 at 12:33
  • 1
    $\begingroup$ This might make the function technically continuous, but it will still be wildly zig-zagging so it will not essentially solve the problem. Unfortunately I don't think there is any way of making a transformation that is smooth in the way that you need it to be (except for very simple distributions). $\endgroup$
    – J. Delaney
    Commented Jun 6, 2022 at 15:12
  • $\begingroup$ I'm giving you the bounty because you were the only that took a crack at this but I am not convinced that there is no good way to do this $\endgroup$
    – kilojoules
    Commented Jun 8, 2022 at 9:21

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.