In-Context Batch Black-box Optimization via Flow Matching

Quoc-Anh Hoang Nguyen; Xinyu Zhang; Daolang Huang; Julien Martinelli; Samuel Kaski

cs.LGOct 14, 2026

In-Context Batch Black-box Optimization via Flow Matching

Authors: Quoc-Anh Hoang Nguyen, Xinyu Zhang, Daolang Huang, Julien Martinelli, Samuel Kaski

Organizations: ELLIS Institute Finland · Department of Computer Science, Aalto University, Finland · Department of Computer Science, University of Manchester

Abstract

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\in\{32,64\}$ and $d\in\{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\in\{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	$\mathbf{2.79\pm 0.01}$	$\mathbf{2.80\pm 0.01}$	$\mathbf{2.79\pm 0.01}$	$\mathbf{2.78\pm 0.01}$
TuRBO	$8.15\pm 0.67$	$9.68\pm 1.57$	$12.17\pm 1.37$	$10.10\pm 1.47$
DiBO	$180.35\pm 3.89$	$156.38\pm 1.97$	$205.71\pm 3.28$	$179.29\pm 1.89$
BEEBO	$106.70\pm 12.28$	$198.17\pm 42.50$	$548.54\pm 112.60$	$789.35\pm 149.53$
qLogEI	$93.69\pm 16.69$	$199.48\pm 23.69$	$350.10\pm 56.05$	$536.13\pm 66.44$
qUCB	$56.85\pm 8.69$	$123.47\pm 14.03$	$436.70\pm 31.04$	$641.55\pm 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.

Synthetic task prior hyperparameters
Episode component	Distribution or value
Search domain $\mathcal{X}_{\tau}$	$[-1,1]^{d}$
Dimension $d_{\tau}$	$[2,3,..,14,15]$
Batch size $q_{\tau}$	$[30,40,50,60,70,80]$
Horizon $H_{\tau}$	$[5,6,7,8,9,10]$
Initial context size $n_{0,\tau}$	$[10,20,30]$

Appendix

Table S2: Hyperparameter settings.

Black-Box Optimization

Black-box optimization finds high-quality solutions when objectives can only be queried, not directly analyzed or differentiated.

See all 176 papers

B3O: Scalable Boltzmann Batch Bayesian Optimization
Jun 29, 2026Maximilian Bloor, Liyuan Xu, Hrvoje Stojic +1Bayesian Optimization Boltzmann Sampling
Imperial College London London, UK · Secondmind Cambridge, UK
Expected Batch Optimal Transport Plans and Consequences for Flow Matching
May 12, 2026Samuel Boïté, Julie Delon, Kimia NadjahiFlow Matching Optimal Transport
1DMA, ENS Paris · Université Paris Cité · DI, ENS Paris, CNRS
Adaptive Candidate Point Thompson Sampling for High-Dimensional Bayesian Optimization
Apr 10, 2026Donney Fan, Geoff PleissBayesian Optimization Thompson Sampling
University of British Columbia Vector Institute

Jun 29, 2026J/K move · Enter open · S save

B3O: Scalable Boltzmann Batch Bayesian Optimization

Maximilian Bloor, Liyuan Xu, Hrvoje Stojic, Victor Picheny

Imperial College London London, UK · Secondmind Cambridge, UK

Bayesian Optimization Boltzmann Sampling

Open paper AskPDF arXiv

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.

Abstract

Black-Box Optimization

Explore similar work

B3O: Scalable Boltzmann Batch Bayesian Optimization

Expected Batch Optimal Transport Plans and Consequences for Flow Matching

Adaptive Candidate Point Thompson Sampling for High-Dimensional Bayesian Optimization