PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers Can Do Bayesian Inference

Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, Frank Hutter

arXiv:2112.10510v7cs.LGstat.ML

TL;DR

Deep learning has not easily delivered Bayesian methods’ benefits of explicit priors and uncertainty, especially for small-data tasks. PFNs train on datasets sampled from a prior to approximate posterior predictions in one forward pass, achieving broad results across probabilistic models and real-world tasks, including large speedups over current methods.

  • Problem

    Deep learning has made major progress with abundant data, but encoding prior knowledge and computing Bayesian posterior predictive distributions remains difficult, especially for small-data problems.

  • Method

    PFNs train Transformers with set-valued datasets sampled from a prior, masking labels so the network learns probabilistic posterior-predictive predictions directly.

  • Results

    PFNs approximate Gaussian-process and Bayesian-neural-network posterior predictive distributions orders of magnitude faster than MCMC with NUTS or SVI with Bayes-by-Backprop, while also performing strongly on tabular and few-shot image tasks.

  • Takeaways & Limitations

    PFNs provide a broadly applicable way to approximate Bayesian inference whenever supervised-learning tasks can be sampled from a prior.

Abstract

from arXiv · show

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 in-context learning in 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 at https://github.com/automl/TransformersCanDoBayesianInference.

1 INTRODUCTION

PFNs use sampled supervised-learning tasks to train deep networks that approximate Bayesian posterior predictions for small-data problems. The approach requires only a sampleable prior and achieves broad empirical results across probabilistic models and real-world tasks.

  • Motivation: Deep learning has advanced with large datasets, but encoding prior knowledge remains challenging for smaller-scale machine-learning tasks.Bayesian inference offers a principled way to specify prior beliefs and represent uncertainty, but posterior predictive distributions are often intractable.
  • Method: PFNs train on datasets sampled from a prior, masking labels so a Transformer learns probabilistic predictions from the remaining set-valued examples.At inference time, the PFN receives an actual dataset and test point and predicts in a single forward propagation.
  • Method: A sampleable prior is the main requirement for PFNs to approximate posterior predictive distributions across a broad set of priors.This includes priors that are difficult to approximate with currently available Bayesian-inference tools.
  • Contributions: PFNs introduce Transformer adaptations for posterior predictive approximation, including a novel predictive distribution for regression tasks.The proposed method is described as simple, cheap, and generally applicable to many priors.
  • Results: PFNs approximate Gaussian-process and Bayesian-neural-network posterior predictive distributions orders of magnitude faster than MCMC with NUTS or SVI with Bayes-by-Backprop.The comparison concerns posterior predictive distribution approximation for these probabilistic models.
  • Results: On real-world tasks, PFNs outperform all baselines on a large small-tabular-dataset benchmark and support few-shot Omniglot classification.The tabular setup uses Bayesian neural networks with priors over architectures, while handwriting priors support the few-shot image task.

2 BACKGROUND

The paper places PFNs within transfer learning, meta-learning, Bayesian posterior prediction, and amortized simulation-based inference. Its distinctive background framing is direct posterior-predictive modeling from sampled tasks rather than explicit posterior approximation.

  • Transfer learning: Large-scale training can transfer knowledge to tasks with limited data, motivating approaches that use abundant data to improve small-dataset performance.Examples include fine-tuning unsupervised models and designing prompts for language models.
  • Meta-learning: Meta-learning generalizes learning methods across datasets, while PFNs use artificially generated datasets to learn Bayesian inference in a single forward pass.The comparison emphasizes both the shared learning-to-learn setting and PFNs’ Bayesian objective.
  • Bayesian prediction: Bayesian posterior prediction models a new output y for input x using a supervised dataset D and a prior over the latent task variable t.Bayes’ theorem defines p(t|D), which supports the posterior predictive distribution p(y|x, D).
  • Approximate inference: Posterior predictive distributions are sometimes available in closed form, but most cases require approximation through methods such as MCMC or variational inference.MCMC methods can be accurate but sometimes very slow, whereas the passage introduces variational methods as another prominent class.
  • Simulation-based inference: Amortized simulation-based inference trains models from prior samples, whereas PFNs model the posterior predictive distribution directly without instantiating the posterior.The paper applies this direct modeling approach to datasets with general priors.

3 PPD APPROXIMATION WITH PFNS

