Debiased Prior-Data Fitted Networks for Amortized Bayesian Semi-Parametric Inference

Valentyn Melnychuk; Emil Javurek; Stefan Feuerriegel

stat.MLOct 20, 2026

Debiased Prior-Data Fitted Networks for Amortized Bayesian Semi-Parametric Inference

Authors: Valentyn Melnychuk, Emil Javurek, Stefan Feuerriegel

Organizations: LMU Munich & Munich Center for Machine Learning, Munich, Germany

Abstract

Prior-data fitted networks (PFNs) have recently gained popularity as an end-to-end approach to amortized Bayesian inference. Yet, their statistical behavior for semi-parametric targets, such as means under missingness or causal effects, remains poorly understood. In this work, we develop theoretical foundations for debiased PFNs to perform inference on finite-dimensional targets in general semi-parametric settings. First, we show that, despite being Bayes-optimal with respect to the target, an idealized PFN can exhibit a posterior concentration rate that implicitly depends on the concentration rates of (latent) nuisance functions and can become slower as the prior becomes more diverse. Then, to speed up the target posterior concentration rates, we introduce two generic debiasing strategies based on efficient influence functions: (i) a label correction and (ii) a prior correction. Finally, we show that the debiased PFNs (i) & (ii) possess concentration rates that asymptotically match a convergence rate given by traditional frequentist bias-corrected estimators. In this way, we establish a novel connection between amortized Bayesian inference with PFNs and frequentist semi-parametric efficiency theory. Under additional assumptions, this allows us to guarantee local and global optimality of PFN-based estimators (namely, the Bernstein-von Mises property and minimax optimality, respectively). Therefore, our work provides guidelines for designing debiased PFNs for semi-parametric inference with principled and provable asymptotic guarantees.

Figures & tables

Figure 1: Overview of our proposed debiasing strategies: (i) label correction and (ii) prior correction (marked with orange ). Compared to the existing approach of Melnychuk et al. (2026) , our debiasing strategies (i) and (ii) are applied at the prior specification & training stages of PFNs and, therefore, do not require post-hoc inference to be performed by a downstream user.

Figure 2: Propagation of

\mu(x)

nuisance errors into the target (based on Example Example B with

d_{x}=1

P(x)=P_{0}(x)=\mathrm{Unif}[0,1]

, and

\pi_{0}(x)=0.2+{0.7}/({1+\exp\{-6(x-0.4)\}})

). Here, (a) a generic posterior,

\Pi_{\mathrm{g}}(\mu\mid\mathcal{D}_{n=30})

, is a Gaussian process with a Matérn kernel (

\beta=0.5

). Then, (b) an oversmoothed posterior,

\Pi_{\mathrm{s}}(\mu\mid\mathcal{D}_{n=30})

, is a Gaussian process with a squared-exponential kernel (

\beta=\infty

), random intercept, and inflated noise level. (c) , (d) The corresponding posteriors of

\psi=\int\mu(x)\mathop{}\!\mathrm{d}{P_{0}}(x)

for growing

n

(each line follows a fixed posterior quantile across

n

) and the bands

\psi_{0}\pm M(P_{0},\mathcal{D}_{n},\Pi)

. In (d) (unlike (c)), the errors of

\mu(x)

barely cancel in

P_{0}(\mu-\mu_{0})

, and the target posterior concentrates at nearly the same rate as

M(P_{0},\mathcal{D}_{n},\Pi)

(i. e., the oversmoothed posterior is rate-close to the worst-case posterior).

Figure 3: Slowing concentration rates for different versions of TabPFN (

\diamond\in\{\mathrm{v2},\mathrm{v2.5},\mathrm{v3}\}

) . Here, we show

\Pi_{0}

-prior-averaged squared

L_{2}

-concentration errors of the conditional outcome mean

\mu(x)

posteriors,

\mathcal{R}_{L_{2}^{2}}(\mu_{\diamond};\Pi_{0},n):=\mathbb{E}_{P_{0}\sim\Pi_{0}}\mathbb{E}_{\mathcal{D}_{n}\sim P_{0}^{\otimes n}}[M(P_{0},\mathcal{D}_{n},\Pi_{\diamond})^{2}]

, depending on dataset sizes

n

and covariate dimensionality

d_{x}

(prior-averaging is based on a sample of

M=24

ground-truth

P_{0}

, each with

D=1

dataset; lower is better). For reference, we also show the error of a frequentist non-parametric estimator, namely, the

\Pi_{0}

-prior-averaged squared

L_{2}

-error of Nadaraya–Watson regression with a Matérn kernel (

\beta=1.0

) (= NW estimator (

\beta=1.0

)). Note that the ground-truth

\Pi_{0}

is based on the adapted prior of the TabPFN

\mathrm{v1}

