Learning State-Tracking from Code Using Linear RNNs

TL;DR

The paper investigates how to learn state-tracking in neural sequence models by recasting permutation composition as code execution through Python REPL traces, enabling next-token prediction as supervision. It shows that linear RNNs with extended eigenvalue spectra in $[-1,1]$, notably DeltaNet, can robustly solve deterministic state-tracking and generalize under sparse supervision, while Transformers fail on these tasks. To address probabilistic and partially observable transitions typical of real code, it formalizes PFSA-SR, a probabilistic finite-state automaton with state reveals, and analyzes two representations for belief tracking: the joint representation, which under partial reveals suffers exponential mass decay, and the marginal representation, which can be stable with adequate reveals but remains vulnerable to adversarial sequences. The results indicate a fundamental barrier for linear RNNs in maintaining stable beliefs without nonlinear renormalization in probabilistic settings, suggesting future work on nonlinear or hybrid architectures and applying the framework to real execution traces. Overall, the work connects deterministic automata benchmarks with code-style state evolution, highlighting both the potential and limits of linear recurrence models for state-tracking in realistic scenarios.

Abstract

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.

Figures (12)

• Figure 1: The State-Tracking Hierarchy. Linear RNNs solve deterministic tasks (blue) but struggle with probabilistic automata (orange). Real-world code often requires belief-state tracking in the orange region.
• Figure 2: Three representations of the permutation tracking task. Left: The shell game analogy showing cups being swapped to track object positions. Center: The sequence-to-sequence modeling approach from merrill2024illusion, where a Transformer/RNN processes input permutations $\sigma_i$ and outputs cumulative states $\sigma_{\leq i} = \prod_{j=1}^i \sigma_j$ at each position. Right: Our code-based representation using Python REPL traces, where variable swaps implement permutations and print statements reveal partial cumulative states for next-token prediction training.
• Figure 3: Transformers require dense state supervision to solve REPL traces. Accuracy drops as reveal spacing increases, with full permutations failing under any sparsity.
• Figure 4: Left: Per reveal position accuracy averaged across 5 seeds, reveal spacing 16 and num commands 128. Only DeltaNet $[-1, 1]$ manages to reliably learn to perform state tracking. Right: Final reveal position accuracy when increasing the reveal spacing and num. commands beyond training regime, 8 and 64 respectively.
• Figure 5: Interpretability Analysis: Distribution of $\beta$ for a sequence of 512 commands with reveal spacing 4. Left:$\beta$ values aggregated per layer across heads. Middle:$\beta$ per head for layer 12; head 3 stands out as the only head in the network to consistently set its $\beta$ values to 2, making it the state-tracking head. Right: Cumulative explained variance of the PCA of the keys. Head 3 again stands out, being explainable with much fewer components.

Theorems & Definitions (1)

• Definition 1: PFSA-SR