Decoupled PFNs: Identifiable Epistemic-Aleatoric Decomposition via Structured Synthetic Priors

TL;DR

This paper introduces decoupled PFNs that explicitly separate epistemic uncertainty from aleatoric noise using synthetic priors, improving decision-making in noisy, heteroscedastic settings.

Contribution

It presents a novel decoupled PFN architecture with separate signal and noise heads, enabling better epistemic uncertainty estimation from synthetic task data.

Findings

Decoupled PFNs outperform baseline models in hyperparameter optimization.

Epistemic-only acquisition reduces failure modes in noisy environments.

Decoupled models achieve the best average rank in synthetic Bayesian optimization.

Abstract

Prior-Fitted Networks (PFNs) amortize Bayesian prediction by meta-learning over a synthetic task prior, but their standard output is a posterior predictive distribution over noisy observations. For sequential decision-making, such as active learning and Bayesian optimization, acquisition should prioritize epistemic uncertainty about the latent signal rather than irreducible aleatoric observation noise. We show that this epistemic--aleatoric split is not identifiable in general from the posterior predictive distribution alone, even when that distribution is known exactly. We then exploit a distinctive advantage of PFNs: because the synthetic data-generating process is under our control, each task can contain an explicit latent signal and noise function, and the generator can provide query-level labels for both the noiseless target and the observation-noise variance. We use these labels to train a decoupled PFN with separate latent-signal and aleatoric heads. The observation-level predictive is induced by convolving the latent signal distribution with the learned noise model. Empirically, epistemic-only acquisition mitigates the failure mode of total-variance exploration in noisy and heteroscedastic settings. In matched comparisons, decoupled models usually improve over tuned observation-level baselines, with the clearest gains in HPO; in broader sweeps, a decoupled model obtains the best average rank in both HPO and synthetic BO.