Transformers Can Do Bayesian Inference

1 INTRODUCTION . In the last decade , supervised machine learning ( ML ) methods using deep learning architectures have made substantial progress on machine learning tasks where a large amount of training data is available ( Vaswani et al. , 2017 ; He et al. , 2016 ; Krizhevsky et al. , 2012 ) . A very important problem in ML is thus to transfer these successes to smaller-scale setups with less data available . In this paper , we propose a way to build models that approximate posteriors with flexible and replaceable priors using deep learning models . It makes specifying a prior as simple as defining a sampling scheme of supervised learning tasks . While the success of deep learning on large datasets can be attributed to the capacity of neural networks to approximate any function , there is still a need for encoding prior knowledge , for example through model architecture ( e.g. , Convolutional Neural Networks ( LeCun et al. , 1989 ) ) or regularizers ( e.g. , data augmentations ( Hendrycks et al. , 2019 ; Cubuk et al. , 2020 ) ) . Otherwise , the no free lunch theorems show that there are no good methods to solve the class of prediction problems ( Wolpert & Macready , 1997 ) . Thus , a large number of specialized algorithms have been developed for different small-scale tasks ( LeCun et al. , 1989 ; Kadra et al. , 2021 ; Chen & Guestrin , 2016 ) . Encoding prior information into a machine learning model can , however , be challenging . A well-defined way to bias a model is to use Bayesian inference . The foundation of Bayesian inference is an assumption on the distribution of the data to appear in a real-world application . This assumption gives rise to a prior belief over the probability for the data to follow a particular model . One might , e.g. , implement a prior , encoding the data as created by a neural network ( Bayesian Neural Networks , ( MacKay , 1992 ) ) , by a polynomial , the likelihood of a Gaussian mixture or random code in a pre-defined programming language ( Solomonoff , 1997 ) . Using Bayesian inference for prediction in supervised learning has the advantages that ( i ) it has a theoretical foundation that makes it valid in setups where the prior p ( t ) fits and ( ii ) can thus better account for the actual likelihood of different events ; ( iii ) it is well calibrated and ( iv ) it is interpretable as the prior describes the expectations of the model . However , retrieving the prediction distribution for a given prior intractable in most cases ( Blei et al. , 2017 ; MacKay , 1992 ) . Figure 1 outlines Prior-Data Fitted Networks ( PFNs ) , for approximating such Bayesian models . We assume a given representative prior distribution over supervised learning tasks ( or functions ) , which provides our inductive bias . To train a PFN , we use supervised learning with set-valued inputs representing entire datasets : we repeatedly sample a meta-train task ( or function ) from the given prior , draw a set of data points and their labels from it , mask one of the labels and learn to make probabilistic predictions for it based on the set-valued input of the rest of the data points . Given an actual real-world dataset , we feed it together with a test point as inputs to the PFN and that outputs its prediction distribution for the test point , conditioned on the dataset . As we will demonstrate , this distribution approximates exact Bayesian posterior prediction . We refer to this step as ( Bayesian ) inference , as opposed to the training of the PFN itself . Our PFNs thus allow us to approximate the posterior predictive distribution for any prior from which we are able to sample data . This is a very weak requirement compared to the standard assumptions of other approximations for Bayesian inference ( Hoffman et al. , 2014 ; 2013 ; Jordan et al. , 1999 ) . This allows a simple approximation of a large set of priors , including priors that are very hard to approximate with currently available tools . We make the following contributions : • We present architectural changes to successfully use Transformers for posterior predictive distribution ( PPD ) approximation , including a novel predictive distribution for regression tasks . The proposed method is simple , cheap and generally applicable to a large set of priors . • We demonstrate that PFNs can approximate the PPD of Gaussian processes and Bayesian neural networks ( BNNs ) orders of magnitude faster than MCMC with NUTS or SVI with Bayes-by-Backprop ( Blundell et al. , 2015 ) . • We demonstrate that PFNs can have an impact on real-world tasks . ( i ) We implement BNNs with priors over architectures on PFNs that enable tuning free predictions in a single forward pass and outperform all baselines on a large benchmark of small tabular datasets . ( ii ) Additionally , we show that simple handwriting priors enable few-shot learning on Omniglot ( Lake et al. , 2015 ) . 2 BACKGROUND . This paper concerns the use of a large amount of data to later yield strong performance on tasks for which only small datasets are available . In multiple machine learning setups , one can see successes based on this approach , such as fine-tuning unsupervised models on image data ( Chen & He , 2021 ; Chen et al. , 2020 ) or large language corpora ( Devlin et al. , 2018 ) . Another way to transfer knowledge from large corpora is to design prompts for language models ( Wang & Komatsuzaki , 2021 ; Brown et al. , 2020 ) . While we use a large amount of data during training , unlike the other approaches , all the data we use is generated artificially . In meta-learning , or learning-to-learn ( Hochreiter et al. , 2001 ; Finn et al. , 2017a ) , one tries to generalize learning methods from a training set of datasets to a validation set of datasets . Similar to our approach , there is a branch of research on building models that can learn to learn in a single forward pass ( Santoro et al. , 2016 ; Hochreiter et al. , 2001 ; Gordon et al. , 2019 ) . More recently ( Conditional ) Neural Processes ( Garnelo et al. , 2018a ; b ; Kim et al. , 2018 ) were proposed for meta-learning . These methods focus on architectural constructions as stochastic processes . The difference , besides the dissimilar model architectures and dissimilar prediction distributions , to previous meta-learning work is that we set up learning problems from which we sample artificial , on-the-fly-generated datasets such that we can show that our method does not only learn to learn , but that it learns to perform Bayesian inference . In this work , we are interested in Bayesian prediction for supervised learning problems , i.e. , we model the output y for a new input x based on a supervised dataset of arbitrary size n , D = { ( xi , yi ) } ni=1 , where yi is the output for xi . To address the supervised learning problem in a Bayesian way , we consider a prior p ( t ) over the latent variable t , the task . We can now use Bayes ' Theorem to define the posterior p ( t|D ) . The crucial distribution for prediction , the posterior predictive distribution ( PPD ) , can then be inferred as p ( y|x , D ) = ∫ t p ( y|x , t ) p ( t|D ) . ( 1 ) In some cases , one can arrive at the PPD p ( y|x , D ) in closed form ( e.g. , for Gaussian Processes , ( Rasmussen & Williams , 2005 ) ) , but in most cases the PPD can only be approximated . Traditionally , one approximates the posterior p ( t|D ) and based on this approximation infers an approximation of the PPD . To approximate the posterior , there are two classes of prominent methods : i ) Markov Chain Monte Carlo ( MCMC ) ( Neal , 1996 ; Andrieu et al. , 2003 ; Welling & Teh , 2011 ) methods , like NUTS ( Hoffman et al. , 2014 ) , allow accurate but sometimes very slow approximations of the posterior ; and ii ) Variational Inference ( VI ) ( Jordan et al. , 1999 ; Wainwright & Jordan , 2008 ; Hoffman et al. , 2013 ) methods approximate the posterior by a tractable distribution , such as a factorized normal distribution . One more line of research to be mentioned is amortized simulation-based inference ( Cranmer et al. , 2020 ; Chan et al. , 2018 ) , in particular neural posterior estimation ( Lueckmann et al. , 2021 ; Papamakarios & Murray , 2016 ) . The objective here is to find a model that approximates the posterior p ( t|X ) by training on samples from p ( t , X ) . Where X can be a dataset D like above , but typically is a single vector.Previous work in simulation-based inference is focused on simulations for which a specific model underlying the data is known ( Lueckmann et al. , 2021 ) . Our method is in so far similar to simulation-based inference as both solely use samples from the prior . In the next section , we will show how we can model the PPD directly , though , without ever instantiating the posterior , and apply this to large datasets with general priors . 3 PPD APPROXIMATION WITH PFNS . Let us consider a parameterized model qθ that can accept a dataset D = { ( xi , yi ) } ni=1 , as well as a query x as input , and which predicts a distribution over possible values of y for the query x . Many neural architectures can be used as such a model ; in this paper , we use a variant of Transformers ( Vaswani et al. , 2017 ) as they are a reliable and powerful architecture . We train this model by cross-entropy over samples drawn from the prior . Our proposed loss , the Prior-Data Negative Log-Likelihood ( Prior-Data NLL ) ` θ is defined as ` θ = ED∪ { x , y } ∼p ( D ) [ − log qθ ( y|x , D ) ] , ( 2 ) Algorithm 1 : Training a PFN model by Fitting Prior-Data Input : A prior distribution over datasets p ( D ) , from which samples can be drawn and the number of samples K to draw Output : A model qθ that will approximate the PPD Initialize the neural network qθ ; for j ← 1 to K do Sample D ∪ { ( xi , yi ) } mi=1 ∼ p ( D ) ; Compute stochastic loss approximation ¯̀θ = ∑m i=1 ( − log qθ ( yi|xi , D ) ) ; Update parameters θ with stochastic gradient descent on∇θ ¯̀θ ; end where D ∪ { x , y } simply is a dataset of size |D|+ 1 sampled from p ( D ) . PFNs are trained to minimize ` θ like outlined in Algorithm 1 : We draw many datasets from our prior and fit our model to predict a hold-out example correctly for these . In contrast to VI and MCMC , we learn to approximate the PPD directly from dataset samples . For all cases we consider , it is simple and cheap to sample from p ( D ) , fulfilling the only requirement our method has . MCMC on the other hand requires access to a non-normalized posterior f ( t|D , x ) ∝ p ( t|D , x ) , and VI requires access to the density values of p ( t , D ) , both of which can be hard to compute . Following similar insights as Gordon et al . ( 2019 ) , we can show that minimizing the Prior-Data NLL ` θ yields an approximation of the PPD in terms of cross-entropy and thus KL-Divergence : Theorem 1 . The proposed objective ` θ is equal to the expectation of the cross-entropy between the PPD and its approximation qθ : ` θ = Ex , D∼p ( D ) [ H ( p ( ·|x , D ) , qθ ( ·|x , D ) ) ] Proof . The above can be shown with a simple derivation . We mark transformations for easy reading . ` θ = − ∫ D , x , y p ( x , y , D ) log qθ ( y|x , D ) = − ∫ D , x p ( x , D ) ∫ y p ( y|x , D ) log qθ ( y|x , D ) ( 3 ) = ∫ D , x p ( x , D ) H ( p ( ·|x , D ) , qθ ( ·|x , D ) ) = Ex , D∼p ( D ) [ H ( p ( ·|x , D ) , qθ ( ·|x , D ) ) ] ( 4 ) Thus , we have an approximation of the PPD . We can formulate this in terms of KL-Divergence , too . Corollary 1.1 . The loss ` θ equals the expected KL-Divergence ED , x [ KL ( p ( ·|x , D ) , qθ ( ·|x , D ) ) ] between p ( ·|x , D ) and qθ ( ·|x , D ) over prior data x , D , up to an additive constant . For the proof , please see Appendix A . In the following , we consider the optimum of the given optimization problem θ∗ ∈ arg minθ ` θ. Corollary 1.2 . If p is a member of the distribution family qθ ( i.e. , there is a θ such that qθ = p ) , we have qθ∗ ( ·|x , D ) = p ( ·|x , D ) for all x , D for which p ( ·|x , D ) is defined . Proof . We assume a θ with qθ = p exists and we know that the cross-entropy is minimized for equal distributions ; thus the optimum θ∗ fulfills qθ∗ ( ·|x , D ) = p ( ·|x , D ) for all x , D where the density of the prior fulfills p ( x , D ) > 0 , which is exactly the subset for which the conditional is defined .

Summary: In this paper the authors show a simple architecture (simple on top of transformer pieces) that can learn to perform supervised learning on new datasets given collections of existing datasets. The model is trained on a collection of supervised learning datasets, and, at test time a novel training dataset and the corresponding test inputs are given and a single forward pass through the network produces the predictive distributions for those test inputs. Several different types of data are tested with close to state-of-the-art results provided, but in much less time.