Improving Bayesian Optimization via Training-Aware Conditional Diffusion Models

2026-06-06

AI Summary

This work addresses the high computational cost of information-based acquisition functions in traditional Bayesian optimization, which rely on posterior sampling from Gaussian processes and struggle to efficiently approximate the distribution of global optima. To overcome this limitation, the authors propose a training-aware conditional diffusion model that directly learns the distribution of global optima and introduces a novel Diffusion-based Maximum Seeking (DMS) acquisition strategy grounded in this model. The proposed approach substantially reduces computational overhead while providing theoretical guarantees on suboptimality. Empirical evaluations across multiple benchmark tasks demonstrate that DMS consistently outperforms standard Bayesian optimization baselines, confirming its effectiveness and superiority in both efficiency and solution quality.

Abstract

Bayesian optimization (BO) is a widely used approach for black-box optimization that uses a Gaussian process (GP) as a surrogate and guides sequential evaluations via an acquisition function, with the ultimate goal of locating the global optimum $\mathbf{x}^{\star}$. To align with this goal, information-based acquisition functions such as Predictive Entropy Search (PES) model $\mathbf{x}^{\star}$ as a random variable and reduce the entropy of its distribution, but approximating this distribution via traditional GP posterior sampling is computationally expensive. To address this limitation, we leverage Conditional Diffusion Models (CDMs) to efficiently approximate the distribution of $\mathbf{x}^{\star}$ and develop BO-inherent training strategies for CDMs. Motivated by the structural properties of the CDM-learned distribution, we further develop an acquisition strategy termed Diffusion-based Mode Seeking (DMS) to guide the sequential evaluation. We establish a sub-optimality guarantee for the CDM-learned distribution and demonstrate through extensive experiments that DMS outperforms standard BO baselines.

Problem

• Bayesian optimization
• acquisition function
• global optimum
• posterior sampling
• computational cost

Innovation

• Conditional Diffusion Models
• Bayesian Optimization
• Acquisition Function
• Diffusion-based Mode Seeking
• Global Optimization