The Bayesian Geometry of Transformer Attention

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

• The paper shows that transformers, leveraging attention mechanisms, achieve exact Bayesian inference with sub-bit calibration error (e.g., 0.003 bits MAE) in controlled wind-tunnel settings.
• It employs discrete bijection tasks and HMM state tracking benchmarks to rigorously validate the necessity of compositional attention over capacity-matched MLPs.
• The study reveals a three-stage mechanism—binding, elimination, refinement—via geometric diagnostics that clarifies transformer-based in-context probabilistic reasoning.

The Bayesian Geometry of Transformer Attention: A Mechanistic Study of Exact Probabilistic Inference

Introduction and Problem Framing

This work rigorously investigates the capacity of transformer architectures to realize exact Bayesian inference in controlled settings, introducing the concept of Bayesian wind tunnels—prediction tasks where the analytic posterior is known in closed form and memorization is provably taboo through combinatorial task construction. The central question addressed is not whether transformer outputs correlate with Bayesian posteriors but whether transformers mechanistically implement the internal computations prescribed by Bayes’ rule, as distinguished from pattern matching or memorized statistical associations.

Two wind-tunnel benchmarks are constructed: a discrete bijection elimination task and a sequential Hidden Markov Model (HMM) state tracking problem. These settings ensure:

1. Known analytic posteriors at every step.
2. Hypothesis spaces large enough to render memorization infeasible.
3. In-context learning is equivalent to performing bona fide probabilistic inference.

Experimental Design and Theoretical Foundations

Theoretical analysis establishes that cross-entropy training on contextual prediction tasks targets the Bayesian posterior predictive distribution in the population-optimum regime. The empirical challenge is to determine which architectures can approximate this target in finite-data, finite-capacity scenarios.

Task 1: Bijection Learning

Each sequence derives from a random bijection $\pi:\{1,\dots,V\}\to\{1,\dots,V\}$, with predictive entropy decaying as hypotheses are sequentially eliminated.

Task 2: HMM State Tracking

Each sequence is constructed with a new HMM, requiring recursive probabilistic state estimation via the forward algorithm.

Model Classes:

Small transformers (2.67–2.68M parameters), with standard attention mechanisms, are compared to capacity-matched MLPs explicitly lacking attention, isolating attention’s necessity.

Transformer Behavioral Fidelity to Bayesian Inference

Bijection Wind Tunnel: Exact Hypothesis Elimination

A transformer is shown to match analytic Bayesian posteriors with a mean absolute entropy error (MAE) of $3\times 10^{-3}$ bits over 2,000 held-out bijections (Figure 1), outperforming the MLP baseline by over three orders of magnitude.

Figure 1: Transformer entropy overlays analytic Bayes, while MLP fails to reduce uncertainty or implement hypothesis elimination.

Per-sequence analysis (Figure 2) demonstrates discrete stepwise hypothesis elimination and confirms that transformer outputs are not merely aggregating statistics but performing Bayesian inference at the sequence level.

Figure 2: Transformer reproduces the staircase entropy trajectory of exact Bayesian elimination, as seen in eight bijection instances.

Layer-wise ablation increases calibration error by more than an order of magnitude (Figure 3), confirming a strictly compositional hierarchy. Head-wise ablation (Figure 4) identifies a single critical Layer-0 attention head vital to forming hypothesis frames.

Figure 3: Removing any transformer layer disrupts compositional Bayesian inference, confirming depth necessity.

Figure 4: Layer-0 hypothesis-frame head is uniquely critical; later heads are partially redundant.

HMM Wind Tunnel: Recursive Bayesian State Tracking

Transformers converge to the recursive forward algorithm with sub-bit precision (MAE $7.5\times 10^{-5}$) at the training horizon (Figure 5). Crucially, the model generalizes gracefully to $1.5\times$ and $2.5\times$ longer sequences without explicit position conditioning, with smooth error growth and no training-boundary discontinuity.

Figure 5: Transformer matches Bayesian entropy trajectories across and beyond training lengths, signaling true algorithmic generalization.

Positional error curves (Figure 6) expose calibration characteristics: slight drift at early positions (uncertain initial state), near-zero error in mid-sequence, and graceful degradation at the tail.

