Prior-Fitted Diffusion Transformers Bring Non-Gaussian Thompson Sampling to Black-Box Optimization

Summary:This work introduces prior-fitted diffusion transformers that perform Thompson sampling in-context at inference time. By moving away from Gaussian-process posteriors, the method can represent non-Gaussian prior structures, including priors over unimodal functions, for continuous global optimization. The models reportedly outperform Gaussian-process-based pipelines across several black-box optimization benchmarks and exhibit scaling behavior comparable to transformers in other modalities. The ap

---

Core Idea: Replacing Gaussian-Process Posterior Inference with Generative Modeling

Thompson sampling for continuous global optimization traditionally relies on a Gaussian process. A GP builds a posterior from observed input-output pairs, samples a candidate function, and selects the function's global minimizer. This is statistically coherent, but exact posterior inference becomes expensive or unavailable for high-dimensional, non-stationary, or strongly non-Gaussian objectives.

This paper introduces prior-fitted diffusion transformers that perform Thompson sampling in-context at inference time. Instead of explicitly computing a GP posterior, the model conditions on previous observations and generates candidate functions through an iterative diffusion process. Sampling a candidate function and then optimizing it preserves the randomized exploration mechanism of Thompson sampling while replacing posterior construction with conditional generation.

This is not simply a neural-network substitute for a GP. It expands the family of function distributions that can be modeled before optimization begins. Non-Gaussian priors can represent unimodal, skewed, locally non-smooth, or otherwise irregular structures that are difficult to capture under Gaussianity. However, the flexibility is still constrained by the finite pretraining distribution and model capacity; it does not provide an exact, unbiased posterior guarantee.

Technical Mechanism: Conditional Diffusion and Candidate Minimization

The model can be viewed as a conditional diffusion process. Starting from noise, it applies a sequence of denoising updates conditioned on observed function values. Each generated sample represents a plausible candidate function rather than a single deterministic prediction. The optimizer then identifies the most promising point on each candidate and evaluates that point using the black-box objective.

This mechanism has two practical advantages. First, multiple candidate functions can be generated in parallel, allowing batched objective evaluations. Second, sampling diversity can preserve multiple basins of attraction and reduce the risk of committing to one incorrect local optimum.

The costs are equally important. Each optimization run may require multiple in-context generations and objective evaluations. Attention over a long observation history generally approaches quadratic sequence-length cost, while diffusion sampling multiplies that inference overhead across denoising steps. When evaluations of the real objective are expensive, a modest number of extra model calls can erase the optimization gain. The approach is therefore most attractive when objective evaluations are costly but can be executed in parallel batches.

Why Non-Gaussian Priors Matter

The limitation of GP-based Thompson sampling is not only computational. Gaussian priors and their posteriors tend to favor smooth responses and locally quadratic structure. They can struggle with objectives that are globally unimodal but strongly non-Gaussian, highly irregular, or sharply curved near the optimum.

The work develops practical non-Gaussian priors for general black-box functions, including priors over unimodal functions. These priors are specified during synthetic pretraining rather than through an analytically tractable posterior. The key mechanism is to learn a sampler for the prior function family, thereby avoiding explicit GP posterior inference.

Prior design nevertheless introduces bias. The closer the pretrained prior family matches the deployment task, the stronger the expected benefit; a mismatch can make the model consistently generate plausible but incorrect minimizers. In practice, the prior family should be treated as a tunable part of the optimizer rather than a fixed assumption.

Scaling Laws and the Capacity–Training Trade-Off

The paper reports scaling behavior similar to that observed in transformers for other modalities and quantifies the trade-off between model size and the number of pretraining iterations. This shifts the question from whether a single network can outperform a GP toward how reliably additional capacity and training translate into optimization gains.

Scaling laws do not guarantee monotonic improvements on every objective. Larger models increase inference memory use, latency, and deployment complexity, while additional pretraining iterations improve prior coverage at greater training cost. A smaller model trained for longer may approximate a larger model but can still create throughput pressure when many optimization requests run concurrently.

Practical comparisons should therefore include:

• improvement in the final objective per black-box evaluation;
• total model calls and wall-clock time per optimization;
• variance across random seeds rather than a single best result;
• whether performance gains survive increased inference and memory costs.

Interpreting the Benchmark Results

The authors report improved performance over Gaussian-process-based pipelines across a set of global optimization benchmarks. This indicates that diffusion-based function sampling can provide stronger optimization behavior on some tasks, but it is not evidence of universal superiority over GP methods.

A rigorous evaluation should specify whether the benchmarks cover high dimensions, discrete constraints, multimodality, noisy observations, and strict evaluation budgets. Because Thompson sampling is stochastic, single-seed results can overstate performance. Comparisons with GP should also equalize the number of objective evaluations, initialization points, parallelism, stopping conditions, and hyperparameter-tuning budgets.

If the evaluation functions resemble the synthetic pretraining distribution, the model may benefit from distributional memorization rather than transferable optimization ability. More convincing evidence would include repeated multi-seed experiments, transfer across prior families, modern Bayesian-optimization baselines with comparable tuning, and the number of black-box evaluations required per successful solution.

Deployment Guidance

The method is best explored first in offline optimization with synthetic objectives whose structure is known. Instead of merely comparing final values, test whether the model recovers the correct global basin. The next stage should include hyperparameter tuning, simulator-driven design, and experimental design with real black-box objectives.

Use three baselines: standard GP-based Thompson sampling, random search, and a mature modern Bayesian-optimization method. Normalize comparisons by objective-evaluation count, but also report GPU time, parallel throughput, and energy use. If each candidate requires a serial experiment on physical equipment, inference overhead must be tightly constrained. When objective evaluations can be batched, the ability of a diffusion transformer to generate several candidate functions becomes more valuable.

The approach is best positioned not as a wholesale replacement for GP, but as an alternative when posterior computation is difficult, the useful prior structure is non-Gaussian, and objective evaluations can be parallelized. Its central contribution is to unify black-box prior modeling, function generation, and randomized exploration within a scalable transformer framework; its main risks are mismatch between pretrained priors and real tasks and the sampling overhead of diffusion inference.

---

Source: ArXiv Machine Learning (cs.LG) (2026-11-13)

— END —

Tags: #Diffusion Transformers Follow #Thompson Sampling Follow #Black-Box Optimization Follow #Gaussian Processes Follow #Prior-Fitting Follow