Title: ANTIFRAGILE AND ROBUST HETEROSCEDASTIC BAYESIAN OPTIMISATION

Abstract: Bayesian Optimisation is an important decision-making tool for high-stakes applications in drug discovery and materials design. An oft-overlooked modelling consideration however is the representation of input-dependent or heteroscedastic aleatoric uncertainty. The cost of misrepresenting this uncertainty as being homoscedastic could be high in drug discovery applications where neglecting heteroscedasticity in high throughput virtual screening could lead to a failed drug discovery program. In this paper, we propose a heteroscedastic Bayesian Optimisation scheme which both represents and optimises aleatoric noise in the suggestions. We consider cases such as drug discovery where we would like to minimise or be robust to aleatoric uncertainty but also applications such as materials discovery where it may be beneficial to maximise or be antifragile to aleatoric uncertainty. Our scheme features a heteroscedastic Gaussian Process (GP) as the surrogate model in conjunction with two acquisition heuristics. First, we extend the augmented expected improvement (AEI) heuristic to the heteroscedastic setting and second, we introduce a new acquisition function, aleatoric-penalised expected improvement (ANPEI) based on a simple scalarisation of the performance and noise objective. Both methods are capable of penalising or promoting aleatoric noise in the suggestions and yield improved performance relative to a naive implementation of homoscedastic Bayesian Optimisation on toy problems as well as a real-world optimisation problem.

1. INTRODUCTION:

Bayesian Optimisation is already being utilised to make decisions in high-stakes applications such as drug discovery [1, 2, 3, 4, 5] , materials discovery [6, 7, 8] , robotics [9] , sensor placement [10] and tissue engineering [11] . In these problems heteroscedastic or input-dependent noise is rarely accounted for and the assumption of homoscedastic noise is often inappropriate. As a case study, heteroscedastic noise is the rule rather than the exception in the majority of scientific datasets. This is the case, not only in experimental datasets, but also in datasets where properties are predicted computationally. We illustrate this for molecular hydration free energies in Figure 1 using the dataset of [12] where there is a distribution of noise values and in general the noise function might be expected to grow in proportion to chemical complexity [13] . The consequences of neglecting heteroscedastic noise are illustrated using a second example in Figure 2 . The homoscedastic model will underestimate noise in certain regions which could induce a Bayesian Optimisation scheme to suggest values possessing large aleatoric noise. In an application such as high-throughput virtual screening the cost of misrepresenting the noise during the screening process could amount to a year wasted in the physical synthesis of a drug [14] .

In materials discovery, we may derive benefit from or be antifragile [15] towards high aleatoric uncertainty. In an application such as the search for high-performing perovskite solar cells, we are faced with an extremely large compositional space, with millions of potential candidates possessing high aleatoric noise for identical reproductions [16] . In this instance we may want to guide search towards a candidate possessing a high photoluminescence quantum efficiency with high aleatoric noise. If the cost of repeating material syntheses is small relative to the cost of the search, the large aleatoric noise will present opportunities to synthesise materials possessing efficiencies far above their mean values.

In this paper we present a heteroscedastic Bayesian Optimisation scheme capable of both representing and optimising aleatoric noise in the suggestions. Our contributions are:

1. The introduction of a novel combination of surrogate model and acquisition function designed to optimise heteroscedastic aleatoric uncertainty to be used in situations where one wants to minimise (be robust to) aleatoric uncertainty as well as situations where one wants to maximise (be antifragile towards) aleatoric uncertainty.

2. A demonstration of our scheme's ability to outperform naive schemes based on homoscedastic Bayesian Optimisation on toy problems in addition to a real-world optimisation problem. 

2. RELATED WORK:

The most similar work to our own is that of [17] where experiments are reported on a heteroscedastic Branin-Hoo function using the variational heteroscedastic GP approach of [18] although to the best of our knowledge this work does not consider sequential evaluations. A modification to EI, expected risk improvement is introduced in [19] and is applied to problems in robotics where robustness to aleatoric noise is desirable. [20, 21] implement heteroscedastic Bayesian Optimisation but don't introduce an acquisition function that penalises aleatoric uncertainty. [22, 23] consider the related problem of safe Bayesian Optimisation through implementing constraints in parameter space. In this instance, the goal of the algorithm is to enforce a performance threshold for each evaluation of the black-box function and so is unrelated to our problem definition. In terms of acquisition functions, [24, 25] propose principled approaches to handling aleatoric noise in the homoscedastic setting that could be extended to the heteroscedastic setting. Our primary focus in this work however, is to highlight that heteroscedasticity in the surrogate model is beneficial and so an examination of a subset of acquisition functions is sufficient for this purpose.

