Transformers and recurrent networks learn group state tracking differently
Tracking States or Tracking Cosets? An Algebraic Account of Learned State Tracking
Machine LearningArtificial Intelligence
Summary
Keeping track of complex states over time is a tricky problem for AI models. The authors study how certain neural networks, like Transformers and recurrent networks, track states represented by mathematical groups. They find that Transformers often learn to recognize groups of states together (called cosets) in a uniform way, while recurrent networks capture more detailed state partitions, including some less obvious ones. This study links the accuracy of the models to the mathematical properties of the groups they work with, revealing new insights about what the networks actually learn internally.
What this means in practice
- •For machine learning engineers: Design neural architectures that better track structured state updates by leveraging coset and subgroup insights to improve model interpretability and accuracy.
- •For robotics software teams: Enhance state estimation in sequential decision-making tasks by applying algebraic analysis of learned internal representations to debug and improve recurrent models.
Authors
Zhiyu Zhang, Yupeng Li
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.