PaperlayerView exact revision on arXiv

Source-linked AI summary

IQS-BO: In-Context Query Selection for Bayesian Optimisation

Luca Geminiani, Nadja Klein

arXiv:2610.01269v1cs.LGstat.ML

TL;DR

Bayesian optimisation repeatedly fits a surrogate and maximises an acquisition function to choose evaluations of expensive black-box objectives. IQS-BO learns candidate selection by supervised training on synthetic priors, with or without a fixed surrogate. It matches or outperforms GP-based and in-context methods, and a richer pre-training prior can improve optimisation, especially early in the evaluation budget.

  • Problem

    Bayesian optimisation targets expensive black-box functions, but typically repeats surrogate fitting and acquisition-function maximisation at each evaluation.

    Source: §1 · p. 1

  • Method

    IQS-BO uses supervised learning on synthetic priors to predict each candidate’s probability of being best and select the highest-probability candidate.

    Source: §1 · p. 2

  • Results

    IQS-BO matches or outperforms GP-based BO and in-context methods, and in batch BO reaches the highest performance while proposing queries more than an order of magnitude faster than GP-based BO.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Adding functions poorly modeled by stationary GPs to the pre-training prior can improve PFN-based optimisation, especially at early evaluations.

    Source: §5 · p. 9

    Models are pre-trained for fixed input dimensions and batch sizes, experiments focus on three and four dimensions, and queries are restricted to a finite candidate set under a myopic one-step objective.

    Source: §5 · p. 9

Abstract

from arXiv · show

Bayesian Optimisation (BO) is a powerful framework for the optimisation of expensive black-box functions, but typically requires refitting a surrogate and maximising an acquisition function at every evaluation step. In-context approaches based on Prior-data Fitted Networks (PFNs) amortise part of this cost by pre-training transformers on functions drawn from synthetic priors. PFNs4BO amortises the surrogate but still relies on a numerically maximised acquisition function, while FIBO performs BO fully in-context by sampling optimiser locations from a learned density, which fixes the decision rule and admits no surrogate. Learned acquisition functions score a finite candidate set with a trained network, but, lacking a label for the query, learn the score by reinforcement learning on previously solved tasks. We propose IQS-BO, a PFN that learns the query decision by supervised learning on synthetic priors. In a single forward pass, IQS-BO predicts the probability that each candidate maximises the objective over the set, and we show that the minimiser of its objective is the posterior probability of this event. The model can be pre-trained without a surrogate for fully in-context BO, or take the predictions of a fixed probabilistic surrogate as additional input, amortising only the decision step. Our method proposes queries at a fraction of the cost of acquisition-based methods, while either matching or outperforming standard BO with Gaussian processes (GPs) and available in-context methods on synthetic and real-world benchmarks. Finally, we propose a mixture prior for pre-training PFNs which combines samples from GPs with functions exhibiting warped inputs, isolated narrow optima, or plateaus that are poorly modeled by stationary kernels common in GP surrogates. We show that pre-training on this prior can lead to improved optimisation performance.

1 INTRODUCTION

Bayesian optimisation repeatedly fits a surrogate and optimises an acquisition function, while in-context methods amortise parts of that process. IQS-BO learns candidate selection from synthetic priors and can operate with or without a fixed surrogate.

1 INTRODUCTION

  • PFNs4BO (M¨uller et al., 2023) amortises surrogate prediction but still numerically maximises an acquisition function, whereas FIBO (de Carvalho et al., 2025) samples from a learned optimiser-location density with a fixed query policy.Learned acquisition methods such as NAP instead train candidate scores from rewards on solved tasks.Source: §1 · pp. 1–2
  • IQS-BO predicts each candidate’s probability of being best and selects the highest-scoring one, with cross-entropy minimised by the posterior probability of that event.Training labels are available because each synthetic-prior function can be evaluated at every candidate.Source: §1 · p. 2
  • IQS-BO can use observations directly for fully in-context optimisation or take predictions from a fixed probabilistic surrogate to amortise query selection.The paper demonstrates the latter with a released PFNs4BO model and a Gaussian process.Source: §1 · p. 2
  • Across synthetic and real benchmarks, IQS-BO matches or outperforms GP-based BO and in-context methods, proposing queries one to three orders of magnitude faster than PFNs4BO and GP-based baselines.In batch optimisation, it also reaches better performance than FIBO at every batch size while maintaining comparable computational cost.Source: §1 · p. 2

1.1 SUMMARY OF MAIN CONTRIBUTIONS

  • A mixture prior adds functions poorly captured by fixed-kernel GPs to GP samples, improving PFN-based optimisation, especially when the initial context is small.The added families include warped, localised, and piecewise-flat functions.Source: §1 · p. 2, §1.1 · p. 3

2 PROBLEM STATEMENT

The problem statement considers maximising an expensive, noisy black-box function from observations accumulated over successive evaluations. Instead of maximising an acquisition function over the full domain, it frames the next query as a choice among sampled candidates.

2 PROBLEM STATEMENT

  • The objective is to maximise an expensive black-box function over [0, 1]^d using noisy evaluations, making each new evaluation costly.Standard BO fits a probabilistic surrogate and maximises an acquisition function over the domain.Source: §2 · p. 3
  • At each round, a stochastic sampler constructs a finite candidate set from the observations, and the learned decision rule selects directly from that set.This replaces maximisation of the acquisition function over the full domain.Source: §2 · p. 3
  • The selected candidate is the one with the largest posterior probability of having the highest function value in the set.The choice is conditioned on both the observed data and the candidate set.Source: §2 · p. 3

3 METHODOLOGY

IQS-BO learns to select queries from synthetic tasks, predicting which candidate is best; it can operate without a surrogate or use a fixed surrogate’s predictions as input.

3 METHODOLOGY

  • IQS-BO matches or outperforms GP-based BO and in-context baselines across synthetic and real-world benchmarks, while proposing queries faster than acquisition-based methods.Its three variants use no surrogate, a PFN surrogate, or a GP surrogate.Source: §3 · p. 3, §3.5 · p. 5

3.1 TRAINING OBJECTIVE

  • Training minimizes expected cross-entropy on synthetic tasks, with the known best candidate providing the label.The function is evaluated without noise at every candidate to compute that label exactly.Source: §3.1 · pp. 3–4

3.2 POSTERIOR-BEST-CANDIDATE INTERPRETATION

  • Under well-specification, the loss minimizer is the posterior probability that each candidate is best, making its highest-probability candidate the Bayes-optimal choice under 0–1 loss.Source: §3.2 · p. 4

3.3 ARCHITECTURE

  • A transformer reads context and candidate tokens and scores each candidate in one forward pass; candidate scores depend only on the context and that candidate.When a fixed surrogate is used, its outputs are added to the token embeddings.Source: §3.3 · p. 4

3.4 PRE-TRAINING

  • Pre-training samples functions from either the HEBO prior or a mixture prior that adds input-warped GPs, local Gaussian bumps, and piecewise-flat plateaus.The mixture includes function types that stationary GP surrogates model poorly.Source: §3.4 · p. 5

3.5 INFERENCE

  • At inference, the model draws a candidate set, predicts each candidate’s probability in one pass, and evaluates the highest-probability candidate without parameter updates.For batches, picks are made sequentially, with earlier picks passed as pending tokens.Source: §3.5 · p. 5, §3.4 · p. 5

4 EMPIRICAL EVALUATION

Across synthetic and real-world problems, IQS-BO performs competitively with GP-based and in-context methods. The evaluation also tests its query-selection rule, pre-training prior, batch performance, and design choices.

4.1 EXPERIMENTAL SETUP

  • The evaluation covers ten 3D problems and five 4D problems, with synthetic and real-world tasks, comparing IQS-BO variants against GP-based BO, in-context methods, and Sobol search.Runs use 90 seeds and compare regret and average rank across initial context sizes.Source: §4.1 · p. 6