PFNs approximate posterior predictive distributions by training a model on datasets sampled from a prior and predicting masked labels. Their Prior-Data NLL corresponds to expected predictive cross-entropy, and at the optimum can recover the exact PPD when representable.

  • Model and objective: PFNs accept a dataset D and query x, then predict a distribution over possible y values for x.The paper uses a Transformer variant trained by cross-entropy on prior-sampled data.
  • Model and objective: Training repeatedly samples datasets from the prior and fits the model to predict a held-out example from the remaining data.This is implemented with stochastic gradient updates over prior-data samples.
  • Comparison with Bayesian approximations: Unlike VI and MCMC, PFNs learn the PPD directly from samples of the prior, requiring only the ability to sample prior datasets.The paper contrasts this with VI's density access and MCMC's non-normalized posterior access requirements.
  • Model and objective: The Prior-Data NLL equals expected cross-entropy between the exact PPD and its approximation, and equals expected KL divergence up to an additive constant.These identities justify using the loss to measure PPD approximation quality.
  • Model and objective: If the exact PPD belongs to the model family, an optimum of the objective recovers it wherever the PPD is defined.This is a representational guarantee conditional on the stated family-membership assumption.

4 ADAPTING THE TRANSFORMER FOR BAYESIAN INFERENCE

The paper adapts Transformer encoders for permutation-invariant dataset inputs and introduces the Riemann Distribution to model continuous predictive distributions.

  • Efficient Architecture: The proposed architecture uses a Transformer encoder without positional encodings, making it invariant to permutations of dataset D.Inputs and queries are linearly projected, and the model returns a PPD for each query based on D and that query.
  • Efficient Architecture: The model samples the numbers of inputs and queries so their total remains a fixed N, supporting varied set-valued inputs.Queries are included for computational efficiency and are the only inputs for which position matters.
  • Riemann Distribution: The Riemann Distribution discretizes continuous outputs into equal-prior-probability buckets represented as neural-network-friendly bar plots.For unbounded-support priors, the outer bars are replaced with appropriately scaled half-normal distributions.

5 POSTERIOR APPROXIMATION STUDIES

PFNs closely approximate Bayesian predictive distributions across tractable and intractable Gaussian-process and neural-network settings. They match or outperform established inference methods while using substantially less computation.

  • Experimental setup: Fixed-hyperparameter GP experiments provide a tractable ground-truth comparison, while hyperparameter-mixed GPs and BNNs test intractable inference.The studies evaluate both predictive accuracy and computational cost.
  • Gaussian Process Approximation: PFN means and confidence intervals for fixed-hyperparameter GPs are virtually indistinguishable from the exact posterior predictive distribution.The approximation remains very close for datasets with up to 2000 examples and multiple features.
  • Gaussian Process Approximation: More than 200× faster than baselines, the PFN approximates mixed-hyperparameter GP PPDs more closely than MLE-II and NUTS.The reported speedups are more than 200× versus MLE-II and 1 000× to 8 000× versus NUTS.
  • Bayesian Neural Networks: 1 000× faster than Bayes-by-Backprop SVI and 10 000× faster than NUTS, PFNs achieve the same BNN inference performance with a fraction of the compute budget.The comparison uses Prior-Data NLL for the BNN prior shared by all solvers.

6 APPLICATION TO TABULAR DATASETS

PFNs use Gaussian-process and Bayesian-neural-network priors for small tabular classification, with the BNN-prior model performing strongly while requiring only a single forward pass for prediction.

  • Method: The BNN prior samples architectures, weights, features, and propagated outputs to represent uncertainty over model architectures.This extends Bayesian treatment beyond network weights to architecture choices.
  • Experimental Results: PFN-GP matches the best evaluated baselines, while PFN-BNN performs strongest overall on the tabular classification comparison.The evaluation uses ROC AUC across subsets of small real-world datasets.
  • Experimental Results: 13 seconds is sufficient to process all 20 datasets combined on an Nvidia Tesla V100, compared with 20 hours for XGBoost.PFNs require only a single forward propagation on a new dataset.
  • Method: PFNs integrate hyperparameter distributions, cross-validation, and training within the model's Bayesian inference step.This enables predictions without separately searching for a single architecture or hyperparameter setting.
  • Experimental Results: PFNs achieve lower Expected Calibration Error than standard BNNs and ensemble baselines such as XGBoost.The authors report this using 100-bin ECE and confidence histograms.

7 APPLICATION TO FEW-SHOT LEARNING

The Omniglot experiment tests whether PFNs can use a simple synthetic handwriting prior and subsequent fine-tuning for few-shot classification.

  • Method: PFNs achieve strong Omniglot performance using a prior that generates handwritten symbols from random straight lines.Each class prototype consists of a random line assortment, with positional noise applied to instances.
  • Evaluation: The evaluation uses 5-shot 5-way classification across tasks formed from Omniglot alphabets.The PFN is trained on the synthetic prior and fine-tuned on 30 training tasks.
  • Results: Table 2 compares meta-learning algorithms by test accuracy, with bold values denoting results within the confidence interval of the best performance.The supplied caption specifies the comparison and highlighting rule but not the numerical entries.

8 CONCLUSION & FUTURE WORK

