Local Latent Space Bayesian Optimization over Structured Inputs

Natalie Maus, Haydn Thomas Jones, Juston Moore, Matt Kusner, John Bradshaw, Jacob R. Gardner

Keywords: Bayesian optimization, ML for molecules

TL;DR: Performing local Bayesian optimization in a jointly trained VAE + GP surrogate improves the performance of latent space Bayesian optimization for molecules by as much as 20x.

Abstract: Bayesian optimization over the latent spaces of deep autoencoder models (DAEs) has recently emerged as a promising new approach for optimizing challenging black-box functions over structured, discrete, hard-to-enumerate search spaces (e.g., molecules). Here the DAE dramatically simplifies the search space by mapping inputs into a continuous latent space where familiar Bayesian optimization tools can be more readily applied. Despite this simplification, the latent space typically remains high-dimensional. Thus, even with a well-suited latent space, these approaches do not necessarily provide a complete solution, but may rather shift the structured optimization problem to a high-dimensional one. In this paper, we propose LOL-BO, which adapts the notion of trust regions explored in recent work on high-dimensional Bayesian optimization to the structured setting. By reformulating the encoder to function as both an encoder for the DAE globally and as a deep kernel for the surrogate model within a trust region, we better align the notion of local optimization in the latent space with local optimization in the input space. LOL-BO achieves as much as 20 times improvement over state-of-the-art latent space Bayesian optimization methods across six real-world benchmarks, demonstrating that improvement in optimization strategies is as important as developing better DAE models.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2201.11872/code)