# Sample-Efficient Optimisation with Probabilistic Transformer Surrogates

arXiv: https://arxiv.org/abs/2205.13902 · HTML: https://mlsanity.com/paper/2205.13902/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 84 references that are in the map.
- First posted on arXiv 27 May 2022.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [VideoBERT: A Joint Model for Video and Language Representation Learning](https://mlsanity.com/paper/1904.01766/) (2019)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Gaussian Processes for Big Data](https://mlsanity.com/paper/1309.6835/) (2013)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Learning Texture Transformer Network for Image Super-Resolution](https://mlsanity.com/paper/2006.04139/) (2020)
- [Hands-on Bayesian Neural Networks -- a Tutorial for Deep Learning Users](https://mlsanity.com/paper/2007.06823/) (2020)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Scalable Bayesian Optimization Using Deep Neural Networks](https://mlsanity.com/paper/1502.05700/) (2015)
- [Attentive Neural Processes](https://mlsanity.com/paper/1901.05761/) (2019)
- [Cross-Modal Self-Attention Network for Referring Image Segmentation](https://mlsanity.com/paper/1904.04745/) (2019)
- [Pitfalls of In-Domain Uncertainty Estimation and Ensembling in Deep Learning](https://mlsanity.com/paper/2002.06470/) (2020)
- [Evaluating Scalable Bayesian Deep Learning Methods for Robust Computer Vision](https://mlsanity.com/paper/1906.01620/) (2019)
- [Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)](https://mlsanity.com/paper/1503.01057/) (2015)
- [Bayesian Optimization is Superior to Random Search for Machine Learning Hyperparameter Tuning: Analysis of the Black-Box Optimization Challenge 2020](https://mlsanity.com/paper/2104.10201/) (2021)
- [High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspaces](https://mlsanity.com/paper/2103.00349/) (2021)
- [Uncertainty Quantification and Deep Ensembles](https://mlsanity.com/paper/2007.08792/) (2020)
- [Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining](https://mlsanity.com/paper/2006.09191/) (2020)
- [Variational Bayesian Optimal Experimental Design](https://mlsanity.com/paper/1903.05480/) (2019)

## Built on by (most cited first, 3 of 3)

- [PFNs4BO: In-Context Learning for Bayesian Optimization](https://mlsanity.com/paper/2305.17535/) (2023)
- [Applications of Gaussian Processes at Extreme Lengthscales: From Molecules to Black Holes](https://mlsanity.com/paper/2303.14291/) (2023)
- [Information-Theoretic Multi-Model Fusion for Target-Oriented Adaptive Sampling in Materials Design](https://mlsanity.com/paper/2602.03319/) (2026)

## Abstract (arXiv)

Faced with problems of increasing complexity, recent research in Bayesian Optimisation (BO) has focused on adapting deep probabilistic models as flexible alternatives to Gaussian Processes (GPs). In a similar vein, this paper investigates the feasibility of employing state-of-the-art probabilistic transformers in BO. Upon further investigation, we observe two drawbacks stemming from their training procedure and loss definition, hindering their direct deployment as proxies in black-box optimisation. First, we notice that these models are trained on uniformly distributed inputs, which impairs predictive accuracy on non-uniform data - a setting arising from any typical BO loop due to exploration-exploitation trade-offs. Second, we realise that training losses (e.g., cross-entropy) only asymptotically guarantee accurate posterior approximations, i.e., after arriving at the global optimum, which generally cannot be ensured. At the stationary points of the loss function, however, we observe a degradation in predictive performance especially in exploratory regions of the input space. To tackle these shortcomings we introduce two components: 1) a BO-tailored training prior supporting non-uniformly distributed points, and 2) a novel approximate posterior regulariser trading-off accuracy and input sensitivity to filter favourable stationary points for improved predictive performance. In a large panel of experiments, we demonstrate, for the first time, that one transformer pre-trained on data sampled from random GP priors produces competitive results on 16 benchmark black-boxes compared to GP-based BO. Since our model is only pre-trained once and used in all tasks without any retraining and/or fine-tuning, we report an order of magnitude time-reduction, while matching and sometimes outperforming GPs.