T RANSFORMERS C AN D O B AYESIAN I NFERENCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Currently, it is hard to reap the beneﬁts of deep learning for Bayesian methods, which allow the explicit speciﬁcation of prior knowledge and accurately capture model uncertainty.
We present Prior-Data Fitted Networks (PFNs) .
PFNs leverage large-scale machine learning techniques to approximate a large set of posteriors.
The only requirement for PFNs to work is the ability to sample from a prior distribution over supervised learning tasks (or functions).
Our method restates the objective of posterior approximation as a supervised classiﬁcation problem with a set-valued input: it repeatedly draws a task (or function) from the 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 a set of samples from a new supervised learning task as input, PFNs make probabilistic predictions for arbitrary other data points in a single forward propagation, having learned to approximate Bayesian inference.
We demonstrate that PFNs can near-perfectly 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 very diverse areas such as Gaussian process regression, Bayesian neural networks, classiﬁcation for small tabular data sets, and few-shot image classiﬁcation, demonstrating the generality of PFNs.
Code and trained PFNs are released under [ANONYMIZED; please see the supplementary material].
I NTRODUCTION
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.
It makes specifying a prior as simple as deﬁning a sampling scheme of supervised learning tasks.
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
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-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, the likelihood of a Gaussian mixture or random code in a pre-deﬁned 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 ) ﬁts 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).
B ACKGROUND
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 ﬁne-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 artiﬁcially.
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 artiﬁcial, on-the-ﬂy-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 = { ( x i , y i ) } ni =1 , where y i is the output for x i .
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 deﬁne the posterior p ( t |D ) .
The crucial distribution for prediction, the posterior predictive distribution (PPD), can then be inferred as
[Equation]
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 ﬁnd 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 speciﬁc 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.
PPD A PPROXIMATION WITH PFN S
Let us consider a parameterized model q θ that can accept a dataset D
= { ( x i , y i ) } 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) (cid:96) θ is deﬁned as
[Equation]
[Equation]
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 ∪ { ( 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 parameters θ with stochastic gradient descent on ∇ θ ¯ (cid:96) θ ; end where D ∪ { x, y } simply is a dataset of size | D | + 1 sampled from p ( D ) .
PFNs are trained to minimize (cid:96) θ like outlined in Algorithm 1: We draw many datasets from our prior and ﬁt 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 ) , fulﬁlling 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 (cid:96) θ yields an approximation of the PPD in terms of cross-entropy and thus KL -Divergence:
Theorem 1.
The proposed objective (cid:96) θ is equal to the expectation of the cross-entropy between the
PPD and its approximation q θ : (cid:96) θ
[Equation]
Proof.
The above can be shown with a simple derivation.
We mark transformations for easy reading.
[Equation]
Thus, we have an approximation of the PPD.
We can formulate this in terms of KL -Divergence, too.
Corollary 1.1.
The loss (cid:96) θ equals the expected KL-Divergence E D,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 θ (cid:96) θ .
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 deﬁned.
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 where the density of the prior fulﬁlls p ( x, D ) > 0 , which is exactly the subset for which the conditional is deﬁned.
[Figure]
A DAPTING THE T RANSFORMER FOR B AYESIAN I NFERENCE
In this section, we propose an efﬁcient architecture based on the Transformer encoder (Vaswani et al., 2017) and a novel regression head for regression problems.
An Efﬁcient Architecture Our architecture, visualized in Figure 2a, is based on a Transformer encoder without positional encodings, which makes it invariant to permutations in the dataset D .
We feed the Transformer with simple linear projections of our inputs and queries.
Our model returns the
PPD for each given query x , only depending on the dataset D and the query itself.
We sample the number of inputs n and queries m randomly such that N = m + n for some ﬁxed N .
See Appendix
D.1 for a more detailed description and Appendix D.3 for an ablation of the permutation invariance.
For additional practical training information see Appendix C.
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 .
PDFs of our distribution are bar plots, see Figure 2b.
We select the boundaries of the buckets B in which we discretize such that they all have the same probability in prior-data: p ( y ∈ b )
= 1 / | B | , ∀ b ∈ B .
We estimate this using a large prior-data sample.
In experiments with priors that need unbounded support, we replace the last bar on each side with an appropriately scaled half-normal distribution, see Figure 2c.
For a more exact deﬁnition, we refer the reader to Appendix D.2.
For an ablation of the Riemann distribution, see Appendix D.3.
Theorem 2.
A Riemann distribution q B with ﬁnite support, where B are the buckets, can model any distribution with a ﬁnite, almost everywhere continuous (Riemann-integrable) density function and full support up to arbitrary precision.
That is, ∀ p ∀ (cid:15) ∃ B .
KL ( q B , p ) ≤ (cid:15) .
For the proof, please see
Appendix B.
[Figure]
P OSTERIOR A PPROXIMATION S TUDIES
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).
Corollary 1.1 shows that the Prior-Data NLL is equal to the KL -Divergence between the exact PPD and the approximation, up to an additive constant.
We thus use it in all following experiments to compare how close methods can match the correct PPD.
For methods that do not approximate the PPD directly, we approximate the integral of the PPD in Equation 1 with a large Monte Carlo integration.
In this section, we report the 95% conﬁdence interval over different sampled datasets of the Prior-Data NLL for each experiment as shaded areas alongside the mean.
For more information of the setup of the following setup we refer the reader to Appendix E.
[Figure]
N
[Figure]
a ti v e L og - L i k e li hood
[Figure]
5. G AUSSIAN P ROCESS A PPROXIMATION
We begin by approximating GPs to study what PFNs 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 to the exact PPD.
To create the prior for this task, we sample the N inputs x i uniformly at random from the unit-cube.
The data-prior of a zero-mean GP with a kernel function k , given the x
i 's, can be described as a multivariate normal distribution N ( 0 , K ) , where K i,j = k ( x i , x j ) .
We sample from this normal distribution to yield a dataset { ( x 1 , y 1 ) , . . .
, ( x N , y N ) } .
As described above, we then train our PFNs to predict the y 's for a held-out subset of the dataset.
The trained PFN can then be used to approximate the PPD for new datasets and arbitrary query points in the unit-cube.
Figure 3 visualizes the predictions of such a PFN with two data points and query points x ranging from 0 to 1 .
We visualize both the estimated PPD as a heat map, as well as conﬁdence bounds, compared to the exact PPD.
The estimated conﬁdence intervals and means are virtually indistinguishable from the ground truth, therefore we marked the most noticeable differences with red arrows.
The model learns on its own, completely automatically, to generate smooth distributions without any explicit knowledge of the positions of the bars in R .
In Figure 7 of the appendix, we show more qualitative examples of the behavior of our approximation together with the exact posterior for a different length scale.
In Figure 8, we additionally compare to Attentive Neural Processes (Kim et al., 2018), showing the clear superiority of our method for Bayesian Inference.
Figure 4a shows that the qualitative results above generalize to datasets with up to 2000 examples with multiple features.
The tractable
GP posterior, which achieves optimal performance by deﬁnition, is approximated very closely, and increasingly better so the more meta-train datasets we train the PFN for.
5. A PPROXIMATING A GP WITH H YPER -P RIORS
A common practice in ﬁelds applying GPs in the real world is to deﬁne distributions over the hyperparameters of these GPs (Rasmussen & Williams, 2006; Snoek et al., 2012), so-called hyperpriors.
Thus, the prior of these models considers a more diverse set of functions, to model the fact that the correlation between data points, the smoothness of functions and the scale of outputs varies across applications.
The downside of using hyper-priors is that it is not possible to compute the PPD of the
GP exactly anymore.
There are two common practices to approximate it anyways, and we evaluate against both.
(i)
Firstly, MLE-II, the most common approximation.
This 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) which also inspired our hyper-priors.
(ii)
Secondly, Markov Chain Monte Carlo (MCMC), a method to sample from any distribution for which one can calculate non-normalized probabilities, that is also frequently applied to sample hyperparameters (e.g., Snoek et al.
Here, we use [Figure] og
[Figure]
i k
[Figure]
li hood
[Figure]
i
[Figure]
the state-of-the-art algorithm NUTS (Hoffman et al., 2014), which 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 1 000 × to 8 000 × faster than NUTS.
Further details can be found in Appendix E.
5. A PPROXIMATING BNN S
Bayesian Neural Networks (BNNs) provide a strong ability to model uncertainty due to their Bayesian nature, compared to traditional neural networks.
Recently, posterior approximation methods, such as stochastic variational inference (SVI) and stochastic gradient MCMC methods (such as NUTS) allow for much faster convergence, which has sparked a lot of interest in the area (Izmailov et al., 2021;
Fortuin et al., 2021).
However, computational feasibility still remains a key issue of current BNNs.
In Figure 5, we show that our PFNs can outperform today's default methods for BNN inference using a fraction of their compute budget (PFNs are 1 000 × faster than Bayes-by-Backprop SVI and 10 000 × faster than NUTS to achieve the same performance).
The ﬁgure shows the Prior-Data NLL of the BNN prior that is used during inference by all solvers.
To generate prior-data for our PFNs, we sampled random weights for a Φ and used normally-distributed i.i.d.
feature vectors
X i to obtain y i :=
Φ( X i ) with D :=
( X 1: n , y 1: n ) .
For the evaluation we sampled 100 datasets, with m =test data points and n = 100 revealed data points per dataset.
A detailed description of baselines as well as further experiments are included in Appendix E.
A PPLICATION TO T ABULAR D ATASETS
We evaluate Transformer-based classiﬁers with priors over Gaussian Processes and Bayesian Neural
Networks on a range of small real-world tabular datasets.
Surprisingly, the simple priors we describe next outperform a large set of baselines, including XGB (Chen & Guestrin, 2016) and Catboost (Prokhorenkova et al., 2018), while greatly surpassing all baselines in uncertainty calibration.
A GP Prior for Tabular Classiﬁcation
We build a classiﬁcation prior based on GPs with a hyperprior, as described in Section 5.2.
Before, we used this prior for regression.
To make it a classiﬁcation prior, we ﬁnd the median over all targets in the dataset y and set the class to one if greater than the median and zero otherwise.
A BNN
Prior over Architectures We generalize our BNN setup to learn a prior over model architectures.
Fitting a neural network requires ﬁnding a suitable architecture, e.g., embedding size, number of layers and activation function.
Commonly, resource-intensive searches need to be employed to ﬁnd suitable architectures (Liu et al., 2018;
Zoph et al., 2018).
The result of this search is not a posterior distribution over architecture choices and their respective parameters, though, but instead
(a) x (1) x (2)
H H
H
Hy (b) x (1) x (2)
H
H
H H
H
Hy x (1) x (2)
H
H
H
Hy x (1) x (2)
H 1 H
H 3 H
H 4 Hy
The network considers a distribu at once only a point-estimate of the architecture choice.
This is true for Bayesian optimisation approaches as well (Kandasamy et al., 2019; Zhou et al., 2019; Ru et al., 2021). Ensembling methods over multiple architectures can yield a rough approximation to a distribution over architectures, but this scales linearly in cost with the number of architectures considered (Zaidi et al., 2020).
Using dropout can be regarded as ensemble learning but is limited to varying connectivity (Hara et al., 2017).
PFNs allow us to be fully Bayesian, not only about the weights of the BNN, but its architecture, too.
By deﬁning a prior probability over the space of model architectures, we remove the need for ﬁnding an optimal neural architecture and instead jointly integrate over the space of architectures and their respective weights in a single Bayesian inference.
This yields a distribution over architectures in the posterior.
The degrees of freedom of model architectures we considered can be found in Appendix
F. Our prior-data sampling algorithm for learning to predict based on a prior over both weights and architectures iterates the following steps:
(1) Sample a model architecture A ∼ p ( A ) (2) Sample model weights for the given architecture: W i,j ∼ p w ( · ) (3) Sample i.i.d.
features x i,f ∼ N (0 , 1) for each data point i and each feature f ( ) Propagate each x i forward through A W yielding { ( x i , A W ( x i )) } Ni =
We binarized targets like for the GP Prior above.
See Appendix F for more details.
Setup of Datasets We used a large collection of tabular datasets from the open-source OpenML
AutoML Benchmark (Gijsbers et al., 2019); we ﬁrst removed datasets with more than one hundred features or missing values, ending up with 20 datasets that represent a diverse set of classiﬁcation problems with numerical and categorical features.
We then simpliﬁed each dataset to a balanced binary classiﬁcation problem, and, for each dataset sampled 20 subsets, each including 100 samples.
Within each subset we provide labels for the ﬁrst 30 samples and evaluate on the remaining samples.
All baselines and PFNs use the same datasets and subsets.
We also deﬁne a set of six unrelated validation datasets used for optimizing the prior distribution over architectures of the PFNs.
This is similar to setting the range of hyperparameters in a cross-validation grid search and can be reused for all similar problems.
See Appendix F for more details.
Setup of Models Our baselines comprise two Gradient Boosted Tree algorithms, XGBoost (Chen & Guestrin, 2016) and Catboost (Prokhorenkova et al., 2018), as well as logistic regression, K-Nearest Neighbors, Gaussian Processes (RBF Kernel) and Bayesian Neural Networks (Bayes-by-Backprop
SVI with Pyro).
Catboost additionally received information on which features are categorical.
We used two PFNs, one PFN for each prior described above.
We were able to apply the same model across different feature dimensions.
We used grid search with 5-fold cross-validation to optimize our baselines' hyperparameters for each dataset, using the hyperparameter spaces described in Table 6 in the appendix.
Experimental Results Table 1 compares the area under the receiver operating curve (ROC AUC) of
PFNs based on BNN and GP priors against the various baselines.
While the PFN-GP demonstrates on-par performance to the best of the evaluated baselines, the PFN-BNN performs strongest in this comparison.
The Wilcoxon signed-rank test (Demˇsar, 2006), a standard metric for comparing classiﬁers across multiple datasets is evaluated on each of the 20 subsets of the 21 evaluation datasets separately and shows a signiﬁcant performance improvement.
While all baselines make use of cross-validation to ﬁnd suitable hyperparameters, the PFNs naturally consider a distribution of hyperparameters and integrate cross-validation and training within the model's Bayesian inference step.
As shown in Table 1, this makes it very fast compared to the baselines: since all it needs to do on a new dataset is a single forward propagation step, when run on a GPU (Nvidia Tesla V100), it requires as little as 13 seconds for all 20 datasets combined.
This is more than 5000 × faster than the 20 hours XGBoost requires, and even faster than K-Nearest Neighbors. [Table]
We can also experimentally conﬁrm that our PFNs make well-adjusted uncertainty predictions.
We ﬁnd that on our data they have lower Expected Calibration Error (ECE, using 100 bins (Naeini et al., 2015)) even than standard BNNs and ensemble methods such as XGBoost, the two baseline methods that best capture uncertainty.
We also provide conﬁdence histograms in Figure 10 of the appendix, visually demonstrating the PFN's strong calibration.
A PPLICATION TO F EW -S HOT L EARNING
As shown in the previous section [Table] PFNs can be trained with priors that are not (easily) possible to use in the traditional Bayesian setups.
In this section [Table] we show another example of this, building a prior for hand-written digits and letters and demonstrating strong performance on Omniglot (Lake et al., 2015) in combination with ﬁne-tuning of the PFN on Omniglot.
As an approximation to hand-written symbols, we create a prior that generates random straight lines.
In each sample from this prior, we consider ﬁve different random line assortments as class prototypes.
For each instance of a class, we apply noise to the positions of the lines.
It required just 55 lines of readable code to create this prior (see Figure 11 in the Appendix).
In Figure 12 of the Appendix we show a sample dataset from this prior.
For our evaluation on Omniglot, we follow the setup of Rothfuss et al.
We perform 5-shot 5-way classiﬁcation, i.e. we consider ﬁve classes and ﬁve examples of each class.
We create one task from the ﬁrst ﬁve classes of each of the 30 train and 20 evaluation alphabets in Omniglot.
We train our PFN on the synthetic prior and then ﬁne-tune it on the 30 training tasks created.
This practice is unlike the previous sections.
Our premise is to show the ﬂexibility of our model: it can even be ﬁne-tuned on similar datasets, if these are available.
We do not use task augmentation, like Rothfuss et al.
but apply translations to each image, as we only use simple linear layers to encode the images in the Transformer.
In Table 2 one can see that for this very hard version of the Omniglot dataset our simple method yields results on par with the state of the art.
For more information on the setup, please see Appendix G.
[Table]
C ONCLUSION & F UTURE W ORK
We presented Prior-Fitted Networks (PFNs), a novel way to efﬁciently approximate the posterior predictive distribution using deep neural networks and showed its capability on a set of diverse tasks and priors.
We expect the following future work to be especially fruitful: (i) Work on ﬁnding novel priors that are now feasible to approximate using PFNs.
(ii) Work on architectures that are well-ﬁt for this task, as we simply used a slight adaption of current Transformer models.
(iii) Work on scaling PFNs to even larger real-world problems.
(iv) Work on using our model for the amortized simulation-based inference setting.
R EPRODUCIBILITY S TATEMENT
Code release In an effort to ensure reproducibility, we release code, including simple notebooks, for all experiment sections, which can be found alongside this submission.
Details of training procedures for PFN and baselines Details shared in the training procedure of all our transformer models can be found in Appendix C. Details for each experiment and subexperiment can be found in Appendices E, F and G. An overview of the hyperparameters used for running our baselines can be found in Table 6.
Complete proof of theoretical results A detailed proof of Corollary 1.1 and Theorem 2 can be found in appendices A and B.
Availability of datasets All datasets used in our experiments are either generated by procedures included in our code submission or freely available with downloading procedures included in the submission as well.
Further detail on the datasets used can be found in Tables 3 and 4.
A P ROOF OF C OROLLARY 1.
Proof.
This can easily be seen by considering the result of Theorem
[Equation]
where C = − E x,D [ H ( ·| x, D )] is a constant that does not depend on θ .
B P ROOF T HEOREM
Proof.
Given an arbitrary distribution p , an arbitrary bound (cid:15) .
For arbitrary minimum (maximum) bucket borders a ( b ), we consider the function b B that maps each position y to its bucket b ∈ B , where B is an arbitrary set of buckets.
We deﬁne the upper bound of p for each bucket
[Equation] [Equation]
Similarly, we deﬁne the lower bound
[Equation]
For each of these we consider their log form, too.
That is ˆ u ( y ) =
log
[Equation]
( y
[Equation]
= sup y
[Equation]
∈ b s ( y ) log p ( y (cid:48) ) and ˆ l ( y ) =
log l ( y ) =
inf y (cid:48)
∈ b s ( y ) log p ( y (cid:48) ) .
We can move the log inside the sup ( inf ), since log is steadily increasing.
Integrals over ˆ u and ˆ l thus handily represent Darboux's upper and lower sums of log p , if log p is
Riemann integrable and ﬁnite.
We assumed that p is Riemann integrable, thus log p is also integrable, since the composition of continuous functions, which log x is for x > 0 , with an almost everywhere continuous function is almost everywhere continuous and any ﬁnite value x has a ﬁnite value log x , if x > 0 .
The requirement x > 0 is given due to the assumed full support of p .
Thus, we can use ∀ (cid:15) (cid:48) ∃ B .
(cid:82) b a ˆ u ( y ) − ˆ l ( y ) dy < (cid:15) (cid:48) .
We deﬁne t = sup y ∈
[ a,b ] p ( y ) , T =
sup y ∈ R p
( y ) and the total probability mass of the interval P = p ( y ∈ [ a, b ]) .
We use (cid:15) (cid:48)
= log 1 /P , which fulﬁlls (cid:15) (cid:48) > 0 since p has full support and thus for any borders a and b we have P < 1 , get B s.t.
[Equation]
[Equation] [Equation]
[Equation]
Let c be a constant s.t. cu ( y ) yields a valid probability distribution with (cid:82) cu ( y
[Equation]
dy = (cid:82)
[Equation]
dy = 1 .
[Equation]
We use c < 1 /P , as 1 /P would be the factor c (cid:48) to make (cid:82) b a c (cid:48) p ( y ) dy
= 1 and c is this factor for u ( y ) , which has u ( y ) ≥ p ( y ) for any y ∈ [ a, b ] .
[Equation] [Equation]
[Equation] [Equation]
− log
[Equation] [Equation]
(cid:90)
[Equation]
−
[Equation]
23) ⇒ KL ( cu ( y )
[Equation]
p ( y ))
< − T log P P − log P (24) ⇒ KL ( cu ( y )
[Equation]
p ( y ))
< − ( T + P )log
[Equation]
Here, cu ( y ) is a valid bar distributions.
Now we want to ﬁnd a and b such that − ( T + 1) log PP <
[Equation]
, which is equivalent to
[Equation]
Here, we can substitute the Lambert W function to arrive at the ﬁnal result, that we should choose a and b such that p ([ a, b ]) ≥ e − W
( (cid:15)T +1 ) , which is always possible as W 0 maps to (0 , ∞ ) for positive inputs, which we have with (cid:15) T +1 .
And we have that e − x < 1 and e − x > 0 for any x ∈ (0 , ∞ ) .
Thus, we can conclude the proof.
We could upper bound the KL -Divergence for any a and b .
Then, based on that bound we found a setting for values a and b such that this upper bounds falls below (cid:15) .
C G ENERAL T RANSFORMER T RAINING D ETAILS
We follow the standard Transformer encoder setup.
In the whole discussion one should note the crucial difference to a standard MLE training: Our optimization problem does not allow for overﬁtting.
An improvement in training loss always means an improvement in the closeness to the exact posterior.
More training just improves performance and calibration.
The only exception of this is the ﬁne-tuning done for the few-shot learning experiment in Section 7.
We use Adam (Kingma & Ba, 2015) and cosine decay (Loshchilov & Hutter, 2016) with warmup.
For all experiments we used a embedding size of 512, only for few-shot classiﬁcation we used 1024.
The only hyper-parameters that we did ﬁne-tune for the Transformer training were the batch size and learning rate.
Although, it usually was enough to choose a large enough batch size while the learning rate was more important.
Learning rates generally where chosen based on training performance, if not speciﬁed differently.
We can do this as we generate new data during training and thus do not need to tune on held-out data.
We saw that search over the set { . 00001 , . 0001 , .
001 } for a learning rate generally yielded good enough results.
Usually performance of the models got better with more training, thus the number of samples used for training generally depends on our compute budget.
For setups models considered we could improve performance with more training.
We stopped when we were satisﬁed with the results.
We did not train longer than a few days on a single standard GPU for any PFN.
[Figure]
D T RANSFORMER A DAPTIONS D. A N E FFICIENT A RCHITECTURE
The ﬁrst step of the model is to encode x
and y ; in all our experiments, we use simple linear layers for this, but one could also use other application-speciﬁc encoding layers.
Our setup is 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 query points together 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 input example to every other input example and allows attention from every query points to every input example.
An example of the model and its attention can be seen in Figure 2a.
During training, we ﬁx the total number of inputs N = n + m , where n is the number of inputs and m is the number of queries, and sample n from a distribution, in order to allow the Transformer to learn to work with different dataset sizes.
Since for smaller n we have more query points to train on, we over-sample larger n , s.t.
the number of query points seen during training for each n is approximately equal.
That is, we sample each n ∈ { 0 , . . .
, N − 1 } with a weight 1 /m = 1 / ( N − n ) .
As prediction head we use a softmax for multi-class classiﬁcation tasks, a sigmoid for binary classiﬁcation tasks and the Riemann distribution for regression.
[Figure]
D. R IEMANN D ISTRIBUTION D EFINITION
Given a set of of buckets B , where each bucket b ∈ B describes an interval such that
[Equation]
(28) and for any b, b (cid:48) ∈ B , we have b ∩ b (cid:48)
= ∅ .
Additionally, we deﬁne the upper-bound (lowerbound) of the lowest (highest) bucket to be l , and a function B ( · ) that maps a y ∈ R to the unique bucket b with y ∈ b .
We assume to have a function w mapping to the width of each bucket as w ( b ) =
max y ∈ b
y − min y ∈ b y .
Lastly, we assume the model gives us a probability p b for each bucket b .
In the ﬁnite case the Riemann 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 = y (cid:48) ) .
For the case of inﬁnite support, we have
[Equation]
Here, the Half-Normal distributions are scaled such that half the probability weight is inside the bounds of the ﬁrst (last) bucket.
D. A BLATION S TUDIES
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 show the impact of both approaches on the Transformer in Figure 9.
Notice 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.
For each setup we searched over a set of the same four learning rates and evaluated the best performing on another sample from the distribution.
[Figure]
E D ETAILS FOR S ECTION
The error intervals that can be seen in the Prior-Data NLL plots are 95% percent conﬁdence intervals over different dataset samples from the prior.
Details for Section 5.1 Here, we generally followed default settings for GPs.
That is, we used the RBF-Kernel with zero-mean, a length scale of 0.6, 1e-4 additional noise on the diagonal of the covariance matrix.
The output scale was set to 1.
Only for the plots in the Figure 7 and Figurethe length scale was decreased to 0.1.
The experiments with more examples in Figure 4a did useinput dimension, while the qualitative visualizations only used one.
For all methods we sampledatasets for evaluation, each of which has a single validation example, for each dataset size.
Details for Section 5.2 In this section, we followed the default settings that can be found in
BoTorch's SingleTaskGP (Balandat et al., 2020).
We only changed the distribution over the diagonal noise, as the default is very high, such that it is hard to see any improvements with any of the methods in NLL with an increasing number of training samples.
The GP is setup as a Mat ´ ern Kernel with ν = 2 .
and we have the following hyper-priors: the noise was sampled from Γ( α = 0 . 0001 , β = 1 . 0) , the output scale from Γ( α = 2 . 0 , β = 0 . 15) and the length scale from Γ( α = 3 . 0 , β = 6 . 0) .
We used the Pyro (Bingham et al., 2018) implementation of NUTS.
Here, we used every sampled hyper-parameter setting as a sample GP.
We averaged the output distribution over all sampled GPs to yield an approximation to the PPD.
We scaled both the warmup and the rollout steps together: with “8 (warmup) steps” we mean 8 warmup steps and then 8 sampling steps.
Only the MCMC experiments were run on CPU, as this was faster than using the GPU.
For all methods we sample 1000 datasets for evaluation, each of which has a single validation example, for each dataset size.
Details for Section 5.3 For SVI (Bayes-by-Backprop) we followed the setup of (Blundell et al., 2015) and used Adam (Kingma & Ba, 2015) as optimizer.
We leveraged Pyro (Bingham et al., 2018) for optimization.
We scaled both the training steps and number of samples averaged together to increase approximation quality.
For MCMC we used the Pyro (Bingham et al., 2018) implementation of NUTS and averaged the output distribution over all sampled BNNs to yield an approximation to the PPD.
We scaled both the warmup and the rollout steps together to increase approximation performance.
We used one MCMC chain.
PFN and SVI are evaluated on a GPU (Nvidia Tesla V100) while MCMC is evaluated on a
CPU since it was faster on CPU than on GPU.
[Figure] F D ETAILS FOR S ECTION
Accomodating a variable number of features Since the evaluated datasets used a variable number of features, the dimensionality of the input is variable.
This can not be accommodated by the ﬁxed size of the linear layer used in our PFNs.
Thus, we ﬁxed a maximum number of features (60) and extended the datasets to this maximum number by appending zeros.
We did this during prior training, sampling the number of features uniformly at random alongside the model architecture and during inference.
We rescale the features by a factor of max features features used during prior training and bayesian inference to keep the mean input scale constant with a varying feature number. Degrees
of freedom for model architectures The evaluated BNN models consider a wide range of model architectures, that is those architectures vary in terms of: number of layers, number of hidden units per layer, layer sparsity, standard deviation of Gaussian noise applied to each unit and activation function applied to each unit.
This list can be easily extended by adapting the sampling procedure.
Tuning hyperparameter ranges for hyper-priors and cross-validation While the Transformerbased methods need no parameter tuning on each dataset individually, they still require setting the prior hyperparameters.
This is similar to setting the range of hyperparameters in a cross-validation grid search and can be reused for all similar problems, e.g., all small-scale tabular classiﬁcation.
We therefore run a hyperparameter search on a second set of validation datasets, listed in Table 4.
The hyperparameters considered in this grid search can be found in 5 In this search, we ﬁnd good settings for all hyper-parameters of our priors for the average over all validation datasets.
This search results in one ﬁnal hyper-parameter setting and only one single Transformer model that is reused for all datasets.
We use the same setup to ﬁnd cross-validation ranges for our baselines.
Standardization of datasets Each subset is standardized to have zero mean and unit variance where the statistics for the standardization are only calculated on the samples considered by the model, thus without any dependence on the test set.
Runtimes of prior ﬁtting step Fitting our transformer with GP prior was done over 100 epochs, taking 2:44h, while ﬁtting the BNN prior took 3:14h for the same number of epochs.
Crucially, this model is reused for all datasets.
G D ETAILS FOR S ECTION
We searched for the hyper-parameters speciﬁed as inputs in Figure 11 with a grid search, where we measure performance on the training tasks as validation set.
We then trained our model as a PFN on samples from this prior.
After this stage we ﬁne-tuned the model on the training tasks of Omniglot.
Here, we applied translations and used a default batch size of 100.
We used a learning rate of 1eto yield the performance shown in the plot after sampling 500,000 samples from the training tasks.
Unlike above, we sample over different samples of our ﬁne-tuned model to yield the conﬁdence interval, to follow the standard reporting on the Omniglot benchmark.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] [Table]
Training Examples
[Table]
[Figure]
[Table]
, scaling the output range from [0, 1] to [max, min] and then rounding the output to the integer range. beta sampler f(a, b) refers to a function that samples from the beta distribution without scaling and rounding [Table] [Table]
[Table]
[Table]
[Table]
[Table]
20 data-subsets are evaluated for each dataset.
The ECE is a calibration measure.
The novel
PFN-BNN performs strongest overall.