$α$ -PFN: Fast Entropy Search via In-Context Learning

Published 5 Jun 2026 in cs.LG | (2606.07134v1)

Abstract: Information-theoretic acquisition functions such as Entropy Search (ES) offer a principled exploration-exploitation framework for Bayesian optimization (BO). However, their practical implementation relies on complicated and slow approximations, i.e., a Monte Carlo estimation of the information gain. This complexity can introduce numerical errors and requires specialized, hand-crafted implementations. We propose a two-stage amortization strategy that learns to approximate entropy search-based acquisition functions using Prior-data Fitted Networks (PFNs) in a single forward pass. A first PFN is trained to be conditioned on information about the optima; second, the $α$ -PFN is trained to predict the expected information gain by training on information gains measured with the first PFN. The $α$ -PFN offers a flexible learned approximation, which replaces the complex heuristic approximations with a single forward pass per candidate, enabling rapid and extensible acquisition evaluation. Empirically, our approach is competitive with state-of-the-art entropy search implementations on synthetic and real-world benchmarks, while accelerating the different entropy search variants across all our experiments, with speed ups over 50x. Source code: https://github.com/automl/AlphaPFN.

Abstract PDF Upgrade to Chat

Authors (7)

Summary

The paper introduces a two-stage PFN approach that bypasses traditional Monte Carlo sampling, achieving over 50x speed-ups in entropy search for Bayesian Optimization.
It utilizes a base PFN for posterior predictive estimation and an α-PFN to directly compute acquisition values via a single transformer forward pass.
Empirical evaluations on synthetic and real-world benchmarks show that the method matches state-of-the-art ES techniques while robustly generalizing to higher dimensions.

Fast Amortization of Entropy Search in Bayesian Optimization via In-Context Prior-Fitted Networks

Introduction and Motivation

$\alpha$ -PFN (" $α$ -PFN: Fast Entropy Search via In-Context Learning" (2606.07134)) addresses the computational bottlenecks inherent in information-theoretic acquisition functions for Bayesian Optimization (BO) such as Entropy Search (ES) and its variants (Predictive Entropy Search, Max-value Entropy Search, Joint Entropy Search). Traditional ES-based BO incurs substantial runtime due to expensive Monte Carlo (MC) or similarly manual, sampling-based estimation of information gain—especially problematic for fully Bayesian Gaussian Processes (GPs) and when optimizing GPs' hyperparameters.

The fundamental contribution of this work is the design of a two-stage, meta-learned approximation of ES-based BO acquisition using Prior-data Fitted Networks (PFNs). By mapping the computation of acquisition values directly to a single forward pass of a transformer trained on millions of GP sample datasets, the authors decouple ES from MC-based inference. This achieves over 50x speed-ups, while empirically matching the performance of state-of-the-art handcrafted ES approximations on synthetic and real-world benchmarks. The PFN-based architecture amortizes the computation across acquisition functions, contexts, and prior settings, supporting flexible and efficient BO with surrogate uncertainty propagation.

Background: Information-Theoretic Acquisition Functions

Classical BO with GPs generally employs acquisition functions such as Expected Improvement (EI), which are myopic and suboptimal on noisy or multi-modal objective landscapes. In contrast, ES-based acquisition functions maximize the expected reduction in uncertainty (entropy) regarding the optimum's location $x^*$ or value $f^*$ after a (potential) observation at $x$ —formally, the mutual information between $x^*$ or $f^*$ and the observation $y$ . ES, MES, and JES involve nested expectations and conditional entropies over posteriors, requiring expensive sampling (e.g., MC or RFF-based function path optimization) for both current context and GP hyperparameters. Fully Bayesian variants require sampling over kernel hyperparameters, further compounding computational overhead.

Prior-data Fitted Networks: Amortized Posterior and Acquisition Inference

PFNs are transformer-based conditional neural processes, trained on massive synthetic datasets sampled from a GP (or other surrogate) prior. A PFN, through in-context learning, outputs the posterior predictive distribution (PPD) conditioned on context points and queried locations, in a single forward pass. This work extends the PFN paradigm in two stages:

Base PFN: Trained to provide the PPD $q(y|x, D, I)$ , optionally conditioned on information $I$ about the optimum (i.e., on $α$ 0, $α$ 1, or both). The base PFN can thus simulate the entropy terms under conditioning required by ES, MES, or JES objectives.
$α$ 2-PFN (Acquisition PFN): Trained on the difference of entropies (using the base PFN)—specifically, on $α$ 3, across massive precomputed datasets where the true optimum is brute-force approximated using Random Fourier Features and global optimization heuristics (see Algorithm~\ref{algo:cluster}).

Rather than sampling $α$ 4 for each evaluation, this meta-learned $α$ 5-PFN directly predicts the distribution of information gain. The mean of this distribution corresponds to the ES-based acquisition function, removing runtime MC estimation at test time.

