Efficient Robust Bayesian Optimization
for Arbitrary Uncertain Inputs

As illustrated in Figure 1(b), we consider an optimization of expensive black-box function: $f(x)$, where $x$ is the design parameter to be tuned. At each iteration $n$, we select a new query point $x_{n}$ according to the optimization heuristics. However, due to the stochastic nature of the process, such as machining error or execution noise, the query point is perturbed to $x_{n}^{\prime}$ before the function evaluation. Moreover, we cannot observe the exact value of $x_{n}^{\prime}$ and only have a vague probability estimation of its value: $P_{x_{n}}$. After the function evaluation, we get a noisy measurement $y=f(x_{n}^{\prime})+\zeta_{n}$, where $\zeta_{n}$ is homogeneous measurement noise sampled from $\mathcal{N}(0,\sigma^{2})$. The goal is to find an optimal design parameter $x^{*}$ that maximizes the expected function value under input uncertainty:

Depending on the specific problem and randomness source, the input distribution $P_{x}$ can be arbitrary in general and even become quite complex sometimes. Here we do not place any additional assumption on them, except assuming we can sample from these input distributions, which can be easily done by approximating it with Bayesian methods and learning a parametric probabilistic model [16]. Additionally, we assume the exact values of $x^{\prime}$ are inaccessible, which is quite common in some real-world applications, particularly in robotics and process control [25].

In this paper, we focus on finding the robust optimum with BO. Each iteration of BO involves two key steps: I) fitting a surrogate model and II) maximizing an acquisition function.

Assume $P_{x}\in\mathcal{P}_{\mathcal{X}}\subset\mathcal{P}$ are a set of distribution densities over $\mathbb{R}^{d}$, representing the distributions of the uncertain inputs. We are interested in building a GP surrogate over the probability space $\mathcal{P}$, which requires to measure the difference between the uncertain inputs.

In this work, we choose MMD as the distance measurement for the uncertain inputs because of its intrinsic connection with distance measurement in RKHS. Given a characteristic kernel $k:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ and associate RKHS $\mathcal{H}_{k}$, define the mean map $\psi:\mathcal{P}\to\mathcal{H}_{k}$ such that $\langle\psi(P),g\rangle=\mathbb{E}_{P}[g],\forall g\in\mathcal{H}_{k}$. The MMD between $P,Q\in\mathcal{P}$ is defined as:

Without any additional assumption on the input distributions, except we can get $m$ samples $\{u_{i}\}_{i=1}^{m},\{v_{i}\}_{i=1}^{m}$ from $P,Q$ respectively, MMD can be empirically estimated as follows [21]:

With the MMD kernel, our surrogate model places a prior $\mathcal{GP}(0,\hat{k}(P_{x},P_{x^{\prime}}))$ and obtain a dataset $\mathcal{D}_{n}=\{(\hat{x}_{i},y_{i})|\hat{x}_{i}\sim P_{x_{i}},i=1,2,...,n)\}$. The posterior is Gaussian with mean and variance:

where ${\bf y}_{n}:=[y_{1},\cdots,y_{n}]^{T}$, $\hat{\bf k}_{n}(P_{\ast}):=[\hat{k}(P_{\ast},P_{x_{1}}),\cdots,\hat{k}(P_{\ast},P_{x_{n}})]^{T}$ and $[\hat{\bf K}_{n}]_{ij}=\hat{k}(P_{x_{i}},P_{x_{j}})$.

To derive the posterior distribution of our robust GP surrogate, it requires estimating the MMD between each pair of inputs. Gretton et al. prove the empirical estimator in Eq. 8 approximates MMD in a bounded and asymptotic way [15]. However, the sampling size $m$ used for estimation greatly affects the approximation error and insufficient sampling leads to a high estimation variance(ref. Figure 3(a)).

Such an MMD estimation variance causes numerical instability of the covariance matrix and propagates into the posterior distribution and acquisition function, rendering the search for optimal query point a challenging task. Figure 3(b) gives an example of MMD-GP posterior with insufficient samples, which produces a noisy acquisition function and impedes the search of optima. Increasing the sampling size can help alleviate this issue. However, the computation and space complexities of the empirical MMD estimator scale quadratically with the sampling size $m$. This leaves us with a dilemma that insufficient sampling results in a highly-varied posterior while a larger sample size can occupy significant GPU memory and reduce the ability for parallel computation.

