💬 NLP

Learning State-Tracking from Code Using Linear RNNs

This paper bridges the gap between state-tracking research and next-token prediction by converting permutation composition into code-based REPL traces, demonstrating that while linear RNNs excel at this task compared to Transformers, they can underperform non-linear RNNs when state reveals are deterministic but actions are not fully observable.

Original authors: Julien Siems, Riccardo Grazzi, Korbinian Pöppel, Kirill Kalinin, Hitesh Ballani, Babak Rahmani

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Julien Siems, Riccardo Grazzi, Korbinian Pöppel, Kirill Kalinin, Hitesh Ballani, Babak Rahmani

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Shell Game" of Code

Imagine you are watching a magician perform the "Shell Game." Three cups are on a table, and a ball is hidden under one. The magician swaps the cups around. Your job is to keep track of where the ball is.

The Ball: A variable in a computer program (like a number stored in a box).
The Swaps: The code instructions that move variables around.
The Goal: To know exactly where the ball is after a long series of swaps.

For a long time, researchers tested AI models on this "Shell Game" using a specific setup: they showed the AI the list of swaps and asked, "Where is the ball now?" The paper argues this is like asking a student to memorize a whole movie script and then recite the ending. It doesn't test how well the AI actually understands the story as it happens.

The New Approach: The "Live Commentary"

The authors changed the test to match how real AI (like the one you are talking to) learns: Next-Token Prediction.

Instead of showing the whole list of swaps at once, they gave the AI a live transcript of a computer program running, line by line.

Line 1: "Move cup A to B."
Line 2: "Look under cup A!" (The computer prints the result).
Line 3: "Swap cup B and C."
Line 4: "Look under cup C!"

The AI has to guess the next word in the transcript. To do this, it must track the state of the cups in its mind as it reads, just like a human following a story.

The Contenders: The "Linear" vs. The "Transformer"

The paper pits two types of AI architectures against each other:

Transformers (The "Photographic Memory"): These are the current champions of AI (like the models behind this chat). They are great at remembering facts and finding patterns if the information is right in front of them.
Linear RNNs (The "Mental Note-Takers"): These are newer, faster models designed to process information sequentially, like reading a book one word at a time.

The Result:

When the "Shell Game" was fully visible (the AI could see every swap and every peek), the Linear RNNs (specifically a type called DeltaNet) were amazing. They could track the ball perfectly, even if the game got very long.
The Transformers struggled. They needed to see the state revealed very frequently to keep up. If the "peeks" were spaced out, they got lost.

The Twist: When the Game Gets "Fuzzy"

The paper then asks: What happens when the game isn't perfectly clear?

In real computer code, things aren't always deterministic. Sometimes, the code makes a random choice, or a variable depends on something the AI can't see (like a hidden environment variable).

The authors created a scenario where the AI has to guess the state based on probabilities (e.g., "There's a 50% chance the ball moved left, and a 50% chance it stayed").

The Problem with Linear RNNs:
The paper discovered a fundamental weakness in Linear RNNs when dealing with this "fuzzy" uncertainty.

The Analogy: Imagine you are trying to keep a stack of papers balanced. Every time you get a new clue (a "reveal"), you have to reorganize the stack.
In a Linear RNN, the math used to update the stack is "linear." It's like a leaky bucket. Every time you get a partial clue, a tiny bit of your "confidence" (mathematical mass) leaks out.
If you get a long series of partial clues without a "full reset" (a clear, total reveal), the confidence in your answer shrinks exponentially. Eventually, the number becomes so small that the computer treats it as zero. The AI forgets everything.

The "Adversarial" Trap:
The authors showed that you can trick these Linear RNNs with a specific sequence of moves:

Mix the cups randomly (creating uncertainty).
Reveal the position of one specific cup (giving a partial clue).
Repeat.

Every time this happens, the Linear RNN loses a little bit of its ability to track the other cups. After enough repetitions, the AI's internal "belief" about where the other cups are vanishes completely, even though a human could still logically deduce the answer.

The Conclusion

Linear RNNs are great at tracking states when the rules are clear and the path is deterministic (like a perfect Shell Game). They can even beat Transformers at this if the training is set up correctly.
Linear RNNs struggle with real-world code where things are probabilistic or partially hidden. Their mathematical structure causes them to "forget" details over time when they have to deal with uncertainty, because they lack a mechanism to "re-normalize" or fix their confidence levels without breaking their linear structure.

In short: Linear RNNs are excellent at following a clear script, but they tend to lose their mind when the script gets fuzzy and random.

Technical Summary: Learning State-Tracking from Code Using Linear RNNs

Problem Statement

State-tracking is a fundamental capability for understanding program execution, game states, and world models. While theoretical work has established a divide where Transformers excel at associative recall and Recurrent Neural Networks (RNNs) perform well at state-tracking, a significant gap exists in how these models are evaluated versus how they are trained.

Most existing benchmarks for state-tracking (e.g., permutation composition) utilize a sequence-to-sequence setup: mapping a sequence of actions directly to a sequence of states. This differs fundamentally from the next-token prediction paradigm used to train modern Large Language Models (LLMs), where the model must predict the next token in a stream that interleaves actions with partial state reveals (e.g., print statements). Furthermore, real-world code execution often involves partial observability and probabilistic transitions (e.g., branching on hidden states or random choices), creating a setting where the true state is not fully observable, requiring the maintenance of a belief distribution rather than a deterministic state.

Methodology

1. Reformulating State-Tracking via Code (REPL Traces)

To bridge the gap between abstract automata benchmarks and LLM training objectives, the authors convert permutation composition tasks into Python REPL traces.

