mlsanityData to 14 Oct 2026

Constrained belief updates explain geometric structures in transformer representations

arXiv 2502.01954 ↗First posted 4 Feb 2025cs.LGcited by 15

Not on the map. Why?

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.

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

11 of 24 references.

References · 11

  1. 2017Attention Is All You Needcited by 60,361
  2. 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,420
  3. 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,165
  4. 2022Toy Models of Superpositioncited by 971
  5. 2023Progress measures for grokking via mechanistic interpretabilitycited by 839
  6. 2023Towards Automated Circuit Discovery for Mechanistic Interpretabilitycited by 724
  7. 2023The Linear Representation Hypothesis and the Geometry of Large Language Modelscited by 716
  8. 2024Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Modelscited by 435
  9. 2024Not All Language Model Features Are One-Dimensionally Linearcited by 212
  10. 2024Transformers represent belief state geometry in their residual streamcited by 77
  11. 2017Nearly Maximally Predictive Features and Their Dimensionscited by 9

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 4; 2026: 11.

15 papers.

Sort:

By area:

Similar, not cited

Close in wording, no citation either way

  1. 2026A Comparative Analysis of Attention versus State-Space Models for In-Context Learning
  2. 2026Fully Interpretable Minimal Transformers: From Geometry to Algorithm
  3. 2025Selective Induction Heads: How Transformers Select Causal Structures In Context
  4. 2026Invariant Learning Dynamics of Transformers in Inductive Reasoning Tasks
  5. 2024Learning interpretable positional encodings in transformers depends on initialization
  6. 2026Trajectory Geometry of Transformer Representations Across Layers
  7. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization
  8. 2026Neural expectation-maximization and meta-models