Variational Last Layers for Bayesian Optimization

Paul Brunzema, Mikkel Jordahn, John Willes, Sebastian Trimpe, Jasper Snoek, James Harrison

Keywords: Bayesian optimization, Bayesian neural networks, Bayesian last layer

TL;DR: We explore variational Bayesian last layers as a surrogate model for Bayesian optimization to optimize functions with complicated input correlations; an area where Gaussian processes tend to struggle.

Abstract: Gaussian Processes (GPs) are widely seen as the state-of-the-art surrogate models for Bayesian optimization (BO) due to their ability to model uncertainty and their performance on tasks where correlations are easily captured, such as those defined by Euclidean metrics. However, the performance of GPs depends on the choice of kernel, and kernel selection for complex correlation structures is often difficult. While Bayesian neural networks are a promising direction for higher capacity surrogate models, they have so far seen limited use due to a combination of cost of use and poor performance. In this paper, we explore the potential of neural networks with variational Bayesian last layers (VBLLs), which offer a simple and computationally lightweight approach to Bayesian uncertainty quantification in neural networks. Our findings suggest that VBLL networks significantly outperform GPs and other BNN architectures on tasks with complicated input correlations, and match the performance of well-tuned GPs on established benchmark tasks. These results highlight their promise as an alternative surrogate model for BO.

Submission Number: 83