Variational Last Layers for Bayesian Optimization

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.