Support-Proximity Augmented Diffusion Estimation for Offline Black-Box Optimization

Offline black-box optimization aims to discover novel designs with high property scores using only a static dataset, a task fundamentally challenged by the out-of-distribution (OOD) extrapolation problem. Existing approaches typically bifurcate into inverse methods, which struggle with the ill-posed nature of mapping scores to designs, and forward methods, which often lack the distributional expressivity to quantify uncertainty effectively. In this work, we propose \textbf{SPADE} (\textbf{S}upport-\textbf{P}roximity \textbf{A}ugmented \textbf{D}iffusion \textbf{E}stimation), a novel framework that reimagines forward surrogate modeling through the lens of conditional generative modeling. SPADE models the forward likelihood $p(y|\boldsymbol{x})$ using a diffusion model, but with two critical enhancements to tailor it for optimization: (1) a \emph{Calibrated Diffusion Estimation} module that enforces global consistency in statistical moments and pairwise rankings, and (2) a \emph{Support-Proximity Regularization} mechanism that implicitly internalizes the data manifold constraint $p(\boldsymbol{x})$ via kNN-based density estimation. Theoretically, we prove that our regularization is first-order equivalent to maximizing a Bayesian posterior with a valid design prior. Empirically, SPADE achieves state-of-the-art performance across Design-Bench tasks and an LLM data mixture optimization benchmark. Our code is available at https://github.com/HarryYoung2018/spade.

Many scientific advances hinge on finding the best design: the best protein for a drug, the best shape for a robot, the best mix of training data for a language model. But real experiments are often slow, expensive, or even dangerous, so researchers must rely on past experimental records to guess what new designs will work. The trouble is that AI models trained on such data tend to be overconfident: they recommend exotic designs that look promising on paper but fail in reality, because the model has never actually seen anything similar. We built SPADE to address this in two ways. First, instead of predicting a single score for each candidate design, it uses a diffusion model to produce a full range of likely outcomes along with how uncertain it is. Second, we taught the model to be humble: when a candidate looks unlike anything in the existing data, it automatically lowers its confidence and predicts a worse score. We prove mathematically that this humility corresponds to a principled statistical rule. Across tasks ranging from DNA sequence design to robot morphology to language-model training, SPADE consistently outperforms 23 competing methods. Building AI that knows when to be cautious is a step toward safer, more reliable scientific discovery.