PaperlayerView exact revision on arXiv

Source-linked AI summary

How Well Do Prior-Data Fitted Networks Amortize Bayesian Inference?

Manuel Hülskamp, Julius Kobialka, Emanuel Sommer, Thomas Nagler, David Rügamer

arXiv:2610.13764v1stat.MLcs.LG

TL;DR

Evidence that PFNs recover their target PPD is limited, and real deployment may shift the task distribution from the pretraining prior. The paper decomposes deployment error and studies amortization and Bayesian model averages over PFN weights. It reports that deep ensembles consistently improve predictions, whereas sampling and variational learning add little.

  • Problem

    Evidence on how well trained PFNs recover the PPD they amortize is thin, and deployment may involve a different task distribution from the pretraining prior.

    Source: §1 · p. 1

  • Method

    The paper decomposes deployment error and studies amortization error, using disagreement across independently trained PFNs and Bayesian model averages over their weights.

    Source: §2.2 · p. 2, §3.2 · p. 4

  • Results

    Deep ensembles consistently improve predictions, whereas sampling and variational learning add little.

    Source: §1 · p. 2

  • Takeaways & Limitations

    PFN disagreement provides a label-free probe of the removable part of amortization error, but does not reveal error shared by ensemble members.

    Source: §3.2 · p. 4

    First-order stochastic optimization provides no general guarantee that a pretrained PFN is close to its target PPD.

    Source: §2.2 · p. 2

Abstract

from arXiv · show

Prior-data fitted networks (PFNs) are trained to approximate the Bayesian posterior predictive distribution (PPD) of a synthetic prior, yet results about the approximation quality are surprisingly thin. Taking an empirical risk minimization perspective, we decompose the deployment error of PFNs into prior, context, and amortization error, and measure the latter against a reference PPD under a realistic prior. Results show that PFNs generally recover the targeted PPD very well. Nevertheless, we find that parts of the residual amortization error persist, especially in \textit{rare} settings under the prior pretraining distribution, even when scaling the model size, which manifests in disagreement between independently pretrained PFNs and leaves uncertainty about the predictor at inference time. As a remedy, we study multiple types of Bayesian model averages and find that these can consistently reduce this uncertainty and improve predictions up to full-size TabICLv2 models.

1 INTRODUCTION

Evidence that PFNs recover the posterior predictive distribution (PPD) they are trained to approximate has been limited, while deployment can involve tasks drawn from a different distribution. The paper studies sources of PFN error and uncertainty, including amortization and Bayesian model averaging over network weights.

1 INTRODUCTION

  • Evidence that trained PFNs closely track the PPD they amortize is thin and has focused more on whether transformers can perform Bayesian inference than on how well they recover it.The cited earlier focus is Müller et al. (2021).Source: §1 · p. 1
  • Pretraining is commonly treated as first-order optimization of a single network, leaving open whether effectively unlimited synthetic data removes uncertainty about its outcome.The paper frames this as a question about amortization uncertainty.Source: §1 · p. 1
  • Real-data deployment may shift the task distribution from the pretraining prior, changing the optimal predictor and complicating error analysis.The paper identifies this shift as an additional challenge when analyzing deployed PFNs.Source: §1 · pp. 1–2
  • The paper decomposes deployment error into prior, context, and amortization error, and studies the amortization error and Bayesian model averages over PFN weights.Its contribution summary reports that deep ensembles consistently improve predictions, whereas sampling and variational learning add little.Source: §1 · p. 2

2 AMORTIZATION ERROR

The section frames PFNs as predictors trained under a prior over tasks and context sizes, then defines deployment risk and amortization error. It measures approximation against an MCMC reference and identifies limits to extending that analysis to realistic priors and shifted deployment tasks.

2.1 Setup

  • A PFN assumes a prior over task and context-size pairs, draws iid contexts from each task, and targets the prior’s posterior predictive distribution.The task is a law over features and labels, and Ψ is the prior’s marginal distribution over context sizes.Source: §2.1 · p. 2

