Learning Source Structure: Bayesian Prediction with Transformers
Authors: Meir Feder, Ruediger Urbanke
Organizations: School of Electrical and Computer Engineering, Tel Aviv University, Tel Aviv, Israel · School of Computer and Communication Sciences, EPFL, Lausanne, Switzerland
A transformer trained on sequences whose source parameters are drawn afresh each time can learn only what the sequences share: the source family and its prior. The Bayesian predictor that knows both, which we call informed Bayes, is the unique minimizer of the training loss and is computable from the family alone; we use it as a quantitative, falsifiable reference. Where earlier work compared transformers with Bayes on sources whose contexts can be learned independently, we build Markov sources with up to 64 symbols in which observations in one context inform prediction in another: product sources with known or hidden encodings, transition rows sharing an unknown base law and depth, and a mixture of three families. Our main result is that informed Bayes predicts, from the family and prior alone, how a trained transformer behaves in four tests that go beyond the learning curve: the loss that remains under a source outside the training family (predicted 2.237 bits, measured 2.309), the doubling of the extra loss when two independent components change rather than one (0.345 predicted, 0.350 measured), the benefit of other contexts when a token appears for the first time, and the depth at which copying an earlier continuation starts to pay. On the learning curves themselves, small transformers match informed Bayes when the shared structure is visible in the tokens, and hidden structure needs more layers and training, with gaps below 0.01 bits per token for the largest models. We explain why the imitation is possible and where it must fail: informed Bayes decomposes into sufficient statistics that attention selects and accumulates and an estimator, fixed by the prior, that the feed-forward network learns; training pays the complexity of informed Bayes under the architecture's prior, and finite-memory universal prediction bounds how closely a bounded network can reach it.
Bayesian accounts of in-context learning face a direct objection: exact posterior predictives for exchangeable data are invariant to task-preserving order, yet transformers change next-token probabilities when the same examples are serialized differently. We show this objection targets a structural invariant rather than the quantity scoring online prediction. For any Bayesian reference, excess prequential code length is exactly cumulative predictive KL. For unordered support sets that must be serialized, the expected regret of a single admissible ordering decomposes into that of the order-averaged predictor plus an order-averaging gain. Exchangeability violations are therefore not binary refutations; they are priced by log loss. We instantiate the theory with KT/Dirichlet finite-alphabet prediction and coarsened Bayesian linear-regression (BLR) predictive distributions. On Qwen2.5-7B/14B, floored candidate distributions at support 256 have one-step excess code lengths of 0.020/0.011 bits for Bernoulli and 0.039/0.022 bits for four-way categorical prediction, with candidate mass above 0.999; coarsened BLR continuations increasingly match the posterior-predictive digit distribution as support grows. A frequentist plug-in baseline sharpens the reading: the predictive distributions sit closer to the Bayesian posterior predictive than to the maximum-likelihood plug-in, by a margin largest at small support, where the plug-in is degenerate, and vanishing as the references converge. Position interventions and a from-scratch ablation localize order sensitivity to the positional encoding, activation patching tests causal use of decoded sufficient statistics, and permutation mixtures quantify the downstream log-loss cost of arbitrary orderings. Transformers need not realize exchangeable posterior predictives for every serialization to be Bayes-competitive prequential predictors.
Sequence modelling requires determining which past tokens are causally relevant from the context and their importance: a process inherent to the attention layers in transformers, yet whose underlying learned mechanisms remain poorly understood. In this work, we formalize the task of estimating token importance as an in-context learning problem by introducing a framework based on Mixture of Transition Distributions, where a latent variable determines the influence of past tokens on the next. The distribution over this latent variable is parameterized by unobserved mixture weights that transformers must learn in-context. We demonstrate that transformers can implement Mirror Descent to learn these weights from the context. Specifically, we give an explicit construction of a three-layer transformer that exactly implements one step of Mirror Descent and prove that the resulting estimator is a first-order approximation of the Bayes-optimal predictor. Corroborating our construction and its learnability via gradient descent, we empirically show that transformers trained from scratch learn solutions consistent with our theory: their predictive distributions, attention patterns, and learned transition matrix closely match the construction, while deeper models achieve performance comparable to multi-step Mirror Descent.
Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that transformers can implement a gradient descent algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of attention depth in enabling the extrapolation abilities of transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.
Gyeonghun Kang, Changwoo J. Lee, Xiang Cheng
Department of Statistical Science, Duke University, Durham, NC, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA