Universal priors: solving empirical Bayes via Bayesian inference and pretraining

Abstract: We theoretically justify the recent empirical finding of [Teh et al., 2025] that a transformer pretrained on synthetically generated data achieves strong performance on empirical Bayes (EB) problems. We take an indirect approach to this question: rather than analyzing the model architecture or training dynamics, we ask why a pretrained Bayes estimator, trained under a prespecified training distribution, can adapt to arbitrary test distributions. Focusing on Poisson EB problems, we identify the existence of universal priors such that training under these priors yields a near-optimal regret bound of $\widetilde{O}(\frac{1}{n})$ uniformly over all test distributions. Our analysis leverages the classical phenomenon of posterior contraction in Bayesian statistics, showing that the pretrained transformer adapts to unknown test distributions precisely through posterior contraction. This perspective also explains the phenomenon of length generalization, in which the test sequence length exceeds the training length, as the model performs Bayesian inference using a generalized posterior.

• The paper introduces a universal prior framework that leverages pretrained transformers to achieve near-minimax regret bounds in empirical Bayes inference.
• It employs prior-on-prior distributions and posterior contraction techniques to justify instance-adaptive performance on Poisson and Gaussian models.
• The work shows that transformer models generalize to longer sequences with regret governed by the effective training sample size, aligning with empirical observations.

Universal Priors and Efficient Empirical Bayes Inference via Pretraining

Introduction and Context

The study addresses a central question in empirical Bayes (EB) estimation: How can deep neural architectures, specifically transformers, pretrained on synthetic data, exhibit strong and near-minimax frequentist performance in empirical Bayes problems, particularly those involving Poisson models? Instead of pursuing architecture-specific or optimization-based explanations, the paper frames the training and deployment of such estimators as hierarchical Bayesian inference with universal priors—specifically, a prior-on-prior (PoP) distribution. This perspective provides rigorous support, grounded in statistical decision theory and posterior contraction, for the empirical success of pretrained transformers in amortized statistical inference.

Formal Setup

The core empirical Bayes problem involves observations $X_i \sim \operatorname{Poi}(\theta_i)$, where $\theta_i$ are i.i.d. latent parameters drawn from an unknown $G_0$ supported on $[0, A]$. The objective is to estimate $\theta^n$ such that the estimator's risk (mean squared error) is uniformly close to the Bayes risk, i.e., has small regret relative to the Bayes estimator with oracle prior knowledge.

The innovative approach, inspired by [teh2025solving], replaces per-problem estimation with a pretrained transformer: training is conducted using data where the priors $G$ themselves are random (drawn from a PoP, e.g., random mixtures of Diracs). Training aims to minimize empirical squared error, and inference proceeds by direct application of the pretrained model to any new test data.

Main Theoretical Results

Existence and Construction of Universal Priors

A central technical contribution is the demonstration that relatively simple PoPs—specifically, mixtures with $k = O(\log n)$ support points—are "universal": that is, pretraining under such PoPs confers nearly minimax regret bounds uniformly over all possible priors $G_0$. The main theorem asserts that, under population risk minimization and sufficient pretraining, the obtained estimator achieves regret at most

$$\sup_{G_0 \in \mathcal{P}([0, A])} \mathrm{Regret}(\widehat{\theta}^n; G_0) \leq C \frac{\log^3 n}{n (\log\log n)^2},$$

where $C$ depends only on the support constraint $A$.

The proof leverages a blending of Bayesian and minimax frequentist decision theory, showing that for any metric entropy regular statistical family, there always exists a least favorable PoP yielding the minimax risk via the Bayes estimator. The universality property is characterized via mass-concentration in neighborhoods (measured by $\chi^2$-divergence) of arbitrary priors $G_0$, echoing classical posterior contraction conditions.

Notably, these universal priors are not engineered in an adversarial sense; random mixtures suffice. The sequence-to-sequence map implemented by the transformer approximates the Bayes rule:

$$\widehat{\theta}^n(X^n) \sim \mathbb{E}_\Pi[\theta^n | X^n],$$