The paper concludes that PFNs efficiently approximate posterior predictive distributions across diverse tasks and priors, while identifying scaling, architecture, and prior design as future directions.

  • Conclusion: PFNs efficiently approximate posterior predictive distributions using deep neural networks across diverse tasks and priors.The conclusion presents this as the paper's demonstrated capability.
  • Future Work: Future work includes novel priors, architectures specialized for PFNs, scaling to larger real-world problems, and amortized simulation-based inference.These directions are explicitly listed by the authors.

A PROOF OF COROLLARY 1.1

The proof establishes that the stated Prior-Data objective recovers the exact posterior predictive distribution under the assumed realizability and support conditions.

  • Proof: If a parameter θ satisfies qθ = p, cross-entropy optimality yields qθ*(·|x,D) = p(·|x,D) wherever p(x,D) > 0.This identifies the optimizer with the exact conditional distribution on its defined support.
  • Proof: The proof constructs bucket-wise upper and lower bounds for the density and their logarithms to control the approximation error.The bounds are connected to Darboux sums under integrability and full-support assumptions.
  • Proof: Choosing interval borders with sufficient probability mass makes the resulting KL-divergence upper bound fall below any prescribed ϵ.The argument uses the Lambert W function to select suitable borders.

D GENERAL TRANSFORMER TRAINING DETAILS

The PFN uses an order-equivariant Transformer to predict labels from set-valued context, with a Riemann distribution for regression outputs and training over varied context sizes.

  • Optimization: Unlike standard MLE training, improving training loss directly improves closeness to the exact posterior, while additional training improves performance and calibration.The stated exception is fine-tuning for the few-shot learning experiment.
  • Transformer architecture: PFN inputs combine encoded (x, y) pairs without positional encodings, making predictions invariant to the ordering of dataset examples.The architecture sums pair encodings before feeding them to the Transformer; application-specific encoders could also be used.
  • Training setup: Training fixes N = n + m and samples n so the model learns to operate across different numbers of inputs and queries.Each n is weighted by 1/m = 1/(N − n), approximately equalizing the number of query points seen for each n.
  • Riemann distribution: The Riemann distribution converts bucket probabilities into a continuous regression density by dividing each bucket probability by its width.This normalization reflects that the model predicts p(y′ ∈ b), whereas the target is a pointwise density p(y = y′).
  • Riemann distribution: For infinite support, the regression density uses Half-Normal tails outside the lowest and highest bucket and the bucket-based density between them.The Half-Normal distributions are scaled so half their probability mass lies within the bounds of the first or last bucket.

E.3 ABLATION STUDIES

The ablations examine the Riemann regression distribution and permutation-invariant architecture, finding a pronounced distribution effect and a sequence-length-dependent equivariance effect.

  • Ablation studies: The Riemann Distribution has a very pronounced effect on Transformer performance when fitting the Gaussian Process prior.The comparison uses the same Gaussian Process as Section 5.1.
  • Ablation studies: Permutation invariance improves with sequence length, but its effect increases only slowly compared with the Riemann Distribution.The ablation removes positional embeddings to make the architecture permutation invariant.

F DETAILS FOR SECTION 5

The experiments use standardized benchmark datasets and varied GP, BNN, and image-task configurations, with calibration and ROC AUC evaluation supporting the reported comparisons.

  • Evaluation: Evaluation uses 95% confidence intervals over different prior-dataset samples, and Table 7 reports ROC AUC and ECE on 30-sample training subsets.Table 7 states that the novel PFN-BNN performs strongest overall.
  • Input dimensionality: Variable-feature datasets are padded to 60 features with zeros, and features are rescaled to keep mean input scale constant as feature count changes.The padding and rescaling are performed during prior training and Bayesian inference.
  • BNN configurations: The evaluated BNN priors vary in layers, hidden units, sparsity, Gaussian noise, and activation functions, with the sampling procedure extensible to additional architectures.These variations define the model-architecture range used in evaluation.
  • Hyperparameter selection: PFN prior hyperparameters are tuned on separate validation datasets, while baseline cross-validation uses the hyperparameter settings listed in the accompanying tables.The tuned prior hyperparameters can be reused across similar small-scale tabular classification problems.
  • Preprocessing: Each dataset subset is standardized using statistics computed only from the samples supplied to the model, without test-set dependence.Standardization gives each subset zero mean and unit variance.
  • Runtime: The GP-prior Transformer takes 2:44h to fit and the BNN-prior Transformer 3:14h over 100 epochs, with each fitted model reused across datasets.The reported fitting times concern the prior-fitting stage rather than per-dataset inference.
  • Few-shot image classification: For Omniglot, the model is fine-tuned on training tasks after prior training, using translations, batch size 100, learning rate 1e-5, and 500,000 sampled training examples.The reported configuration follows grid-search selection of the Figure 11 prior inputs.
  • Datasets: The tabular evaluation uses OpenML AutoML Benchmark datasets filtered to at most 100 features and no missing values.The dataset inventory is provided in Table 3.
Loading 2112.10510v7…