( Hollmann et al., 2023 ) (

\mu_{0}(\cdot)\sim\Pi_{\mathrm{v1}}(\mu)\text{ and }P_{0}(y\mid x)=N(\mu_{0}(x),\sigma^{2}=0.3^{2})

) and is assumed to be in support of the priors of the higher versions.

Figure 4: Debiasing strategies for the mean value (based on Example Example A with

Z\in[0,1]

and

\psi(P)=\mathbb{E}_{P}[Z]

). (a) Posterior draws of

P

under a non-parametric prior

\Pi

(a logistic Gaussian process with a squared-exponential kernel (

\beta=\infty

P(z)\propto\exp\{f(z)\}

) and the corresponding target labels

\psi(P)

. (b) (i) Label correction: The posterior of

P

is unchanged, but the label of every draw is replaced with

K=1

sample Bayesian bootstrap (shown with arrows). (c) (ii) Prior correction: Posterior draws

\tilde{P}

under the corrected prior

\widetilde{\Pi}

(black), each the exponential tilt of an untilted draw

P

(gray, vertical arrows) that moves its label from

\psi(P)

\psi(\tilde{P})

(horizontal arrows), (d) - (f) The resulting posteriors of the target for growing

n

(each line follows a fixed posterior quantile across

n

) and the bands

\psi_{0}\pm F_{\psi}(P_{0},\mathcal{D}_{n},\star)/\sqrt{12}

, where

\star\in\{\Pi,\widetilde{\Pi}\}

. Without a correction, the uncorrected posterior draws pull

\psi

towards

1/2

. Both corrections remove this first-order error: The posteriors after the corrections are better centered and concentrate at a faster rate around

\psi_{0}

Figure 5:

L_{2}

risk on the evaluation prior

\check{\Pi}

(based on Example Example C with

w\equiv 1

; axes are log-scaled). Shown: mean

\pm

\check{\Pi}

-averaged

L_{2}

concentration errors (solid lines) and the corresponding power-law regression fits (dotted lines) depending on dataset sizes

n

and covariate dimensionality

d_{x}

(over

M=200

ground-truth

\check{P}

, each with

D=200

datasets; lower is better). Here,

\diamond\in\{\psi,\tilde{\psi}\}

are the target posteriors from the standard and debiased PFNs, respectively.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Prior-induced confounding bias of the training targets for the ATE (based on Example Example C with

w\equiv 1

; see Remark 2 ). Here, we show the densities of the standardized degree of observed confounding

\Delta

(= differences between the training targets,

\diamond\in\{\psi(P),\tilde{\psi}(P),\psi(\tilde{P})\}

, and the unadjusted treated–control contrast,

m(P):=\mathbb{E}_{P}[Y\mid A=1]-\mathbb{E}_{P}[Y\mid A=0]

, namely,

\Delta:=\{\diamond-m\}/\mathrm{sd}[Y]

) depending on the covariate dimensionality

d_{x}

. The densities are estimated over

M=512

draws of ground-truth

P\sim\Pi

from the original prior or their debiased counterparts: (i) label correction,

\tilde{\psi}(P)

, with

D=1

dataset

\mathcal{D}_{n}

and

K=4

Bayesian bootstrap draw per

P

; and (ii) prior correction,

\psi(\tilde{P})

with

\tilde{P}=\operatorname{LF}_{\psi}(P;\lambda)

and

\lambda\sim N(0,\sigma_{n}^{2}=n^{-1/2})

(both for

n=256

;

m

and

\mathrm{sd}[Y]

are taken under

\tilde{P}

for (ii)). For reference, we also show

\Delta

under the worst-case evaluation prior

\check{\Pi}

Figure 7: Total variation, averaged over the training and evaluation priors,

\Pi

and

\check{\Pi}

, respectively (based on Example Example A with

w\equiv 1

). Shown: mean

\pm

3SE

\Pi/\check{\Pi}

-averaged total variation between the PFN posterior and the asymptotic normal distribution given by the efficient frequentist estimator (A-IPTW w/ ground-truth nuisance functions),

\mathcal{R}_{\mathrm{TV}}(\diamond;\star,n)

, depending on dataset sizes

n

and covariate dimensionality

d_{x}

(over

M=200

ground-truth

\check{P}

, each with

D=200

datasets; lower is better; axes are log-scaled). Here,

\diamond\in\{\psi,\tilde{\psi}\}

are the target posteriors from the standard and debiased PFNs, respectively.

Figure 8:

L_{2}

risk on the evaluation prior

\check{\Pi}

(based on Example Example A ). Shown: mean

\pm

3SE

\check{\Pi}

-averaged

L_{2}

concentration errors (solid lines) and the corresponding power-law regression fits (dotted lines) depending on dataset sizes

n

(over

M=200

ground-truth

\check{P}

, each with

D=200

datasets; lower is better; axes are log-scaled). Here,

\diamond\in\{\psi,\tilde{\psi}\}

are the target posteriors from the standard and debiased PFNs, respectively.

Figure 9: Total variation, averaged over the training and evaluation priors,

\Pi

and

\check{\Pi}

, respectively (based on Example Example A ). Shown: mean

\pm

3SE

\Pi/\check{\Pi}

-averaged total variation between the PFN posterior and the asymptotic normal distribution given by the efficient frequentist estimator (sample mean),

\mathcal{R}_{\mathrm{TV}}(\diamond;\star,n)

, depending on dataset sizes

n

(over

M=200

ground-truth

\check{P}

, each with

D=200

datasets; lower is better; axes are log-scaled). Here,

\diamond\in\{\psi,\tilde{\psi}\}

are the target posteriors from the standard and debiased PFNs, respectively.

Prior-Data Fitted Networks

Prior-Data Fitted Networks (PFNs) learn in-context prediction by training on synthetic datasets sampled from designed priors.

See all 45 papers

Oct 9, 2026

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

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 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.

Manuel Hülskamp, Julius Kobialka, Emanuel Sommer +2

LMU Munich and Munich Center for Machine Learning

Bayesian Model Averaging Prior-Data Fitted Networks Amortized Inference Bayesian Inference

Ask PDF

Mar 12, 2026

Frequentist Consistency of Prior-Data Fitted Networks for Causal Inference

Foundation models based on prior-data fitted networks (PFNs) have shown strong empirical performance in causal inference by framing the task as an in-context learning problem. However, it is unclear whether PFN-based causal estimators provide uncertainty quantification that is consistent with classical frequentist estimators. In this work, we address this gap by analyzing the frequentist consistency of PFN-based estimators for the average treatment effect (ATE). (1) We show that existing PFNs, when interpreted as Bayesian ATE estimators, can exhibit prior-induced confounding bias: the prior is not asymptotically overwritten by data, which, in turn, prevents frequentist consistency. (2) As a remedy, we suggest employing a calibration procedure based on a one-step posterior correction (OSPC). We show that the OSPC helps to restore frequentist consistency and can yield a semi-parametric Bernstein-von Mises theorem for calibrated PFNs (i.e., both the calibrated PFN-based estimators and the classical semi-parametric efficient estimators converge in distribution with growing data size). (3) Finally, we implement OSPC through tailoring martingale posteriors on top of the PFNs. In this way, we are able to recover functional nuisance posteriors from PFNs, required by the OSPC. In multiple (semi-)synthetic experiments, PFNs calibrated with our martingale posterior OSPC produce ATE uncertainty that (i) asymptotically matches frequentist uncertainty and (ii) is well calibrated in finite samples in comparison to other Bayesian ATE estimators.

Valentyn Melnychuk, Vahid Balazadeh, Stefan Feuerriegel +1

LMU Munich & Munich Center for Machine Learning (MCML), Munich, Germany · University of Toronto & Vector Institute, Toronto, Canada

Causal Effect Estimation Prior-Data Fitted Networks Bayesian Inference Semiparametric Inference Probability Calibration

Ask PDF

Oct 10, 2026

Characterizing the Separability Gap in Prior-Data Fitted Networks

Prior-data fitted networks (PFNs) are transformers pretrained on simulated datasets to approximate Bayesian posterior predictive distributions. Standard PFN theory assumes a separable task prior, under which the feature distribution and the mechanism generating labels from features are independent. We show that common PFN priors are non-separable, even in simple settings. Consequently, standard pretraining targets a predictive distribution based on a task posterior updated using both the labeled context and the query features. We characterize its difference from the predictive based on the context posterior alone and show that this difference vanishes under posterior concentration. Under misspecification (i.e., when the true joint distribution lies outside the prior's support), however, prior non-separability can favor better feature fit at the expense of conditional label fit. We illustrate the consequences for feature attribution, Bayesian optimization, and martingale checks of Bayesian behavior. The same dependence also enables PFNs to learn from additional unlabeled observations. We demonstrate this with SemiTabICL, a PFN explicitly pretrained to use these observations, which improves accuracy over a matched supervised baseline. Our findings highlight the need for explicit rules on unlabeled data use in PFN evaluation and establish the dependence between feature and label distributions as a consequential prior design choice.

Julius Kobialka, Manuel Hülskamp, Thomas Nagler +1

Department of Statistics, LMU Munich · Munich Center for Machine Learning

Prior-Data Fitted Networks Bayesian Inference

Ask PDF

Abstract

Prior-Data Fitted Networks

Explore similar work

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

Frequentist Consistency of Prior-Data Fitted Networks for Causal Inference

Characterizing the Separability Gap in Prior-Data Fitted Networks