Meta-Learning to Perform Bayesian Inference in a single Forward Propagation
Abstract
Currently, it is hard to reap the beneﬁts of deep learning for Bayesian methods.
We present Prior-Data Fitted Networks (PFNs) , a method that allows to employ large-scale machine learning techniques to approximate a large set of posteriors.
The only requirement for PFNs is the ability to sample from a prior distribution over supervised learning tasks (or functions).
The method repeatedly draws a task (or function) from this prior, draws a set of data points and their labels from it, masks one of the labels and learns to make probabilistic predictions for it based on the set-valued input of the rest of the data points.
Presented with samples from a new supervised learning task as input, it can then make probabilistic predictions for arbitrary other data points in a single forward propagation, effectively having learned to perform Bayesian inference.
We demonstrate that PFNs can nearperfectly mimic Gaussian processes and also enable efﬁcient Bayesian inference for intractable problems, with over 200-fold speedups in multiple setups compared to current methods.
We obtain strong results in such diverse areas as Gaussian process regression and Bayesian neural networks, demonstrating the generality of
PFNs.
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 ﬂexible and replaceable priors using deep learning models.
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 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).
In particular, this is true for ML applications to small datasets; thus, a large number of specialized algorithms have been developed for different tasks (LeCun et al., 1989; Kadra et al., 2021).
Encoding prior information into such a model can, however, be challenging. [Figure] Train [Figure]
by minimizing
[Figure] A well-deﬁned 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 or the likelihood of a Gaussian mixture.
In practice, however, it is often very difﬁcult or intractable to retrieve the Posterior Predictive Distribution (PPD) for a given prior exactly.
In this paper, we use large-scale neural networks to approximate Bayesian models.
We outline our method of Prior-data Fitted Networks (PFNs) in Figure 1.
We assume a given representative prior distribution over supervised learning tasks (or functions), which provides our inductive bias.
To train a PFN, we repeatedly sample a meta-train task (or function) from this 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 (meta-test) dataset, we feed it and a test point as inputs to the PFN and the network outputs its prediction for the test point, conditioned on the dataset.
As we will demonstrate, this probabilistic prediction approximates exact
Bayesian posterior prediction.
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 assumption that unnormalized posterior probabilities can be computed, which is made by currently common methods, such as Markov Chain Monte Carlo (MCMC) Neal (1996); Andrieu et al.
; Welling & Teh (2011); Hoffman et al.
and
Variational Inference (VI) (Jordan et al., 1999; Wainwright & Jordan, 2008; Hoffman et al., 2013).
This weak requirement of being able to sample 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 introduce the technique of Prior-data Fitted Networks (PFNs); • We show that via meta learning on data sets sampled from a dataset prior, PFNs yield an approximation technique for the posterior predictive distribution (PPD);
• We present architectural changes to successfully use Transformers for PPD approximation, including a novel predictive distribution for regression tasks;
• We demonstrate the successful application of PFNs of in approximating the PPD for Gaussian processes and Bayesian neural networks.
Background on Bayesian Inference
Bayesian inference relies on priors for real world applications.
By combining the knowledge from a general prior distribution with the data from a particular task, one obtains a posterior distribution, which can be used to predict outputs for new data points from the task at hand.
In this work, we are interested in supervised learning problems, i.e., we model the relationship between inputs x and outputs y .
Generally, we consider predictions based on datasets of arbitrary size n , D
= { ( x i , y i ) | i = 1 , . . . , n } , where y i is the output for x i .
There are two kinds of Bayesian models commonly used: i)
In parametric models e.g. Bayesian Neural Networks (BNNs we model a distribution p
( t
over functions or tasks explicitly, where
t
is
the random variable
representing
the
task. In parametric models , t can be represented explicitly, e.g., as a vector, and either represents a conditional distribution p ( y | x, t ) or a joint distribution p ( x, y | t ) over the data.
Thus , we can generate a prior over datasets
p ( D ) = p ( t ) (cid:81) ni =1 p ( x i , y i | t ) , where in the case of the conditional
deﬁnition p ( x i y i | t ) factors into p ( y i | x i , t ) and a a simple task-independent distribution p ( x i ) .
ii) Non-parametric models , e g., Gaussian Processes (GPs), can be viewed similarly to the above, but here we cannot explicitly instantiate t
Thus we directly model , that means without the indirection over the task t , both the prior on data sets p ( D ) , as well as the PPD p ( y | x, D
For both models, the prior distribution (PD) p ( D ) is easily accessible.
For all cases we consider, it is simple and cheap to sample from p(D), fulﬁlling the only requirement our method has.
In practice, using Bayesian inference has the advantages that (i) it has a theoretical foundation that makes it valid in setups which the prior p ( t ) ﬁts, (ii) it can thus better account for the actual likelihood of different events; (iii) it is well calibrated and (iv) it is very interpretable as the prior describes the expectation of the model.
PPD Approximation with Prior-Data Fitting
Let us consider a parameterized model q θ that can accept a set D
= { ( x i , y i ) | i = 1 , . . . , n } , as well as a query x as input, and which predicts a distribution over possible values of y for the query x .
Many current neural network models 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 Data
Prior Negative Log-Likelihood (Prior-Data NLL) (cid:96) is deﬁned as
[Equation]
where D ∪ { x, y } is simply a dataset of size | D | + 1 sampled from p ( D ) .
We minimize the Prior-Data NLL using a stochastic approximation of (cid:96) θ , as we show in Algorithm 1.
That means, we draw many samples of datasets from our prior and ﬁt our model to predict a hold-out example correctly for these.
Algorithm 1: Training the model by Fitting Prior-Data
Result:
A model q θ that will approximate the PPD
Deﬁne data prior p ( D ) ; initialize neural network q θ ; while not converged do
Sample D ∪ { ( x i , y i ) } mi =1 ∼ p ( D ) ;
Compute stochastic loss approximation ¯ (cid:96) θ = (cid:80) mi =1 ( − log q θ ( y i | x i , D )) ;
Update model θ with gradient ∇ θ ¯ (cid:96) θ ; end
We can show that minimizing (cid:96) yields an approximation to the true PPD in a similar way as in Gordon et al.
Theorem 1.
The proposed objective (cid:96) θ is equal to the expectation of the cross-entropy E x,D [ H ( p ( ·| x, D ) , q θ ( ·| x, D ))] between the PPD and its approximation.
Proof.
The above can be shown with a simple derivation.
We mark transformations for easy reading.
[Equation]
We see that by optimizing (cid:96) θ , we optimize for similarity in terms of cross-entropy between the model and the PPD across the support of p ( x, D ) .
The proof above implies a minimization in terms of KL -Divergence, too.
Corollary 1.1.
The loss (cid:96) θ is up to a constant, equal to the expected KL-Divergence of p ( ·| x, D ) with q θ ( ·| x, D ) over prior data x, D .
For the proof, please refer to Appendix A.
In the following, we consider the optimum of the given optimization problem θ ∗ = arg min θ (cid:96) θ .
We consider the optimal result of our process for the case where q θ describes a distribution family that can exactly approximate the posterior, that means there is a θ such that q θ = p .
Corollary 1.2.
If q θ is a distribution family that can exactly approximate the posterior, we have q θ ∗ ( ·| x, D ) =
p ( ·| x, D ) for all x, D with p ( x, D ) > 0 .
Proof.
We assume a θ with q θ = p exists
and we know that the cross-entropy is minimized for equal distributions; thus the optimum θ ∗ fulﬁlls q θ ∗ ( ·| x, D ) =
p ( ·| x, D ) for all x, D with p ( x, D ) > 0 .
Adapting the Transformer for Bayesian Inference
We use a standard transformer (Vaswani et al., 2017) with only a few changes to its architecture to ﬁt our problem better.
Transformers were developed with sequence modelling and prediction in mind, but in our application we want to map a set D
= { ( x i , y i ) }
Ki =1 of data points ( x i , y i ) and a query x to a distribution q θ ( y | x, D ) over y , which in many cases should be continuous.
We propose an efﬁcient architecture that allows batch-wise training and prediction, as well as a novel regression head for continuous y .
4.1 An Efﬁcient Architecture
The ﬁrst step in the architecture is to encode x and y ; in all our experiments, we use simple linear layers for this, to project from the given dimension to the dimension of the transformer, but one could also use other application-speciﬁc encoding layers.
Our setup should be equivariant to the ordering of the input dataset’s examples.
To achieve this, we remove positional encodings and feed ( x, y ) pairs together to the transformer as a sum of their encodings; see Appendix D for an ablation of the impact of the equivariance to the ordering.
In practice, we feed multiple query points to the transformer for efﬁciency reasons; these are the only inputs for which position matters, and we thus use the output at their positions as prediction for their corresponding PPD over y .
Since different query points should not depend on each other, we use an attention mask that allows attention from every ( x, y ) pair to every other ( x, y ) pair and allows attention from all query points to the ( x, y ) pairs.
The mask M we use can be characterized by the number of input examples m and the number of query points n , we write M m,n .
An example for
[Figure]
−− 1 030 .0 .y
D e n s it y [Figure] −− 1 030 .0 .y
q
D e n s it
y such a matrix is
[Equation] [Equation] We predict the distribution of y from the output of the last time step. During training, we ﬁx the total size M = m + n and sample m from a distribution, in order to allow the Transformer to learn to work with different dataset sizes.
Since for smaller m we have more query points, we over-sample larger m , s.t.
the number of query points seen during training for each m is equal during training.
That is, we sample each m ∈ { 0 , . . .
, M − 1 } with a weight 1 /n =
1 / ( M − m ) .
4.2 Riemann Distribution
The modelling of continuous distributions is hard with neural networks.
To yield strong performance in modelling PPDs, we developed a distribution that works particularly well with neural networks.
Based on the knowledge that neural networks work well for classiﬁcation and inspired by discretizations in distributional reinforcement learning (Bellemare et al., 2017), we created a discretized continuous distribution that we call Riemann Distribution , as it is based on a similar idea to approximation as Riemann integrals.
PDFs of our distribution are bar plots, see Figure 2a.
The problem remaining is how to select the boundaries of the buckets.
To make the problem balanced for the Transformer, we simply select buckets such that the likelihood of each bucket b has the same prior predictive probability p ( y ∈ b )
= 1 / | B | , ∀ b ∈ B .
We approximate this by taking a large sample of y values from our prior and adjusting the buckets such that an equal number of these values fall in each bucket.
The distributions described above are very powerful; e.g., in contrast to most parametric distributions they can also represent multimodal distributions.
Their downside, however, is that they have ﬁnite support: the probability for y to be greater than the highest border is 0.
This is a problem whenever one tries to model distributions that have unbounded support.
A simple trick that worked well for us is to replace the last bar on each side with an appropriately scaled half-normal distribution.
We visualize a resulting PDF in Figure 2b.
For an exact mathematical deﬁnition of the distributions, we refer the reader to Appendix C. For an ablation of the Riemann distribution, see Appendix D.
Posterior Approximation Studies
In our ﬁrst set of experiments, we study the capability of PFNs to perform Bayesian inference for the tractable case of Gaussian Processes (GPs) with ﬁxed hyperparameters (where we can compare to ground truth data; Section 5.1) and the intractable cases of GPs with unknown hyperparameters (Section 5.2) and Bayesian Neural Networks (BNNs; Section 5.3).
Based on Theorem 1, we present a novel way to evaluate methods approximating PPDs and Posteriors, the Prior-Data NLL itself.
We simply take the negative log-likelihood of a method on data generated from the prior over datasets.
We can do this, of course, for our models, but we can do this as easily with other well-known methods.
We simply treat samples from the prior as training and test sets 1 .x y 0 .0 .0 .0 . 1 .x For additional comparisons, see for these methods and evaluate the negative log-likelihood of the approximated distributions on the predicted outputs for the test sets.
As we show in Corollary 1.1, the metric we evaluate is directly representative of the KL divergence between the true PPD and the approximation.
There only is a constant difference.
If a method only approximates the posterior and not the PPD directly, we approximate PPD with a large MC integration.
5.1 Gaussian Process approximation
We begin by approximating GPs as a study as to what the transformers are capable of; GPs with ﬁxed hyperparameters are convenient for this purpose since the tractability of their PPD allows us to assess how close our approximation is.
In this case, our meta-train datasets are functions f sampled from a GP with the given ﬁxed hyperparameters.
Since GPs describe conditional distributions given an input x , we sample n + 1 inputs x i uniformly at random from the unit-cube and evaluate y i = f
( x i ) .
As usual, we then train our PFNs to predict y n +1 from the dataset { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } and the query point x n +1 .
The trained PFN can then be used to compute the PPD for new datasets and arbitrary query points.
Figure 3 visualizes the predictions of such a PFN for a small dataset with two data points and query points x ranging from 0 to 1 .
In Figure 3 (left), one can still see the very tiny boxes in which the Riemann distribution is discretized.
In Figure 3 (right), we compare these predictions to those of a ground-truth GP with the given hyperparameters; both the mean function and the plotted 95% conﬁdence interval are virtually indistinguishable from this ground truth.
The model also learns on its own, completely automatically, to generate smooth distributions without any explicit knowledge of the positions of the bars in R .
In Figure 6 in the appendix we show more qualitative examples of the behavior of our approximation together with the exact posterior.
Figure 4a shows that the qualitative results above generalize to much large datasets with multiple inputs.
The
GP Posterior achieves optimal performance by deﬁnition, and we can approximate it very closely, and better so with longer training.
Additionally, there does not seem to be a trend towards worse approximation with larger datasets.
5.2 Approximating A GP with Hyper-Priors
A common practice in ﬁelds applying GPs in the real world is to deﬁne hyper-priors over the hyperparameters of GPs Rasmussen & Williams (2006); Snoek et al.
Thus, the prior of these models considers a more diverse set of functions.
The correlation between data points, the smoothness of the function and the scale of outputs can be variable.
This allows the user to apply
GPs to a larger set of tasks.
The downside is that it is not possible to compute the PPD of the GP with hyper-priors exactly.
There are two common practices to approximate it anyways and we evaluate against both.
(i)
Firstly, MLE-II is the most common, which is a simple special case of variational inference.
Here, suitable hyperparameters are found by maximum a posteriori (MAP) estimation (e.g., Rasmussen & Williams (2006)).
We use the ﬁtting setup of BoTorch (Balandat et al., 2020) from which our Hyper-Priors for all models are inspired as well.
(ii)
Secondly, Markov Chain Monte Carlo (MCMC), a method to sample from any distribution for which one can calculate non-normalized probabilities, is also frequently applied to sample hyperparameters (e.g. Snoek et al.
Here
(a)
[Figure] [Figure]
N e g a ti v e L og - L
i k
e li hood (500Kdatasets) (1Mdatasets) (
[Figure]
[Figure] [Figure] [Figure]
−
[Figure]
N e g a ti v e L og - L
i
k
e li hood
[Figure]
we compare against doing so by the NUTS method (Hoffman et al., 2014), a state-of-the-art method that uses gradients to facilitate the Hamiltonian Monte Carlo algorithm.
We plot the performance of the different methods in Figure 4b.
Our model can approximate the PPD on the prior much closer than both of these methods, and its inference is more than 200 × faster than MLE-II and 1000 × to 8 000 × faster than NUTS.
5.3 Approximating the BNN PPD
Bayesian Neural Networks (BNNs) provide an inherent ability to model uncertainty and strong regularization by inherent ensembling.
Recently, approximation methods such as stochastic variational inference or stochastic gradient Markov Chain Monte Carlo methods allow for much faster convergence, which has sparked a great deal of interest in the area (Izmailov et al., 2021; Fortuin et al., 2021).
However, these methods still converge too slowly for many practical applications.
In Figure 5, we show that our method can outperform today's default methods using a much smaller compute budget.
The ﬁgure shows the approximation accuracy on synthetic datasets generated using the BNN prior that is used during inference by all solvers.
To generate these synthetic datasets, we sampled random weights for a BNN BNN and used normally-distributed i.i.d.
features X to obtain y :=
BNN ( X ) with D :=
( X 1: n , y 1: n ) .
Conclusion & Future Work
We present a novel way to approximate the posterior predictive distribution very efﬁciently using deep neural networks.
We show it’s capability on a set of diverse tasks and priors.
This opens the possibility for a multitude of ground breaking future work.
We expect especially the following to be fruitful: (i) Work on ﬁnding novel priors that are cheap to sample for particular application areas.
(ii)
Work on architectures that are well-ﬁt for this task, as we simply used a slight adaption of a current model.
A Proof of Corollary 1.
Proof.
This can easily be seen by considering the result of Theorem[-
-]
[Equation]
[Equation]
11) where C = − E x, D [ H ( ·| x, D ]] .
[h] 0 .
− + y 0 .
− +  0 .
− + x y 0 .
[-e (cid:15) 
-]
− + x
C Riemann Distribution Deﬁnition
Given a set of of buckets B , where each bucket b ∈ B describes an interval such that (cid:83) b ∈ B b =
[ a, b ] and for any b, b (cid:48)
∈ B we have b ∩ b (cid:48)
= ∅ .
Additionally, we deﬁne the upper-bound (lower-bound) of the lowest (highest) bucket to be l ( f ).
Additionally, we use a mapping B ( · ) that maps a y ∈ R to the unique bucket b with y ∈ b and the width of each bucket to be deﬁned as w ( b ) =
max y ∈ b
y − min y ∈ b y .
The model gives us a probability for each bucket p b .
In the ﬁnite case the distribution is deﬁned as
[Equation]
We normalize with the width of the bucket since p b describes the probability p ( y (cid:48) ∈ b ) , but we are interested in the probability p ( y (cid:48) =
y ) .
For the case of inﬁnite support given probabilities p b for each bucket b , we have
[Equation]
otherwise
[Equation]
Here, the Half-Normal distributions are scaled such that half the probability weight is in the ﬁrst (last) bucket.
D Ablation
Above, we introduce a novel way of doing regression with neural networks that departs from the tradition of using normal distributions.
Additionally, we remove the default positional embeddings from the Transformer.
We shot the impact of both on the Transformer in Figure 7.
We can see that, while the effect of the Riemann Distribution is very pronounced, the effect of making our architecture permutation invariant only slowly increases with sequence length.
[Figure]