Bayesian optimization improves discrete design with guided candidate generation

Acquisition Function-Guided Generation for Bayesian Optimization over Discrete Spaces

Machine Learning

Summary

Finding the best options from a set of discrete choices can be hard because you can't use traditional smooth math tricks. The authors created a method that works directly with these choices and uses a clever way to generate promising candidates by guiding a sampler with an objective to optimize. They use advanced models called diffusion and flow to do this efficiently, achieving competitive or better results on molecular design tasks. Their approach keeps domain knowledge easy to include and helps find good solutions in complex discrete spaces.

What this means in practice

  • For chemical engineers: Design molecules with better properties by generating and testing new candidates efficiently using acquisition function guided generation.
  • For machine learning engineers: Build algorithms that optimize complex discrete design problems while incorporating domain-specific kernels in the model.

Authors

Tommaso Ancilli, Thomas Wedenig, Henry Moss, Robert Peharz

Abstract

Bayesian optimization (BO) over discrete spaces is challenging due to missing gradients and the sparsity of valid solutions. Latent-space BO sidesteps both issues by fitting a surrogate in the latent space of a generative model, which however substantially complicates the direct incorporation of domain knowledge that is naturally expressed in the original discrete space. In this work, we instead combine a pretrained generative model with a surrogate defined directly in data space, where inductive biases such as domain-specific kernels can be readily incorporated. Rather than maximizing the acquisition function over the discrete design space, our method samples candidates from the generative model under a distribution tilted by the acquisition function. To recover useful gradients for sampling, we embed this tilted distribution into a diffusion bridge that yields a smooth, posterior-averaged acquisition function. This leads to \emph{acquisition function guided generation} (AFGG), which we instantiate with a categorical diffusion model and an autoregressive normalizing flow. Across a broad suite of molecular benchmarks, the diffusion model variant performs on par with leading BO methods, while the flow model variant outperforms state-of-the-art methods in most cases and remains highly competitive otherwise. Our implementation is available at https://anonymous.4open.science/r/afgg.