Regime-Adaptive Bayesian Optimization via Dirichlet Process Mixtures of Gaussian Processes
Abstract
Standard Bayesian Optimization (BO) assumes uniform smoothness across the search space—an assumption violated in multi-regime problems such as molecular conformation search through distinct energy basins or drug discovery across heterogeneous molecular scaffolds. A single GP either oversmooths sharp transitions or hallucinates noise in smooth regions, yielding miscalibrated uncertainty. We propose RAMBO, a Dirichlet Process Mixture of Gaussian Processes that automatically discovers latent regimes during optimization, each modeled by an independent GP with locally-optimized hyperparameters. We derive collapsed Gibbs sampling that analytically marginalizes latent functions for efficient inference, and introduce adaptive concentration parameter scheduling for coarse-to-fine regime discovery. Our acquisition functions decompose uncertainty into intra-regime and inter-regime components. Experiments on synthetic benchmarks and real-world applications—including molecular conformer optimization, virtual screening for drug discovery, and fusion reactor design—demonstrate consistent improvements over state-of-the-art baselines on multi-regime objectives. Code is available at https://github.com/AnthonyZhangYan/RAMBO.
Lay Summary
Many real-world design problems — finding a low-energy shape of a molecule, screening drug candidates, or shaping a fusion reactor — boil down to searching a vast space where every attempt is expensive. Bayesian optimization is a popular automated search method, but it relies on a single statistical model that assumes the landscape behaves the same way everywhere. Real scientific landscapes do not: they are patchworks of qualitatively different "regimes" separated by sharp transitions (for example, distinct conformations of a molecule, or different magnetic topologies of a plasma). We introduce RAMBO, a search algorithm that automatically discovers these regimes from the data it has already collected, models each one with its own statistical surrogate, and uses that finer-grained picture to decide where to look next. Across synthetic test problems and three scientific applications — molecular conformer optimization, drug screening against a cancer-related protein, and stellarator fusion reactor design — RAMBO reaches better solutions with fewer expensive evaluations than current state-of-the-art methods. This makes computational scientific design faster, cheaper, and more accessible.