3. BACKGROUND:

Bayesian Optimisation features a surrogate model for the black-box objective. The surrogate model should maintain calibrated uncertainty estimates in order to guide the acquisition of new data points. GPs are often a popular choice of surrogate model in BO applications because of their ability to represent uncertainty.

3.1. GAUSSIAN PROCESSES:

In the terminology of stochastic processes we may formally define a GP as follows: Definition 1. A Gaussian Process is a collection of random variables, any finite number of which have a joint Gaussian distribution.

The random variables consist of function values f (x) at different locations x within the design space.

The GP is characterised by a mean function

EQUATION

and a covariance function

EQUATION

The process is written as follows

EQUATION

In our experiments, the prior mean function will be set to the data mean after normalisation of the inputs and standardisation of the outputs. As such, m(x) = 0 will be assumed henceforth. The covariance function computes the pairwise covariance between two random variables (function values). The covariance between a pair of output values f (x) and f (x ) is a function of an input pair x and x . As such, the kernel encodes smoothness assumptions about the latent function being modelled; similarity in input space yields outputs that are close [26] . The most widely-utilised kernel is the squared exponential (SE) kernel

k SQE (x, x ) = σ 2 f • exp -x -x 2 2 2 (4)

where σ 2 f is the signal amplitude hyperparameter (vertical lengthscale) and is the (horizontal) lengthscale hyperparameter. We use the squared exponential kernel in all experiments. For further information on Gaussian Processes the reader is referred to [27] .

3.2. BAYESIAN OPTIMISATION:



Problem Statement:

The global optimisation problem is defined as

EQUATION

where x * is the global optimiser of a black-box function f : X → Y. X is the design space and is typically a compact subset of R d . What makes this optimisation problem practically relevant in applications are the following properties:

1. Black-box Objective: We do not have the analytic form of f . We can however evaluate f pointwise anywhere in the design space X .

2. Expensive Evaluations: Choosing an input location x and evaluating f (x) takes a very long time.

3. Noise: The evaluation of a given x is a noisy process. In addition, this noise may vary across X , making the underlying process heteroscedastic.

We have a dataset of past observations D = {(x i , t i )} n i=1 consisting of observations of the black-box function f and fit a surrogate model to these datapoints. We then leverage the predictive mean as well as the uncertainty estimates of the surrogate model to guide the acquisition of the next data point x n+1 according to a heuristic known as an acquisition function.

4. HETEROSCEDASTIC BAYESIAN OPTIMISATION:

We wish to perform Bayesian Optimisation whilst optimising input-dependent aleatoric noise. In order to represent input-dependent aleatoric noise, a heteroscedastic surrogate model is required. We take the most likely heteroscedastic GP approach of [28] , adopting the notation presented there for consistency. We have a dataset D = {(x i , t i )} n i=1 in which the target values t i have been generated according to t i = f (x i ) + i . We assume independent Gaussian noise terms i ∼ N (0, σ i ) with variances given by σ i = r(x i ). In the heteroscedastic setting r is typically a non-constant function over the input domain x. In order to perform Bayesian Optimisation, we wish to model the predictive distribution P (t * | x * 1 , . . . , x * q ) at the query points x * 1 , . . . , x * q . Placing a GP prior on f and taking r(x) as the assumed noise rate function, the predictive distribution is multivariate Gaussian N (µ * , Σ * ) with mean

EQUATION

and covariance matrix

EQUATION

where

K ∈ R n×n , K ij = k(x i , x j ), K * ∈ R q×n , K * ij = k(x * i , x j ), K * * ∈ R q×q , K * * ij = k(x * i , x * j ), t = (t 1 , t 2 , . . . , t n ) T , R = diag(r) with r = (r(x 1 ), r(x 2 ), . . . , r(x n )) T , and R * = diag(r * ) with r * = (r(x * 1 ), r(x * 2 ), . . . , r(x * q )) T .