2.2 Risk Under a Task Distribution

  • Under a task distribution, the corresponding PPD uniquely minimizes risk, and its risk is the lowest attainable; amortization error measures a PFN’s expected KL distance from its target PPD.In practice, stochastic first-order optimization provides no general guarantee that the trained network is close to that target.Source: §2.2 · p. 2
  • Deployment error includes prior error when deployment tasks differ from the training prior, while context error remains even for the exact PPD.Only the amortization term is controlled by pretraining.Source: §2.2 · p. 2

2.3 Measuring Amortization Error

  • Directly measuring amortization error requires the true PPD, which is intractable for current tabular priors, and distribution shift further complicates analysis on real deployments.Earlier evaluations were low-dimensional, used no current PFN architecture at scale, and considered only small context sizes.Source: §2.3 · pp. 2–3
  • The PFN approximates the MCMC reference PPD very well across context sizes, while the reference still dominates and the gap depends mostly on dataset probability under the prior.The figure compares aggregate NLL for a classification task using Gaussian and TabICL prior feature distributions; deviations are noted at small context sizes for uniform Ψ and large context sizes for loguniform Ψ.Source: Figure 2 · p. 3, §2.3 · p. 3
  • A realistic reference prior retains covariate structure from an SCM prior but uses a Bayesian neural network labeling mechanism, making its PPD tractable for MCMC comparison.The covariate generator and neural-network weights are drawn independently, leaving the posterior over weights unchanged after conditioning on a query covariate.Source: §2.3 · p. 3

3 AMORTIZATION ERROR UNDER DEPLOYMENT

The section separates PFN deployment error into components with different sources and remedies, then shows how disagreement between independently trained members reveals a removable part of amortization error. It also characterizes how deployment shift changes the interpretation of that disagreement.

3.1 Amortization Error Without a Reference

  • Without a tractable reference PPD, amortization error is analyzed as estimation, approximation, and optimization error.Estimation error is expected to decrease at a 1/M rate; approximation depends on hypothesis-class capacity, while optimization error is not guaranteed to fall with model size or training duration.Source: §3.1 · pp. 3–4

3.2 Shared & Removable Amortization Error

  • Averaging PFN predictions improves expected log loss whenever members disagree under the task distribution.The mixture’s expected log loss is never worse than that of the average member, with equality only when members agree almost everywhere.Source: §3.2 · p. 4
  • Across PFNs, average amortization error splits into shared error and a removable, member-varying part measurable as the mixture’s log-loss gain over its average member.Member disagreement metrics provide label-free probes of the removable part, but are blind to shared error.Source: §3.2 · p. 4

3.3 Decomposing the Deployment Error

  • Deployment excess risk separates into coverage, transport, off-support, context, and amortization terms; only amortization is controlled by pretraining.Coverage, amortization, and context terms are non-negative, while transport and off-support terms are sign-indefinite.Source: §3.3 · p. 5
  • Coverage, transport, and off-support terms describe how the training prior relates to deployment, but quantifying these terms for PFNs remains open.Coverage vanishes only when every deployment task can be expressed by the prior; transport reflects whether deployment overweights tasks on which the PFN performs poorly.Source: §3.3 · p. 5

4 BAYESIAN MODEL AVERAGES OVER PFN WEIGHTS

The section motivates Bayesian averaging of PFN weights as a remedy for uncertainty that persists across independently pretrained networks. It describes a tempered weight posterior and reports that averaging can improve predictions when its members disagree.

4 BAYESIAN MODEL AVERAGES OVER PFN WEIGHTS

  • When residual amortization error persists and is partly removable, a Bayesian model average over network weights is the proposed remedy.The authors frame this as addressing uncertainty in pretraining outcomes and additional uncertainty at deployment.Source: §4 · p. 5

A Tempered Posterior Over PFN Weights

  • A tempered posterior over PFN weights uses β as an effective sample size to decouple posterior concentration from the number of pretraining datasets M.Its predictions average S samples from an approximation of that posterior; the paper compares three approximations.Source: A Tempered Posterior Over PFN Weights · p. 6

4 BAYESIAN MODEL AVERAGES OVER PFN WEIGHTS

  • Averaging members never worsens expected log loss and strictly improves it when they disagree; positive dAU JSD provides a label-free strict-gain certificate for full-support outcomes.On the expressible part, averaging also reduces the removable amortization error.Source: A Tempered Posterior Over PFN Weights · p. 6

5 EXPERIMENTAL STUDY

The experiments test persistent PFN disagreement across pretraining-data and capacity scales, then compare Bayesian averaging methods from smaller models through full-size TabICLv2. Deep ensembles deliver the strongest and most consistent gains, including at full scale.

5.1 Setup

  • Experiments evaluate TabICLv2 at four scales and nanoTabPFN at two, using synthetic TabICL-prior tasks and real-world TabArena and TALENT datasets.The models are pretrained on synthetic tasks with up to 1024 rows, 100 features, and 10 classes.Source: §5.1 · p. 6

5.2 Disagreement Along Scaling Axes

  • Increasing pretraining data reduces estimation error at the expected rate, but disagreement flattens beyond about 10^4 datasets and a small share remains.Across the larger training scale, disagreement falls from 9.1% to 7.6% for TabICLv2 and from 6.7% to 3.7% for nanoTabPFN.Source: §5.2 · pp. 6–7

Deep Ensembles

  • Deep ensembles improve ACC, ROC AUC, and LPPD over their average member across model and context sizes, with 1.7–3.1% LPPD gains within and 1.8–2.7% beyond the pretraining range.At 1000 rows, an ensemble of one size roughly matches a single member of the next larger size.Source: Figure 4 · p. 7, Deep Ensembles · p. 7

5.2 Disagreement Along Scaling Axes

  • Increasing model size from 1.5M to 9M parameters reduces TabArena d_AU JSD by 41% and disagreement from 7.5% to 5.9%, but does not eliminate residual disagreement.Contexts beyond the pretraining range do not reduce the disagreement rate either.Source: §5.2 · p. 7
  • The largest averaging gains plausibly occur on inputs rare under the prior, where more pretraining or larger models barely reduce uncertainty.The difference to the MCMC reference and the averaging gain both concentrate on a small share of test inputs.Source: §5.2 · p. 7

Deep Ensembles

  • Within-basin methods recover some removable error but less than deep ensembles: IVON gains 0.2–2.5% within the pretraining range, while pSMILE adds up to 0.42%.IVON gains shrink with capacity and saturate by S ≈ 8; beyond the pretraining range, additional IVON samples do not help, and adaSGHMC provides no reliable additional gains.Source: Deep Ensembles · pp. 7–8

5.4 Deep Ensembles of Full-Size TabICLv2

  • At every training stage, the eight-member ensemble beats its average member on all 38 TabArena classification and 13 regression datasets, and beats the best single member on classification.On TabArena classification, its Elo of 1552 is within the confidence interval of the official checkpoint at 1575.Source: Figure 5 · p. 8, §5.4 · p. 8

Deep Ensembles

  • Plain deep ensembles deliver most of the attainable gain at matched inference cost, without tempering or sampler hyperparameters.The members train independently.Source: Deep Ensembles · p. 8

6 DISCUSSION

PFNs closely track their target posterior predictive distribution, with small deviations concentrated on prior-rare inputs. Independently pretrained ensembles address persistent disagreement; the analysis is limited by its reference and model coverage.

6 DISCUSSION

  • PFN deviations from the target posterior predictive distribution are small and concentrated on inputs rare under the prior.Residual disagreement between independently pretrained models is also concentrated on few inputs and is barely reduced by more pretraining data or capacity.Source: §6 · p. 8
  • The analysis measures amortization gaps against finite-budget MCMC and empirically validates claims only for TabICLv2 and nanoTabPFN.Some studies also use smaller models.Source: §6 · p. 8

AI Use Statement

  • The authors used generative AI to assist with implementation, proof drafting and review, manuscript presentation, and literature exploration, and report reviewing AI-assisted work.Source: AI Use Statement · p. 9

CHECKLIST

The checklist records reproducibility and theoretical-reporting items, while the appendices provide related-work context and proof assumptions. The related work distinguishes predictor-level amortization uncertainty from uncertainty conditional on a fixed PFN.

CHECKLIST

  • The checklist points to mathematical-setting explanations and assumption statements, with proofs provided in Appendix B.Source: CHECKLIST · p. 13
  • The checklist reports code and instructions for reproducing experiments, training details, metric and error-bar definitions, and computing infrastructure.Source: Appendix B · p. 13

A RELATED WORK

  • Related work traces PFNs to Müller et al. (2021) and identifies TabICL (Qu et al., 2025, 2026) as an open-source alternative for tabular foundation models.Source: A RELATED WORK · p. 14
  • The related work separates uncertainty conditional on a fixed predictor from amortization uncertainty about the predictor itself, which can reflect stochastic pretraining.Source: A RELATED WORK · p. 14
  • Prior PFN studies mainly establish whether models implement Bayesian inference, while evidence on how closely they recover their intended posterior predictive distribution remains limited.The paper revisits and extends earlier experiments, noting limitations in released code from Müller et al. (2021).Source: A RELATED WORK · p. 14

B PROOFS

  • Appendix B states finiteness, projection, and support assumptions and defines the task-and-context laws used in the proofs.Source: B PROOFS · p. 15

B.1 Proof of Proposition 1

The proofs establish the prior PPD as the unique risk-minimizing predictor, decompose deployment and ensemble error, and characterize which disagreement averaging can remove.

B.1 Proof of Proposition 1

  • The prior PPD uniquely minimizes risk: every predictor has at least its risk, with equality only when it matches the PPD almost surely.The proof applies the KL inequality and its equality condition.Source: §B.1 · p. 15

B.2 Formal Statement and Proof of Proposition 2

  • Deployment error separates into coverage, amortization, context, transport, and off-support terms, distinguishing prior mismatch from approximation and deployment effects.The amortization term is the only one directly targeted by pretraining; context error depends on the prior, while transport and off-support terms also depend on deployment.Source: Figure 6 · p. 16, §B.2 · pp. 16–17

B.3 Proof of Lemma 1

  • Ensemble diversity gain is nonnegative and vanishes exactly when all member predictions agree almost everywhere under the task distribution.The result follows from strict concavity of the logarithm.Source: §B.3 · p. 17

B.4 Proof of Corollary 1

  • The mixture’s amortization error is the part shared by all members; averaging removes their disagreement component.The decomposition applies pointwise to predictions and then averages over the prior.Source: §B.4 · p. 17

B.5 Ensemble gain under the deployment distribution

  • Under deployment, averaging gain includes the members’ prediction-diversity gain plus terms reflecting deployment reweighting and off-support behavior.The gain under the deployment distribution is related to its prior-average value through the change in task weighting.Source: §B.5 · p. 18

C A TRACTABLE SYNTHETIC PRIOR

The section evaluates PFNs against an MCMC reference for a tractable BNN prior, including shifted inputs and regression. PFNs track the reference closely overall, but errors concentrate in difficult or sparsely covered settings, and the reference itself has limitations.

