# Constrained belief updates explain geometric structures in transformer representations

arXiv: https://arxiv.org/abs/2502.01954 · HTML: https://mlsanity.com/paper/2502.01954/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 24 references that are in the map.
- First posted on arXiv 4 Feb 2025.

## Builds on (most cited first, 11 of 11)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models](https://mlsanity.com/paper/2403.19647/) (2024)
- [Not All Language Model Features Are One-Dimensionally Linear](https://mlsanity.com/paper/2405.14860/) (2024)
- [Transformers represent belief state geometry in their residual stream](https://mlsanity.com/paper/2405.15943/) (2024)
- [Nearly Maximally Predictive Features and Their Dimensions](https://mlsanity.com/paper/1702.08565/) (2017)

## Built on by (most cited first, 15 of 15)

- [Are LLM Belief Updates Consistent with Bayes' Theorem?](https://mlsanity.com/paper/2507.17951/) (2025)
- [Transformers learn factored representations](https://mlsanity.com/paper/2602.02385/) (2026)
- [Neural networks leverage nominally quantum and post-quantum representations](https://mlsanity.com/paper/2507.07432/) (2025)
- [Next-token pretraining implies in-context learning](https://mlsanity.com/paper/2505.18373/) (2025)
- [From monoliths to modules: Decomposing transducers for efficient world modelling](https://mlsanity.com/paper/2512.02193/) (2025)
- [The Topological Trouble With Transformers](https://mlsanity.com/paper/2604.17121/) (2026)
- [Finding Belief Geometries with Sparse Autoencoders](https://mlsanity.com/paper/2604.02685/) (2026)
- [Predictive Statistics Shape Emergent World Representations of Grid Walkers](https://mlsanity.com/paper/2603.16689/) (2026)
- [Erased, Rerouted, or Rescaled? Post-Training and the Causal Quotient of a Language Model's Belief State](https://mlsanity.com/paper/2610.05292/) (2026)
- [Structure and Scale in Simplicial Sequence Modelling](https://mlsanity.com/paper/2606.01302/) (2026)
- [Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization](https://mlsanity.com/paper/2606.30813/) (2026)
- [Elicitation without Backpropagation: Steering Model Behavior by Optimizing the Latent Posterior](https://mlsanity.com/paper/2607.18804/) (2026)
- [World models of environment, agent and joint agent-environment systems](https://mlsanity.com/paper/2608.20401/) (2026)
- [Small transformers track Bayesian evidence for latent common causes via a context-invariant mechanism](https://mlsanity.com/paper/2609.35161/) (2026)
- [Bayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow](https://mlsanity.com/paper/2610.00679/) (2026)

## Abstract (arXiv)

What computational structures emerge in transformers trained on next-token prediction? In this work, we provide evidence that transformers implement constrained Bayesian belief updating -- a parallelized version of partial Bayesian inference shaped by architectural constraints. We integrate the model-agnostic theory of optimal prediction with mechanistic interpretability to analyze transformers trained on a tractable family of hidden Markov models that generate rich geometric patterns in neural activations. Our primary analysis focuses on single-layer transformers, revealing how the first attention layer implements these constrained updates, with extensions to multi-layer architectures demonstrating how subsequent layers refine these representations. We find that attention carries out an algorithm with a natural interpretation in the probability simplex, and create representations with distinctive geometric structure. We show how both the algorithmic behavior and the underlying geometry of these representations can be theoretically predicted in detail -- including the attention pattern, OV-vectors, and embedding vectors -- by modifying the equations for optimal future token predictions to account for the architectural constraints of attention. Our approach provides a principled lens on how architectural constraints shape the implementation of optimal prediction, revealing why transformers develop specific intermediate geometric structures.