Transformers Represent Belief State Geometry in their Residual Stream
Adam S. Shai, Lucas Teixeira, Alexander Gietelink Oldenziel, Sarah Marzen, Paul M. Riechers
Abstract
What computational structure are we building into large language models when we train them on next-token prediction? Here, we present evidence that this structure is given by the meta-dynamics of belief updating over hidden states of the data-generating process. Leveraging the theory of optimal prediction, we anticipate and then find that belief states are linearly represented in the residual stream of transformers, even in cases where the predicted belief state geometry has highly nontrivial fractal structure. We investigate cases where the belief state geometry is represented in the final residual stream or distributed across the residual streams of multiple layers, providing a framework to explain these observations. Furthermore we demonstrate that the inferred belief states contain information about the entire future, beyond the local next-token prediction that the transformers are explicitly trained on. Our work provides a general framework connecting the structure of training data to the geometric structure of activations inside transformers.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 1474fd8e-44fd-42c6-a422-4e2642418984Cited by top-tier papers16
- Symmetries in language statistics shape the geometry of model representationsDhruva Karkada, Daniel Korchinski, Andres Nava, Matthieu Wyart et al.ICML 2026 · 15 citations
- Not All Language Model Features Are One-Dimensionally LinearJoshua Engels, Eric J. Michaud, Isaac Liao, Wes Gurnee et al.ICLR 2025 · 3 citations
- Spectral Filters, Dark Signals, and Attention SinksNicola CanceddaACL 2024 · 3 citations
- Transformers learn factored representationsAdam Shai, Loren Amdahl-Culleton, Casper Christensen, Henry R Bigelow et al.ICML 2026 · 2 citations
- Transformers Learn Latent Mixture Models In-Context via Mirror DescentFrancesco D'Angelo, Nicolas FlammarionICLR 2026 · 2 citations
Builds on4
- Better & Faster Large Language Models via Multi-token PredictionFabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz et al.ICML 2024 · 286 citations
- The Pitfalls of Next-Token PredictionGregor Bachmann, Vaishnavh NagarajanICML 2024 · 163 citations
- Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic TaskKenneth Li, Aspen K. Hopkins, David Bau, Fernanda B. Viégas et al.ICLR 2023 · 60 citations
- Transformers Learn Shortcuts to AutomataBingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy et al.ICLR 2023 · 11 citations
Related papers
- Constrained Belief Updates Explain Geometric Structures in Transformer RepresentationsMateusz Piotrowski, Paul M. Riechers, Daniel Filan, Adam S. ShaiICML 2025
- How Transformers Represent Hierarchies: A Local-to-Global MechanismZhiling Zhou, Tianhao Wang, Zhuoran YangICML 2026
- Provable Long-Range Benefits of Next-Token PredictionXinyuan Cao, Santosh S. VempalaSTOC 2026
- Residual Stream Analysis with Multi-Layer SAEsTim Lawson, Lucy Farnik, Conor J. Houghton, Laurence AitchisonICLR 2025
- Lines of Thought in Large Language ModelsRaphaël Sarfati, Toni J. B. Liu, Nicolas Boullé, Christopher J. EarlsICLR 2025
