Return of the Latent Space COWBOYS: Re-thinking the use of VAEs for Bayesian Optimisation of Structured Spaces

Henry Moss, Sebastian W. Ober, Tom Diethe

TL;DR: Don't do Bayesian optimisation in the latent space of a VAE ....

Abstract: Bayesian optimisation in the latent space of a VAE is a powerful framework for optimisation tasks over complex structured domains, such as the space of valid molecules. However, existing approaches tightly couple the surrogate and generative models, which can lead to suboptimal performance when the latent space is not tailored to specific tasks, which in turn has led to the proposal of increasingly sophisticated algorithms. In this work, we explore a new direction, instead proposing a decoupled approach that trains a generative model and a GP surrogate separately, then combines them via a simple yet principled Bayesian update rule. This separation allows each component to focus on its strengths— structure generation from the VAE and predictive modelling by the GP. We show that our decoupled approach improves our ability to identify high-potential candidates in molecular optimisation problems under constrained evaluation budgets.

Lay Summary: The rise of Generative AI, capable of producing novel images, molecules, and engineered structures, holds the ability to fundamentally redefine how experiments are conceived, conducted, and iterated. Trained on historical experiment logs or libraries of valid designs (e.g. synthesisable molecules), these models distil the scientific community's tacit knowledge and permit generation of a rich, diverse set of informative designs on demand. Yet, while generative models have become routine in image synthesis, data augmentation, and creative content, we still lack a rigorous approach to use generative models to aid scientists and innovators in running their experimental programmes. In this work we propose one such approach.

Primary Area: Probabilistic Methods->Gaussian Processes

Keywords: Bayesian Optimisation

Submission Number: 3247