The Bayesian Geometry of Transformer Attention
Naman Aggarwal, Siddhartha R. Dalal, Vishal Misra
TL;DR
This work demonstrates that transformers can implement exact Bayesian posterior inference in controlled wind-tunnel tasks where memorization is infeasible and the true posterior is known. By analyzing two settings—bijection elimination and HMM state tracking—the authors uncover a unified three-stage mechanism: layer-0 creates an orthogonal hypothesis frame, middle layers perform sequential Bayesian elimination via progressive QK sharpening, and late layers refine posterior precision along a low-dimensional value manifold while keeping routing stable. They show that Transformers achieve near-perfect calibration of predictive entropy with respect to Bayes, while capacity-matched MLPs fail to implement the necessary content-addressable routing and hierarchical refinement. The findings provide mechanistic interpretability for Bayesian reasoning in transformers, offer a rigorous lower bound for reasoning in LLMs, and propose Bayesian wind tunnels as a principled tool to connect small verifiable systems to more complex language models. This work lays the groundwork for extending such wind-tunnel diagnostics to natural-language tasks and larger architectures, enabling principled evaluation of inference capabilities beyond surface performance metrics.
Abstract
Transformers often appear to perform Bayesian reasoning in context, but verifying this rigorously has been impossible: natural data lack analytic posteriors, and large models conflate reasoning with memorization. We address this by constructing \emph{Bayesian wind tunnels} -- controlled environments where the true posterior is known in closed form and memorization is provably impossible. In these settings, small transformers reproduce Bayesian posteriors with $10^{-3}$-$10^{-4}$ bit accuracy, while capacity-matched MLPs fail by orders of magnitude, establishing a clear architectural separation. Across two tasks -- bijection elimination and Hidden Markov Model (HMM) state tracking -- we find that transformers implement Bayesian inference through a consistent geometric mechanism: residual streams serve as the belief substrate, feed-forward networks perform the posterior update, and attention provides content-addressable routing. Geometric diagnostics reveal orthogonal key bases, progressive query-key alignment, and a low-dimensional value manifold parameterized by posterior entropy. During training this manifold unfurls while attention patterns remain stable, a \emph{frame-precision dissociation} predicted by recent gradient analyses. Taken together, these results demonstrate that hierarchical attention realizes Bayesian inference by geometric design, explaining both the necessity of attention and the failure of flat architectures. Bayesian wind tunnels provide a foundation for mechanistically connecting small, verifiable systems to reasoning phenomena observed in large language models.