Modeling arbitrary uncertain inputs: We demonstrate MMDGP’s capabilities by employing an RKHS function as the black-box function and randomly selecting 10 samples from its input domain. Various types of input randomness are introduced into the observation and produce training datasets of $\mathcal{D}=\{(x_{i},f(x_{i}+\delta_{i}))|\delta_{i}\sim P_{x_{i}}\}_{i=1}^{10}$ with different $P_{x}$ configurations. Figure 2(a) and 2(b) compare the modeling results of a conventional GP and MMDGP under a Gaussian input uncertainty $P_{x}=\mathcal{N}(0,0.01^{2})$. We observe that the GP model appears to overfit the observed samples without recognizing the input uncertainty, whereas MMDGP properly incorporates the input randomness into its posterior.

To further examine our model’s ability under complex input uncertainty, we design the input distribution to follow a beta distribution with input-dependent variance: $P_{x}=beta(\alpha=0.5,\beta=0.5,\sigma=0.9(\sin{4\pi x}+1))$. The MMDGP posterior is shown in Figure 2(c). As the input variance $\sigma$ changes along $x$, inputs from the left and right around a given location $x_{i}$ yield different MMD distances, resulting in an asymmetric posterior (e.g., around $x=0.05$ and $x=0.42$). This suggests that MMDGP can precisely model the multimodality and asymmetry of the input uncertainty.

Moreover, we evaluated MMDGP using a step-changing Chi-squared distribution $P_{x}=\chi^{2}(g(x),\sigma=0.01)$, where $g(x)=0.5$ if $x\in[0,0.6]$, and $g(x)=7.0$ otherwise. This abrupt change in $g(x)$ significantly alters the input distribution from a sharply peaked distribution to a flat one with a long tail. Figure 2(d) illustrates that our model can accurately capture this distribution shape variation, as evidenced by the sudden posterior change around $x=0.6$. This demonstrates our model can thoroughly quantify the characteristics of complex input uncertainties.

Comparing with the other surrogate models: We also compare our model with the other surrogate models under the step-changing Chi-squared input distribution. The results are reported in Figure 7 and they demonstrate our model outperforms obviously under such a complex input uncertainty (see Appendix D.1 for more details)

Estimation variance of MMD: We first examine the variance of MMD estimation by employing two beta distributions $P=beta(\alpha=0.4,\beta=0.2,\sigma=0.1)$ and $Q=beta(\alpha=0.4,\beta=0.2,\sigma=0.1)+c$, where $c$ is an offset value. Figure 3(a) shows the empirical MMDs computed via Eq. 8 with varying sampling sizes as $Q$ moves away from $P$. We find that a sampling size of 20 is inadequate, leading to high estimation variance, and increasing the sampling size to 100 stabilizes the estimation.

We further utilize this beta distribution $P$ as the input distribution and derive the MMDGP posterior via empirical estimator in Figure 3(b). Note that the MMD instability caused by inadequate sampling subsequently engenders a fluctuating posterior and culminates in a noisy acquisition function, which prevents the acquisition optimizer (e.g., L-BFGS-B in this experiment) from identifying the optima. Although Figure 3(c) shows that this issue can be mitigated by using more samples during empirical MMD estimation, it is crucial to note that a larger sampling size significantly increases GPU memory usage because of its quadratic space complexity of $O(MNm^{2})$ ($M$ and $N$ are the sample number of training and testing, $m$ is the sampling size for MMD estimation). This limitation severely hinders parallel inference for multiple samples and slows the overall speed of posterior computation.

Table 1 summarizes the inference time of MMDGP posteriors at 512 samples with different sampling sizes. We find that, for beta distribution defined in this experiment, the Nyström MMD estimator with a sampling size of 100 and sub-sampling size of 10 already delivers a comparable result to the empirical estimator with 100 samples (as seen in the acquisition plot of Figure 3(d)). Also, the inference time is reduced from 8.117 to 0.78 seconds by enabling parallel computation for more samples. For the cases that require much more samples for MMD estimation (e.g., the input distribution is quite complex or high-dimensional), this Nyström-based acceleration can have a more pronounced impact.

Effect of Nyström estimator on optimization: To investigate the effect of Nyström estimator on optimization, we also perform an ablation study in Appendix D.2, the results in Figure 8 suggest that Nyström estimator slightly degrades the optimization performance but greatly improves the inference efficiency.