4.3 CONTROLLED COMPARISONS AND EFFECT OF THE PRIOR

  • IQS-BO-GP outperforms expected improvement on the same GP and candidate set at every context size, and the advantage grows with more observations.The comparison isolates how candidates are scored; the candidate design also contributes to IQS-BO-GP’s gain over GP-MAP.Source: §4.3 · pp. 7–8, Figure 3 · p. 8
  • The mixture prior improves IQS-BO, IQS-BO-PFN, and PFNs4BO in most runs, especially early; its benefit is smaller on GP-well-specified problems and negligible for FIBO.The authors relate the gains to exposure during pre-training to varied functions, including local bumps and plateaus.Source: §4.3 · p. 8

4.2 MAIN RESULTS

  • Across prior settings, IQS-BO matches or outperforms GP-based BO and improves on in-context baselines; in batch BO it reaches the highest GAP at every tested batch size while proposing queries faster than GP-based BO.For mixture-prior models, all three variants outperform competing learned methods; with HEBO or GP-based priors, IQS-BO often achieves the best or near-best final regret as context grows.Source: Table 1 · p. 6, §4.2 · p. 7, Figure 2 · p. 7

4.4 BATCH BAYESIAN OPTIMISATION

  • IQS-BO and IQS-BO-PFN achieve the highest GAP at every tested batch size, while IQS-BO is almost unaffected by increasing batch size and proposes queries more than an order of magnitude faster than GP-based BO.The models are pre-trained at the deployed batch size, learning to account for earlier picks and diversify the batch.Source: §4.4 · pp. 8–9, Figure 4 · p. 9

4.5 ABLATION STUDY

  • Candidate composition matters most: replacing near-incumbent perturbations with space-filling candidates degrades performance in most runs, while sequential batch selection outperforms selecting all points in one pass.Fewer candidates hurt, whereas more candidates modestly improve final regret and do not improve average regret.Source: §4.5 · p. 9, Table 2 · p. 9

5 CONCLUSION

IQS-BO selects queries by predicting which candidate is best, and its reported comparisons favor it over in-context methods on the same prior while matching or outperforming GP-based BO. The appendix distinguishes methods by what they learn, their training signal, and how they propose queries.

5 CONCLUSION

  • IQS-BO matches or outperforms GP-based BO and in-context methods trained on the same prior; in batch BO it reaches the highest performance and proposes queries over an order of magnitude faster than GP-based BO.Enriching the prior with functions poorly modeled by stationary GPs can further improve performance, especially at early evaluations.Source: §5 · p. 9
  • The conclusion identifies fixed input dimension and batch size, experiments in three and four dimensions, finite candidate sets, and myopic one-step decisions as boundaries for future work.Proposed extensions include models spanning dimensions and batch sizes, jointly learned candidate sets, and objectives beyond one-step selection.Source: §5 · p. 9

A COMPARISON TO IN-CONTEXT AND LEARNED ACQUISITION METHODS

A.1 IN-CONTEXT METHODS

  • PFNs4BO amortizes surrogate prediction but still numerically maximizes a specified acquisition function at every round; FIBO samples optimiser locations, fixing its decision rule and precluding surrogate input.PFNs4BO's acquisition function must have a closed form on the PFN output; FIBO's batches are independent draws.Source: §A.1 · p. 13

A.2 LEARNED ACQUISITION METHODS

  • NAP, TAMO, MetaBO, FSAF and BOFormer learn query scores or policies from rewards, while FSAF also requires repeated GP fitting and task-specific gradient updates.NAP and TAMO operate in-context; MetaBO, FSAF and BOFormer receive task information through a refitted GP.Source: §A.1 · p. 13, §A.2 · p. 14

A.3 RELATION TO OUR METHOD

IQS-BO differs from prior query-selection methods through supervised training on synthetic priors and optional use of a fixed surrogate. Its implementation uses candidate and pending tokens, sequential batch selection, and a mixture prior spanning function families.

A.3 RELATION TO OUR METHOD

  • IQS-BO uses supervised classification on synthetic-prior labels, needs no solved target tasks, and can run with or without a fixed surrogate.With a surrogate, only query selection is amortised; without one, the full BO loop runs in-context.Source: §A.3 · p. 14