Where the amortization error concentrates

  • 60–66% of KL at n = 64 and over 80% at n ≥1024 comes from the 10% highest-KL inputs, which lie near the decision boundary.At n = 4096, the 20% of inputs with the smallest predicted class-1 probability carry 80% of KL; the 10% highest-KL tasks carry at most a third.Source: C A TRACTABLE SYNTHETIC PRIOR · p. 19
  • At n = 4096, PFN KL to the reference exceeds a single chain’s by 5.5 (11), and PFN predictions fall outside all 24 chain predictions for over a third (two thirds) of inputs.The larger parenthetical gaps and fractions correspond to log-uniform row counts; the gap grows with context size and the PFN becomes underconfident at large n.Source: C A TRACTABLE SYNTHETIC PRIOR · pp. 19–20
  • For n ≥64, PFN and reference calibration slopes are both within 0.03 of 1 against the oracle, so aggregate miscalibration does not explain the amortization error.Calibration can be close to ideal even when a predictor remains far from the PPD.Source: C A TRACTABLE SYNTHETIC PRIOR · p. 20

Testing the PFN under deployment distribution shifts

  • Under covariate shift, the PFN–MCLMC NLL gap grows with feature scale for every n ≥64, and at n = 4096 exceeds MCLMC’s excess NLL over the oracle.Pretraining input-distribution choices matter more under this shift.Source: C A TRACTABLE SYNTHETIC PRIOR · p. 20

Synthetic prior amortization gap for regression

  • PFNs come very close to the MCMC reference across the pretraining context-size range, while residual amortization error concentrates near decision boundaries and can exceed finite-budget reference error several-fold.The boundary is also where chains disagree most, limiting pointwise attribution of the discrepancy.Source: C A TRACTABLE SYNTHETIC PRIOR · pp. 20–21
  • In regression, the PFN nearly matches MILE on CRPS, while quantile-derived density makes its NLL gap larger; low-noise tasks drive intermediate-context gaps, and MILE can become over-dispersed at large n.The 999-quantile density is sensitive to differences between neighboring predicted quantiles, inflating NLL while barely affecting CRPS; at the largest contexts, the PFN can outperform the reference.Source: C A TRACTABLE SYNTHETIC PRIOR · pp. 20–21, Figure 9 · p. 21

D AMORTIZATION DISAGREEMENT METRICS

The section defines disagreement measures across independently trained PFNs and distinguishes distributional disagreement from disagreement in predicted decisions.

D AMORTIZATION DISAGREEMENT METRICS

  • Amortization disagreement metrics: label-free probes of between-member predictive variation, using distributional JSD and decision-level argmax disagreement.The JSD captures distributional disagreement, while the pairwise rate registers only prediction changes in the argmax class.Source: D AMORTIZATION DISAGREEMENT METRICS · pp. 21–22

E.1 Implementation and data

The appendix documents implementation and evaluation choices and extends analyses of pretraining-data scaling, model capacity, and Bayesian model averages.

E.1 Implementation and data

  • Evaluation uses 255 TALENT datasets for reported results and 30 in-scope TabArena classification datasets, with five-fold cross-validation and five random subsamples.The TALENT-tiny core's 45 datasets are validation data only; full-size models use all 38 TabArena classification and 13 regression datasets.Source: §E.1 · p. 22

E.2 Bayesian model averages

  • Deep-ensemble members use disjoint pretraining data; SG-MCMC chains warm-start from ensemble members, while IVON averages draws from a diagonal-Gaussian variational posterior.The IVON comparison includes its posterior mean as the S=0 baseline.Source: §E.2 · pp. 22–23, Table 7 · p. 27

E.3 Full-size TabICLv2

  • Full-size TabICLv2 ensembles use eight members trained through three stages, extending contexts to 60,000 rows in the final stage.The recipe uses bfloat16, QK-normalization, and bounded QASSMax scaling.Source: §E.3 · pp. 23–24, Table 5 · p. 24

F.1 Pretraining data scaling

  • The tiny-model data-scaling analysis plots estimation error and member disagreement against the number of pretraining datasets M.Figure 11 adds synthetic-test disagreement for tiny and small models; neither caption states the outcome.Source: §F.1 · p. 24, Figure 10 · p. 25, Figure 11 · p. 25

