Bayesian Optimization via Continual Variational Last Layer Training

TL;DR

This paper introduces a novel online training method for variational Bayesian last layer neural networks, which outperform Gaussian processes and other Bayesian neural networks on complex correlation tasks, offering a scalable alternative for Bayesian optimization.

Contribution

The paper presents a new online training algorithm for variational Bayesian last layer models, connecting them to Gaussian process conditioning, and demonstrating improved performance on complex tasks.

Findings

VBLL networks outperform GPs on complex correlation tasks.

VBLLs match GPs on benchmark tasks.

Proposed method is efficient and scalable for online training.

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) and their ability to be efficiently updated online. However, the performance of GPs depends on the choice of kernel, and kernel selection for complex correlation structures is often difficult or must be made bespoke. While Bayesian neural networks (BNNs) are a promising direction for higher capacity surrogate models, they have so far seen limited use due to poor performance on some problem types. In this paper, we propose an approach which shows competitive performance on many problem types, including some that BNNs typically struggle with. We build on variational Bayesian last layers (VBLLs), and connect training of these models to exact conditioning in GPs. We exploit this connection to develop an efficient online training algorithm that interleaves conditioning and optimization. Our findings suggest that VBLL networks significantly outperform GPs and other BNN architectures on tasks with complex input correlations, and match the performance of well-tuned GPs on established benchmark tasks.