PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Context Batch Black-box Optimization via Flow Matching

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

arXiv:2610.20138v2cs.LG

TL;DR

Expensive evaluations and few affordable rounds make batch selection important, while standard Bayesian optimization can require difficult joint searches as batch size grows. FMTR learns a history-conditioned continuous proposal from refined Thompson-sampling targets, and reports strong benchmark performance at lower batch-selection cost.

  • Problem

    Expensive evaluations make effective batch construction important, while joint Bayesian-optimization acquisitions become difficult to optimize as batch size grows.

    Source: §1 · p. 1

  • Method

    FMTR learns a continuous query proposal conditioned on history and remaining budget, training it on horizon-refined Thompson-sampling argmaxes.

    Source: §1 · p. 3

  • Results

    Across synthetic and real-world benchmarks, FMTR achieves strong optimization performance at substantially lower batch-selection cost.

    Source: §1 · p. 3

  • Takeaways & Limitations

    FMTR samples requested-size batches with fixed parameters, without task-specific surrogate fitting or acquisition-function optimization.

    Source: §1 · p. 3

    FMTR samples batch members independently given history and remaining budget, without explicitly modeling within-batch redundancy.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

FMTR targets costly batch selection by learning a continuous proposal distribution conditioned on optimization history and remaining budget. It uses horizon-dependent refinement of Thompson-sampling targets and reports strong performance at lower batch-selection cost.

1 INTRODUCTION

  • Batch selection must cover promising and underexplored regions without redundancy, while expensive evaluations make effective use of each round important.Joint Bayesian-optimization acquisitions become difficult to optimize as batch size grows because they involve qd variables.Source: §1 · p. 1
  • Existing end-to-end in-context optimizers commonly map histories and finite candidate sets to single queries, rather than sampling continuous locations for batches.Source: §1 · p. 2
  • FMTR models a continuous query-location proposal conditioned on observed history and remaining budget, and samples requested-size batches with fixed parameters.Figure 1 positions FMTR as directly modeling a query distribution for end-to-end batch selection.Source: Figure 1 · p. 2, §1 · p. 3
  • Training targets refine Thompson-sampling argmaxes toward high-value regions according to the remaining horizon, preserving more early exploration and increasing late exploitation.The authors present this as an alternative to trajectory-level reinforcement learning.Source: §1 · p. 3
  • Across synthetic and real-world benchmarks, FMTR achieves strong optimization performance at substantially lower batch-selection cost.Source: §1 · p. 3

2 BACKGROUND

The background frames optimization as selecting large batches over a small number of rounds, then introduces parallel Thompson sampling and flow matching as relevant constructions.

2 BACKGROUND

  • The setting focuses on large batches and a small horizon, with each batch evaluated in parallel before the history is updated.Source: §2 · p. 3
  • Thompson sampling draws plausible objective functions conditioned on observed history and selects their maximizers; parallel TS forms batches from independent posterior paths.Different paths can place their maximizers in distinct regions that remain plausible under the posterior.Source: §2 · p. 3
  • Flow matching learns a time-dependent velocity field that transforms a source distribution into a target, with conditional training using straight-line paths between source and target points.At inference, samples from the source are transported by solving the learned field's ODE from t = 0 to t = 1.Source: §2 · p. 3

3 METHODOLOGY

FMTR learns a continuous proposal over query locations from synthetic optimization tasks and histories, then samples variable-size batches with fixed parameters. Its training targets refine Thompson-sampling locations according to the remaining horizon, and closed-loop rollouts expose the model to deployment-like histories.

3.1 PRETRAINING TASKS AND HISTORIES

  • Synthetic pretraining tasks combine randomized differentiable random Fourier feature GP paths with quadratic trends to create landscapes with clearer basins and optima.Training states include histories encountered during optimization, with varying context sizes and remaining horizons.Source: §3.1 · p. 4

3.2 QUERY TARGET DISTRIBUTION

  • Thompson-sampling targets represent uncertainty over plausible optimizer locations, using posterior paths built from the episode’s known kernel and current history.Each path is maximized independently; the known kernel avoids estimating kernel parameters from sparse histories.Source: §3.2 · p. 4
  • Thompson-sampling locations are refined using the known synthetic objective, with more objective-guided updates as fewer optimization rounds remain.Stochastic perturbations are retained during refinement, while value weighting further favors high-value regions as the horizon shrinks.Source: §3.2 · p. 5

3.3 MODEL ARCHITECTURE AND TRAINING OBJECTIVE

  • The proposal conditions on observed history and remaining-round fraction, using a transformer history encoder and a continuous conditional flow decoder.Query states are processed independently given history, and the encoded history is reused during generation.Source: §3.3 · p. 5
  • Flow matching trains on samples from the refined target, approximating its value-tilted objective with self-normalized importance sampling.The velocity head models a Gaussian mixture because an intermediate flow state can correspond to different source-target pairs and target velocities.Source: §3.3 · pp. 5–6