Mechanism: Instead of mapping a sequence of permutations $\sigma_1, \dots, \sigma_m$ to cumulative states, the model processes a sequence of Python commands (variable assignments, swaps) interleaved with print statements that reveal partial state information.
Objective: The model is trained via next-token prediction to predict the next token in the trace, which may be a command or a revealed state value. This forces the model to learn state-tracking implicitly through the prediction task.

2. Architectural Comparison

The study compares three architectural classes on these tasks:

Transformers: Standard attention-based models.
Linear RNNs (DeltaNet [0, 1]): Linear RNNs restricted to eigenvalues in $[0, 1]$ , known to be incapable of complex state-tracking like parity or permutation composition.
Linear RNNs (DeltaNet [-1, 1]): Linear RNNs with extended eigenvalue ranges $[-1, 1]$ , theoretically capable of solving deterministic permutation composition.

3. Theoretical Framework: PFSA-SR

To analyze the difficulty of state-tracking in real code, the authors introduce the Probabilistic Finite-State Automaton with State Reveals (PFSA-SR).

Definition: A tuple $(Q, \Sigma, \delta, \rho, q_0)$ where $\delta$ is a probabilistic transition kernel and $\rho$ is a reveal function mapping inputs to subsets of states.
Dynamics: The system evolves via probabilistic transitions, but the observer receives "hard" constraints (reveals) that prune the set of possible states (support pruning) rather than providing soft likelihoods.
Belief Update: Tracking requires maintaining a belief distribution $b_t$ over states. The update involves a reveal step (zeroing out inconsistent states and renormalizing) followed by a transition step. Crucially, the renormalization step $f(x) = x/\|x\|_1$ is non-linear.

4. Representation Analysis

The paper investigates two representations for belief states in linear RNNs:

Joint Representation: Explicitly enumerates all $n!$ permutations. The authors argue this is numerically unstable for linear RNNs because the input injection term $B(x_t)$ cannot adaptively replenish the probability mass of surviving configurations after a partial reveal, leading to exponential norm decay.
Marginal Representation: Tracks the marginal probability of each position independently (an $n \times n$ doubly stochastic matrix). While this allows $B(x_t)$ to directly inject certainty for revealed entries, the authors show it is still susceptible to exponential decay under adversarial reveal sequences where specific unrevealed entries are never directly replenished.

Key Results

1. Supervision Density and Architecture

Transformers: Require dense state supervision to learn state-tracking. As the spacing between print statements (reveals) increases, Transformer performance drops rapidly. They fail to generalize to full permutations even with dense supervision.
Linear RNNs (DeltaNet [-1, 1]): Excel in this setting. They learn to track states reliably even with sparse supervision and demonstrate robust length extrapolation (generalizing to sequences longer than training).
Linear RNNs (DeltaNet [0, 1]): Fail to solve the task, confirming that the extended eigenvalue range $[-1, 1]$ is necessary for complex state-tracking.

2. Interpretability of State-Tracking

Analysis of the trained DeltaNet [-1, 1] reveals that state-tracking is localized to specific components:

Eigenvalue Extension: At least one attention head in a specific layer (e.g., Layer 12, Head 3) learns to set its decay parameter $\beta$ to 2, resulting in an eigenvalue of $-1$.
Subspace Structure: This specific head operates in a significantly lower-dimensional linear subspace compared to other heads, consistent with the theoretical requirements for tracking permutations.
Intervention: Scaling the $\beta$ range of this specific head back to $[0, 1]$ causes immediate performance collapse, while scaling other heads has no effect, confirming this head as the dedicated state-tracking mechanism.

3. Limits of Linear RNNs in Probabilistic Settings

The paper demonstrates that while linear RNNs solve deterministic tasks, they face inherent barriers in probabilistic settings (PFSA-SR):

Exponential Decay: Under the joint representation, partial reveals cause the unnormalized state norm to decay exponentially ( $1/2 \to 1/4 \to 1/8 \dots$ ) because the linear recurrence cannot perform the necessary non-linear renormalization.
Adversarial Vulnerability: Even with the marginal representation, an adversary can construct a sequence of probabilistic swaps and selective reveals that causes the probability mass of unrevealed entries to decay exponentially until it vanishes in finite precision.
Conclusion: Stable belief maintenance in probabilistic settings with partial observability appears infeasible for linear RNNs without either non-linear renormalization in the recurrence or frequent full-state resets.

Significance and Claims

The paper claims to provide a more realistic testbed for state-tracking that aligns with the next-token prediction objective of language models, moving beyond synthetic sequence-to-sequence benchmarks.

Bridging the Gap: It demonstrates that linear RNNs capable of state-tracking (specifically DeltaNet with extended eigenvalues) can learn these capabilities via next-token prediction on code-like traces, whereas Transformers cannot.
Theoretical Barrier: It identifies a fundamental limitation of linear RNNs in real-world code scenarios. Real code often involves probabilistic transitions and partial observability (e.g., hidden branching, random choices). The authors show that tracking states in such environments requires non-linear renormalization, which linear RNNs cannot perform internally. This suggests that the "complexity gap" observed in real-world tasks may stem from the probabilistic nature of code execution rather than just sequence length or complexity.
Benchmarking Implications: The authors suggest that benchmarking architectures solely on Deterministic Finite-State Automata (DFA) tasks obscures the potential of non-linear RNNs and the limitations of linear RNNs. The proposed PFSA-SR framework and the associated failure modes (exponential norm decay) offer a theoretical explanation for why standard linear recurrent models struggle when uncertainty reasoning is required.

The paper concludes modestly, noting that its experiments focus on synthetic permutation groups and that future work is needed to extend these findings to real-world execution traces and to investigate hybrid architectures that interleave linear recurrence with periodic non-linear normalization.