Research map / Architectures and efficiency / Attention and transformers
(How) Do Language Models Track State?
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.
Studies how transformer LMs fine-tuned on permutation composition (S3/S5 word problems) implement state tracking, finding they consistently learn either an associative-scan mechanism (AA) or a parity-then-associative mechanism (PAA), and that intermediate parity training can steer which one is learned.
Problem
How transformer language models track unobserved state, using permutation composition (word problems on S3 and S5) as a model system for state tracking in general.
Approach
Combines theoretically derived algorithm signatures (sequential, parallel, associative, parity-associative) with activation patching, probing, attention analysis, and training dynamics to show LMs consistently learn either an associative-scan mechanism or a parity-heuristic-plus-associative-scan mechanism, and that intermediate tasks (parity prediction, topic-model NTP) can steer which is learned.
Key results
- Models consistently fall into one of two mechanism clusters (AA or PAA) across architectures, sizes, and seeds, identifiable from generalization curves after 10k training steps.
- AA-type models generalize better (state and parity cutoff lengths equal) and converge faster to lower loss than PAA-type models, whose parity cutoff is longer than state cutoff.
- Training Pythia models first on parity prediction converts pre-trained Pythia's AA-type mechanism to PAA-type; topic-model NTP pretraining from scratch always yields AA-type.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| learn an associative algorithm rather than a parity-associative algorithm | language models trained on permutation composition · S3 and S5, held-out sequence lengths | state generalization | Helped |
| learn an associative algorithm rather than a parity-associative algorithm | language models trained on permutation composition · S3 and S5, training | convergence speed | Helped |
| train from scratch with a topic-modeling next-token-prediction task before S3 training | randomly initialized language models · S3 | AA mechanism emergence | Helped |
| train on state parity before the S3 objective | GPT-2 and Pythia trained from scratch · S3 | PAA mechanism emergence | Helped |
| train on state parity before the S3 objective | Pile-pretrained Pythia models · S3 | PAA mechanism emergence | Helped |
| fine-tune on random S3 next-token-prediction documents | language models · S3 | ratio of AA and PAA mechanisms | No clear effect |
| use PAA rather than AA | language models trained on permutation composition · S3 and S5, held-out sequence lengths | state accuracy cutoff length | Hurt |
| use PAA rather than AA | language models trained on permutation composition · S3 and S5, held-out sequence lengths | parity versus state accuracy cutoff length | Mixed |
| pre-train Pythia on the Pile before natural-language permutation training | Pythia · Natural-language S3 permutation composition | parity representation depth | Helped |
| use pre-trained rather than non-pre-trained Pythia | Pythia · Natural-language S3 permutation composition | state computation depth | Helped |
Surprising findings
- Whether a model learns AA or PAA is determined by architecture and initialization rather than model size, and the choice is essentially fixed by these factors plus data order.
- Parity-head attention heads appear only in PAA models, and no parity heads are found in AA models.
Limitations
- Experiments focus on small-scale models (up to roughly 1B parameters) and synthetic permutation tasks.
- Signatures are necessary but not sufficient conditions for the proposed algorithms; the exact mechanism may be more complex.
- Natural-language experiments are preliminary, and token-position alignment for patching is imperfect.
Experimental setup
Datasets: synthetic S3 permutation composition sequences synthetic S5 permutation composition sequences The Pile topic-model-generated S3 documents natural language swap/rotate permutation sequences
Benchmarks: synthetic S3 permutation composition synthetic S5 permutation composition natural language permutation composition
Metrics: state accuracy parity accuracy normalized logit difference probe accuracy cutoff length parity head score
Baselines: sequential algorithm parallel algorithm
Builds on: Pythia GPT-2 activation patching associative scan theory of Transformers
Closest papers by meaning
- The Illusion of State in State-Space Models (ICML 2024 · similarity 0.91)
- The Expressive Capacity of State Space Models: A Formal Language Perspective (NeurIPS 2024 · similarity 0.89)
- Repeat After Me: Transformers are Better than State Space Models at Copying (ICML 2024 · similarity 0.87)
- Rethinking State Tracking in Recurrent Models Through Error Control Dynamics (NeurIPS 2026 · similarity 0.87)
- On the Expressiveness and Length Generalization of Selective State Space Models on Regular Languages (AAAI 2025 · similarity 0.87)
- Do LLMs Build World Representations? Probing Through the Lens of State Abstraction (NeurIPS 2024 · 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.