Scalable First-Order Bayesian Optimization via Structured Automatic Differentiation

TL;DR

This paper introduces a scalable approach to first-order Bayesian Optimization that leverages structured automatic differentiation to efficiently incorporate gradient and Hessian information, enabling high-dimensional optimization.

Contribution

It develops a structure-aware automatic differentiation method that exploits kernel matrix structures for scalable gradient and Hessian computations in Bayesian Optimization.

Findings

Achieves $ ext{O}(n^2d)$ matrix-vector multiply for gradient observations.

Enables automatic extension to complex kernels like neural networks and spectral mixtures.

Scales first-order Bayesian Optimization to high-dimensional problems.

Abstract

Bayesian Optimization (BO) has shown great promise for the global optimization of functions that are expensive to evaluate, but despite many successes, standard approaches can struggle in high dimensions. To improve the performance of BO, prior work suggested incorporating gradient information into a Gaussian process surrogate of the objective, giving rise to kernel matrices of size $nd \times nd$ for $n$ observations in $d$ dimensions. Na\"ively multiplying with (resp. inverting) these matrices requires $\mathcal{O}(n^2d^2)$ (resp. $\mathcal{O}(n^3d^3$)) operations, which becomes infeasible for moderate dimensions and sample sizes. Here, we observe that a wide range of kernels gives rise to structured matrices, enabling an exact $\mathcal{O}(n^2d)$ matrix-vector multiply for gradient observations and $\mathcal{O}(n^2d^2)$ for Hessian observations. Beyond canonical kernel classes, we derive a programmatic approach to leveraging this type of structure for transformations and combinations of the discussed kernel classes, which constitutes a structure-aware automatic differentiation algorithm. Our methods apply to virtually all canonical kernels and automatically extend to complex kernels, like the neural network, radial basis function network, and spectral mixture kernels without any additional derivations, enabling flexible, problem-dependent modeling while scaling first-order BO to high $d$.