[Paper Review] Achieving Robustness to Aleatoric Uncertainty with Heteroscedastic Bayesian Optimisation

Ryan‐Rhys Griffiths, Alexander A. Aldrick+3 more|arXiv (Cornell University)|Oct 17, 2019

Advanced Multi-Objective Optimization AlgorithmsComputer Science79 references30 citations

TL;DR

This paper proposes a heteroscedastic Bayesian optimisation framework that models and minimises input-dependent aleatoric uncertainty using a heteroscedastic Gaussian process surrogate and two novel acquisition functions: Heteroscedastic Augmented Expected Improvement (HAEI) and Aleatoric Noise-Penalised Expected Improvement (ANPEI). The method outperforms homoscedastic Bayesian optimisation and random sampling on synthetic and real-world scientific datasets by explicitly penalising high-noise regions, improving robustness in materials and drug discovery applications.

ABSTRACT

Bayesian optimisation is a sample-efficient search methodology that holds great promise for accelerating drug and materials discovery programs. A frequently-overlooked modelling consideration in Bayesian optimisation strategies however, is the representation of heteroscedastic aleatoric uncertainty. In many practical applications it is desirable to identify inputs with low aleatoric noise, an example of which might be a material composition which consistently displays robust properties in response to a noisy fabrication process. In this paper, we propose a heteroscedastic Bayesian optimisation scheme capable of representing and minimising aleatoric noise across the input space. Our scheme employs a heteroscedastic Gaussian process (GP) surrogate model in conjunction with two straightforward adaptations of existing acquisition functions. First, we extend the augmented expected improvement (AEI) heuristic to the heteroscedastic setting and second, we introduce the aleatoric noise-penalised expected improvement (ANPEI) heuristic. Both methodologies are capable of penalising aleatoric noise in the suggestions and yield improved performance relative to homoscedastic Bayesian optimisation and random sampling on toy problems as well as on two real-world scientific datasets. Code is available at: \url{https://github.com/Ryan-Rhys/Heteroscedastic-BO}

Motivation & Objective

• To address the lack of robustness in Bayesian optimisation when faced with heteroscedastic aleatoric uncertainty in real-world scientific applications.
• To develop a surrogate model that explicitly captures input-dependent noise across the design space.
• To introduce acquisition functions that penalise high-aleatoric-noise regions to prioritise robust, low-noise solutions.
• To demonstrate improved performance over homoscedastic BO and random sampling on both synthetic and real-world datasets with heteroscedastic noise.
• To provide an open-source implementation for reproducibility and broader adoption in scientific optimisation.

Proposed method

• Employs a heteroscedastic Gaussian process (GP) surrogate model that jointly models the mean and variance of the black-box function.
• Adapts the augmented expected improvement (AEI) heuristic to the heteroscedastic setting, resulting in the Heteroscedastic AEI (HAEI) acquisition function.
• Introduces the Aleatoric Noise-Penalised Expected Improvement (ANPEI) acquisition function, which explicitly penalises predictions with high predicted noise variance.
• Uses the predictive variance from the heteroscedastic GP to guide exploration toward low-noise, high-performance regions of the input space.
• Applies variational inference for approximate posterior inference in the heteroscedastic GP, enabling scalable and differentiable optimisation.
• Employs a two-stage optimisation strategy: first optimise the acquisition function, then evaluate the black-box function at the selected point.

Experimental results

Research questions

• RQ1 Can a heteroscedastic GP surrogate model improve robustness in Bayesian optimisation when aleatoric noise varies across the input space?
• RQ2 Does the ANPEI acquisition function outperform homoscedastic BO and random sampling in identifying low-noise, high-performance inputs?
• RQ3 How does the proposed method perform on real-world scientific datasets with known heteroscedastic noise, such as FreeSolv and soil phosphorus?
• RQ4 What is the impact of using a heteroscedastic surrogate when the true noise structure is unknown or mischaracterised?
• RQ5 Is the performance gain from heteroscedastic modelling robust across different noise regimes, including homoscedastic and noiseless settings?

Key findings

• The ANPEI acquisition function significantly outperforms homoscedastic Bayesian optimisation and random sampling on synthetic heteroscedastic problems, achieving faster convergence to low-noise, high-performance regions.
• On the FreeSolv dataset, the proposed method successfully identifies molecular compositions with lower computational noise, improving robustness in virtual screening.
• On the soil phosphorus dataset, the heteroscedastic GP model better captures input-dependent noise compared to homoscedastic models, avoiding overconfidence in high-noise regions.
• The ablation study shows that using a heteroscedastic surrogate does not degrade performance on homoscedastic or noiseless tasks, indicating robustness to model misspecification.
• The open-source implementation on GitHub enables reproducibility and facilitates adoption in materials and drug discovery pipelines.