B.1 TOKENS AND INPUTS

  • Observation tokens attend to observations, while candidate tokens attend to observations; each candidate score depends only on the data and that candidate.With a surrogate, its predictions are added to token embeddings; pending points are supplied at fantasy values, and the GP is conditioned on them without refitting.Source: §B.1 · pp. 15–16
  • For batches, selected points become pending tokens with fantasy values; later candidate scores condition on those pending points.IQS-BO uses zero as the fantasy value, while surrogate variants use the surrogate posterior mean.Source: §B.1 · pp. 15–16

B.2 NETWORK AND DEPLOYMENT

  • The deployed encoder has four layers, width 512, eight attention heads, and 12.88M trainable parameters for IQS-BO; the surrogate variants have 12.89M and 13.15M.The surrogate-based IQS-BO-PFN count excludes the frozen PFN's 26.8M parameters; deployment uses 512 candidates.Source: §B.2 · p. 16

B.3 BATCH SELECTION

  • The batch loss trains later selections only when earlier picks have not selected the best candidate, teaching the network to diversify the batch.For batches larger than two, losses are counted only while no earlier pick is best, with up to four random slots averaged.Source: §B.3 · p. 16

C PRE-TRAINING DETAILS

  • The pre-training setup generates tasks from a prior, samples contexts and candidates, and labels the best candidate by evaluating the noise-free function.For GP surrogates, hyperparameters are fit to the sampled context during pre-training.Source: §B.3 · p. 16

C.1 PRIORS

  • The mixture prior combines stationary, warped, bumps, and plateau families to represent smooth functions, varying effective lengthscales, narrow optima, and sharp steps.The warped family transforms inputs; bumps create narrow optima in nearly flat regions, while plateaus create piecewise-flat functions.Source: §C.1 · p. 17

C.2 CANDIDATE SAMPLING

  • Candidate sets combine 40% scrambled Sobol points, 30% uniform points, and 30% perturbations around the three best observations.Space-filling points support exploration, while perturbations allow refinement using moves spanning coarse and fine scales.Source: §C.2 · p. 17

C.3 TRAINING CONFIGURATIONS AND COST

The paper reports training configurations and costs for the models, with IQS-BO models trained on a single H100 GPU. Its training-time table reports H100 hours for batch size two.

C.3 TRAINING CONFIGURATIONS AND COST

  • All IQS-BO models were trained on a single NVIDIA H100 GPU; Table 9 reports training time as median per-epoch time excluding the first epoch, multiplied by epochs, excluding evaluation and data generation.Source: §C.3 · p. 18
  • NAP training time is reported for the full PPO run rather than by epochs, and its 3D and 4D models use different configurations.The paper attributes the difference in NAP training times to those different configurations.Source: §C.3 · p. 18
  • PFNs4BO, FIBO, and NAP use their own training recipes rather than the shared IQS-BO settings.The appendix provides the baseline configurations separately.Source: §C.3 · p. 18

D.1 GP-MAP AND GP-FB

The appendix specifies implementations for GP-based BO, PFNs4BO, FIBO, and NAP. It also reports a comparison between the authors’ FIBO implementation and the paper’s reported performance.

D.1 GP-MAP AND GP-FB

  • The GP-MAP and GP-FB baselines maximize qLogNoisyExpectedImprovement jointly over the batch, differing in how they handle GP hyperparameters.GP-MAP fits hyperparameters by MAP each round; GP-FB samples their posterior with NUTS and averages the acquisition over samples.Source: §D.1 · pp. 18–19, Table 9 · p. 18

D.2 PFNS4BO

  • PFNs4BO uses released HEBO and BNN checkpoints, transforms observed values, and maximizes expected improvement; batches use the kriging-believer rule.The released code proposes one point per round and does not support batch optimization.Source: §D.2 · p. 19

D.3 FIBO

  • 0.90 GAP after 58 evaluations: the authors’ FIBO implementation, versus 0.93 after 200 evaluations in the paper on its prior at d = 3.The implementation reports that it beats random search on every test function.Source: §D.3 · p. 19

