# Learning State-Tracking from Code Using Linear RNNs

arXiv: https://arxiv.org/abs/2602.14814 · HTML: https://mlsanity.com/paper/2602.14814/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 70 references that are in the map.
- First posted on arXiv 16 Feb 2026.

## Builds on (most cited first, 30 of 46)

- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [xLSTM: Extended Long Short-Term Memory](https://mlsanity.com/paper/2405.04517/) (2024)
- [CRUXEval: A Benchmark for Code Reasoning, Understanding and Execution](https://mlsanity.com/paper/2401.03065/) (2024)
- [Theoretical Limitations of Self-Attention in Neural Sequence Models](https://mlsanity.com/paper/1906.06755/) (2019)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Kimi Linear: An Expressive, Efficient Attention Architecture](https://mlsanity.com/paper/2510.26692/) (2025)
- [The Parallelism Tradeoff: Limitations of Log-Precision Transformers](https://mlsanity.com/paper/2207.00729/) (2022)
- [Transformer Quality in Linear Time](https://mlsanity.com/paper/2202.10447/) (2022)
- [The Illusion of State in State-Space Models](https://mlsanity.com/paper/2404.08819/) (2024)
- [RWKV-7 "Goose" with Expressive Dynamic State Evolution](https://mlsanity.com/paper/2503.14456/) (2025)
- [Parallelizing Linear Recurrent Neural Nets Over Sequence Length](https://mlsanity.com/paper/1709.04057/) (2017)
- [What Formal Languages Can Transformers Express? A Survey](https://mlsanity.com/paper/2311.00208/) (2023)
- [Evaluating the World Model Implicit in a Generative Model](https://mlsanity.com/paper/2406.03689/) (2024)
- [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://mlsanity.com/paper/2411.12537/) (2024)
- [CWM: An Open-Weights LLM for Research on Code Generation with World Models](https://mlsanity.com/paper/2510.02387/) (2025)
- [DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products](https://mlsanity.com/paper/2502.10297/) (2025)
- [Chess as a Testbed for Language Model State Tracking](https://mlsanity.com/paper/2102.13249/) (2021)
- [A Formal Hierarchy of RNN Architectures](https://mlsanity.com/paper/2004.08500/) (2020)
- [Learning Associative Inference Using Fast Weight Memory](https://mlsanity.com/paper/2011.07831/) (2020)
- [What Has a Foundation Model Found? Using Inductive Bias to Probe for World Models](https://mlsanity.com/paper/2507.06952/) (2025)

## Built on by (most cited first, 6 of 6)

- [Olmo Hybrid: From Theory to Practice and Back](https://mlsanity.com/paper/2604.03444/) (2026)
- [Can Video World Models Track Unobserved World States?](https://mlsanity.com/paper/2608.30692/) (2026)
- [On Subquadratic Architectures: From Applications to Principles](https://mlsanity.com/paper/2606.12364/) (2026)
- [How Wrong Can a Good Predictor Be? Diverging Updates with Vanishing Predictive KL](https://mlsanity.com/paper/2609.11132/) (2026)
- [The Automaton Underneath: The Additive Input Pathway Is a Parasitic Attractor for State Tracking in Householder Linear RNN](https://mlsanity.com/paper/2609.18966/) (2026)
- [Switching Linear Attention](https://mlsanity.com/paper/2609.39034/) (2026)

## Abstract (arXiv)

Over the last years, state-tracking tasks, particularly permutation composition, have become a testbed to understand the limits of sequence models architectures like Transformers and RNNs (linear and non-linear). However, these are often sequence-to-sequence tasks: learning to map actions (permutations) to states, which is incompatible with the next-token prediction setting commonly used to train language models. We address this gap by converting permutation composition into code via REPL traces that interleave state-reveals through prints and variable transformations. We show that linear RNNs capable of state-tracking excel also in this setting, while Transformers still fail. Motivated by this representation, we investigate why tracking states in code is generally difficult: actions are not always fully observable. We frame this as tracking the state of a probabilistic finite-state automaton with deterministic state reveals and show that linear RNNs can be worse than non-linear RNNs at tracking states in this setup.
