In-Context Batch Black-box Optimization via Flow Matching
Batch black-box optimization requires selecting many parallel queries across only a few expensive rounds, yet conventional Bayesian optimization suffers from costly joint acquisition maximization and per-task surrogate refitting. This paper proposes Flow Matching for Thompson-Sampling Refinement (FMTR), an amortized approach that offline-learns a continuous, history-conditioned proposal distribution over query locations from which variable-size batches can be directly sampled at test time with fixed parameters. By marrying in-context adaptation to a generative flow model, FMTR aims to eliminate online acquisition optimization while preserving sample efficiency.
FMTR presents a well-motivated amortization strategy for large-batch, few-round optimization that successfully combines in-context learning with continuous generative modeling. The method convincingly reduces online computational cost relative to joint-acquisition baselines, and the horizon-dependent Langevin refinement of Thompson-sampling targets offers a principled supervised alternative to trajectory-level reinforcement learning. Nevertheless, the paper openly acknowledges limitations that cap its edge: as noted in Section 6, "FMTR samples batch members independently given the history and remaining budget, without explicitly modeling within-batch redundancy."
The construction of the training target is the paper's strongest asset. The authors start with Thompson-sampling argmaxes to preserve posterior-driven exploration, refine them with a horizon-dependent Langevin pushforward toward higher values, and apply a Gibbs tilt to concentrate mass on high-value regions, yielding a clear supervised objective that avoids unstable long-horizon RL. On synthetic benchmarks "FMTR achieves the best final average rank for each function, consistently across tested dimensions," and the wall-clock comparison shows orders-of-magnitude faster batch selection than diffusion or BoTorch-based competitors. Ablations further demonstrate that removing either the Langevin refinement or the distributional tilt degrades performance, confirming both components are necessary.
The primary concern is the fragility of transfer from the synthetic pretraining prior to real-world tasks, magnified by the independent-sampling simplification. On real-world benchmarks "its advantage is less uniform, with several baselines closely matching its performance, such as TuRBO, which even obtains higher mean values on PM2.5," indicating that the RFF+trend prior does not reliably capture the structure of practical objectives. Because batch members are sampled independently given the history, there is no explicit mechanism to enforce within-batch diversity or prevent redundant evaluations, a gap the authors admit could "further improve complementarity between evaluations" if joint proposals were used. Additionally, the value-weighted target is approximated via self-normalized importance sampling over $M=200$ samples; this can suffer from high variance when the Gibbs tilt strongly reweights the refined base distribution, making the learned objective sensitive to the temperature schedule.
The evidence strongly supports the headline claim of fast, rank-competitive batch selection on synthetic functions, but it is weaker on real-world generalization. The baseline set---TuRBO, BEEBO, DiBO, GIT-BO, qLogEI, qTS-Matheron, and qUCB---is appropriate, yet the real-world evaluation spans only four tasks and shows overlapping confidence intervals where FMTR does not uniformly dominate. The critique of DiBO that "accurately learning the acquisition from the limited observations available during optimization can be difficult, potentially hindering its ability to identify promising regions" fairly highlights the benefit of offline amortization, though the same limited-data challenge applies to any pretrained model at deployment. The paper would benefit from calibrated uncertainty estimates or statistical tests beyond average ranks, especially given the small sample of real-world tasks.
Experimental details are unusually thorough: hyperparameters are listed in Table S2, pretraining task generation is fully specified in Appendix B, and the DAgger-style closed-loop pretraining algorithm is given step-by-step. However, the code is not yet accessible; the reproducibility statement only promises that "the code will be made publicly available upon acceptance." Until then, independent reproduction requires re-implementing the differentiable RFF objective sampler, the pathwise GP posterior construction via Matheron's rule, the horizon-dependent Langevin refinement, and the bucketed replay buffer with context-length stratification---a substantial engineering barrier. Deployment also relies on integrating 100-step ODE solves per query with a GMM velocity head, and minor deviations in solver step size or noise scheduling could affect empirical outcomes in ways that are not exhaustively ablated.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.