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

TL;DR

This paper proposes a decoupled approach for Bayesian optimisation in the latent space of VAEs, separating the generative and surrogate models to improve molecular optimisation performance.

Contribution

It introduces a novel decoupled framework combining a VAE and Gaussian Process with a Bayesian update, enhancing optimisation in structured spaces.

Findings

Improved candidate identification in molecular optimisation tasks.

Decoupled approach outperforms coupled models under limited evaluation budgets.

Enhanced flexibility and robustness in structured domain optimisation.

Abstract

Bayesian optimisation in the latent space of a Variational AutoEncoder (VAE) is a powerful framework for optimisation tasks over complex structured domains, such as the space of scientifically interesting 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 Gaussian Process (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.