Research map / Architectures and efficiency / Attention and transformers
Constrained Belief Updates Explain Geometric Structures in Transformer Representations
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 single-layer transformers trained on Mess3 HMM sequences implement a constrained (parallelized, attention-based) Bayesian belief update, predicting attention patterns, OV vectors, and embedding geometry via a spectral analysis of the HMM transition matrix.
Problem
Optimal next-token prediction requires recursive Bayesian belief updating, but transformer architectures compute in parallel via attention; it is unclear how transformers reconcile this and why intermediate residual-stream activations form fractal geometries that differ from the final belief-state g
Approach
Model the first attention layer as implementing a constrained belief update: each past token contributes an independent correction to the stationary belief, Pr(S_d|Z_s) - Pr(S_d), summed in parallel. A spectral decomposition of the HMM transition matrix then predicts attention weights (decaying as (1-3x)^n), OV vectors, and embeddings, and multiple heads are needed when eigenvalues are negative.
Key results
- Predicted attention pattern decays as (1-3x)^n and OV vectors and embeddings align closely with trained single-layer transformer activations.
- For x > 1/3 (negative eigenvalue), a single head fails to capture belief geometry while two heads converge to the predicted solution.
- In 4-layer transformers, first-layer post-attention activations fit the constrained belief geometry better, while post-MLP activations fit the full belief geometry better.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add a second attention head instead of using a single head (ablation) | single-layer transformer · Mess3 with x=0.5 | belief geometry fit | Helped |
| use a single attention head rather than two (ablation) | single-layer transformer · Mess3 with x>1/3, including x=0.5 | belief geometry fit | Hurt |
| add a second attention head to capture oscillatory updates (ablation) | single-layer transformer · Mess3 with negative transition eigenvalue | KL divergence to optimal probabilities | Helped |
| add subsequent transformer layers after the first attention layer | 4-layer transformer · Mess3 | fit to full Bayesian belief geometry | Helped |
Surprising findings
- Intermediate (post-attention) residual stream forms fractal structure distinct from the final belief geometry.
- Negative eigenvalues of the transition matrix require two attention heads, since attention weights must be non-negative.
Limitations
- Experiments limited to small 1- and 4-layer transformers on synthetic Mess3 HMM data, far from natural language.
- Mapping from theory to MLP output is not characterized mathematically; deviations appear for alpha outside [0.2, 0.6].
- Discrepancy in sign of predicted embeddings for the first two positions remains unexplained.
Experimental setup
Datasets: Mess3 HMM-generated sequences
Benchmarks: Mess3 HMM belief-state geometry
Metrics: MSE to constrained belief geometry MSE to full belief geometry validation KL divergence explained variance ratio
Baselines: optimal full Bayesian belief updating
Builds on: computational mechanics belief state geometry transformer circuits framework
Closest papers by meaning
- Transformers Represent Belief State Geometry in their Residual Stream (NeurIPS 2024 · similarity 0.91)
- Predictive Geometry of Hidden Trajectories in Transformers (NeurIPS 2026 · similarity 0.87)
- Transformers on Markov data: Constant depth suffices (NeurIPS 2024 · similarity 0.87)
- Looking Beyond the Top-1: Transformers Determine Top Tokens in Order (ICML 2025 · similarity 0.87)
- Selective Induction Heads: How Transformers Select Causal Structures in Context (ICLR 2025 · similarity 0.87)
- How Transformers Learn Structured Data: Insights From Hierarchical Filtering (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.