Appendix D.2 derives this approximation, while lastly, Section 3.4 describes history collection.

  • At deployment, FMTR transforms the requested number of independent Gaussian samples through the learned flow, with fixed parameters and shared encoded history.Batch size determines how many samples are generated.Source: Appendix D.2 · p. 6

3.4 CLOSED-LOOP PRETRAINING

  • Closed-loop pretraining mixes target-distribution and current-model batches, with the target share decreasing across rollout iterations as histories increasingly reflect FMTR’s proposals.Replay is bucketed by context size so frequently occurring history lengths do not dominate minibatches.Source: §3.4 · p. 6

4 RELATED WORK

  • Unlike batch BO methods that retain task-specific surrogates and online search, FMTR learns a fixed-parameter continuous sampler rather than a policy over supplied candidate sets.The related-work discussion also contrasts FMTR with generative BO methods that update task-specific generators as observations accumulate.Source: §4 · pp. 6–7

5 EXPERIMENTS

FMTR performs strongly on synthetic and real-world benchmarks, with less uniform gains across real-world tasks. Ablations find benefits from the trend component, Langevin refinement, value weighting, and horizon-dependent temperature scheduling.

5.1 RESULTS

  • FMTR ranks best finally on all three synthetic functions across dimensions, outperforms qTS, and has the lowest reported runtime.Evaluations use batch size q = 64; runs start from 10 observations and proceed for H = 10 rounds, with means and 95% confidence intervals reported over 15 runs.Source: Figure 2 · p. 7, §5.1 · pp. 8–9, §5 · p. 8
  • FMTR achieves the best final average rank across real-world tasks for both metrics, though its advantage varies and TuRBO has higher mean values on PM2.5.The metrics are normalized best value and Top 10 normalized best value; the latter measures identification of multiple high-value candidates.Source: Figure 3 · p. 8, §5.1 · p. 9

5.2 ABLATION STUDIES

  • Including the random trend substantially improves performance over the RFF-only pretraining variant.The trend adds smoother global structure while preserving the RFF component’s local variability.Source: Figure 4 · p. 8, §5.2 · p. 9
  • Removing either Langevin refinement or distributional tilt causes similar degradation, while removing both yields the weakest FMTR variant.The doubly ablated model still outperforms qTS, indicating an additional benefit from amortized pretraining.Source: Figure 4 · p. 8, §5.2 · p. 9
  • All three temperature schedules outperform variants without distributional tilt, while scheduler A achieves the lowest NLR in later rounds.The results indicate that the annealing profile affects the strength and timing of exploitation.Source: Figure 4 · p. 8, §5.2 · p. 9

6 CONCLUSION

  • A single pretrained model remains competitive from d = 2 to 15; within-batch redundancy is unmodeled, and extensions beyond D = 15 and to discrete domains remain future work.Joint proposals could improve complementarity, but would require structured pretraining targets and an explicit batch-level objective.Source: §6 · pp. 9–10

REPRODUCIBILITY STATEMENT

  • The paper identifies locations for experimental details, hyperparameters, training algorithms, baseline settings, and benchmark descriptions.The code is stated to be made publicly available upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

AI DISCLOSURE

  • The authors report using generative AI for methodology feedback, implementation, brainstorming, searching, and editing, and state that they reviewed AI-assisted work.They also state that LLM-generated code was verified and tested by the authors.Source: AI DISCLOSURE · p. 10

A.1 MORE EXPERIMENTS ON SYNTHETIC FUNCTIONS

Figures S1–S3 report synthetic-function results for batch sizes 16, 32, and 64, using log simple regret across 15 runs and average ranks across dimensions.

A.1 MORE EXPERIMENTS ON SYNTHETIC FUNCTIONS

  • For q = 16, Figure S1 reports log simple regret over 15 runs and average rank across dimensions for each function.Curves show means with 95% confidence intervals.Source: Figure S1 · p. 14
  • For q = 32, Figure S2 reports log simple regret over 15 runs and average rank across dimensions for each function.Curves show means with 95% confidence intervals.Source: Figure S2 · p. 15
  • For q = 64, Figure S3 reports log simple regret over 15 runs and average rank across dimensions for each function.Curves show means with 95% confidence intervals.Source: Figure S3 · p. 15

A.2 FURTHER DETAIL ON ABLATION STUDIES

The appendix explains the normalized regret metric and gives additional ablation and rollout visualizations, but the supplied passages do not report specific ablation outcomes.