Figure 6: Absolute entropy error per position across sequence lengths highlights transformer stability and smooth generalization.

Sequence-specific dynamics (Figure 7) show that transformers track fluctuations in uncertainty due to emissions and transitions.

Figure 7: Transformer captures sequence-specific rises and falls in uncertainty, reflecting exact stochastic belief tracking.

Semantic invariance under hidden-state relabeling (Figure 8) confirms that transformer computations are structurally probabilistic, invariant to arbitrary relabelings.

Figure 8: Entropy error is identical before/after hidden-state permutation, confirming semantic invariance of transformer’s Bayesian inference.

Disabling late-layer attention collapses the model’s ability to generalize at extended lengths (Figure 9), establishing that late-layer attention is vital for stable rollout.

Figure 9: No-late-attention transformer diverges dramatically with sequence length, confirming necessity for recursive stability.

MLP controls (Figure 10) fail to learn recursive Bayesian updates, with flat ~0.4 bit error, evidencing the architectural indispensability of attention.

Figure 10: Transformer sustains low error while MLP maintains high, position-independent error, failing Bayesian update replication.

Multi-seed robustness (Figures 11, 12) further verifies that transformer inference is not sensitive to initial conditions or training noise.

Figure 11: Transformer’s error curves across five seeds nearly overlap, supporting robust algorithmic learning.

Figure 12: Single-seed trajectory matches multi-seed average, ruling out initialization artifacts.

Mechanistic Insights: Geometric Implementation of Bayesian Updates

Layer 0 instantiates an orthogonal key basis encoding the hypothesis space (Figure 13); this frame remains stable during training, while downstream value representations refine posterior precision.

Figure 13: Layer-0 keys are nearly orthogonal, forming a hypothesis basis central to Bayesian computation.

Progressive query-key alignment (Figure 14) concentrates attention on feasible hypotheses as depth increases, paralleling analytic Bayesian conditioning.

Figure 14: Early layers diffuse attention; later layers concentrate sharply on feasible hypotheses via Q–K alignment.

Block-wise probe deltas (Figure 15) show that feed-forward networks implement the numerical Bayesian update, while attention serves content-addressable routing.

Figure 15: FFN blocks reduce entropy prediction loss, confirming their role in belief updating.

Value-manifold diagnostics (Figures 16, 17) reveal that training unfurls a low-dimensional manifold parameterized by posterior entropy, enabling fine-grained belief state encoding.

Figure 16: PCA of attention outputs colored by entropy demonstrates value-manifold refinement during training.

Figure 17: Calibration improves with value-manifold unfurling, especially at late positions.

A three-stage mechanism emerges (binding, elimination, refinement), illustrated architecturally (Figure 14): Layer 0 constructs the hypothesis frame, mid layers perform sequential elimination, and late layers refine manifold precision.

Implications, Limitations, and Theoretical Connections

These findings empirically validate that transformers, by geometric design, achieve exact Bayesian inference when memorization is de facto impossible. Attention supplies content-addressable hypothesis retrieval, compositional refinement via depth, and stable routing for long-horizon updates.

The failure of capacity-matched MLPs corroborates that attention is architecturally necessary in contextual Bayesian structure learning. The geometric mechanisms—orthogonal key bases, progressive Q–K sharpening, and value-manifold unfolding—suggest a structural underpinning behind algorithmic generalization in LLMs, which cannot be ascribed to scale alone.

Relevant limitations include task simplicity, modest latent dimensionality, and the need for extension to more complex probabilistic structures and higher-dimensional ambient spaces. The route toward mechanistic analysis in natural-language wind tunnels is proposed, foreshadowing research into LLM interpretation using the established geometric diagnostics.

Conclusion

Transformers, in stringent wind-tunnel settings, perform exact Bayesian inference with sub-bit calibration error, leveraging compositional depth, attention-furnished routing, and geometric encoding in residual streams. MLPs fail by large margins across all metrics and controls, affirming the necessity of hierarchical attention for probabilistic reasoning. The architectural geometry of transformers carves a Bayesian manifold that can be directly probed, providing foundational insight into the algorithmic interpretation of transformer-based in-context learning and a template for future analysis in high-dimensional and naturalistic domains.