$\alpha$-PFN: Fast Entropy Search via In-Context Learning

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.

Many important tasks involve finding the best option by trial and error, where each attempt is slow or costly. This happens when tuning complex software systems, but also when designing experiments in science and engineering. In all of these cases, the goal is to reach a good choice in as few attempts as possible. One principled approach, known as "entropy search", decides what to try next by asking which experiment would reveal the most about where the best option lies. The difficulty is that these methods are hard to compute and depend on slow, complex approximations. In this work, we trained a machine learning model to learn this computation instead, which removes the need for hand-built approximations. On both synthetic and real-world tasks, our approach reaches a quality similar to existing methods while running more than 50 times faster. This makes this type of optimization more practical and easier to use without large computing resources.