TRANSFORMERS CAN DO BAYESIAN INFERENCE
Year: 2021

Abstract:
Currently, it is hard to reap the benefits of deep learning for Bayesian methods, which allow the explicit specification 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 classification 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 efficient 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, classification for small tabular data sets, and few-shot image classification, demonstrating the generality of PFNs. Code and trained PFNs are released under [ANONYMIZED; please see the supplementary material].

heading: 1 INTRODUCTION
text: 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).

heading: 2 BACKGROUND
text: 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.

heading: 3 PPD APPROXIMATION WITH PFNS
text: 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.

heading: 4 ADAPTING THE TRANSFORMER FOR BAYESIAN INFERENCE
text: In this section, we propose an efficient architecture based on the Transformer encoder (Vaswani et al., 2017) and a novel regression head for regression problems.
An Efficient 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 fixed 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 classification 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 definition, we refer the reader to Appendix D.2. For an ablation of the Riemann distribution, see Appendix D.3.
Theorem 2. A Riemann distribution qB with finite support, where B are the buckets, can model any distribution with a finite, almost everywhere continuous (Riemann-integrable) density function and full support up to arbitrary precision. That is, ∀p∀ ∃B.KL(qB, p) ≤ . For the proof, please see Appendix B.

heading: 5 POSTERIOR APPROXIMATION STUDIES
text: In our first set of experiments, we study the capability of PFNs to perform Bayesian inference for the tractable case of Gaussian Processes (GPs) with fixed 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% confidence 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.

heading: 5.1 GAUSSIAN PROCESS APPROXIMATION
text: We begin by approximating GPs to study what PFNs are capable of; GPs with fixed 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 xi uniformly at random from the unit-cube. The data-prior of a zero-mean GP with a kernel function k, given the xi's, can be described as a multivariate normal distribution N (0,K), where Ki,j = k(xi, xj). We sample from this normal distribution to yield a dataset {(x1, y1), . . . , (xN , yN )}. 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 confidence bounds, compared to the exact PPD. The estimated confidence 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 definition, is approximated very closely, and increasingly better so the more meta-train datasets we train the PFN for.

heading: 5.2 APPROXIMATING A GP WITH HYPER-PRIORS
text: A common practice in fields applying GPs in the real world is to define 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 fitting 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. (2012)). Here, we use
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.

heading: 5.3 APPROXIMATING BNNS
text: 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 figure 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 Xi to obtain yi := Φ(Xi) with D := (X1:n, y1:n). For the evaluation we sampled 100 datasets, with m = 200 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.

heading: 6 APPLICATION TO TABULAR DATASETS
text: We evaluate Transformer-based classifiers 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 Classification We build a classification prior based on GPs with a hyperprior, as described in Section 5.2. Before, we used this prior for regression. To make it a classification prior, we find 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 finding a suitable architecture, e.g., embedding size, number of layers and activation function. Commonly, resource-intensive searches need to be employed to find 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)
H1 H2 H3 H4
y
(b)
x(1) x(2)
H1 H2 H3
H1 H2 H3
y
x(1) x(2)
H1 H2 H3 H4
y
x(1) x(2)
H1 H2 H3H2 H4 H3
y
Figure 6: (a) A BNN. The inputs x are mapped to the output y. (b) A prior over BNN architectures. The network considers a distribution over different model architectures 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 defining a prior probability over the space of model architectures, we remove the need for finding 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: Wi,j ∼ pw(·) (3) Sample i.i.d. features xi,f ∼ N (0, 1) for each data point i and each feature f (4) Propagate each xi forward through AW yielding {(xi, AW (xi))}Ni=1 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 first removed datasets with more than one hundred features or missing values, ending up with 20 datasets that represent a diverse set of classification problems with numerical and categorical features. We then simplified each dataset to a balanced binary classification problem, and, for each dataset sampled 20 subsets, each including 100 samples. Within each subset we provide labels for the first 30 samples and evaluate on the remaining samples. All baselines and PFNs use the same datasets and subsets. We also define 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 (Demšar, 2006), a standard metric for comparing classifiers across multiple datasets is evaluated on each of the 20 subsets of the 21 evaluation datasets separately and shows a significant performance improvement. While all baselines make use of cross-validation to find 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.
We can also experimentally confirm that our PFNs make well-adjusted uncertainty predictions. We find 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 confidence histograms in Figure 10 of the appendix, visually demonstrating the PFN's strong calibration.