A.2 FURTHER DETAIL ON ABLATION STUDIES

  • Normalized Log Simple Regret rescales log simple regret to [0, 1] per function and dimensionality before averaging across functions.The normalization uses the minimum and maximum log simple regret across methods evaluated in Section 5.2.Source: §A.2 · p. 16
  • Figure S4 shows log simple regret for batch sizes 32 and 64, averaged over 15 randomized runs with 95% confidence intervals.Rows 1–3 correspond to batch size 32 and rows 4–6 to batch size 64; lower values are better.Source: Figure S4 · p. 16, §A.2 · p. 16

A.3 VISUALIZATION ON PROCESS

  • Figure S5 follows a q = 16, H = 10 rollout as Thompson-sampling argmaxes are refined and FMTR proposals concentrate on high-value regions.The panels show iterations 1, 3, 5, 7, 9, and 10, including history, target locations, refinements, and learned proposals.Source: Figure S5 · p. 17, §A.3 · p. 17

A.4 MORE EXPERIMENTS ON REAL-WORLD APPLICATIONS

Figure S6 presents real-world-function results for batch sizes 32 and 64, including best-value and Top 10 normalized-value summaries.

A.4 MORE EXPERIMENTS ON REAL-WORLD APPLICATIONS

  • Figure S6 reports normalized best value and Top 10 normalized best value for real-world functions at q ∈ {32, 64}.Means and 95% confidence intervals use 15 random-initialization runs; the last column gives average rank across tasks for each batch size.Source: Figure S6 · p. 18

A.5 EXAMPLES OF TRAINING TASKS

The appendices illustrate training tasks and specify the synthetic prior, rollout collection, posterior-path sampling, finite-step target refinement, and value weighting used in FMTR.

A.5 EXAMPLES OF TRAINING TASKS

A.6 RUNTIME COMPARISON

  • On Ackley with q = 64 and H = 10, FMTR has the fastest execution time among the evaluated baselines across 15 independent runs.Table S1 reports wall-clock time over all 10 iterations in seconds.Source: §A.6 · p. 19, Table S1 · p. 19

B.1 SYNTHETIC OBJECTIVE PRIOR

  • The synthetic objective combines a random-Fourier-feature approximation to a sampled GP with a randomized quadratic trend that adds smooth global structure.The explicit representation makes function values and gradients available for target construction without task-specific kernel estimation.Source: §B.1 · p. 19

B.2 SAMPLING KERNEL AND TREND FUNCTION

  • Synthetic tasks randomize Matérn kernel families, lengthscale structures, rotations, trend locations and shapes, and relative component scales.The sampling choices include heterogeneous effective dimensionality and optima that may lie near the search-space boundary.Source: §B.2 · p. 20

B.3 CLOSED-LOOP PRETRAINING

  • Closed-loop pretraining increasingly collects histories from FMTR proposals while periodically augmenting the replay buffer with updated-model rollouts.Training minibatches sample a non-empty context-size bucket uniformly, then draw records from that bucket.Source: §B.3 · p. 20

C.1 PATHWISE GP POSTERIOR SAMPLING

  • Pathwise posterior samples combine an RFF prior path with a history-dependent correction, enabling differentiable paths and gradient-based maximization.The episode supplies the kernel and hyperparameters, avoiding task-specific estimation during target construction.Source: §C.1 · pp. 21–23

C.2 STOCHASTIC REFINEMENT OF THOMPSON-SAMPLING TARGETS

  • Targets use exactly K horizon-dependent refinement updates, retaining the Thompson-sampling initialization rather than assuming a stationary Langevin distribution.Reflection keeps refined points inside domain bounds, and other finite refinement procedures may suit different task families.Source: §C.2 · pp. 23–24

C.3 VALUE-WEIGHTED TARGET DISTRIBUTION

  • Value weighting adjusts the basin masses of the finite-step refinement reference law toward higher objective values, but cannot add mass to regions the reference law excludes.The temperature affects both value weighting and, in FMTR, the refinement budget that shapes the reference law.Source: §C.3 · p. 24

D MODEL ARCHITECTURE AND TRAINING DETAILS

The appendix details FMTR’s conditional proposal architecture, from variable-dimensional history and flow-state representations to its horizon-conditioned decoder and mixture output. It also explains how the decoder uses observed-history memory during generation.

