IQS-BO: In-Context Query Selection for Bayesian Optimisation
IQS-BO tackles the per-iteration cost of Bayesian optimisation by replacing surrogate fitting and acquisition-function maximisation with a single transformer forward pass. The core idea is to cast query selection as supervised classification: given a finite candidate set, the model predicts which candidate maximises the objective, trained with cross-entropy on synthetic priors where the true best candidate is known for free. The method can run fully in-context or on top of a fixed surrogate, and it introduces a mixture prior of GP and non-stationary functions to improve optimisation performance.
The paper presents a compelling and cleanly derived in-context BO method. Proposition 1 rigorously justifies the cross-entropy objective, and the empirical evaluation shows strong performance against GP-based and in-context baselines in 3D and 4D, with orders-of-magnitude speedups. However, the method inherits key limitations of in-context approaches: it is restricted to fixed input dimensions and batch sizes, relies on a hand-crafted finite candidate set, and assumes a one-step (myopic) decision horizon.
The supervised-learning formulation is elegant and avoids the instability of reinforcement learning for acquisition optimisation. Proposition 1 provides a crisp Bayesian justification: the cross-entropy minimiser is exactly the posterior probability of each candidate being the best. The empirical speedup is dramatic—orders of magnitude—and the controlled ablation (EI-CAND) cleanly isolates the value of the learned selector from the surrogate and candidate design. The proposed mixture prior, which augments GP samples with warped, localised and piecewise-flat functions, is a thoughtful addition that yields consistent early-round improvements across PFN-based methods.
The strongest concerns are architectural rigidity and dependence on candidate design. Models are pre-trained for a fixed dimension and batch size, and the experiments are limited to 3D and 4D—scalability to higher dimensions is neither demonstrated nor argued for. The candidate set is a heuristic mix of Sobol points, uniform samples, and local perturbations; the ablation shows that replacing perturbations with purely space-filling candidates raises regret, meaning the method is not domain-agnostic in its proposal mechanism.
Proposition 1 also assumes well-specification, which is not guaranteed by the finite transformer. The theoretical claim that the network therefore approximates the Bayes-optimal rule rests on expressivity and optimisation assumptions that are not verified. The objective is myopic by construction, and extending it to non-myopic or multi-step lookahead is left entirely to future work.
The evidence is solid for the studied low-dimensional regime. The benchmarks cover a reasonable range of synthetic and emulated real-world landscapes, and the comparisons to GP-MAP, GP-FB, PFNs4BO, FIBO and NAP are thorough, with win rates, ranks and confidence intervals reported. Two caveats weaken the comparison. First, the NAP baseline required a custom training script by the authors, which may not reflect its best possible performance. Second, the claim that PFNs4BO gains less from the mixture prior because expected improvement reduces the posterior to mean and variance is speculative and not directly tested.
The real-world evidence is also limited: only three chemistry emulators from Olympus plus one constructed bioprocess surface are used, leaving open how the method behaves on truly noisy, high-stakes experimental data.
The paper provides extensive hyperparameters and implementation details in the appendices—including exact training recipes, architecture specifications, and baseline configurations—which is commendable. However, the reproducibility statement claims all information is in the paper but makes no mention of a public code repository or released model weights. Without publicly available code, independent reproduction would require reimplementing the transformer, priors, and candidate sampler from scratch.
The reliance on a custom reimplementation of FIBO and a self-written training loop for NAP further raises the risk of baseline drift. The paper does not report random seeds, and pretrained models are large (up to ~13M parameters plus a frozen PFN), so reproducing the exact computational cost claims would be difficult without the authors' infrastructure.
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.
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.