Experiment setup: To experimentally validate AIRBO’s performance, we implement our algorithm ¹¹1The code will be available on https://github.com/huawei-noah/HEBO, and more implementation details can be found in Appendix C.1. based on BoTorch [2] and employ a linear combination of multiple rational quadratic kernels [6] to compute the MMD as Eq. 9. We compare our algorithm with several baselines: 1) uGP-UCB [26] is a closely related work that employs an integral kernel to model the various input distributions. It has a quadratic inference complexity of $O(MNm^{2})$, where $M$ and $N$ are the sample numbers of the training and testing set, and $m$ indicates the sampling size of the integral kernel. 3)GP-UCB is the standard GP with UCB acquisition, which represents a broad range of existing methods that focus on non-robust optimization. 3) SKL-UCB employs symmetric Kullback-Leibler divergence to measure the distance between the uncertain inputs [20]. Its closed-form solution only exists if the input distributions are the Gaussians. 4) ERBF-UCB is the robust GP with the expected Radial Basis Function kernel proposed in [13]. It computes the expected kernel under input distribution using the Gaussian integrals. Assuming the input distributions are sub-Gaussians, this method can efficiently find the robust optimum. Since all the methods use UCB acquisition, we simply distinguish them by their surrogate names in the following tests.

1D RKHS function: We begin the optimization evaluation with an RKHS function that is widely used in previous BO works [1, 24, 10]. Figure 4(a) shows its exact global optimum resides at $x=0.892$ while the robust optimum is around $x=0.08$ when the inputs follow a Gaussian distribution $\mathcal{N}(0,0.01^{2})$. According to Figure 4(b), all the robust BO methods work well with Gaussian uncertain inputs and efficiently identify the robust optimum, but the GP-UCB stagnates at a local optimum due to its neglect of input uncertainty. Also, we notice the regret of our method decrease slightly slower than uGP works in this low-dimensional and Gaussian-input case, but later cases with higher dimension and more complex distribution show our method is more stable and efficient.

1D double-peak function: To test with more complex input uncertainty, we design a blackbox function with double peaks and set the input distribution to be a multi-modal distribution $P_{x}=beta(\alpha=0.4,\beta=0.2,\sigma=0.1)$. Figure 4(c) shows the blackbox function (black solid line) and the corresponding function expectations estimated numerically via sampling from the input distribution (i.e., the colored lines). Note the true robust optimum is around $x=0.251$ under the beta distribution, but an erroneous location at $x=0.352$ may be determined if the input uncertainty is incorrectly presumed to be Gaussian. This explains the results in Figure 4(d): the performance of SKL-UCB and ERBF-UCB are sub-optimal due to their misidentification of inputs as Gaussian variables, while our method accurately quantifies the input uncertainty and outperforms the others.

10D bumped-bowl function: we also extend our evaluation to a 10D bumped-bowl function [27] under a concatenated circular distribution. Figure 9 demonstrates AIRBO scales efficiently to high dimension and outperforms the others under complex input uncertainty(see Appendix D.3).

Robust robot pushing: To evaluate AIRBO in a real-world problem, we employ a robust robot pushing benchmark from [31], in which a ball is placed at the origin point of a 2D space and a robot learns to push it to a predefined target location $(g_{x},g_{y})$. This benchmark takes a 3-dimensional input $(r_{x},r_{y},r_{t})$, where $r_{x},r_{y}\in[-5,+5]$ are the 2D coordinates of the initial robot location and $r_{t}\in[0,30]$ controls the push duration. We set four targets in separate quadrants, i.e., $g1=(-3,-3),g_{2}=(-3,3),g_{3}=(4.3,4.3)$, and a “twin” target at $g_{3}^{\prime}=(5.1,3.0)$, and describe the input uncertainty via a two-component Gaussian Mixture Model (defined in Appendix D.4). Following [7, 10], this blackbox benchmark outputs the minimum distance to these 4 targets under squared and linear distances: $loss=\min(d^{2}(g_{1},l),d(g_{2},l),d(g_{3},l),d(g_{3}^{\prime},l))$, where $d(g_{i},l)$ is the Euclidean distance between the ball’s ending location $l$ and the $i$-th target. This produces a loss landscape as shown in Figure 5(a). Note that $g_{2}$ is a more robust target than $g_{1}$ because of its linear-form distance while pushing the ball to quadrant I is the best choice as the targets, $g_{3}$ and $g_{3}^{\prime}$, match the dual-mode pattern of the input uncertainty. According to Figure 5(b), our method obviously outperforms the others because it efficiently quantifies the multimodal input uncertainty. This can be further evidenced by the push configurations found by different algorithms in Figure 6, in which each dot represents the robot’s initial location and its color represents the push duration. We find that AIRBO successfully recognizes the targets in quadrant I as an optimal choice and frequently pushes from quadrant III to quadrant I. Moreover, the pushes started close to the origin can easily go far away under input variation, so our method learns to push the ball from a corner with a long push duration, which is more robust in this case.