The most likely heteroscedastic GP algorithm [28] executes the following steps:

1. Estimate a homoscedastic GP, G 1 on the dataset D = {(x i , t i )} n i=1 2.

Given G 1 , we estimate the empirical noise levels for the training data using var

[t i , G 1 (x i , D)] = 0.5 (t i -E[x]) 2 forming a new dataset D = {(x i , z i )} n i=1 .

A note on the form of this variance estimator is give in Appendix B. The Bayesian Optimisation problem may be framed as

EQUATION

where the black-box objective f , to be minimised has the form

EQUATION

s(x) is, in this instance, the true noise rate function. We investigate extensions of the expected improvement [29] acquisition criterion, the principal form of which may be written in terms of the targets t and the incumbent best objective function value, η, found so far as

EQUATION

where p(t | x) is the posterior predictive marginal density of the objective function evaluated at x. (η -t) + ≡ max (0, η -t) is the improvement over the incumbent best objective function value η.

Evaluations of the objective are noisy in all of the problems we consider and so we use expected improvement with plug-in [30] , the plug-in value being the GP predictive mean [31] .

We propose two extensions to the expected improvement criterion. The first is an extension of the augmented expected improvement criterion

EQUATION

of [32] where σ n is the fixed aleatoric noise level. We extend AEI to the heteroscedastic setting by exchanging the fixed aleaotric noise level with the input-dependent one:

EQUATION

where r(x) is the predicted aleatoric uncertainty at input x under the most likely heteroscedastic GP and var[t * ] is the predictive variance of the heteroscedastic GP incorporating both aleatoric and epistemic components of the uncertainty. We also propose a simple modification to the expected improvement acquisition function that explicitly penalises regions of the input space with large aleatoric uncertainty. We call this acquisition function aleatoric noise-penalised expected improvement (ANPEI) and denote it

EQUATION

where α is a scalarisation constant which we set to 0.5 for the experiments in this paper. In the multiobjective optimisation setting a particular value of α will correspond to a point on the Pareto frontier. We use both the modification to AEI (het-AEI) and ANPEI acquisition function in conjunction with the most likely heteroscedastic GP surrogate model in the experiments that follow.

For the experiments on antifragility to aleatoric uncertainty we propose the antifragile augmented expected improvement (AAEI) acquisition function

EQUATION

)

which has the effect of promoting solutions with high aleatoric noise.

5. EXPERIMENTS ON ROBUSTNESS TO ALEATORIC UNCERTAINTY:



5.1. IMPLEMENTATION:

Experiments were run using a custom implementation of Gaussian Process regression and most likely heteroscedastic Gaussian Process regression. The code is available at https://anonymous. 4open.science/r/3361287c-6879-4b38-9153-5c9491271200/. The squared exponential kernel was chosen as the covariance function for both the homoscedastic GP as well as G 1 and G 2 of the most likely heteroscedastic GP. The lengthscales, , of the homoscedastic GP were initialised to 1.0 for each input dimension across all toy problems after standardisation of the output values following the recommendation of [33] . The signal amplitude σ 2 f was initialised to a value of 2.5. The lengthscale, , of G 2 of the most likely heteroscedastic GP [28] was initialised to 1.0, the initial noise level of G 2 was set to 1.0. The EM-like procedure required to train the most likely heteroscedastic GP was run for 10 iterations and the sample size required to construct the variance estimator producing the auxiliary dataset was 100. Hyperparameter values were obtained by optimising the marginal likelihood using the scipy implementation of the L-BFGS optimiser. The objective function in all cases is the principal objective g(x) minus one standard deviation of the ground truth noise function s(x).

5.2. 1D TOY OBJECTIVE WITH LINEAR NOISE RATE FUNCTION:

Referring to Equation 9 from section 4, in the first experiment we take a one-dimensional sin wave

EQUATION

