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: This paper presents the posterior inference framework named Prior-Data Fitted Networks. The theoretical background is to approximate the posterior predictive distribution (PPD) with a proposal distribution and optimize the KL divergence between the proposal and the PPD. The architecture used as approximation networks is adapting from a transformer encoder, by tweaking the attention mask to output the estimation of $y$ for queries alone. In addition, the authors mainly discussed two application cases, Gaussian process (GP) and Bayesian neural networks (BNN). In the experiment, the paper demonstrates the proposed method can be successful in few-shot learning.