Figure 1: The base PFN models the PPD for a candidate query, optionally conditioned on true optimum information ( $α$ 6, $α$ 7, or both) to enable analytic entropy difference computation for the acquisition.

Figure 2: Left—base PFN is trained on multiple conditional regimes; Middle—inducing PES acquisition as a meta-learning objective; Right—at inference, a single forward pass of the $α$ 8-PFN yields the acquisition function for BO.

Training Pipeline and Generalization

Training utilizes synthetic datasets from GP priors, for dimensions $α$ 9 and context sizes up to $x^*$ 0. Maximization of GP sample paths (for ground-truth $x^*$ 1) is performed via efficient global and local search ensembles. To prevent distributional shift due to clustering of context points in BO traces, a trace generation heuristic models exploration-exploitation in the context distribution; this is crucial for high-dimensional generalization.

Figure 3: Clustered trace generation matches BO context distribution and is critical for performance beyond toy settings.

Fully Bayesian coverage is provided by sampling GP hyperparameters, allowing the base PFN and $x^*$ 2-PFN to properly integrate posterior uncertainty, in contrast to the per-hyperparameter averaging typical of GP-MCMC approaches.

Empirical Results

Synthetic Function Benchmarks

Evaluation on classic benchmarks (Branin, Hartmann, Ackley) demonstrates that $x^*$ 3-PFN's inference regret is highly competitive with conventional GP-based ES, PES, MES, and JES—in many cases, nearly indistinguishable. Variance across traces and dimensionality is well-handled, indicating robust amortization and generalization.

Figure 4: PFN-based ES variants yield low inference regret, matching GP-based MC approximations.

Real-World HPO Benchmarks

On surrogate HPO datasets (LCBench, HPO-B)—with higher dimensionality and discrete domains—the $x^*$ 4-PFN consistently achieves predicted-best accuracy and mean regret similar to or better than GP-MCMC-based ES, across ES, PES, and JES variants.

Out-of-Distribution and Robustness

In OOD experiments with higher observation noise, $x^*$ 5-PFN precision degrades at parity with non-amortized methods, confirming faithful approximation rather than overconfident extrapolation.

Figure 5: In high-noise settings, $x^*$ 6-PFN’s performance and degradation mirror GP-based ES, indicating robust generalization.

Ablation and Qualitative Analysis

The dependence on clustered context distributions is significant; training with uniform contexts yields poor BO performance in higher dimensions, confirming the necessity of synthetic trace generation.

Figure 6: Context trace ablation underscores the importance of matching the true BO context distribution during PFN training.

Qualitative acquisition comparisons in 1D demonstrate close correlation between MC-estimated and direct PFN-acquisition mappings, with PFN-based approaches showing less variance in regions of high acquisition uncertainty.

Figure 7: JES acquisition function estimates by GP, base-PFN + MC sampling, and $x^*$ 7-PFN for a 1D function; PFN-based estimation exhibits reduced variance.

Runtime Comparison

$x^*$ 8-PFN achieves substantial practical runtime savings. Across all experimental settings, speedups range from 2x up to 72x, with the typical reduction 10x-50x for challenging settings (e.g., fully Bayesian HPO-B).

Theoretical Analysis

The authors prove that the $x^*$ 9-PFN’s mean corresponds to the required expectation of the information gain as in PES/MES/JES acquisition functions, leveraging the law of the unconscious statistician and KL-divergence minimization under cross-entropy loss. This justifies using the learned $f^*$ 0-PFN for direct single-pass acquisition evaluation without any MC or sampling stage at test time.

Discussion, Limitations, and Future Directions

The results demonstrate that amortized, PFN-based acquisition function prediction matches the empirical utility of classic Monte Carlo or sampling-based methods, at a fraction of runtime and compute. The flexibility of PFN-based amortization supports rapid evaluation on unseen tasks, domain shifts, higher-dimensional spaces, and arbitrary priors.

However, retraining distinct $f^*$ 1-PFNs is required per prior family or class of surrogates; broader-coverage, adaptive prior families and domain adaptation methods (such as distribution transformers) might resolve this. Current scaling is evaluated up to 16D at test time; leveraging recent PFN scaling results to 500D is a promising avenue for further extension. There is also clear scope for amortizing additional sophisticated acquisition functions, e.g., batch/parallel and constrained settings, or BO for active experimentation.

Conclusion

The $f^*$ 2-PFN framework provides a practical, flexible, and highly efficient amortized solution for Bayesian optimization acquisition in both synthetic and real-world contexts. Its single-pass, meta-learned architecture supplants sampling-based MC approximations without sacrificing informativeness or quality of acquisition. This paradigm invites investigation into foundation models for optimization acquisition, adaptive prior generalization, and large-scale, real-time active learning.