where the expectation is over the hierarchical model $G \sim \Pi$, $\theta^n \sim G^{\otimes n}$, $X^n \sim \prod_i \operatorname{Poi}(\theta_i)$.

Posterior Contraction and Adaptation to Test Priors

The paper’s main technical mechanism is posterior contraction: regardless of the actual $G_0$, the Bayes estimator under the PoP posterior, given sufficient data, concentrates around the true Bayes estimator for $G_0$. The conditional mean map is thus nearly instance-adaptive. Formally, for any $G_0$ and $X^n \sim f_{G_0}^{\otimes n}$,

$$H^2(f_{G_0}, \Pi_{X_n | X^{n-1}}) = O\left(\frac{\log^2 n}{n \log\log n} + B_n/n \right)$$

with high probability, where $H$ is the Hellinger distance and $B_n$ is the universal prior mass rate.

Length Generalization

The length generalization phenomenon is rigorously justified. Transformers trained on length $n$ sequences generalize to inference on longer sequences—yet, the regret saturates at the rate corresponding to the training length, not the longer test length. This is formalized by showing that when the transformer is applied at test time to sequences of length $n' > n$, it approximates Bayes inference under an $\alpha$-posterior, where

$\alpha = \frac{n}{n'}.$

Posterior contraction still applies, but the rate is determined by the effective training sample size. This aligns with empirical observations that regret decreases with length but stops improving for $n' \gg n$.

Extensions

• Subexponential Prior Tails: The framework is generalized to subexponential (effectively unbounded support) priors, yielding a regret rate of $O(\log^4 n / n)$.
• Gaussian EB: The methodology is extended to the Gaussian means problem, providing near-optimal rates up to logarithmic factors.
• Functionals of $\theta$: The method covers estimating functionals $g(\theta)$, with regret rates that match known minimax rates for polynomial functions.

Numerical Evidence

Empirical studies substantiate the theoretical claims:

• Pretrained transformers consistently outperform NPMLE-based estimators, both when test priors are sampled from neural or multinomial PoPs.
• The Bayes estimator under the hierarchical model matches the pretrained transformer's predictions up to negligible error—confirming that learned inference is approximating Bayesian inference under the training PoP.
• For out-of-distribution test sequence lengths, the outputs of the transformer closely match $\alpha$-posterior hierarchical Bayes, with $\alpha$ empirically matching the predicted training/test ratio.

Implications and Theoretical Significance

This work provides a formal, non-architecture-specific justification for the empirically observed effectiveness of pretraining large models for amortized statistical inference. By characterizing the role of universal priors, the authors establish that neural estimators—when trained appropriately—can achieve near-optimal instance-adaptive performance at test time, even under profound prior uncertainty and without explicit access to test-time prior data.

The perspective bridges Bayesian and frequentist paradigms, showing that amortized inference (in the sense of TabPFN/transformer-based EB) can be theoretically supported within a well-defined minimax framework, provided sufficient coverage of the prior class during pretraining.

Potential Future Directions

• Determining minimax-universal PoPs with polynomial (rather than superpolynomial) pretraining complexity.
• Explicit design of PoPs for other hierarchical or nonparametric models.
• Architectural adjustment to further enhance length generalization, e.g., function classes or regularization schemes that mimic the true infinite-sample Bayes estimator.
• Analysis of sample complexity and generalization for finite-batch pretraining.
• Extension to more complex observation and prior models, such as hierarchical or dependent latent structures common in real-world datasets.

Conclusion

The paper establishes that hierarchical Bayes estimation with universal priors, efficiently implemented via transformer-based pretraining, yields near-optimal empirical Bayes estimators with provable small minimax regret. This framework not only explains high performance across arbitrary test priors but also rigorously characterizes length generalization as approximate inference under fractional posteriors. The results unify and extend amortized inference paradigms and provide explicit theoretical foundations for deploying large pretrained models in diverse statistical estimation problems (2602.15136).