Advanced Monte Carlo for Acquisition Sampling in Bayesian Optimization

Abstract

Optimizing complex systems usually involves costly and time-consuming experiments, where selecting the experiments to perform is fundamental. Bayesian optimization (BO) has proved to be a suitable optimization method in these situations thanks to its sample efficiency and principled way of learning from previous data, but it typically requires that experiments are sequentially performed. Fully distributed BO addresses the need for efficient parallel and asynchronous active search, especially where traditional centralized BO faces limitations concerning privacy in federated learning and resource utilization in high-performance computing settings. Boltzmann sampling is an embarrassingly parallel method that enables fully distributed BO using Monte Carlo sampling. However, it also requires sampling from a continuous acquisition function, which can be challenging even for advanced Monte Carlo methods due to its highly multimodal nature, constrained search space, and possibly numerically unstable values. We introduce a simplified version of Boltzmann sampling, and we analyze multiple Markov chain Monte Carlo (MCMC) methods with a numerically improved log EI implementation for acquisition sampling. Our experiments suggest that by introducing gradient information during MCMC sampling, methods such as the MALA or CyclicalSGLD improve acquisition sampling efficiency. Interestingly, a mixture of proposals for the Metropolis–Hastings approach proves to be effective despite its simplicity.

1. Introduction

Solving engineering and scientific problems often requires the use of complex computer models, setting up expensive prototypes, resource-intensive experimentation, or tuning the best parameters of a deep learning model. Identifying the best configuration for these processes requires evaluating various options with a trial-and-error approach. Given the high cost and time associated with these tasks, it is essential to achieve high sample efficiency. Although traditional optimization methods and evolutionary algorithms can be effective in finding the optimal configuration, they may not always be as sample-efficient as certain expensive contexts require. Carefully selecting which experiments to conduct is crucial. Bayesian optimization (BO) [1,2,3] is a method that helps streamline the process of selecting the most promising experiments and identifying optimal solutions as efficiently as possible.

BO builds a probabilistic surrogate model that approximates the unknown target function from observations. Then, an acquisition function based on the surrogate model is used to find which point should be observed next, balancing exploration and exploitation to achieve sample efficiency. This process is intrinsically sequential, and the standard BO formulation is limited to single experiment evaluations.

Alternatively, many costly problems can be performed cheaper and more efficiently using a parallel BO strategy that enables parallel experimentation, such as chemical experiments in a wet lab for drug discovery [4], multi-agent experiments [5], and autoML in a cluster [6]. Parallel BO can be achieved with a batch setting [7,8], where multiple experiments can be proposed to be carried out in parallel, including asynchronous execution [9]. These BO parallel approaches are centralized, with a single central node typically carrying out the bulk of the BO computation of updating the surrogate model and proposing the batch of experiments, which will be evaluated by many worker nodes. This hinders maintaining privacy in federated learning [10] and resource utilization when experiments are carried out with high-performance computing [11]. To solve this, BO can be adapted to work in an asynchronous decentralized fashion [12,13]. This fully distributed workflow enables the experiment to be carried out even if there is a partial failure in the distributed network since data is shared asynchronously and new experiments can be proposed independently. Figure 1 compares them.

A simple approach to achieving a fully distributed BO is replicating the BO algorithm across multiple nodes, sharing data asynchronously and, to prevent different nodes from suggesting the same exact experiments, relying on sampling approaches such as Thompson sampling [4] and Boltzmann sampling [13] with different random seeds. In particular, Boltzmann sampling has some interesting properties, such as being able to operate with existing well-understood and tested acquisition functions and offering better robustness to modeling inaccuracies [14]. Boltzmann sampling also requires setting a temperature parameter that depends on the acquisition function values. However, there is high variability in the acquisition values, since they can change considerably as the optimization progresses. This means that, in some cases, it is not feasible to select a temperature parameter that performs well throughout all the optimization steps. Since direct sampling is not possible due to not having access to the explicit form of the quantity we are sampling from, Boltzmann sampling relies on using Markov chain Monte Carlo (MCMC) methods [15], which enables efficient sampling in higher dimensions of unnormalized probabilities.

In this paper, we propose an alternative formulation of Boltzmann sampling by instead proposing to sample directly from the acquisition function, which avoids having to set a temperature parameter and has better performance across benchmarks. We analyze different MCMC methods found across the literature for performing acquisition sampling and evaluating the resulting Bayesian optimization performance. In particular, we study the main challenges of applying MCMC to acquisition sampling: (1) we deal with a constrained domain defined by the search space of BO, (2) numerical instabilities of the acquisition function prevent proper, accurate representation of the target density, and (3) acquisition functions are highly multimodal, and sampling them is challenging for MCMC methods.

1.1. Bayesian Optimization

Bayesian optimization [16] is a sequential optimization method that aims to optimize expensive functions in a sample-efficient manner. It employs a probabilistic surrogate model $p\left(f\right)$ that learns the characteristics of the function that we aim to optimize, capturing the uncertainty and providing predictive distributions, which is then used by an acquisition function $\alpha (\mathbf{x},p(f\left)\right)$ that rates the next queries. Formally, it aims to optimize an unknown objective function $f:\mathcal{X}\to \mathbb{R}$ over a domain $\mathcal{X}\subset {\mathbb{R}}^d$, by carefully selecting the queries of the function to efficiently identify the optima ${\mathbf{x}}^{*}$. At iteration t, for readability, all previously observed values as $y_{1:t}$ at queried points ${\mathbf{x}}_{1:t}$ are used to construct a probabilistic surrogate model ${\mathcal{M}}_t=p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})$. For readability, our notation will omit the subscript for observed data, defining $\mathbf{X}={\mathbf{x}}_{1:t}$ and $\mathbf{y}=y_{1:t}$, with a set of points and values defined as $D_{1:t}=\{{\mathbf{x}}_{1:t},y_{1:t}\}$.

This probabilistic surrogate model frequently takes the form of a Gaussian process (GP). It is a distribution over functions with a continuous domain, which is defined by $GP\left(\mu \right(·),k(·,·\left)\right)$, where $\mu :\mathcal{X}\to \mathbb{R}$ is the mean function and $k:\mathcal{X}\times \mathcal{X}\to \mathbb{R}$ is a covariance function or kernel that has kernel hyperparameters $\mathit{\theta }$. For the experiments, we assume $\mu (·)=0$, and the kernel k is the Matérn kernel $5/2$, defined as

$$k_{\mathrm{Matérn}}(\mathbf{x},{\mathbf{x}}^{\prime })={\sigma }_f\left(1+{\displaystyle \frac{\sqrt{5}|\mathbf{x}-{\mathbf{x}}^{\prime }|}{\mathsf{\ell }}}+{\displaystyle \frac{5{(\mathbf{x}-{\mathbf{x}}^{\prime })}^2}{3{\mathsf{\ell }}^2}}\right)exp\left(-{\displaystyle \frac{\sqrt{5}|\mathbf{x}-{\mathbf{x}}^{\prime }|}{\mathsf{\ell }}}\right)$$

(1)

with parameter ℓ as the lengthscale and the scaler parameter ${\sigma }_f$ as the amplitude parameter. These are the hyperparameters of the GP that we will need to adjust $\mathit{\theta }=\{\mathsf{\ell },{\sigma }_f\}$. Note that all methods in this paper are independent of mean and kernel choice.

Having decided both the mean $\mu (·)$ and kernel $k(·,·)$ functions, we can now use the GP posterior model to make predictions at query points ${\mathbf{x}}_q$, which are normally distributed ${\mathbf{y}}_q\sim \mathcal{N}(\mu \left({\mathbf{x}}_q\right),{\sigma }^2\left({\mathbf{x}}_q\right))$. This predictive distribution can be computed in closed form:

$$\begin{array}{cc}\hfill \mu \left({\mathbf{x}}_q\right)& =\mu \left({\mathbf{x}}_q\right)+\mathbf{k}{\left({\mathbf{x}}_q\right)}^T{\mathbf{K}}^{-1}\mathbf{y}\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill {\sigma }^2\left({\mathbf{x}}_q\right)& =k({\mathbf{x}}_q,{\mathbf{x}}_q)-k{({\mathbf{x}}_q,\mathbf{x})}^T{\mathbf{K}}^{-1}k({\mathbf{x}}_q,\mathbf{x})\hfill \end{array}$$

(3)

where the covariance matrix of the observed data points is $\mathbf{K}=k(\mathbf{X},\mathbf{X})$ for noiseless observations, or, when modeling the noise in the observations $y=f\left(\mathbf{x}\right)+ϵ$, the kernel becomes $\mathbf{K}=k(\mathbf{X},\mathbf{X})+{\sigma }_n^{2}I$, with ${\sigma }_n^2$ being the Gaussian likelihood noise variance.

With the GP as our model $\mathcal{M}$, we can now define the overall algorithm for Bayesian optimization, as shown in Algorithm 1. First, it requires some initial data. Then, it involves sequentially updating the model $\mathcal{M}$ with the current observations, choosing a query with a selection strategy and performing the observation of said query. This is repeated until the budget T is exhausted, in which it returns the best-observed value $\{{\mathbf{x}}^{*},y^{*}\}$. This is a valid recipe for sequential and batched parallel selection strategies, including asynchronous evaluations [9], which can be done by fantasizing about possible values of pending observations.

Algorithm 1 Bayesian Optimization
Require: Budget T
1: $D_{1:p}$ ← Query p initial points based on LHS
2: for $t=p\dots T$ do
3:    ${\mathcal{M}}_t$ ← Update surrogate model with $D_{1:t}$
4:    ${\mathbf{x}}_q$ ← Selection strategy using ${\mathcal{M}}_t$	▹ See Section 2
5:    $D_{1:t+1}$ ← $D_{1:t}\cup \{{\mathbf{x}}_q,f\left({\mathbf{x}}_q\right)+ϵ\}$	▹ Observe query ${\mathbf{x}}_q$
6: end for
7: return $\leftarrow \{{\mathbf{x}}^{*},y^{*}\}=D_i\mathrm{s}.\mathrm{t}.i=arg{max}_i{\mathbf{y}}_i$

Updating the model involves estimating the Gaussian process hyperparameters $\mathit{\theta }$, typically by using the log marginal likelihood. The marginal likelihood or evidence $p\left(\mathbf{y}\right|\mathbf{X},\mathit{\theta })$ is the probability of the observed data given the input locations $\mathbf{X}$ and the hyperparameters $\mathit{\theta }$. Log marginal likelihood is expressed as

$$logp\left(\mathbf{y}\right|\mathbf{X},\mathit{\theta })=-\underset{\mathrm{Data}\mathrm{fit}}{\underbrace{{\displaystyle \frac{1}{2}}{\left(\mathbf{y}-\mu \left(\mathbf{X}\right)\right)}^T{\mathbf{K}}^{-1}\left(\mathbf{y}-\mu \left(\mathbf{X}\right)\right)}}-\underset{\begin{array}{c}\mathrm{Complexity}\\ \mathrm{penalty}\end{array}}{\underbrace{{\displaystyle \frac{1}{2}}log\left|\mathbf{K}\right|}}-\underset{\mathrm{Normalization}}{\underbrace{{\displaystyle \frac{n}{2}}log2\pi }}$$

(4)

with the data term assessing how the model explains the observed data; the complexity penalty term discourages complexity that might overfit the data and normalization constant to produce a valid probability distribution. We will employ an empirical Bayes approach, where the hyperparameters are set to the most likely value, that is the maximization of the log marginal likelihood. This optimization can be done with numerical methods such as gradient descent or gradient-free methods such as L-BFGS.

Bayesian optimization requires some initial data $D_{1:p}$ so that we can start making predictions and select the next queries. Therefore, the optimization is initialized with p evaluations from a low discrepancy sequence, such as a Latin Hypercube sampling or a Sobol sequence.

1.2. Distributed BO

Most parallel versions of Bayesian optimization use a form of batch BO, where in each step, a set of query points are generated following a specific acquisition method that guarantees diversity and avoids repetitions on the query set [7,9,17,18]. As shown in Figure 1, these methods rely on a central node that has access to all the information and guarantees the diversity of query points. Furthermore, batch BO requires that queries are generated synchronously, which might be a problem in applications where different experiments might vary in time [9]. Instead, here we propose a fully distributed BO method as shown in Figure 1, which does not have any centralized node; all computations can be performed independently at each node within the distributed system, where diversity is built by construction. An advantage of this structure is that queries can be generated asynchronously as soon as new resources are available to perform experiments.

The BO algorithm from Algorithm 1 can be modified to a fully distributed system, as shown in Algorithm 2, where it shows the operations of a single node of the distributed system. Compared to the sequential BO algorithm, the only changes required are that the evaluation of initialization data is distributed across nodes, that nodes broadcast and receive their latest observations, and, most importantly, that the selection strategy is compatible with being fully distributed, that is, it ensures that the nodes are capable of observing different queries even if they have exactly the same data $D_{1:t}$ and model $\mathcal{M}$, such as in Thompson sampling [4,12] or Boltzmann sampling [13,14].

Algorithm 2 Distributed Bayesian Optimization Node
Require: Budget T
1: $D_{1:p/N}$ ← Query $p/N$ initial points based on LHS
2: Broadcast ← $D_{1:p/N}$	▹ To all other available nodes
3: for $t=p\dots T$ do
4:    $D_{1:t+1}$ ← $D_{1:t}\cup \{{\mathbf{x}}_{\mathrm{other}},{\mathbf{y}}_{\mathrm{other}}\}$	▹ Collect other node data if available
5:    ${\mathcal{M}}_t$ ← Update surrogate model with $D_{1:t}$
6:    ${\mathbf{x}}_q$ ← Selection strategy using ${\mathcal{M}}_t$
7:    $D_{1:t+1}$ ← $D_{1:t}\cup \{{\mathbf{x}}_q,f\left({\mathbf{x}}_q\right)+ϵ\}$
8:    Broadcast ← $\{{\mathbf{x}}_q,f\left({\mathbf{x}}_q\right)+ϵ\}$	▹ To all other available nodes
9: end for
10: Sync data with other nodes	▹ So that all nodes return the same optimum
11: return $\leftarrow \{{\mathbf{x}}^{*},{\mathbf{y}}^{*}\}=D_i\mathrm{s}.\mathrm{t}.i=arg{max}_i{\mathbf{y}}_i$

Figure 1 shows how the nodes only need to broadcast their queries and observed values $\{{\mathbf{x}}_i,y_i\}$ to have access to all the information available at each moment, which requires minimal communication bandwidth. However, each node can perform optimization independently. Therefore, communication can be asynchronous and resilient to network failures, as the order of queries and observations does not affect the process, and nodes can keep working when isolated from the rest of the network if needed.

2. Selection Strategies

In this section, we present several acquisition function methods to generate new queries. Traditionally, acquisition maximization has been used in Bayesian optimization as the optimal decision method. However, we have seen how it can be difficult to parallelize and suffer from model bias. In contrast, acquisition sampling methods can be trivially parallelized, and they are more robust to model mismatch [14].

2.1. Acquisition Maximization

Traditionally, query selection in BO involves greedily optimizing an acquisition function $\alpha (·)$ to select the next query ${\mathbf{x}}_{t+1}\in \mathcal{X}$, formally:

$${\mathbf{x}}_{t+1}=arg\underset{\mathbf{x}\in \mathcal{X}}{max}\alpha \left(\mathbf{x},p\left(f\right|{\mathbf{y}}_{1:t},{\mathbf{x}}_{1:t})\right)$$

(5)

Acquisition functions are designed to trade off exploitation and exploration. Exploitation involves selecting queries that improve upon the current best or that have a high predictive value $\mu \left({\mathbf{x}}_q\right)$. Exploration involves selecting queries that reduce model uncertainty or that have high uncertainties ${\sigma }^2\left({\mathbf{x}}_q\right)$, respectively.

Common acquisition functions are simple utility functions, such as Upper Confidence Bound [19] and expected improvement (EI) [2]. Other acquisition functions follow a more information theoretic approach about the optimum for query selection, such as Entropy Search [20], Predictive Entropy Search [21], and Max-value Entropy Search [22], and typically involve more compute-intensive calculations.

2.1.1. Expected Improvement

In this paper, without the loss of generality to the choice of acquisition function, we will focus our analysis on EI [2], a popular choice for acquisition function since it provides good BO efficiency in practice while involving relatively fast computations since it relies on the posterior predictive $\mu$ and ${\sigma }^2$ to rate a query point ${\mathbf{x}}_q\in \mathcal{X}$:

$$E{I}_t\left(\mathbf{x}\right)={\mathbb{E}}_{p\left({\mathbf{y}}_{t+1}\right|{\mathbf{y}}_{1:t},{\mathbf{x}}_{1:t})}\left[max(0,{\mathbf{y}}_{t+1}-{\rho }_t)\right]=\mathrm{h}\left({\displaystyle \frac{\mu \left(\mathbf{x}\right)-{\rho }_t}{\sigma \left(\mathbf{x}\right)}}\right)\sigma \left(\mathbf{x}\right)$$

(6)

where ${\rho }_t=max({\mathbf{y}}_1,\dots ,{\mathbf{y}}_t)$ is the incumbent optimum at current iteration t and $\mathrm{h}\left(\mathbf{z}\right)=\varphi \left(\mathbf{z}\right)+\mathbf{z}\Phi \left(\mathbf{z}\right)$, with $\varphi$ and $\Phi$ as the standard normal density and cumulative distribution function, respectively.

Acquisition function optimization is still challenging due to acquisition functions being non-convex. Furthermore, as more data are acquired into the GP, $E{I}_t\left(\mathbf{x}\right)$ becomes highly multimodal, which can be seen in Figure 2.

This is a challenge since optimization can be prone to get stuck in local minima. In practice, BO implementations circumvent this by using a multi-restart optimization approach, that is, multiple optimizations are performed from different initial points $\mathbf{x}\in \mathcal{X}$ with the hope of one reaching the mode of the global optimum.

Theoretically, since we use a GP, $EI\left(\mathbf{x}\right)$ is computed with a Gaussian expectation, $p\left(\mathbf{y}\right|\mathbf{x})\sim \mathcal{N}(\mu \left({\mathbf{x}}_q\right),{\sigma }^2\left({\mathbf{x}}_q\right))$, which has unbounded support, meaning that $EI\left(\mathbf{x}\right)$ values and gradients are non-zero—although there are some edge cases, such as evaluating at data points locations with zero GP noise for stationary kernels—which is a useful property to help guide the optimization towards higher $EI\left(\mathbf{x}\right)$ values. However, in practice, when $\mathbf{z}=(\mu \left(\mathbf{x}\right)-{\rho }_t)/\sigma \left(\mathbf{x}\right)$ is small, meaning that the current model is very confident that there will be marginally small chances of improvement at $\mathbf{x}$, the $EI\left(\mathbf{x}\right)$, EI values can become numerically zero. This is a problem when computing the logarithm, as can be seen in Figure 3.

To avoid this, we suggest using the log expected improvement function $LogEI\left(\mathbf{x}\right)$ as presented in [23], which introduces a numerically stable implementation to the logarithm of $EI\left(\mathbf{x}\right)$. As shown in Figure 3, this is more accurate compared to naively computing $log\left(EI\right(\mathbf{x}\left)\right)$. Since the optimum of $log\left(EI\right(\mathbf{x}\left)\right)$ is the same as that of the $EI\left(\mathbf{x}\right)$, we can use it interchangeably to find the next query. $LogEI\left(\mathbf{x}\right)$ can be computed as

$$LogEI\left(\mathbf{x}\right)=log\left(h\left({\displaystyle \frac{\mu \left(\mathbf{x}\right)-{\rho }_t}{\sigma \left(\mathbf{x}\right)}}\right)\right)+log\left(\sigma \left(\mathbf{x}\right)\right)$$

(7)

where $log\left(h\right(·\left)\right)$ can be computed in a numerically stable way with

$$log\left(h\left(z\right)\right)=\left\{\begin{array}{cc}log\left(\varphi \right(z)+z\Phi (z\left)\right)\hfill & z>-1\hfill \\ -z^2/2-c_1+\log 1\mathrm{mexp}(log(\mathrm{erfcx}(-z/\sqrt{2})\left|z\right|)+c_2\hfill & 1/\sqrt{ϵ}<z\le -1\hfill \\ -z^2/2-c_1-2log\left(\right|z\left|\right)\hfill & z\le -1/\sqrt{ϵ}\hfill \end{array}\right.$$

(8)

where $c_1=log\left(2\pi \right)/2$ and $c_2=log(\pi /2)/2$ and $ϵ$ is the numerical precision. $\log 1\mathrm{mexp}$ and $\mathrm{erfcx}$ are numerically stable implementations of $log(1-exp(\mathbf{x}\left)\right)$ and the scaled complementary error function $exp\left({\mathbf{z}}^2\right)\mathrm{erfc}\left(\mathbf{z}\right)$, respectively, and can be found in scientific software libraries, such as the Tensorflow Probability library [24] or SciPy [25].

2.1.2. Thompson Sampling

An alternative approach involves working directly with function values, such as sampling from the posterior distribution of the Gaussian process, drawing i.i.d. samples from the distribution of $\widehat{f}\sim p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})$, and using the maximizer as the next query point:

$${\mathbf{x}}_{t+1}=arg\underset{\mathbf{x}\in \mathcal{X}}{max}\widehat{f}\left(x\right)$$

(9)

This approach is known as Thompson sampling (TS), a strategy for selecting the next action in a multi-armed bandit problem but, in this case, adapted to the BO setting [4]. Similarly to other acquisition functions, TS is capable of leveraging the exploration and exploitation trade-off. The randomness in the drawn function means that points with high average value (exploitation) and points with high uncertainty (exploration) can yield high values and be the maximizers.

When sampling, $\widehat{f}$ is infinite-dimensional; thus, generating exact samples is unfeasible. For computing exact samples, we can use a discretization over the domain $\mathcal{X}$ and evaluate samples at those points. However, sampling from a multivariate Gaussian distribution has $O\left(N^3\right)$ computational complexity due to requiring a Cholesky decomposition, which limits the granularity of the discretization. It is beyond the scope of this paper, but note that there are existing approaches to reduce the computation complexity, such as sparse GP approximations [26], leveraging the use of random Fourier features [27] to approximate kernel computations or enabling continuous samples using pathwise conditioning [28].

One of the advantages of TS is that we can parallelize the BO algorithm by drawing multiple functions $\widehat{f}$ from the GP, which provides numerous maximizers that we can use to acquire parallel next queries [4].

2.2. Acquisition Sampling

Selection strategies such as the acquisition functions presented in Section 2.1.1 are spatially greedy as they only select the point that maximizes the acquisition function and ignores all information provided by the acquisition function across the optimization domain $\mathcal{X}$. This can be seen in Figure 4. Moreover, despite acquisition functions being designed to trade off exploration and exploitation to efficiently carry, they rely on the assumption that the surrogate model is good enough to predict both the expected function value and its uncertainty accurately. Since, in practice, surrogate model parameters are learned as the BO progresses, model inaccuracies are common, especially at the beginning of the BO, which might impact the theoretical optimality of some acquisition functions.

2.2.1. Boltzmann Sampling

One method that aims to use all the information provided by the acquisition function across the domain and improves model robustness is Boltzmann sampling [13], a sampling approach that uses the Boltzmann policy of the acquisition function:

$$p\left({\mathbf{x}}_{t+1}\right|y_{1:t},{\mathbf{x}}_{1:t})={\displaystyle \frac{e^{{\beta }_t\alpha \left({\mathbf{x}}_{t+1},p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})\right)}}{{\int }_{\mathbf{x}\in \mathcal{X}}{e}^{{\beta }_t\alpha \left(\mathbf{x},p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})\right)}d\mathbf{x}}}$$

(10)

This defines a probability distribution for the next query ${\mathbf{x}}_{t+1}\sim p\left({\mathbf{x}}_{t+1}\right|y_{1:t},{\mathbf{x}}_{1:t})$. Based on temperature ${\beta }_t$, it allows the model to explore even if the model is completely biased and improves robustness to model inaccuracies [14]. As we will see in Section 3, some methods can work without normalizing the probabilities, which simplifies sampling the next query.

2.2.2. Direct Acquisition Sampling

Boltzmann sampling requires tuning a parameter for the temperature ${\beta }_t$. In practice, this means that this temperature must work throughout the entire Bayesian optimization. This can be challenging due to the variability of the resulting acquisition function as the Bayesian optimization progresses. If we use the EI as an example, at the beginning of BO, due to the sparsity of the data, we can expect higher improvements and thus higher EI values. As BO observes and gets closer to the global optimum, these improvements become smaller. This shows the variability of the EI. Inspired by Boltzmann, we propose carrying the sampling directly on the acquisition function $\alpha (·)$ as an alternative, to avoid tuning ${\beta }_t$, so that the sampling probability becomes

$$p\left({\mathbf{x}}_{t+1}\right|y_{1:t},{\mathbf{x}}_{1:t})={\displaystyle \frac{\alpha \left({\mathbf{x}}_{t+1},p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})\right)-b}{{\int }_{\mathbf{x}\in \mathcal{X}}\alpha \left(\mathbf{x},p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})\right)d\mathbf{x}}}$$

(11)

with b being a bias to ensure strictly positive p values to enable sampling when using non-negative acquisition functions. For the EI acquisition function that we will be using, we only need to set $b=0$.

This formulation has good practical performance, as we will show in our experiments, without having to tune the ${\beta }_t$ parameter. Additionally, a nice property of this formulation is that since we no longer have an exponential, we can introduce the $LogEI\left(\mathbf{x}\right)$ formulation, which improves the numerical stability of $EI\left(\mathbf{x}\right)$. This is possible due to sampling methods, such as MCMC, often working with ratios of probabilities, for which, to avoid numerical precision errors, computations are typically carried out using the logarithm of the probability [29]. This improves upon the original Boltzmann sampling [13] with the added benefits of the $LogEI\left(\mathbf{x}\right)$ computations. Our target log-density probability that we will sample from becomes

$$logp\left({\mathbf{x}}_{t+1}\right|y_{1:t},{\mathbf{x}}_{1:t})\propto LogEI\left({\mathbf{x}}_{t+1},p\left(f\right|y_{1:t},{\mathbf{x}}_{1:t})\right)$$

(12)

As we will see next, some sampling methods do not require the normalization of these log probability values.

3. MCMC for Acquisition Sampling

Sampling from an unknown probability typically is performed with a Markov chain Monte Carlo (MCMC) method [15]. The idea behind it is to generate a Markov chain $({\mathbf{x}}_0,{\mathbf{x}}_1,\dots ,{\mathbf{x}}_n)$ such that the stationary distribution of the chain is the target distribution $p\left(\mathbf{x}\right)$, that is, as $n\to \infty$, it guarantees that ${\mathbf{x}}_n\sim p\left(\mathbf{x}\right)$. In our case, the target distribution is the next query probability $p\left(\mathbf{x}\right)=p\left({\mathbf{x}}_{t+1}\right|y_{1:t},{\mathbf{x}}_{1:t})$, as defined in Equations (10) and (11).

MCMC implementations often operate with log densities directly to prevent floating point precision errors. This means that, for acquisition sampling, we can seamlessly introduce $LogEI\left(\mathbf{x}\right)$, which helps in preventing numerically zero $EI\left(\mathbf{x}\right)$ values and gradients.

Despite the diverse number of MCMC methods in the literature, sampling an acquisition function for BO presents a unique setting that is not typically found in MCMC applications. First, most acquisition functions are highly multimodal, with low probability areas between modes, as showcased in Figure 2. Second, sampling is performed in a constrained domain because BO operates within fixed bounds, typically a bounding box. In this work, we study different MCMC methods in the acquisition sampling setting and their impact on the resulting performance of the Bayesian optimization method.

3.1. Metropolis–Hastings

Metropolis–Hastings (MH) [30] is a classical MCMC method that performs a random walk combined with an acceptance–rejection sampling scheme. Given a current state ${\mathbf{x}}_n$, we want to generate ${\mathbf{x}}_{n+1}$ such that it satisfies the MCMC property. For that, it first generates a candidate ${\mathbf{x}}^{\prime }$ from a proposal distribution, denoted as $q\left({\mathbf{x}}^{\prime }\right|{\mathbf{x}}_n)$. Then a ratio of acceptance,

$$r_a={\displaystyle \frac{p\left({\mathbf{x}}^{\prime }\right)q\left({\mathbf{x}}_n\right|{\mathbf{x}}^{\prime })}{p\left({\mathbf{x}}_n\right)q\left({\mathbf{x}}^{\prime }\right|{\mathbf{x}}_n)}}$$

(13)

is used to accept or reject the candidate ${\mathbf{x}}^{\prime }$. For $u\in [0,1]$, generated using a uniform, it accepts the candidate ${\mathbf{x}}_{n+1}={\mathbf{x}}^{\prime }$ if $u\le r_a$, otherwise, it rejects it by setting ${\mathbf{x}}_{n+1}={\mathbf{x}}_n$. Note that, if the proposal q is symmetric, this simplifies to $r_a=p\left({\mathbf{x}}^{\prime }\right)/p\left({\mathbf{x}}_n\right)$.

For applying MH in a constrained domain, the general approach is to automatically reject the new candidate $x^{\prime }$ if it lies outside the domain $\mathcal{X}$, that is, we assume that the target probability outside the domain is zero. For dealing with the high multimodality of $\alpha (·)$, we need to carefully select a proposal distribution q, since it needs to propose big transitions to jump between modes and smaller transitions so that the acceptance rate is high. If needed, q can be defined as mixture distribution, that enables more complex and adapted distributions.

3.2. Metropolis-Adjusted Langevin Algorithm

The Metropolis-adjusted Langevin algorithm (MALA) [31] takes steps using a discretized Langevin diffusion, defined by the stochastic differential equation

$$d{\mathbf{x}}_t={\displaystyle \frac{1}{2}}{\nabla }_{x}log\left(p\left({\mathbf{x}}_t\right)\right)dt+d{B}_t$$

(14)

where B denotes a d-dimensional Brownian motion, where in our case d corresponds to the problem dimensionality of the inputs. For discrete-time MCMC, a first-order Euler discretization is used:

$${\mathbf{x}}_{n+1}={\mathbf{x}}_n+{\displaystyle \frac{{ϵ}^2}{2}}{\nabla }_{\mathbf{x}}log\left(p\left({\mathbf{x}}_n\right)\right)+ϵu^n$$

(15)

where $u\sim \mathcal{N}(0,I)$ and $ϵ$ is the integration step size. To correct first-order integration errors from discretization, the MH acceptance $r_a$ is used. In this case, the proposal distribution is $q\left({\mathbf{x}}_{n+1}\right|{\mathbf{x}}_n)\sim \mathcal{N}(\mu ,{\sigma }^2)$, where $\mu ={\mathbf{x}}_n+{\displaystyle \frac{{ϵ}^2}{2}}{\nabla }_{\mathbf{x}}log\left(p\left({\mathbf{x}}_n\right)\right)$ and ${\sigma }^2={ϵ}^2$.

3.3. Hamiltonian Monte Carlo

Hamiltonian Monte Carlo (HMC) [32] is a particular case of MH that uses the gradients of the density function to improve the acceptance ratio. It uses an approximate Hamiltonian dynamics simulation with leapfrog numerical integration alongside a Metropolis acceptance step. The idea is to simulate the trajectory of a particle in a potential field defined by the target distribution. It introduces an auxiliary momentum variable $\rho$ to the state and draws from a joint density

$$p(\rho ,\mathbf{x})=p\left(\rho \right|\mathbf{x}\left)p\right(\mathbf{x})$$

(16)

that defines the following Hamiltonian

$$H(\rho ,\mathbf{x})=-logp\left(\rho \right|\mathbf{x})-logp(\mathbf{x})$$

(17)

which represents the total energy of the system H, with the resulting two log terms often described as the kinetic and potential energy of the system, respectively.

New state candidate $({\rho }^{\prime },{\mathbf{x}}^{\prime })$ is generated by first drawing momentum ${\rho }^{\prime }\sim \mathcal{N}(0,M)$, where M is the mass matrix, and then $({\rho }^{\prime },{\mathbf{x}}_n)$ is run under Hamilton’s equations. To solve these motion equations, a symplectic integrator is needed, a numerical integrator for Hamiltonian systems that preserves the Hamiltonian energy. Furthermore, being a particular case of MH, it needs a time-reversible integrator. We will use a leapfrog integrator [15], which is a second-order time-reversible and symplectic integrator. Given a time discretization step size $ϵ$, it simulates the exact trajectory as

$$\begin{array}{cc}\hfill {\rho }_{t+ϵ/2}& ={\rho }_t-{\displaystyle \frac{ϵ}{2}}\nabla logp\left({\mathbf{x}}_t\right)\hfill \end{array}$$

(18)

$$\begin{array}{cc}\hfill {\mathbf{x}}_{t+ϵ}& ={\mathbf{x}}_t+ϵ{\rho }_{t+ϵ/2}\hfill \end{array}$$

(19)

$$\begin{array}{cc}\hfill {\rho }_{t+ϵ}& ={\rho }_{t+ϵ/2}-{\displaystyle \frac{ϵ}{2}}\nabla logp\left({\mathbf{x}}_{t+ϵ}\right)\hfill \end{array}$$

(20)

which updates the momentum for half a step ${\rho }_{t+ϵ/2}$, then performs a full step update of position ${\mathbf{x}}_{t+ϵ}$ using the updated momentum and, finally, the momentum final half step ${\rho }_{t+ϵ}$ is updated using the new position. At the end of the leapfrog integration steps, it generates a candidate $({\rho }^{\prime },{\mathbf{x}}^{\prime })$ that is subjected to an MH acceptance such that $r_a=exp(H(\rho ,{\mathbf{x}}_n)-H({\rho }^{\prime },{\mathbf{x}}^{\prime }))$.

3.4. No U-Turn Sampler

No U-turn sampler (NUTS) [33] is an HMC method that performs multiple steps informed by first-order gradient information but, as opposed to HMC, it uses a “no U-turn” criterion when sampling trajectories, that is, if the trajectory loops back and begins retracing itself, it stops expanding the trajectory. This improves efficiency by avoiding unnecessary computations and enables an adaptive trajectory length.

3.5. Cyclical SGLD

Stochastic gradient Langevin dynamics (SGLD) is a sampling method that combines stochastic gradient descent (SGD) and Langevin dynamics to achieve faster chain convergence. It uses stochastic gradient descent to approximate the gradient calculation; where there is no closed form, it is estimated from a dataset. In this case, we could use SGD to compute the gradient of the GP with a minibatch of the dataset, which could result in a reduced computation load. However, for comparison, we have decided to use the closed-form derivative of the full GP as in the other methods. Thus, the method becomes similar to the MALA. Another advantage of SGLD is that it uses a schedule on the step size such that ${\sum }_{n=1}^{\infty }{ϵ}_n=\infty$ and ${\sum }_{n=1}^{\infty }{ϵ}_n^2<\infty$, which guarantees convergence in probability without the need for the MH acceptance $r_a$, as in the MALA.

Furthermore, for multimodal distributions, SGLD can be extended to allow the chain to move between different modes while maintaining the same convergence capabilities. This is called cyclical stochastic gradient Langevin descent (CyclicalSGLD) [34], which extends SGLD by introducing warm restart cycles. For each cycle, it restarts the step size to an initial value and then lowers the step size as discussed before. This process alternates between exploration and sampling stages within each cycle, as shown in Figure 5. Exploration involves stochastic gradient descent (SGD), with a larger step size allowing jumping between modes of the distribution. The sampling stage involves SGLD with a lower step size to properly characterize the local density of the target distribution.

4. Results

4.1. Methodology

We will evaluate the methods presented in previous sections: expected improvement maximization (MaxEI), Thompson sampling (TS), and direct acquisition sampling (AS). The AS is broken down into different MCMC sampling methods that we evaluate. The analysis will be performed for parallel queries similar to [13], that is, we perform multiple queries per GP update for TS and AS methods. We have not included Boltzmann sampling (BS) in the overall comparison as the performance was consistently worse than AS for all the experiments, as shown in Section 4.2.3. Both AS and BS use the expected improvement as the underlying acquisition function.