F.3 Disagreement concentration

  • Pointwise ensemble gain is the mixture's log-probability advantage over members' average log-probability, and the analysis tracks its concentration among highest-gain test points.Figure 12 compares tiny and large TabICLv2 ensembles with a benchmark of gain spread evenly across points.Source: §F.3 · p. 26, Figure 12 · p. 26

F.4 Deep ensembles; F.5 Variational sampling via IVON; F.6 SG-MCMC sampling

  • The appendix also presents nanoTabPFN deep-ensemble scaling on TabArena and IVON sampling comparisons on TALENT classification datasets.The IVON table compares averaged posterior samples with the posterior mean; the SG-MCMC section compares two samplers with the deep ensemble.Source: §F.4 · p. 26, Figure 13 · p. 27, §F.5 · p. 27

F.6 SG-MCMC sampling

  • Despite extensive tuning, no adaSGHMC setting both moved the sampler and avoided performance loss versus the deep ensemble.The comparison used 30 TabArena classification datasets.Source: §F.6 · p. 27

F.7 Full-size TabICLv2 deep ensembles

Full-size TabICLv2 deep ensembles improve on individual members across TALENT and TabArena, and the gains persist through training stages and context sizes. Independently trained members can be run concurrently, trading more total compute and inference cost for shorter wall-clock training time.

Setup

  • The study trains eight TabICLv2-like members and gives each the same inference-time ensemble of eight preprocessing and permutation configurations.It restricts the full-scale analysis to deep ensembles and extends it to regression.Source: §F.7 · p. 28

Results

  • 0.8% LPPD gain at every training stage on 170 TALENT classification datasets; ensembles also beat the average member on all 38 TabArena classification datasets.On TabArena, ensembles beat the best member on 27 datasets; at stage 3, the ensemble significantly beats the after-the-fact best member on all three metrics.Source: §F.7 · p. 28, Table 10 · p. 29
  • 0.7–0.9% ensemble gain persists across context-size bins and training stages, even as average-member LPPD improves 6.0% on contexts beyond 10,240 rows from stage 1 to stage 3.Further training improves members without reconciling their differences.Source: §F.7 · p. 28, Table 11 · p. 29
  • Giving members distinct preprocessing configurations rather than sharing the same eight leaves the ensemble essentially unchanged, pointing to independently pretrained weights rather than preprocessing randomness as the gain source.TabArena Elo is 1555 versus 1552.Source: §F.7 · p. 28

Ensembling as parallel pretraining

  • The eight-member ensemble reaches TabArena Elo 1552, within the official checkpoint's 1575 confidence interval, using about 0.37× its wall-clock time but 3× its GPU hours and 8× its inference cost.Members train concurrently on separate nodes without communicating.Source: §F.7 · p. 28, Table 9 · p. 29

F.8 Ensemble size analysis

The appendix examines how ensemble size affects full-size TabICLv2 performance and compares member disagreement within the prior's support with disagreement on shifted or out-of-scope datasets.

F.8 Ensemble size analysis

  • Four members achieve about 85% of the full eight-member gain on TALENT; adding members improves LPPD and nRMSE but quickly saturates and mainly reduces variation across sub-ensembles.The analysis evaluates every possible sub-ensemble from K=1 to 8.Source: §F.8 · p. 30, Figure 14 · p. 30

F.9 Disagreement beyond the prior

  • Disagreement comparisons distinguish in-support TabArena datasets from BeyondArena datasets with temporal or grouped splits or tasks outside the prior's scope.Figure 15 reports empirical CDFs for d AU JSD and pairwise argmax disagreement; the caption does not state which group disagrees more.Source: Figure 15 · p. 30, §F.9 · p. 30, Table 12 · p. 31

Related papers

Loading 2610.13764v1…