with noise rate function s(x) = 0.25x. These functions as well as the black-box objective f (x) are shown in Figure 3 . The Bayesian Optimisation problem is designed such that the first maximum in 3(a) is to be preferred as samples from this region of the input space will have a smaller noise rate. The black-box objective in 3(c) illustrates this trade-off. In 6(a) we compare the performance of a Bayesian Optimisation scheme involving a vanilla GP in conjunction with the EI acquisition function with the most likely heteroscedastic GP in conjunction with the ANPEI and het-AEI acquisition functions. The experiment is designed to contrast the performance of a standard Bayesian Optimisation scheme against our approach in a situation where minimising aleatoric noise is desirable.

5.3. BRANIN-HOO WITH NON-LINEAR NOISE RATE FUNCTION:

In the second experiment we consider the Branin-Hoo function as g(x) with a non-linear noise rate function given by s(x) = 1.4x 1 2 + 0.3x 2 . Given that this example is a minimisation problem, the black-box objective consists of the sum of the Branin-Hoo function and the noise rate function. Contour plots of the functions are shown in Figure 4 . A comparison, in terms of the best objective function values found, between the vanilla GP and EI acquisition function with the most likely heteroscedastic GP and ANPEI and AEI acquisition functions is given in 6(b).

5.4. OPTIMISING THE PHOSPHORUS FRACTION OF SOIL:

In this real-world problem we apply heteroscedastic Bayesian optimisation to the problem of optimising the phosphorus fraction of soil. Soil phosphorus is an essential nutrient for plant growth and is widely used as a fertilizer in agriculture. While the amount of arable land worldwide is declining, global population is expanding and so is food demand. As such, understanding the availability of plant nutrients that increase crop yield is essential. To this end, [34] have curated a dataset on soil phosphorus, relating phosphorus content to variables including soil particle size, total nitrogen, organic carbon and bulk density. In this experiment, we study the relationship between bulk soil density and the phosphorus fraction, the goal being to minimise the phosphorus content of soil subject to heteroscedastic noise. We provide evidence that there is heteroscedasticity in the problem by comparing the fits of a homoscedastic GP and the most likely heteroscedastic GP in Figure 2 and provide a predictive performance comparison based on negative log predictive density values in the appendix. In this problem, we do not have access to a continuous-valued black-box function or a ground truth noise function. As such, the surrogate models were initialised with a subset of the data and the query locations selected by Bayesian Optimisation were mapped to the closest data points in the heldout data. The following kernel smoothing procedure was used to generate pseudo ground truth noise values:

1. Fit a homoscedastic GP to the full dataset. 2. At each point x i , compute the corresponding square error s 2 i = (y i -µ(x i )) 2 . 3. Estimate variances by computing a moving average of the squared errors, where the relative weight of each s 2 i was assigned with a Gaussian kernel.

The performances of homoscedastic Bayesian Optimisation using EI and AEI and heteroscedastic Bayesian Optimisation using ANPEI and AEI are compared in 5(c).

5.5. DISCUSSION:

In all robustness experiments, the most likely heteroscedastic GP and ANPEI combination/heteroscedastic GP and het-AEI combination outperform the homoscedastic GP and EI. The fact that the homoscedastic GP has no knowledge of the heteroscedasticity of the noise rate function puts it at a serious disadvantage. In the first sin wave problem, designed to highlight this point, the heteroscedastic Bayesian Optimisation scheme consistently and preferentially finds the first maximum as that which minimises aleatoric noise. In contrast the homoscedastic GP, finds it impossible to differentiate between the two maxima. The experiments provide strong evidence that modelling heteroscedasticity in Bayesian Optimisation is a more flexible approach to assuming homoscedastic noise.

6 EXPERIMENTS ON ANTIFRAGILITY TO ALEATORIC UNCERTAINTY

The implementational details remain the same as for the robustness experiments save for the fact that the acquisition functions have changed such as to promote antifragility to aleatoric uncertainty. This was achieved by changing the sign in the expression for ANPEI and using the AAEI acquisition function in place of het-AEI. The baseline homoscedastic GP and EI combination remains the same. The homoscedastic GP is incapable of representing input-dependent aleatoric uncertainty and so is unable to incorporate information about this quantity in any acquisition function that is used in conjunction with it.

The results of Figure 6 demonstrate that heteroscedastic Bayesian Optimisation is capable of seeking out solutions possessing high aleatoric noise, a desirable property in materials discovery applications.