All methods are implemented using Python 3.10 [36], Jax 0.4.27 [29], and Blackjax 1.2.3 [35] for MCMC implementations and tinygp 0.3.0 [37] for Gaussian processes. Methods relying on gradients use automatic differentiation and parallel queries are computed using automatic vectorization.

All methods use a Gaussian process (GP) with a Matérn kernel with $\nu =5/2$ and additive Gaussian noise ${\sigma }_n=1\times {10}^{-3}$.

Initial hyperparameters of the GP are the amplitude scaler ${\sigma }_f=0.1$, lengthscale $l=1.0$, and mean $\mu =0$, which are then updated by optimizing the log marginal likelihood of the GP. Starting from 10 initial points, each method performs 150 observations. All methods use common random numbers, and experiments are repeated 10 times. Internally, the input domain for each benchmark is normalized to ${[0,1]}^d\in \mathbb{R}$, which alleviates the need to choose different parameters for initial PG and method parameters across the different functions.

MaxEI and AS methods use the $LogEI\left(\mathbf{x}\right)$ numerically stable computation [23]. TS and AS will propose a batch of five queries since its main application is in parallelization and following the experimental setup of [13].

For AS methods, we will evaluate each of the MCMC algorithms introduced within Section 2.2.1. All methods use chains of 4000 length. Note that, in AS methods, the chain length is not equivalent to the number of acquisition evaluations since methods can use multiple chains or perform multiple evaluations for the leapfrog integrator.

The specific parameters for each method are detailed next:

• Hamiltonian Monte Carlo (HMC) uses a step size of $1\times {10}^{-2}$, with five integration steps and identity mass matrix $M=I^D$.
• The Metropolis-adjusted Langevin algorithm (MALA) uses a step size of $1\times {10}^{-2}$.
• No U-turn sampler (NUTS) uses a step size of $1\times {10}^{-2}$ and identity mass matrix $M=I^D$.
• CyclicalSGLD uses 30 cycles, alternating exploration and sampling, with a 25–75% ratio, respectively, an initial step size of $1\times {10}^{-3}$, and an SGD learning rate of $1\times {10}^{-4}$.
• Mixture distribution Metropolis–Hastings (MMH) uses a proposal that combines different sizes of normal distributions:

  $$q\left({\mathbf{x}}_n\right)={\displaystyle \frac{1}{4}}\left(\mathcal{N}({\mathbf{x}}_n,0.01)+\mathcal{N}({\mathbf{x}}_n,0.1)+\mathcal{N}({\mathbf{x}}_n,0.3)+{\mathcal{U}}^d(0,1)\right)$$

  (21)

  which can be sampled following a stratified sampling strategy:

  $${\mathbf{x}}_n\sim \left\{\begin{array}{cc}\mathcal{N}({\mathbf{x}}_n,0.01)\hfill & u<0.25\hfill \\ \mathcal{N}({\mathbf{x}}_n,0.1)\hfill & 0.25<u\le 0.5\hfill \\ \mathcal{N}({\mathbf{x}}_n,0.3)\hfill & 0.5<u\le 0.75\hfill \\ {\mathcal{U}}^d(0,1)\hfill & 0.75<u\le 1.0\hfill \end{array}\right.$$

  (22)

  where $u\sim \mathcal{U}(0,1)$ is the selecting variable. This proposal achieves a good balance between modeling the locality of the density and jumping between different modes of the acquisition function, considering that the input space is normalized.

To acquire N queries from MCMC, there are two practical approaches found in the literature. One involves maintaining a single long chain and uniformly selecting N samples from the chain [38]. The other approach involves running multiple independent chains [39], where the chains can be used for convergence analysis, and allowing the chains to be run in parallel. For our implementation, we will rely on the second approach by running N independent chains and taking a single sample after a burn-in period of 4000 samples in the chain. The reason is that we can use automatic vectorization to parallelize both the chains and the computations of the target probabilities.

To help prevent any MCMC method from going out of bounds, we drew inspiration from [40], which proposes an unbounded BO with bounds dynamically increasing as needed. To reflect the current boundary in the acquisition function, they put a smooth weighting function over the acquisition function that tends to zero as it approaches the current boundary. In a similar fashion, we propose to smooth the bounds by using

$$w\left(\mathbf{x}\right)=aexp{\displaystyle \frac{-{\mathbf{x}}^2}{2c^2}}$$

(23)

with parameters $a=1$ and $c=1\times {10}^{-2}$; $w\left(\mathbf{x}\right)$ can be used to smooth a step function using the bounds of the optimization search, which maintains derivability and helps the gradient to point back towards the search domain, which is useful to prevent MCMC methods from going out of bounds. Figure 6 shows a step function and its equivalent bump function using $w\left(\mathbf{x}\right)$.

4.2. Experiments

Experiments are carried out using benchmark functions typically used in Bayesian optimization [41]. This allows us to know the global optima to better assess the performance of different selection methods and, in particular, the different MCMC methods. The figures show the progress of the incumbent, that is, the current best point found at each optimization step t. Since experiments are run multiple times, in the plots, we show the average and a shaded region corresponding to two times the standard deviation of the incumbents across the 10 repetitions. For ease of visualization, we also show the logarithm of the incumbent with respect to the global minima, helping discern the best-performing methods visually. In addition, the legend of each plot is sorted in terms of the performance ranking.

In the following subsections, we study experiments performed on benchmark functions, such as a smooth valley-shaped 2D function and multiple N-dimensional functions, and a real robot problem of an autonomous rover finding the optimal trajectory in uneven terrain.

4.2.1. Benchmark Functions

The Rosenbrock function [42] is a 2D problem where the optimum lies within a very plain valley. Finding this valley with exploration is easy; however, identifying the global solution is harder. It has the following equation:

$$f_{\mathrm{Rosenbrock}}\left(\mathbf{x}\right)=100{(x_2-x_1^2)}^2+{(x_1-1)}^2$$

(24)

with the function being evaluated within the domain $\{x_1,x_2\}\in \{[-0.5,3],[-1.5,2]\}$. As shown in Figure 7, it is easy to explore with BO and arrive at the valley, as shown by all methods finding good values very quickly. However, if we look at the logarithmic, it shows the improvements during the exploitation, with MMH finding more refined solutions.

We also evaluate using N-dimensional benchmark functions, that is, functions that are also parametrized by the number of dimensions d. This allows us to study the behavior of different methods as dimensionality increases. In particular, we will focus on functions showcasing many local minima, such as Ackley and Alpine.

The Ackley function [42] is a function that contains many local minima. It has the following equation:

$$f_{\mathrm{Ackley}}\left(\mathbf{x}\right)=-aexp\left(-b\sqrt{{\displaystyle \frac{1}{d}}\sum _{i=1}^d{x}_i^2}\right)-exp\left({\displaystyle \frac{1}{d}}\sum _{i=1}^{d}cosc{x}_i\right)+a+exp\left(1\right)$$

(25)

with recommended variables $a=20$, $b=0.2$, and $c=2\pi$ and evaluated within the domain $x_i\in [-32.768,32.768]$, and we evaluated it for the 3D, 5D, and 10D cases. Results in Figure 8 show that as dimensions increase, the performance of MaxEI and TS deteriorates compared to AS. Within the AS methods, MMH, CyclicalSGLD, and MALA methods show better performance.

We also evaluate the Alpine family of functions [41], which are two functions that also have many local minima. They have the following equations:

$$\begin{array}{cc}\hfill f_{\mathrm{Alpine}1}\left(\mathbf{x}\right)& =\sum _{i=1}^d\left|x_{i}sin\left(x_i\right)+0.1x_i\right|\hfill \end{array}$$

(26)

$$\begin{array}{cc}\hfill f_{\mathrm{Alpine}2}\left(\mathbf{x}\right)& =\prod _{i=1}^d\sqrt{x_i}sin\left(x_i\right)\hfill \end{array}$$

(27)

with the search space being $x_i\in [-10,10]$ and $x_i\in [1,10]$ respectively. Alpine1 evaluation at 5D and 10D in Figure 9 shows how MaxEI performs better. For AS methods, only CyclicalSGLD matches the performance of MaxEI. MMH is the second, followed by the MALA. However, it seems that methods are stuck in suboptimal regions. This can be seen in Alpine2, evaluated at 5D and 10D, shown in Figure 10, which has a smaller search space and is able to reach better incumbent values. In this case, the best-performing method is CyclicalSGLD, followed by the MALA and MMH.

4.2.2. Rover Trajectory Problem

In this section, we cover the problem of finding the lowest-cost trajectory of a rover in a surface map between an initial point and an endpoint, where the cost represents possible terrain irregularities or obstacles that we wish to avoid. This setting is based on optimizing rover trajectories from Wang et al. [43], and we follow the changes introduced in [14]. We define a parametrized trajectory with a B-spline that starts fixed at the starting position, ends fixed at the goal position, and, as parameters, uses two control points that will define the curvature of the B-spline. The overall dimensionality of the problem we are solving is 4D since we have to optimize the two control points that are 2D. Trajectory cost is integrated along the trajectory at a fixed interval using 1000 steps. We use two different maps, as shown in the picture inside Figure 11, with both maps sharing the same start and goal but differing on the costs (e.g., different obstacles). Results are also shown in Figure 11, where it shows that MaxEI has some trouble finding the optimal zero-cost trajectory. Meanwhile, AS and TS are capable of finding such a trajectory. In particular, MMH and the MALA are the AS methods that are capable of solving the problem.

4.2.3. Ablation on Boltzmann Sampling

Here we compare Boltzmann sampling (BS) with different temperatures to direct acquisition sampling (AS). We compared them using the same MMH sampling method, which we chose since it was the best-performing MCMC method overall. The results in Figure 12 showcase how difficult is to choose a good fixed temperature for BS. We can clearly see that there is a point in the BO wherein the BS methods underperform, which aligns with our hypothesis that a single fixed temperature is not capable of working properly throughout the optimization, whereas our proposed approach AS circumvents this limitation. We only show results on the Ackley function, but the results were consistent in the other benchmarks.

5. Discussion

The proposed experiments focused on studying the different MCMC methods when sampling from the acquisition function. Overall, out of the methods that we evaluate, CyclicalSGLD, MMH, and the MALA are the only methods that consistently rank on top.

The most underperforming methods, NUTS and HMC, suffer from one of the challenges of acquisition sampling, that is, the variability in the function that we sample as the BO progresses and its relationship to the MCMC parameters. In particular, if the mass matrix M is a poor approximation of the target distribution covariance, it will be required to compensate for other parameters to perform a more fine-grained integration (i.e., smaller step size $ϵ$ and a higher number of integration steps) to maintain numerical precision. Even when actively tuning these parameters, it is still challenging [44]. In our setup, failing to choose good parameters across all Bayesian optimization steps leads to poor performance. Out of the best-performing methods, CyclicalSGLD and the MALA share that they use the gradient information, helping to improve MCMC efficiency, and the parameters are not as critical for chain convergence or hard to choose. Finally, the most interesting result is the MMH. Metropolis–Hastings is a classic algorithm that is also quite simple compared to other methods since it only has a single parameter: the proposal distribution. In our experiments, we opted for a mixture of distributions as a proposal, with a uniform distribution allowing global jumps to change between multiple modes of the acquisition function and multiple Gaussian processes that enable traversing the locality of the probability. Experiments show that, even on higher dimensions, this mixture approach is performant.

Regarding the bounded problems, our bump function helps to keep the chain within the search space of BO. However, with high gradients, the chain might jump out of bounds directly and reject many states of the MCMC chain. As a takeaway, we propose MMH as the better alternative, since we can reject jumps out-of-bounds, avoiding the problems associated with the previous methods.

Future research of this work will involve the bounded setting of MCMC methods, which, in the context of MCMC, is quite limited in the literature, especially if we focus on highly multimodal densities. In particular, research will include modifying good-performing methods, such as the MALA and CyclicalSGLD, to include bounds or complement with methods that actively deal with bounded domains, such as hit-and-run methods [45]. Additionally, automatic tuning of the MCMC parameters using the initial N samples as burn-in could be used to improve some of the methods’ overall performance or adapt the parameters using an inner loop of Bayesian optimization [44].