Research map / Architectures and efficiency / RNNs and normalization
Transformers Represent Belief State Geometry in their Residual Stream
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that small transformers trained on next-token prediction from HMM-generated data linearly represent the belief-state (mixed-state) geometry of the data-generating process in their residual stream, including fractal geometry.
Problem
What computational structure do transformers learn when trained on next-token prediction? The paper asks whether optimal Bayesian belief-state geometry from computational mechanics is linearly represented in transformer activations.
Approach
Uses the mixed-state presentation from computational mechanics to predict belief-state geometry for HMM-generated data, then regresses residual stream activations linearly onto ground-truth belief states to test for that geometry; finds it is linearly embedded, sometimes spread across layers.
Key results
- Belief-state simplex (fractal) for Mess3 is linearly recoverable from final-layer residual stream (2D subspace of 64-dim activations), with MSE decreasing over training.
- For RRXOR (36 belief states, 4-simplex), belief geometry is not in the final layer or any single layer but is recovered from concatenated layers; pairwise distance correlation with next-token probabilities is R^2=0.31 versus 0.95 with belief geometry.
- Lower normalized validation loss correlates with lower MSE of belief-state representation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Concatenate residual-stream activations across layers (ablation) | Transformer trained on RRXOR · RRXOR | belief-state geometry representation | Helped |
| Use activations after rather than before final LayerNorm (ablation) | Transformer trained on Mess3 · Mess3 | belief-state geometry regression MSE: 0.0004 -> 0.0003 | Helped |
| Shuffle input-to-belief-state correspondences (ablation) | Linear regression of residual activations onto belief states · Mess3 | belief-state geometry representation | Hurt |
| Fit the regression using only 20% of input-activation pairs (ablation) | Transformer trained on Mess3 · Mess3 | belief-state geometry regression MSE | No clear effect |
| Train the transformer on the HMM-generated process | Transformer · Mess3 | belief-state geometry representation | Helped |
| Train the transformer on the HMM-generated process | Transformer · RRXOR | belief-state geometry representation | Helped |
| Represent belief states rather than using next-token predictions as the pairwise-distance account | Transformer trained on RRXOR · RRXOR | pairwise-distance correlation R²: 0.31 -> 0.95 | Helped |
Surprising findings
- For RRXOR, belief-state geometry is absent from every individual layer and only appears when layers are concatenated.
- Hand-constructed transformers solving RRXOR do not linearly represent belief geometry, whereas SGD-trained ones do.
Limitations
- Experiments limited to small HMMs (3-5 states, 2-3 token vocabularies) and small transformers, far from natural language scale.
- For realistic settings the belief simplex dimension may exceed residual stream dimension, requiring compressed representation.
- Does not test whether the representations are used for Bayesian updating circuits.
Experimental setup
Datasets: Mess3 HMM generated sequences RRXOR HMM generated sequences
Benchmarks: Mess3 process RRXOR process
Metrics: mean squared error of linear regression R^2 of pairwise distances
Baselines: next-token probability geometry
Builds on: computational mechanics epsilon-machine mixed-state presentation
Closest papers by meaning
- Constrained Belief Updates Explain Geometric Structures in Transformer Representations (ICML 2025 · similarity 0.91)
- Next-Latent Prediction Transformers Learn Compact World Models (NeurIPS 2026 · similarity 0.88)
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.87)
- Dynamics of the Transformer Residual Stream: Coupling Spectral Geometry to Network Topology (NeurIPS 2026 · similarity 0.87)
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.87)
- Looking Beyond the Top-1: Transformers Determine Top Tokens in Order (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.