The scalarisation heuristic used in the modified ANPEI acquisition function outperforms the AAEI acquisition function on the toy sin function example suggesting that this approach may be both valid and beneficial in its simplicity in certain scenarios.

7. CONCLUSION AND FUTURE WORK:

We have presented an approach for performing Bayesian Optimisation with the explicit goal of optimising aleatoric uncertainty in the suggestions. We posit that such an approach can prove useful for the natural sciences in the search for molecules and materials that are robust and antifragile to experimental noise and in future work we plan to apply our approach to molecular property optimisation [4] . We demonstrate concrete improvements on one and two-dimensional toy problems as well as a real-world optimisation problem and contribute an open-source implementation of the most likely heteroscedastic GP as a surrogate model for Bayesian Optimisation.

:



Figures and Tables:

Figure 1: Fig. (a) shows the frequency histogram of the calculated uncertainties against the error magnitude (kcal/mol) for the FreeSolv hydration energy dataset ([12]). Fig. (b) shows a similar plot for the experimental uncertainties.

Figure 2: Comparison of Homoscedastic and Heteroscedastic GP Fits to the Soil Phosphorus Fraction Dataset.

Estimate a second GP, G 2 on D . 4. Estimate a combined GP, G 3 on D using G 2 to predict the logarithmic noise levels r i . 5. If not converged, set G 3 to G 1 and repeat.

Figure 3: Toy 1D Problem. The toy objective in a) is corrupted with heteroscedastic noise according to the function in b). The combined objective, which when optimised maximises the sin wave subject to the minimisation of aleatoric noise, is given in c) and is obtained by subtracting the noise function from the 1D sinusoid.

Figure 4: Toy 2D Problem. The Branin-Hoo objective function in a) is corrupted by the heteroscedastic noise function in b) s(x 1 , x 2 ) = 1.4x 1 2 + 0.3x 2 . The black-box objective function c) is obtained by summing the functions in a) and b). The sum is required to penalise regions of large aleatoric noise because the objective is being minimised.

Figure 5: Results of heteroscedastic and homoscedastic Bayesian Optimisation on the 3 robustness problems considered. Error bars are computed using 10 random initialisations. The first problem is a maximisation problem whereas the second and third are minimisation problems.

Peer review (iclr2020):
The paper considers the heterogeneous noise in Bayesian optimisation. The paper utilised the existing heterogeneous Gaussian process to model the surrogate function and proposes two acquisition functions from the Expected improvement to deal with such heterogeneous noise. The proposed acquisition function is heuristic and straightforward from the existing ones, ie. augmented EI and EI. I would not count this as much in terms of novelty. The idea of dealing with heterogeneous noise in BO is interesting. The related background in Gaussian process is standard and could be omitted. The experimental section is weak. The experiments are demonstrated using low dimensional functions (1-2 dim). The experiment needs to compare with Random baseline. Under the noise setting, the Random approach performs very competitive to BO. The y-axis in the experiment, the paper has considered the objective function value + noise. The reviewer suspects that the high/low performance is due to the high level of noise (?) rather than the better objective function value. The released source code includes the Lidar, Scallop, Silverman datasets, but not the PHOSPHORUS soil. I would encourage the authors to check the source code before releasing. The paper focuses on demonstrating the heteroscedasticity in the surrogate model using [1]. The reviewer is wondering what is the performance for other heteroscedastic GP approaches, such as [2,3,4] ? Minor points: Page 4, step number 2, what is and how can we get it in a new dataset D’. [1] Kersting, Kristian, et al. *Most likely heteroscedastic Gaussian process regression.* Proceedings of the 24th international conference on Machine learning. ACM, 2007. [2] Le, Quoc V., Alex J. Smola, and Stéphane Canu. *Heteroscedastic Gaussian process regression.* Proceedings of the 22nd international conference on Machine learning. ACM, 2005. [3] Binois, M., Gramacy, R. B., & Ludkovski, M. (2018). Practical heteroscedastic gaussian process modeling for large simulation experiments. Journal of Computational and Graphical Statistics, 27(4), 808-821. [4] Lázaro-Gredilla, M., & Titsias, M. K. (2011, June). Variational Heteroscedastic Gaussian Process Regression. In ICML (pp. 841-848).