Organizations: ELLIS Institute Finland · Department of Computer Science, Aalto University, Finland · Department of Computer Science, University of Manchester
Batch black-box optimization is important for expensive systems where many candidates can be evaluated in parallel, yet only a few rounds of evaluation are affordable. In this regime, effective batch construction must identify multiple promising regions while avoiding redundant evaluations, and conventional Bayesian optimization can require increasingly difficult batch-selection procedures as the batch size grows. We introduce Flow Matching for Thompson-Sampling Refinement (FMTR), an amortized approach that learns offline a continuous, history-conditioned proposal distribution from which variable-size batches can be sampled. FMTR is trained offline by conditional flow matching on Thompson-sampling argmaxes refined toward high-value regions via a horizon-dependent Langevin pushforward map, adapting the exploration-exploitation tradeoff over the optimization horizon. At test time, the model parameters remain fixed and batch generation requires neither task-specific surrogate refitting nor numerical acquisition optimization. FMTR achieves strong performance against batch BO and in-context baselines on synthetic and real-world benchmarks, while substantially reducing the computational cost of batch selection.
Figures & tables
Figure 1: Overview of FMTR and its positioning relative to related methods. FMTR directly models a query distribution for end-to-end batch selection. (Top) Pretraining and deployment with closed-loop rollouts for data generation.
Figure 2: Experiments on synthetic functions. Log simple regret over 15 runs, shown as mean with 95% CI. The last column reports the average rank across dimensions for each function.
Figure 3: Experiments on real-world functions . Top row: normalized best value. Bottom row: Top 10 normalized best value. Mean with 95% CI computed across 15 runs with random initializations. Higher is better. Last column is the average rank across all tasks for each batch size.
Figure 4: Ablation studies. Performance is normalized and averaged on the Ackley, Levy, and Rosenbrock functions with q∈{32,64} and d∈{3,5,10,15} . (a) Effect of the pretraining dataset composition. (b) Effect of the Langevin refinement and Distributional Tilt. (c) Effect of temperature scheduler. These results represent averages, individual results can be found in Figure S4 .
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure S1: Experiments on synthetic functions with q=16 . Log simple regret over 15 runs, shown as mean with 95% CI. The last column reports the average rank across dimensions for each function.
Figure S2: Experiments on synthetic functions with q=32 . Log simple regret over 15 runs, shown as mean with 95% CI. Last column reports average rank across dimensions for each function.
Figure S3: Experiments on synthetic functions with q=64 . Log simple regret over 15 runs, shown as mean with 95% CI. Last column reports average rank across dimensions for each function.
Figure S4: More detail on Ablation Studies . Rows 1–3: Log simple regret with batch size 32. Rows 4–6: Log simple regret with batch size 64. Mean with 95% CI computed across 15 runs with random initializations. Lower is better.
Figure S5: Evolution of FMTR over a complete optimization rollout. Example two-dimensional trajectory with q=16 and H=10 , shown at iterations 1,3,5,7,9, and 10 . From left to right: the accumulated context with the previous batch highlighted, TS argmaxes conditioned on the current history, their horizon-dependent Langevin refinements, and samples from the learned FMTR proposal. The background shows the objective. As observations accumulate and the horizon shrinks, the target and FMTR proposal concentrate on high-value regions.
Figure S6: Experiments on real-world functions with q∈{32,64} . Top 2 rows: normalized best value. Bottom 2 rows: Top 10 normalized best value. Mean with 95% CI computed across 15 runs with random initializations. Higher is better. Last column is the average rank across all tasks for each batch size.
Figure S7: Example of training samples used during pretraining with input dimension d=2 .
Method
d=3
d=5
d=10
d=15
FMTR
2.79±0.01
2.80±0.01
2.79±0.01
2.78±0.01
TuRBO
8.15±0.67
9.68±1.57
12.17±1.37
10.10±1.47
DiBO
180.35±3.89
156.38±1.97
205.71±3.28
179.29±1.89
BEEBO
106.70±12.28
198.17±42.50
548.54±112.60
789.35±149.53
qLogEI
93.69±16.69
199.48±23.69
350.10±56.05
536.13±66.44
qUCB
56.85±8.69
123.47±14.03
436.70±31.04
641.55±22.02
Appendix
Table S1: Runtime comparison. Wall-clock time is measured over all H=10 iterations with q=64 . The measurement unit is seconds.
Modern engineering workflows increasingly rely on massive parallel simulation, driving the need for scalable, large-batch Bayesian Optimization (BO). Existing batch BO methods, however, incur large computational cost or rely on approximations that erode batch diversity. We propose B3O (Boltzmann Batch Bayesian Optimization), a framework that reframes batch generation as a pure sampling problem: drawing samples directly from the Boltzmann distribution defined by the acquisition function avoids the bottlenecks of existing large-batch methods. Theoretically, we prove that queries sampled from this distribution incur only negligible additional regret. Empirically, B3O outperforms existing batch BO methods on standard synthetic benchmarks and adapts robustly across complex applied tasks, including multi-objective electrode design and mixed-variable race car configuration.
Maximilian Bloor, Liyuan Xu, Hrvoje Stojic +1
Imperial College London London, UK · Secondmind Cambridge, UK
Solving optimal transport (OT) on random minibatches is a common surrogate for exact OT in large-scale learning. In flow matching (FM), this surrogate is used to obtain OT-like couplings that can straighten probability paths and reduce numerical integration cost. Yet, the population-level coupling induced by repeated minibatch OT remains only partially understood. We formalize this coupling as the expected batch OT plan πk, obtained by averaging empirical OT plans over independent minibatches of size k. We then establish its large-batch consistency and, in the semidiscrete case relevant to generative modeling, derive rates for both the transport-cost bias and the convergence of πk to the OT plan. For FM, this yields a population coupling whose induced velocity field is regular enough to define a unique flow from the source to the discrete target. We finally quantify how OT batch size interacts with numerical integration in a tractable two-atom model and in synthetic and image experiments.
Samuel Boïté, Julie Delon, Kimia Nadjahi
1DMA, ENS Paris · Université Paris Cité · DI, ENS Paris, CNRS
In Bayesian optimization, Thompson sampling selects the evaluation point by sampling from the posterior distribution over the objective function maximizer. Because this sampling problem is intractable for Gaussian process (GP) surrogates, the posterior distribution is typically restricted to fixed discretizations (i.e., candidate points) that become exponentially sparse as dimensionality increases. While previous works aim to increase candidate point density through scalable GP approximations, our orthogonal approach increases density by adaptively reducing the search space during sampling. Specifically, we introduce Adaptive Candidate Thompson Sampling (ACTS), which generates candidate points in subspaces guided by the gradient of a surrogate model sample. ACTS is a simple drop-in replacement for existing TS methods -- including those that use trust regions or other local approximations -- producing better samples of maxima and improved optimization across synthetic and real-world benchmarks.