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

2025-07-05

AI Summary

In Bayesian optimization (BO) over complex structured spaces—such as molecular spaces—conventional variational autoencoder (VAE)–Gaussian process (GP) coupling leads to poor generalization and necessitates task-specific architectural customization due to strong interdependence between the VAE’s latent space and the surrogate model. To address this, we propose a decoupled BO framework that independently trains the VAE generative model and the GP surrogate, coordinating them via Bayesian updating: the VAE learns transferable structural representations, while the GP models uncertainty of the objective function in the fixed latent space. This decoupling eliminates reliance on task-specific latent geometries, enhancing both optimization efficiency and stability under limited evaluation budgets. On multiple molecular optimization benchmarks, our method identifies higher-performing candidate molecules with fewer evaluations than state-of-the-art coupled VAE–GP approaches.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Optimizing structured spaces using VAEs and Bayesian methods
• Decoupling generative models from surrogate models for better performance
• Improving molecular optimization under constrained evaluation budgets

Innovation

• Decouples VAE and Gaussian Process surrogate models
• Uses principled Bayesian update rule for combination
• Enhances molecular optimization under budget constraints