Tracking States or Tracking Cosets? An Algebraic Account of Learned State Tracking
cs.LG, cs.AI
Submitted: 2026-09-24
Updated: 2026-09-24
Comments: 69 pages including appendices; 9 pages of main text
License: http://creativecommons.org/licenses/by/4.0/
The gist: State tracking requires composing a sequence of updates, but accuracy alone does not reveal what a model has learned.
Terminology
Abstract
State tracking requires composing a sequence of updates, but accuracy alone does not reveal what a model has learned. We study neural networks trained to predict the running product of group elements. We identify quotient solutions in Transformers, where models recover the quotient class while predicting nearly uniformly among its members. The reciprocal of class size predicts partial accuracy without a fitted parameter, extending parity-based accounts to non-parity quotients. Our baseline Transformers' predictions change little under prefix reordering beyond the exact-tracking frontier. We prove that, for finite groups under uniform i.i.d. full-group inputs, optimal order-blind exact accuracy converges to the reciprocal of abelianization class size as prefix length grows, consistent with the observed abelianization plateaus. Sequential updates permit more: any partition into right cosets of a subgroup, normal or not, survives sequential updates. In our census of standard Transformers, every recovered coset partition comes from a normal subgroup, whereas parameter-matched recurrent networks pass through both normal and non-normal right-coset stages during training. On A 5, we identify low-dimensional subspaces of the recurrent state that encode non-normal cosets. In the three-dimensional cases, coset mean vectors form approximate dodecahedra, and swapping the state components in these subspaces transfers the donor's coset state through a shared input suffix. Our results connect partial accuracy, learning stages, and internal computation through the subgroup cosets that models learn to track.
Sources
- Hultman numbers, polygon gluings and matrix integrals
- GPT-NeoX-20B: An Open-Source Autoregressive Language Model
- A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations
- Rethinking State Tracking in Recurrent Models Through Error Control Dynamics
- Neural Networks and the Chomsky Hierarchy
- Learning Dynamics of Chain-of-Thought State Tracking in a Solvable Transformer Model
- Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks
- A Held-Out Transition-Pair Falsifier for Long-Horizon Non-Abelian State Tracking
- (How) Do Language Models Track State?
- Transformers Learn Shortcuts to Automata
- Decoupled Weight Decay Regularization
- Sequential Group Composition: A Window into the Mechanics of Deep Learning
- A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers
- The Illusion of State in State-Space Models
- Recirculation
- Progress measures for grokking via mechanistic interpretability
- Show Your Work: Scratchpads for Intermediate Computation with Language Models
- Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
- Marginals Before Conditionals
- Learning words in groups: fusion algebras, tensor ranks and grokking
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks