Technical Summary: Regime-Adaptive Bayesian Optimization via Dirichlet Process Mixtures of Gaussian Processes (RAMBO)

1. Problem Statement

Standard Bayesian Optimization (BO) relies on Gaussian Process (GP) surrogates that assume uniform smoothness and stationary noise characteristics across the entire search space. This assumption is frequently violated in scientific design problems characterized by multi-regime structures, where the objective function consists of distinct, locally coherent regions separated by sharp boundaries rather than gradual transitions.

The paper identifies three primary domains where this limitation is critical:

• Molecular Conformation Search: Rotatable bonds create distinct energy basins separated by torsional barriers; each basin is locally smooth, but the global landscape comprises hundreds of such basins with incommensurable curvature.
• Drug Discovery: Chemical landscapes are fragmented across molecular scaffolds, where different chemical families exhibit fundamentally different structure-activity relationships (SAR).
• Fusion Reactor Design: Plasma geometry variations traverse qualitatively different stability regimes, with abrupt transitions between stable and unstable configurations.

In these scenarios, a single global GP either oversmooths sharp transitions or hallucinates noise in smooth regions, leading to miscalibrated uncertainty. Existing non-stationary kernel approaches (e.g., input-dependent length scales or deep GPs) typically model smoothly varying hyperparameters and require specifying the functional form of variation a priori, failing to capture the discrete, abrupt heterogeneity of real scientific landscapes.

2. Methodology: RAMBO

The authors propose RAMBO (Regime-Adaptive Mixture Bayesian Optimization), which replaces the monolithic GP surrogate with a Dirichlet Process Mixture of Gaussian Processes (DPMM-GP). This non-parametric Bayesian framework adaptively partitions the search space into an unknown number of regimes inferred directly from data.

2.1 Generative Model

The objective function is modeled as a countable mixture of independent GPs. The generative process involves:

1. Stick-Breaking Construction: Mixture weights are generated via $\beta_k \sim \text{Beta}(1, \alpha)$, where $\alpha$ is the concentration parameter.
2. Regime Assignment: Each observation $i$ is assigned a latent regime $z_i$ drawn from a Categorical distribution based on the weights.
3. Local Modeling: Each regime $k$ is modeled by an independent GP with locally optimized hyperparameters $\theta_k = \{\sigma^2_{f,k}, \ell_k, \sigma^2_{n,k}\}$ (signal variance, length scale, and noise variance).
4. Priors: Hyperparameters follow Inverse-Gamma priors to ensure numerical stability and finite moments.

2.2 Inference via Collapsed Gibbs Sampling

To address the intractability of optimizing the discrete regime assignments and the trans-dimensional parameter space, the authors derive a collapsed Gibbs sampler.

• Analytical Marginalization: The latent function values $f$ are analytically marginalized out, reducing the state space to only the discrete assignments $z$ and hyperparameters $\Theta$. This significantly improves mixing efficiency compared to methods that sample $f$ directly (e.g., HMC).
• Sampling Assignments: Regime assignments are updated based on the Chinese Restaurant Process (CRP) prior and the conditional GP likelihood.
• Hyperparameter Updates: Hyperparameters for active regimes are updated via empirical Bayes (maximizing marginal log-likelihood using Adam) or Metropolis-Hastings steps.

2.3 Spatially-Modulated Predictive Distribution

A key innovation is the handling of the predictive distribution. Standard DPMMs aggregate regimes based on global popularity ($\pi_k$), which is input-independent. For multi-regime landscapes, this averages over irrelevant regimes at a specific test point.

• Proposed Solution: The authors introduce spatially-modulated predictive weights $w_k(x^*)$. These weights combine the global popularity ($\pi_k$) with a spatial confidence term derived from the GP posterior variance ($\sigma^{-1}_{*,k}(x^*)$).
• Justification: This weighting scheme is justified by two perspectives: (i) expected posterior responsibility under the component's own predictive, and (ii) Jeffreys' scale-invariant reference measure. It suppresses regimes that are uncertain at $x^*$ without introducing new learnable gating parameters.

2.4 Adaptive Concentration Parameter Scheduling

The concentration parameter $\alpha$ controls the propensity to create new regimes. The paper introduces an adaptive scheduling strategy (Log-Sqrt Schedule) where $\alpha_t = \alpha_0 \cdot \frac{\sqrt{t}}{\log(t+e)}$.

• Rationale: Early in optimization, data is sparse; a small $\alpha$ prevents premature fragmentation. As data accumulates, $\alpha$ increases to allow the discovery of fine-grained regime structures. This mirrors the exploration-exploitation tradeoff but operates at the model complexity level.

2.5 Acquisition Functions

The framework derives a closed-form Expected Improvement (EI) for the DPMM-GP posterior. The acquisition value is a weighted sum of the EI of each constituent GP component, where weights are the spatially-modulated $w_k(x)$. This naturally decomposes uncertainty into:

• Intra-regime variance: Aleatoric uncertainty within a specific regime.
• Inter-regime disagreement: Epistemic uncertainty arising from disagreement between different regimes.

The paper also outlines extensions to other acquisition functions (UCB, Thompson Sampling, MES, KG, PES) using the derived mixture moments.

3. Key Contributions

1. DPMM-GP Surrogate: Development of a complete DPMM-GP surrogate for BO that uses collapsed Gibbs sampling to analytically marginalize latent functions, improving inference efficiency.
2. Adaptive $\alpha$-Scheduling: Introduction of a dynamic scheduling mechanism for the concentration parameter to balance model parsimony and expressiveness throughout the optimization process.
3. Regime-Aware Acquisition: Derivation of closed-form Expected Improvement that decomposes uncertainty into intra-regime and inter-regime components, enabling robust handling of multi-regime objectives.
4. Spatial Modulation: A predictive modeling choice that introduces spatial dependence into the mixture weights without additional gating parameters, ensuring predictions are relevant to the local search space.

4. Experimental Results

The authors evaluate RAMBO on synthetic benchmarks and real-world scientific applications, comparing it against state-of-the-art baselines including standard SGP, TuRBO, SAASBO, BAxUS, ALEBO, HEBO, and others.

• Synthetic Benchmarks: On the Levy and Schwefel functions (6D and 10D), RAMBO with adaptive scheduling consistently matches or outperforms baselines. It converges significantly faster on the rugged Levy function and avoids getting trapped in suboptimal basins on the deceptive Schwefel landscape.
• Molecular Conformer Optimization (12D): RAMBO achieves a 39.73% improvement in energy reduction compared to the best baseline (SAASBO) after 200 iterations, efficiently navigating between distinct rotameric basins.
• Virtual Drug Screening (50D): In optimizing docking scores for a cancer-related protein, RAMBO achieves a 4.06% improvement in predicted binding affinity over TuRBO, successfully partitioning chemical space into scaffold-specific regimes.
• Fusion Reactor Design (80D): RAMBO discovers 3–5 regimes corresponding to distinct magnetic topologies, achieving confinement quality values approximately 51.55% higher than the best competing method. It avoids the failure modes of embedding-based methods (which assume low-rank structure) and trust-region methods (which plateau at stability boundaries).

5. Significance and Claims

The paper claims that RAMBO addresses a fundamental limitation of standard BO: the inability to model discrete heterogeneity and abrupt transitions in scientific design problems. By replacing stationary surrogates with a non-parametric mixture model, RAMBO captures the "patchwork" nature of real-world landscapes where distinct regimes exhibit locally coherent but globally incommensurable properties.

The authors emphasize that their approach does not require specifying the functional form of non-stationarity in advance. Instead, the regime boundaries and the number of regimes are inferred directly from the data. The consistent improvements across diverse, high-dimensional, and multi-modal benchmarks suggest that RAMBO provides a robust framework for accelerating scientific discovery in domains where function evaluations are expensive and the underlying physics or chemistry exhibit complex, non-stationary structures.

The work is presented as a method to reduce the number of costly simulations or experiments required to reach high-quality solutions, thereby lowering the resource and energy footprint of computational science workflows. The authors note that while the method accelerates optimization, it produces point recommendations under uncertainty, and downstream decision-makers should treat acquisition values as inputs to a deliberative process rather than ground-truth rankings, particularly in safety-critical applications.