D.4 NAP

  • NAP is trained with PPO on the GP task family of MetaBO and uses a regret-based reward plus an auxiliary predictive loss.Its 3D and 4D protocols differ, including episode length, initial points, and candidate counts.Source: §D.4 · pp. 19–20

E BENCHMARK PROBLEMS

The benchmarks combine analytic landscapes, a constructed bioprocess surface, Olympus emulator campaigns, and GP-prior draws in three and four dimensions. Optima and regret references are checked or recomputed, and the appendix provides additional results and ablations.

E.1 PROBLEMS

  • The benchmark set spans analytic functions, a constructed bioprocess surface, Olympus campaigns, and GP-prior draws across three and four dimensions.Problems are mapped to the unit cube and maximised; minimisation objectives are negated.Source: §E.1 · p. 20
  • The bioprocess surface is unimodal but strongly non-stationary, with about 0.3% of its domain within 10% of its optimum.Its optimum is f*=1 at x*=(0.10, 0.176, 0.35).Source: §E.1 · p. 20

E.2 OPTIMA

  • Regret uses recomputed global optima; closed-form optima are exact, while other optima are found with multi-start L-BFGS-B and independently confirmed.GP-prior draws are filtered for interior, uniformly spread optima and checked against 2 × 10^6 Sobol points.Source: §E.2 · p. 21

F ADDITIONAL RESULTS AND ABLATIONS

  • The appendix extends the main evaluation with per-problem results, paired win rates, prior effects, batch optimisation, and additional ablations.It covers the 3D and 4D problems and compares expected improvement as well as pre-training priors.Source: F ADDITIONAL RESULTS AND ABLATIONS · p. 21

F.1 ADDITIONAL RESULTS ON THE 3D AND 4D PROBLEMS

At smaller initial context sizes, the method ordering is unchanged from ninit=40, and mixture-prior IQS-BO variants beat GP and in-context baselines in most paired runs. Per-problem outcomes vary: IQS-BO excels on several irregular landscapes, while well-specified smooth GP problems favor GP baselines.

Smaller context sizes and paired win rates

  • At ninit=10 and 20, mixture-prior IQS-BO variants beat GP-MAP and GP-FB in about 80% of runs and every in-context baseline in more than 83%.The method ordering matches ninit=40; with the HEBO prior, IQS-BO and IQS-BO-GP beat GP-MAP in 65–75% of runs, while IQS-BO-PFN ties at ninit=10.Source: §F.1 · p. 21

Per-problem results

  • In 3D, IQS-BO has lower regret than all baselines early and through the run on Ackley, Camel, Levy, and Olympus colours.Rosenbrock is an exception: GP baselines finish with lower regret where the Matérn GP is well specified.Source: §F.1 · p. 21
  • In 4D, mixture-prior IQS-BO variants reach lower final regret than GP-MAP on Ackley, Levy, and Rosenbrock.GP-prior draws and benzylation favor GP baselines; on GP-prior draws, IQS-BO-GP improves on GP-MAP in 55–66% of paired runs.Source: §F.1 · p. 21
  • On benzylation, GP-MAP and GP-FB can query exactly at the optimum, while IQS-BO and other in-context baselines stall at small non-zero regret.This is attributed to the smooth response surface and gradient-based acquisition maximisation.Source: §F.1 · pp. 21–22

F.2 LEARNED QUERY SELECTION AGAINST EXPECTED IMPROVEMENT

The section reports paired comparisons of IQS-BO against expected improvement and a PFN-surrogate baseline. With the HEBO prior, IQS-BO-GP is better than EI-CAND at final regret but weaker early; IQS-BO-PFN ends with lower regret than PFNs4BO-HEBO in most runs.

PFN surrogate

  • The appendix reports per-iteration average regret for each 3D problem (ninit=10, 20, 40; 90 seeds; 95% confidence intervals), curves for five 4D problems, and mean normalized final regret per 4D problem at ninit=40.Source: Figure 7 · p. 23, Figure 8 · p. 24, Table 14 · p. 25

Dependence on the prior

  • IQS-BO-GP with the HEBO prior beats EI-CAND in 54–58% of runs on Rfinal but wins only 30–47% on RAUC.Thus it reaches better final regret but is slower early in the run; the clear advantage reported for Section 4.3 uses the mixture prior.Source: §F.2 · p. 24