heading: 7 APPLICATION TO FEW-SHOT LEARNING
text: As shown in the previous section, PFNs can be trained with priors that are not (easily) possible to use in the traditional Bayesian setups. In this section, 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 fine-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 five 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. (2021). We perform 5-shot 5-way classification, i.e. we consider five classes and five examples of each class. We create one task from the first five classes of each of the 30 train and 20 evaluation alphabets in Omniglot. We train our PFN on the synthetic prior and then fine-tune it on the 30 training tasks created. This practice is unlike the previous sections. Our premise is to show the flexibility of our model: it can even be fine-tuned on similar datasets, if these are available. We do not use task augmentation, like Rothfuss et al. (2021), 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.

heading: 8 CONCLUSION & FUTURE WORK
text: We presented Prior-Fitted Networks (PFNs), a novel way to efficiently 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 finding novel priors that are now feasible to approximate using PFNs. (ii) Work on architectures that are well-fit 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.

heading: 9 REPRODUCIBILITY STATEMENT
text: 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.

heading: A PROOF OF COROLLARY 1.1
text: Proof. This can easily be seen by considering the result of Theorem 1
Ex,D[KL(p(·|x,D), qθ(·|x,D))] (5) = −Ex,D [∫
y
p(y|x,D) log qθ(y|x,D) p(y|x,D)
] (6)
= −Ex,D [∫
y
p(y|x,D) log qθ(y|x,D) ] + Ex,D [∫
y
p(y|x,D) log p(y|x,D) ]
(7)
= Ex,D[H(p(·|x,D), qθ(·|x,D))]− Ex,D[H(p(·|x,D))]] (8) = `θ + C, (9)
where C = −Ex,D[H(·|x,D)] is a constant that does not depend on θ.

heading: B PROOF THEOREM 2
text: Proof. Given an arbitrary distribution p, an arbitrary bound . For arbitrary minimum (maximum) bucket borders a (b), we consider the function bB that maps each position y to its bucket b ∈ B, where B is an arbitrary set of buckets. We define the upper bound of p for each bucket
u(y) = sup y′∈bs(y)
p(y′). (10)
Similarly, we define the lower bound
l(y) = inf y′∈bs(y)
p(y′). (11)
For each of these we consider their log form, too. That is û(y) = log u(y) = supy′∈bs(y) log p(y ′) and l̂(y) = log l(y) = infy′∈bs(y) log p(y ′). 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 finite. 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 finite value x has a finite value log x, if x > 0. The requirement x > 0 is given due to the assumed full support of p. Thus, we can use ∀ ′∃B. ∫ b a û(y) − l̂(y)dy < ′. We define t = supy∈[a,b] p(y), T = supy∈R p(y) and the total probability mass of the interval P = p(y ∈ [a, b]). We use ′ = log 1/P , which fulfills ′ > 0 since p has full support and thus for any borders a and b we have P < 1, get B s.t.
∫ b a û(y)− l̂(y)dy < log 1 P
(12)
⇒ ∫ b a t(û(y)− l̂(y))dy < t log 1/P (13)
⇒ ∫ b a t log u(y) l(y) dy < t log 1/P (14)
⇒ ∫ b a u(y) log u(y) l(y) dy < t log 1/P (15)
Let c be a constant s.t. cu(y) yields a valid probability distribution with ∫ cu(y)dy = ∫ b a cu(y)dy = 1.
⇒ ∫ b a cu(y) log u(y) l(y) dy < ct log 1/P (16)
⇒ ∫ b a cu(y) log cu(y) l(y) dy < ct log 1/P + ∫ b a cu(y) log c (17)
⇒ ∫ b a cu(y) log cu(y) l(y) dy < ct log 1/P + log c (18)
We use c < 1/P , as 1/P would be the factor c′ to make ∫ b a c′p(y)dy = 1 and c is this factor for u(y), which has u(y) ≥ p(y) for any y ∈ [a, b].
⇒ ∫ b a cu(y) log cu(y) l(y) dy < ct log 1/P + log 1/P (19)
⇒ ∫ b a cu(y) log cu(y) l(y) dy < (ct+ 1) log 1/P (20)
⇒ ∫ b a cu(y) log cu(y) l(y) dy < −t logP P − logP (21)
⇒ ∫ b a cu(y) log cu(y) l(y) dy < −T logP P − logP (22)
⇒ ∫ b a cu(y) log cu(y) p(y) dy < −T logP P − logP (23)
⇒ KL(cu(y), p(y)) < −T logP P − logP (24)
⇒ KL(cu(y), p(y)) < −(T + P ) logP P
(25)
⇒ KL(cu(y), p(y)) < −(T + 1) logP P . (26)
Here, cu(y) is a valid bar distributions. Now we want to find a and b such that −(T + 1) logPP < , which is equivalent to
1 P log 1/P < T + 1 (27)
Here, we can substitute the Lambert W function to arrive at the final result, that we should choose a and b such that p([a, b]) ≥ e−W0( T+1 ), which is always possible as W0 maps to (0,∞) for positive inputs, which we have with 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 .

heading: C GENERAL TRANSFORMER TRAINING DETAILS
text: 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 overfitting. 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 fine-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 classification we used 1024. The only hyper-parameters that we did fine-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 specified 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, .0003, .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 satisfied with the results. We did not train longer than a few days on a single standard GPU for any PFN.

heading: D TRANSFORMER ADAPTIONS
text: D.1 AN EFFICIENT ARCHITECTURE
The first 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-specific 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.3 for an ablation of the impact of the equivariance to the ordering. In practice, we feed query points together for efficiency 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 fix the total number of inputsN = n+m, where n is the number of inputs andm 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 classification tasks, a sigmoid for binary classification tasks and the Riemann distribution for regression.
D.2 RIEMANN DISTRIBUTION DEFINITION
Given a set of of buckets B, where each bucket b ∈ B describes an interval such that⋃ b∈B b = [a, b] (28)
and for any b, b′ ∈ B, we have b ∩ b′ = ∅. Additionally, we define 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) = maxy∈b y−miny∈b y. Lastly, we assume the model gives us a probability pb for each bucket b.
In the finite case the Riemann distribution is defined as
p(y) = pb(y)/w(b(y)). (29)
We normalize with the width of the bucket since pb describes the probability p(y′ ∈ b), but we are interested in the probability p(y = y′).
For the case of infinite support, we have
f(yq) =  2 · fN (yq − l), if yq ≥ l 2 · fN (f − yq), if yq ≥ f pb(yq)/b(yq), otherwise.
(30)
Here, the Half-Normal distributions are scaled such that half the probability weight is inside the bounds of the first (last) bucket.
D.3 ABLATION STUDIES
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.

heading: E DETAILS FOR SECTION 5
text: The error intervals that can be seen in the Prior-Data NLL plots are 95% percent confidence 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 Figure 8 the length scale was decreased to 0.1. The experiments with more examples in Figure 4a did use 5 input dimension, while the qualitative visualizations only used one. 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.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 Matérn Kernel with ν = 2.5 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.
(a) Bayesian Neural Network
(b) PFN-BNN
(c) XGBoost

heading: F DETAILS FOR SECTION 6
text: 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 fixed size of the linear layer used in our PFNs. Thus, we fixed 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 featuresfeatures 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 classification. 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 find good settings for all hyper-parameters of our priors for the average over all validation datasets. This search results in one final hyper-parameter setting and only one single Transformer model that is reused for all datasets. We use the same setup to find 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 fitting step Fitting our transformer with GP prior was done over 100 epochs, taking 2:44h, while fitting the BNN prior took 3:14h for the same number of epochs. Crucially, this model is reused for all datasets.

heading: G DETAILS FOR SECTION 7
text: We searched for the hyper-parameters specified 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 fine-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 1e-5 to 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 fine-tuned model to yield the confidence interval, to follow the standard reporting on the Omniglot benchmark.
d e f p e n c i l s t r o k e p r i o r ( n u m c l a s s e s =2 , s i z e =28 , m i n m a x s t r o k e s = ( 1 , 3 ) , min max len = ( 5 / 2 8 , 2 0 / 2 8 ) , m i n m a x s t a r t = ( 2 / 2 8 , 2 5 / 2 8 ) , min max width = ( 1 / 2 8 , 4 / 2 8 ) , m a x o f f s e t = 4 / 2 8 , m a x t a r g e t o f f s e t = 2 / 2 8 ) :

Review summaries:
This paper describes a novel way to do bayesian inference with deep learning models by employing a meta-learning approach. This is done by creating a synthetic prior over datasets, training a deep learning model on samples from this prior, which results in an approximation of the posterior predictive distribution over the datasets. The authors demonstrate their approach on two synthetic tasks: approximating gaussian processes and bayesian neural networks, then on two real world tasks: the OpenML tabular datasets and the few-shot learning Omniglot dataset.

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.

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.