D.1 INPUT REPRESENTATION AND CONDITIONAL DECODER

  • FMTR zero-pads inputs to its supported maximum dimension and pairs each location with an active-dimension mask; padded coordinates stay fixed at zero.The active d-dimensional coordinates are used for readability in Section 3.3.Source: §D.1 · p. 24
  • History observations are embedded by location and objective value, then a pre-normalized transformer with a learned global token produces contextualized observation tokens as memory.Source: §D.1 · p. 24
  • History memory is permutation-equivariant to observation order and is computed once per optimization round, then reused throughout flow generation.No positional encodings are used for observations.Source: §D.1 · p. 25
  • The flow-state embedding adds a time- and horizon-conditioned term to the state embedding, and that term vanishes at flow time t = 1.The state shifts from being dominated by the Gaussian base sample near t = 0 toward a query location near t = 1.Source: §D.1 · p. 25
  • The decoder uses cross-attention to history memory without self-attention among flow tokens; its AdaLN-Zero residual gates initialize to zero, making each block initially an identity on its query token.Flow time and remaining-round fraction condition the block’s scale, shift, and gating parameters; memory is normalized but not flow-time modulated.Source: §D.1 · p. 25
  • The final decoder representation predicts a Gaussian-mixture velocity distribution with C logits, mean vectors, and diagonal scale vectors; scales have a positive lower bound for numerical stability.Mixture weights are the softmax of the logits.Source: §D.1 · p. 25

D.2 TRAINING OBJECTIVE

The training objective uses importance weighting to approximate flow matching under a value-tilted target from samples of the refined proposal distribution. The accompanying appendix also specifies compute resources, domain handling, and the baseline methods used for evaluation.

D.2 TRAINING OBJECTIVE

  • FMTR’s weighted flow-matching loss approximates training under a value-tilted target by reweighting samples from the refined proposal distribution with self-normalized importance sampling.The tilt is proportional to exp(fτ(x)/T(h, H)); finite-sample self-normalized estimates are generally biased but consistent under standard integrability conditions.Source: §D.2 · pp. 26–27
  • Refinement defines the proposal target p1, while value-based weighting approximates training under the tilted target; objective values are used for pretraining supervision, not as inputs to pθ(x | D, α).Source: §D.2 · p. 27

D.3 COMPUTATIONAL RESOURCES

  • Training took approximately 19 hours per run on one NVIDIA H100 80GB HBM3 GPU, and evaluation used Tesla V100-SXM2-32GB GPUs.Source: §D.3 · p. 27

E EXPERIMENTAL DETAILS

  • For evaluation, observed query points are mapped to the normalized domain [−1, 1]d for FMTR generation, then mapped back to the original domain for function evaluation.Source: E EXPERIMENTAL DETAILS · p. 27

E.1 BASELINES

  • The baseline set spans local trust-region GP search, task-fitted diffusion proposals, pretrained-surrogate search, and GP-based batch acquisitions, with both joint and independent-path batch rules.TuRBO uses five trust regions; BEEBO uses T = 2 before switching to T = 0 in the final round.Source: §E.1 · pp. 27–28

E.2 SYNTHETIC BENCHMARK FUNCTIONS

The synthetic evaluation covers six standard benchmark functions spanning multimodal and complementary nonconvex optimization landscapes.

E.2 SYNTHETIC BENCHMARK FUNCTIONS

  • The synthetic suite contains six standard benchmark functions with different optimization landscapes.Source: §E.2 · p. 28
  • Ackley, Levy, and Rastrigin are the suite’s highly multimodal functions.Source: §E.2 · p. 28
  • Rosenbrock, Dixon–Price, and Styblinski–Tang provide complementary nonconvex landscapes, with x = (x1, . . . , xd) and d denoting input dimensionality.Source: §E.2 · p. 28

E.3 REAL-WORLD BENCHMARK TASKS

The real-world evaluation uses four data-derived optimization benchmarks, covering oil-well depth, XGBoost hyperparameter tuning, asteroid diameter, and air quality. The discrete tasks are interpolated onto continuous domains using an RBF interpolator.

E.3 REAL-WORLD BENCHMARK TASKS

  • Four data-derived optimization benchmarks from Hsieh et al. (2021) are considered.Source: §E.3 · p. 29
  • Oil-4D searches latitude and longitude coordinates of well surfaces and bottoms for wells with greatest drilled depth.Source: §E.3 · p. 29
  • hpobenchXGB-6D associates six XGBoost hyperparameters, including learning rate and regularization, with classification performance and favors higher-performing configurations.Source: §E.3 · p. 29
  • Asteroid-12D represents candidates with twelve numerical asteroid descriptors and seeks those associated with the largest diameters.Source: §E.3 · p. 29
  • pm25-14D uses fourteen inputs, including sulfur dioxide, carbon monoxide, and ozone measurements from northern Taiwan air-quality observations collected in 2015.Its objective is to identify conditions associated with the highest PM2.5 concentration.Source: §E.3 · p. 29
  • The discrete tasks are interpolated onto continuous domains using a linear-kernel radial basis function interpolator with 32 nearest neighbors.Source: §E.3 · p. 29

Related papers

Loading 2610.20138v2…