PFN surrogate

  • IQS-BO-PFN ends with lower regret than PFNs4BO-HEBO in 60%, 67%, and 75% of runs at ninit=10, 20, and 40, respectively.It uses the released PFNs4BO-HEBO model unchanged and differs only in query selection; gains are largest where expected improvement on the PFN posterior does worst against GP-MAP.Source: §F.2 · pp. 24–25

Dependence on the prior

  • Table 15 reports paired win rates of the first method against the second across ten 3D problems, with 900 paired runs per cell.EI-CAND requires no pre-training, so its GP-MAP comparison is the same under both priors and is shown once.Source: Table 15 · p. 25

F.3 EFFECT OF THE PRE-TRAINING PRIOR

The mixture prior improves PFN-based optimisation, with its largest gains early in the run, and does not generally hurt on problems where GP surrogates are well specified. In batch optimisation, IQS-BO maintains high GAP across batch sizes and proposes batches far faster than GP-MAP at q=20.

F.3 EFFECT OF THE PRE-TRAINING PRIOR

  • For all three IQS-BO variants, mixture-prior win rates are higher on RAUC than Rfinal, and regret curves show the gap opening in early rounds and persisting.The authors relate early gains to greater reliance on pre-training before enough observations accumulate.Source: §F.3 · p. 25, Table 16 · p. 26, Figure 10 · p. 27
  • On problems with well-specified GP surrogates, mixture-prior IQS-BO and IQS-BO-PFN win about half or more of runs on Rosenbrock and GP-prior draws.At ninit=40 on benzylation, HEBO-prior IQS-BO and IQS-BO-GP win 61% and 68% of paired runs; the mixture prior lowers performance in a few specified model/problem cases.Source: §F.3 · p. 25, Table 17 · p. 27

F.4 BATCH BAYESIAN OPTIMISATION

  • IQS-BO's mean GAP stays between 0.91 and 0.93 across batch sizes, and its margin to the best baseline grows from 0.05 at q=10 to 0.08 at q=50.The batch-size ablation attributes this robustness to pre-training at the deployed batch size.Source: §F.4 · p. 27
  • At q=20, IQS-BO and IQS-BO-PFN are more than an order of magnitude faster than GP-MAP for a single batch proposal.Only FIBO is faster, drawing batch points independently without conditioning on earlier draws.Source: §F.4 · p. 27

F.5 ADDITIONAL ABLATIONS

The ablations show that candidate composition matters across scoring rules, while batch-size-specific pretraining and sequential selection improve batch performance under the reported comparisons.

Candidate composition with other scoring rules

  • Perturbed candidates let IQS-BO, IQS-BO-GP, and EI-CAND beat GP-MAP in 77–81% of paired runs, while space-filling candidates alone leave all three below GP-MAP.The passage reports individual ranges of 77–80%, 78–81%, and 67–71% with the default set, versus 43–45%, 36–40%, and 27–33% with Sobol and uniform candidates only.Source: §F.5 · p. 27

Additional ablations

  • Table 18 reports GAP after 200 evaluations per problem and averaged over seven problems, plus seconds spent proposing queries over a complete run.The caption specifies means over 90 seeds and 95% confidence-interval half-widths in units of 10^-3.Source: Table 18 · p. 28

Batch size and batch construction

  • At q = 10, 20, and 50, batch-size-pretrained IQS-BO beats the q = 2 checkpoint in 62%, 91%, and 98% of paired runs, respectively.The q = 2 checkpoint's mean GAP falls from 0.88 to 0.63 and 0.44, while the batch-size-pretrained checkpoint stays above 0.90; Figure 11 plots GAP per evaluation by problem and batch size.Source: §F.5 · p. 28, Figure 11 · p. 29
  • At q = 10, sequential selection outperforms top-q selection in 68% of paired runs.Top-q selection takes the q highest-scoring candidates from one forward pass without pending points.Source: §F.5 · p. 28

Related papers

Loading 2610.01269v1…