1 Introduction↩︎

State-tracking is fundamental across many domains: In order to understand a program state, models must track variable states during code execution [1], board configurations in game-playing [2], [3], and environment representations in world-modeling [4], [5]. Theoretical work has established a divide between associative recall, where transformers excel, and state-tracking, where recurrent neural networks perform well [6], [7]. Recently, linear RNNs like Mamba [8], [9] and DeltaNet [10], [11] were introduced which allow parallelization across the sequence length. While early linear RNNs were incapable of complex state-tracking [12], [13], extending the eigenvalues of the state-transition matrix from \([0, 1]\) to \([-1, 1]\) [14]–[16], or introducing recurrent fixed-point self-iteration over linear RNNs enables solving permutation composition  [17]. Yet empirical gains on real-world tasks remain modest [14], [18]. Concurrently, work on parallelizing nonlinear RNNs has emerged [19]–[21], but benchmarking these architectures solely on tasks expressible as Deterministic Finite-State Automata (DFA) obscures their potential, since linear RNNs already solve such tasks efficiently [16].

To understand how state-tracking can be learned by language models, we study it through the lens of next-token prediction. Current benchmarks typically use a sequence-to-sequence setup—mapping sequences of actions to states, which deviates significantly from the objective used to train language models. In this work, we address this gap and present the following contributions:

• State-Tracking via Next-Token Prediction: We convert permutation composition into code via Python REPL traces that interleave variable transformations with partial state-reveals (print statements). This creates a realistic setting where state-tracking must be learned via next-token prediction rather than sequence-to-sequence supervision.

• Linear RNNs vs. Transformers: We show that linear RNNs capable of state-tracking (specifically DeltaNet with extended eigenvalues) excel in this setting even with sparse supervision, whereas Transformers fail to generalize.

• Probabilistic State-Tracking Limits of linear RNNs: We investigate why tracking states in real code is strictly harder than permutations: actions are often not fully observable. We frame this as tracking the state of a Probabilistic Finite-State Automaton with State Reveals (PFSA-SR) and show concrete adversarial sequences under which the natural linear RNN representations of belief states suffer exponential norm decay.

2 Related Work↩︎

2.0.0.1 Deterministic state tracking.

Prior work has largely studied state tracking in deterministic settings, regular languages and finite-state automaton (FSA) emulation, using length generalization as the primary stress test [22]–[25]. These studies reveal sharp, architecture-dependent generalization gaps across the Chomsky hierarchy [24] and characterize when sequence models can implement finite-state computation under bounded depth and precision [12], [26]. However, architectural improvements on synthetic benchmarks translate to only modest gains in language-model evaluations [14], [15], [18]. Moreover, widely used setups such as group-word problems [25] differ substantially from code execution, where supervision is sparse and state evolution often stochastic. Moreover, the standard group-word problem benchmarks are sequence-to-sequence: inputs are sequence of actions and outputs are sequences of states. This differs significantly from the next-token prediction paradigm used to train modern language models. We address this gap by studying probabilistic state tracking in code-like data, in particular REPL traces, which interleave actions with partial state reveals in the same sequence. Recent work in reinforcement learning under partial observability [27], [28] and language modeling [29] empirically confirms that standard linear recurrent models struggle when uncertainty reasoning is required, motivating explicit probabilistic mechanisms such as Kalman filtering. Our PFSA-SR framework provides a theoretical explanation for these empirical observations.

2.0.0.2 Transformers and linear RNN expressivity.

Transformer expressivity analyses failure modes on state-tracking tasks including shortcut learning and poor length extrapolation [12], [22]–[24], [30], [31]. In parallel, work on efficient recurrent and SSM-style architectures has explored structured transition parameterizations trading efficiency for expressive power [10], [16], [32]–[34], showing that richer spectra or factorizations unlock new state-tracking behaviors without changing asymptotic cost [14], [15]. These works primarily address deterministic transitions, abstracting away belief evolution under uncertainty or sparse supervision.

2.0.0.3 Probabilistic automata and partial observability.

Classical probabilistic finite-state models—weighted automata, IO-HMMs, POMDPs—formalize uncertainty via belief-state updates under partial observations [35]–[38]. Recent work characterizes recurrent models as restricted probabilistic finite-state systems [39]–[41], but focuses on representational capacity of nonlinear RNNs. Our PFSA-SR instead analyzes the failure mode of linear RNNs under partial observability induced by hard, support-pruning reveals (e.g., asserts, prints) that deterministically eliminate inconsistent trajectories, enabling direct comparison of linear and nonlinear dynamics for belief-state updates. Also note the concurrent work by [42] investigating nearly deterministic state-tracking.

3 Background↩︎

3.0.0.1 Linear RNNs.

Linear RNNs process input sequences through stacked layers. Each layer transforms input vectors \({\boldsymbol{x}}_1, \dots, {\boldsymbol{x}}_t\in\mathbb{R}^l\) into outputs \(\hat{{\boldsymbol{y}}}_1,\dots,\hat{{\boldsymbol{y}}}_t\in\mathbb{R}^p\) via the linear recurrence \[\label{eq:linrnn} \begin{align} {\boldsymbol{H}}_i = {\boldsymbol{A}}({\boldsymbol{x}}_i) {\boldsymbol{H}}_{i-1} + {\boldsymbol{B}}({\boldsymbol{x}}_i),\qquad \hat{{\boldsymbol{y}}}_i = \mathrm{dec}({\boldsymbol{H}}_i, {\boldsymbol{x}}_i)\qquad \text{for } i \in \{1, \dots, t\} \end{align}\tag{1}\] where \({\boldsymbol{H}}_0 \in \mathbb{R}^{n \times d}\) is the initial state, \({\boldsymbol{A}}: \mathbb{R}^l \to \mathbb{R}^{n \times n}\) produces the state-transition matrix, \({\boldsymbol{B}}: \mathbb{R}^l \to \mathbb{R}^{n \times d}\) injects new information, and \(\mathrm{dec}: \mathbb{R}^{n \times d} \times \mathbb{R}^l \to \mathbb{R}^p\) generates outputs. These functions are learned, with \(\mathrm{dec}\) typically containing a feedforward network. Different linear RNN variants differ in their implementations of \({\boldsymbol{A}}\), \({\boldsymbol{B}}\), and \(\mathrm{dec}\).

We focus on DeltaNet [32], [43], recently shown to be parallelizable across the sequence length [10], [11]. DeltaNet parameterizes the recurrence as \({\boldsymbol{A}}({\boldsymbol{x}}_i) = {\boldsymbol{I}}-\beta_i {\boldsymbol{k}}_i{\boldsymbol{k}}_i^\top,\quad{\boldsymbol{B}}({\boldsymbol{x}}_i)=\beta_i {\boldsymbol{k}}_i{\boldsymbol{v}}_i^\top, \quad\mathrm{dec}({\boldsymbol{H}}_i, {\boldsymbol{x}}_i)=\psi({\boldsymbol{H}}_i^\top{\boldsymbol{q}}_i)\) where \(\beta_i \in [0,1]\) and \({\boldsymbol{q}}_i, {\boldsymbol{k}}_i \in \mathbb{R}^n\) (with \(\left\lVert{\boldsymbol{q}}_i\right\rVert =\left\lVert{\boldsymbol{k}}_i\right\rVert = 1\)), \({\boldsymbol{v}}_i \in \mathbb{R}^d\) are learned functions of \({\boldsymbol{x}}_i\). Here \({\boldsymbol{A}}({\boldsymbol{x}}_i)\) is a Householder transformation [44] with eigenvalues 1 (multiplicity \(n-1\)) and \(1-\beta_i\) (multiplicity 1). Restricting eigenvalues to \([0,1]\) limits state-tracking capabilities, but extending to \([-1,1]\) enables complex tracking behaviors like parity and permutation composition [14]. The linear recurrence enables parallelization across sequence length [45]–[49].

4 From Permutation Groups to Variable Tracking in Code↩︎

We translate the sequence-to-sequence setup for group-word problems from [12] to next-token prediction. This setup has been influential in understanding the state-tracking abilities of recurrent model [14], [15], [17], [50] and parallels the shell game where cups containing objects are shuffled.

Sequence-to-sequence modeling: For permutation group \(S_n\), we sample input permutations \(\sigma_{\textrm{IN}} = [\sigma_1, \ldots, \sigma_{m}]\). At position \(i\in [1,\dots,m]\), the model predicts the cumulative state \(\sigma_{\leq i}=\prod_{j=1}^i\sigma_j\) with labels \(\sigma_{\textrm{OUT}}=[\sigma_{\leq 1}, \dots, \sigma_{\leq m}]\). This provides dense supervision by computing loss at every position.

Next-token prediction (NTP): We adapt permutation groups for NTP using Python REPL traces [51], [52] that demonstrate variable shuffling. 2 shows an example for \(S_3\). We interleave commands with print statements revealing partial states rather than only end of sequence states, providing denser signal while remaining realistic (similar to logging during execution). Full reveals would trivialize the task by allowing the model to ignore prior permutations.

4.1 Experiments↩︎

Transformers require dense state supervision for state-tracking. To examine how supervision density affects state-tracking in pretrained Transformer LLMs, we finetune a Qwen3-0.6B-Base [53] model using standard NTP on Python REPL traces of 64 commands and \(S_5\). We consider both elementary swaps and full permutations, and vary the reveal spacing from 1 to 4, thereby progressively reducing the density of explicit state information available during training. As shown in 3, Transformers rely critically on dense supervision: as reveal spacing increases and state information becomes sparser, the model rapidly loses the ability to track permutations. For full permutations the model does not learn to solve the task.

Architecture Determines Whether State-Tracking Extrapolates. We train DeltaNet \([-1,1]\), a state-tracking capable architecture, and two non-state-tracking capable architectures DeltaNet \([0,1]\), and Transformer models (all with 280M parameters) on a curriculum of 15,000 REPL traces from \(S_5\), with 8, 16, 32, and 64 commands and reveal spacings of 1, 2, 4, and 8. We use full permutations rather than elementary swaps, to get a clearer signal on the difference between transformer models and linear RNNs.

The length extrapolation results in 4 (averaged across five seeds) show a clear architectural separation: DeltaNet \([-1,1]\) learns to state-track perfectly and extrapolates reliably across all seeds, whereas the Transformer does not learn to perform state-tracking even when trained from scratch on these sequences.

Interpretability Analysis. [14] showed that to learn state-tracking, at least one head must learn to set \(\beta\) values to 2 in one of the layers (to obtain an eigenvalue of -1) and find appropriate keys in the right subspace. This was previously demonstrated for a single layer DeltaProduct\(_2\) model by [15] for the \(S_4\) group. We verify whether DeltaNet\([-1,1]\) learns to perform state-tracking using NTP by leveraging its extended eigenvalue range. Therefore, we pass sequences of 512 commands with spacing 4 through the model and retrieve the \(\beta\) values using NNsight [54]. Figure 5 shows the distribution of \(\beta\) across layers and heads. We observe that many \(\beta\) values exceed 1, which results in negative eigenvalues for the generalized Householder. Layer 12 stands out, with at least one head consistently estimating most \(\beta\) values at 2 (Fig. 5, middle). Furthermore, PCA analysis reveals that the keys of this head lie in a significantly lower-dimensional linear subspace than the keys of any other head in the network. We hypothesize that this head has learned to be the state-tracking head of the network. Results for the remaining layers are in 7.2.

To confirm that head 3 in layer 12 is the state-tracking head, we perform an intervention where we gradually scale \(\beta\) to restrict the eigenvalue range from \([-1, 1]\) to \([0, 1]\) for each head independently with results shown in Fig. 6. We evaluate on sequences with 5 variables, reveal spacing 2, and 64 commands. Scaling any head except head 3 leaves the state prediction accuracy unchanged. However, scaling head 3 causes significant performance degradation, confirming that it alone enables state-tracking for the model.

5 Why State-Tracking Fails on Real Code: From Deterministic to Probabilistic Automata with State Reveals↩︎

The preceding experiments establish a clear recipe for learning state-tracking via next-token prediction: reveal intermediate states through print statements (for learnability) and use an architecture with state-tracking capabilities. DeltaNet\([-1,1]\) satisfies both criteria and generalizes reliably on synthetic permutation traces (§[subsec:motivation95arch]). However, this success relies on an idealization that real code violates: every transition is fully observable. In our REPL traces, the model sees exactly which variables are swapped at each step. Real code execution rarely gives such transparency.

Figure 7 illustrates three representative sources of transition uncertainty that arise when language models execute code. (A) Unobservable branching occurs when code branches on conditions inaccessible to the model, whether from explicit randomness, hidden environment state, or external API calls. (B) Approximation arises when models skip expensive operations like loops or recursion, trading exact computation for efficiency but introducing uncertainty about intermediate states. (C) Inverse inference tasks, such as CRUXEval-Input [55] or symbolic execution [56], require recovering inputs from outputs, which often admits multiple valid solutions.

In all three cases, a model cannot deterministically compute the next state but must instead maintain a distribution over possible states. This probabilistic setting differs fundamentally from the deterministic permutation tracking studied in prior work on linear RNNs [12]–[14], [26]. While probabilistic state-tracking has been analyzed for nonlinear RNNs [39]–[41], its implications for linear RNNs remain unexplored.

5.1 Probabilistic Finite-State Automata with State Reveals↩︎

To analyze this formally, we introduce a model that explicitly represents probabilistic transitions (e.g. conditional array swaps) coupled with partial state reveals (e.g. outputs of print statements of the array elements). This formulation treats the input sequence as a series of operations where partial state observation precedes state evolution.

The consistency constraint is the defining feature of the reveal mechanism: since the environment must choose \(\sigma_t\) so that \(q_t \in \rho(\sigma_t)\), the observer (who sees only \(\sigma_t\)) can eliminate all states outside \(\rho(\sigma_t)\) from the current belief. This models code-level observations such as print statements and asserts, which constrain but do not fully determine the latent state. Two special cases are worth noting: (1) reveal-only symbols, where \(\delta(q,\sigma)\) places all mass on \(q\) for every \(q\) (i.e. the state does not change), and (2) transition-only symbols, where \(\rho(\sigma) = \mathcal{Q}\) (i.e. the reveal is vacuous).

Belief Update. If we only observe the sequence of symbols, the uncertainty in the transitions prevents us from directly tracking the state. What we can model is instead the conditional probability distribution of the state given the input history, denoted by \(p(q_t \mid \sigma_1, \dots, \sigma_t)\), which we refer to as the belief \(b_t\). If we know \(\delta\) and \(\rho\), we can update the belief in two stages. First, the reveal step conditions the current belief \(b_t\) on the reveal \(\rho(\sigma_t)\), since we are guaranteed that \(q_t \in \rho(\sigma_t)\). The intermediate belief \(b'_t\) is computed by zeroing out states inconsistent with \(\rho(\sigma_t)\) and renormalizing. Second, the transition step computes the next belief \(b_{t+1}\) by propagating the corrected belief \(b'_t\) through the transition kernel: \[b'_{t}(q) = b_{t}(q) \cdot \mathbb{I}[q \in \rho(\sigma_t)] / \textstyle\sum_{k \in \rho(\sigma_t)} b_{t}(k) \qquad b_{t+1}(q') = \textstyle\sum_{q \in \mathcal{Q}} b'_t(q) \cdot \delta(q, \sigma_t)(q')\] Let \(m = |\mathcal{Q}|\) be the number of states. We represent the belief \(b_t\) as a column vector in the unit simplex \(\Delta_m \subset \mathbb{R}^m\). For each input \(\sigma \in \Sigma\), we define the reveal matrix \(Z_\sigma \in \mathbb{R}^{m \times m}\) as a diagonal matrix where \((Z_\sigma)_{ii} = 1\) if state \(i \in \rho(\sigma)\) and \(0\) otherwise. We define the transition matrix \(T_\sigma \in \mathbb{R}^{m \times m}\) as a column-stochastic matrix where \((T_\sigma)_{ij} = \delta(j, \sigma)(i)\) represents the probability of transitioning to state \(i\) from state \(j\). The full belief update for input \(\sigma_t\) is: \[\label{eq:beliefupdate} b_{t+1} = f(T_{\sigma_t} Z_{\sigma_t} b_t), \quad \text{where } f(x) = x / \|x\|_1\tag{2}\] This model captures two key aspects of code execution. First, partial observability: we receive reveals that prune possible states (e.g., through print statements) rather than observing the full state. Second, probabilistic transitions arise from operations with random choices or external inputs.

Concrete example: unobservable branching as a PFSA-SR. To illustrate how code execution maps to a PFSA-SR, consider the unobservable branching scenario from Figure 7 (A). Suppose a program maintains a list [a, b] and executes if random() < 0.5: a, b = b, a. We model this as a PFSA-SR with \(\mathcal{Q} = S_2 = \{(a,b),\, (b,a)\}\), two states corresponding to the two possible orderings. The conditional swap is a transition-only symbol \(\sigma_{\mathrm{swap}}\) with \(\rho(\sigma_{\mathrm{swap}}) = \mathcal{Q}\) (no information revealed) and \(\delta(q, \sigma_{\mathrm{swap}})\) assigning probability \(1/2\) to each state. A subsequent print(a) revealing that a equals its original value is a reveal-only symbol \(\sigma_{\mathrm{rev}}\) with \(\rho(\sigma_{\mathrm{rev}}) = \{(a,b)\}\) (only the identity ordering is consistent) and \(\delta(q, \sigma_{\mathrm{rev}})(q) = 1\) (no state change). After observing \(\sigma_{\mathrm{swap}}\) followed by \(\sigma_{\mathrm{rev}}\), the belief collapses from \([1/2,\, 1/2]\) back to \([1,\, 0]\) via the reveal’s support pruning and renormalization.

Connections to existing models. PFSA-SR subsumes and relates to several classical models. It reduces to a DFA when \(\delta\) is deterministic and the reveals are not informative, i.e. \(\rho (\sigma) = \mathcal{Q}\) for every \(\sigma \in \Sigma\). If instead of hard observations we receive soft observations rather than hard constraints, the model is closely related to IO-HMMs [36] or POMDP [37], [38] observation models. If reveals are absent and \(\delta\) is linear, we recover a Weighted Finite Automaton (WFA) viewpoint used in linear RNN theory. Compared to POMDPs and IO-HMMs, the essential difference here is the emphasis on hard, support-pruning reveals that eliminate some states rather than likelihood-weighted observations.

5.2 Two Representations for Probabilistic Permutation Tracking↩︎

To ground the PFSA-SR in a concrete but difficult problem, we analyze permutation tracking under uncertainty. We want to track the configurations of a list of \(n\) distinct elements after a sequence of probabilistic permutations, each a convex combination of elements from group \(S_n\). We consider two representations for the belief state, which differ in their computational complexity and stability under linear recurrence.

5.2.1 Joint Representation↩︎

The joint representation explicitly enumerates all possible automaton states. For permutation tracking over \(S_n\), the state space contains \(n!\) configurations (one for each permutation of the list). The belief \(b_t \in \mathbb{R}^{n!}\) is a probability vector over these configurations. Note that this is much less efficient than the deterministic case, where we can track the state using a linear RNN with an \(n-1\)-dimensional state: each permutation in \(S_n\) can be represented by a \(n{-}1 \times n{-}1\) permutation matrix.

A key limitation of the joint representation in a linear recurrence is the role of the input injection term \({\boldsymbol{B}}({\boldsymbol{x}}_t)\) (Eq. 1 ). In the marginal representation (discussed below), a reveal directly replenishes decayed entries via \({\boldsymbol{B}}\). In the joint representation, however, a reveal of the form “position \(i\) contains element \(j\)” is consistent with \((n{-}1)!\) of the \(n!\) configurations, and the subset that survives depends on the current belief \(b_t\), not just the input \({\boldsymbol{x}}_t\). Since \({\boldsymbol{B}}\) is a function of the input alone, it cannot adaptively identify which configurations to replenish. We return to this point in the stability analysis (§5.3), where we show it implies that the joint representation is always numerically unstable under repeated partial reveals. We provide a detailed worked example in 8.1.

5.2.2 Marginal Representation↩︎

As a more tractable alternative, instead of tracking the joint probability of all positions, we track the marginal probability of each position independently. Consider a list of elements starting from \((1,2,\dots, n)\). The belief state is an \(n \times n\) matrix \({\boldsymbol{H}}_t\) where entry \(({\boldsymbol{H}}_t)_{ij}\) represents the probability that position \(i\) contains element \(j\). The starting state is \({\boldsymbol{H}}_0 = I\). To align with standard state-transition notation (left multiplication), rows (\(i\)) represent positions in memory and columns (\(j\)) represent elements/variables.

This matrix must satisfy two constraints derived from the definition of a permutation: (1) each position contains exactly one element, \(\sum_{j} ({\boldsymbol{H}}_t)_{ij} = 1\) for all \(i\); and (2) each element exists at exactly one position, \(\sum_{i} ({\boldsymbol{H}}_t)_{ij} = 1\) for all \(j\). Matrices satisfying these properties form the Birkhoff polytope of doubly stochastic matrices. Tracking the distribution over \(S_n\) corresponds to moving a point within this polytope. By the Birkhoff–von Neumann theorem [57], [58], any doubly stochastic matrix can be expressed as a convex combination of permutation matrices.

The marginal representation admits a natural bi-linear recurrence: \[{\boldsymbol{H}}_t = {\boldsymbol{A}}_l({\boldsymbol{x}}_t){\boldsymbol{H}}_{t-1}{\boldsymbol{A}}_r({\boldsymbol{x}}_t) + {\boldsymbol{B}}({\boldsymbol{x}}_t),\quad\hat{{\boldsymbol{y}}}_t = \mathrm{dec}({\boldsymbol{H}}_t, {\boldsymbol{x}}_t) \quad \text{where } t \in \{1, \dots, T\}\] Using the identity (sometimes referred to as Roth’s column lemma) \(\mathrm{vec}({\boldsymbol{B}}{\boldsymbol{X}}{\boldsymbol{A}}^\top)=({\boldsymbol{A}}\otimes {\boldsymbol{B}})\mathrm{vec}({\boldsymbol{X}})\) [59], [60], we can see that this remains an affine update on the vectorized state: \(\text{vec}({\boldsymbol{H}}_{t}) = ({\boldsymbol{A}}_r({\boldsymbol{x}}_t)^\top \otimes {\boldsymbol{A}}_l({\boldsymbol{x}}_t)) \text{vec}({\boldsymbol{H}}_{t-1}) + \text{vec}({\boldsymbol{B}}({\boldsymbol{x}}_t))\). We note that a similar two-sided recurrence structure is present in the open-source implementation of Kimi Delta Attention (KDA) [61] (code), although the recurrence was not explicitly described in the paper.

For representing the marginal distribution during probabilistic state-tracking, we parameterize the recurrence as \({\boldsymbol{A}}_l ({\boldsymbol{x}}_t) = {\boldsymbol{P}}_{s} ({\boldsymbol{x}}_t){\boldsymbol{D}}_l({\boldsymbol{x}}_t)\) and \({\boldsymbol{A}}_r({\boldsymbol{x}}_t) = {\boldsymbol{D}}_r({\boldsymbol{x}}_t)\), where \({\boldsymbol{D}}_l({\boldsymbol{x}}_t), {\boldsymbol{D}}_r({\boldsymbol{x}}_t)\) are diagonal matrices and \({\boldsymbol{P}}_s({\boldsymbol{x}}_t)\) is a (possibly stochastic) permutation matrix.

Probabilistic Transitions (Mixing). When the input \({\boldsymbol{x}}_t\) encodes a shuffle (e.g., “swap contents of Position \(a\) and Position \(b\)”), we permute the rows of \({\boldsymbol{H}}_t\). If the shuffle is probabilistic, we apply a linear mixture of permutation matrices: \[{\boldsymbol{H}}_{t+1} = {\boldsymbol{P}}_s({\boldsymbol{x}}_t) {\boldsymbol{H}}_t,\qquad {\boldsymbol{P}}_s({\boldsymbol{x}}_t) = \sum_{i=1}^{n!} c_i {\boldsymbol{P}}_i,\qquad \sum_{i=1}^{n!} c_i = 1\] where \({\boldsymbol{P}}_s\) acts on the Positions (rows) and \({\boldsymbol{P}}_i\) are permutation matrices. The other components are set to \({\boldsymbol{D}}_l({\boldsymbol{x}}_t) = {\boldsymbol{D}}_r({\boldsymbol{x}}_t) = {\boldsymbol{I}}\) and \({\boldsymbol{B}}({\boldsymbol{x}}_t) = \mathbf{0}\). Note that \({\boldsymbol{P}}_s({\boldsymbol{x}}_t)\) is doubly stochastic by definition, so if \({\boldsymbol{H}}_{t}\) is doubly stochastic, \({\boldsymbol{H}}_{t+1}\) will be as well.

State Reveals. If \({\boldsymbol{x}}_t\) encodes the constraint “Position \(i\) contains Element \(j\)”, we enforce: (1) entry \((i, j)\) is confirmed (set to \(1\)); (2) position \(i\) cannot contain any other element (zero out row \(i\) except column \(j\)); (3) element \(j\) cannot be in any other position (zero out column \(j\) except row \(i\)). This translates into: \[{\boldsymbol{P}}_s({\boldsymbol{x}}_t) = {\boldsymbol{I}},\quad {\boldsymbol{D}}_l({\boldsymbol{x}}_t) = {\boldsymbol{I}}- {\boldsymbol{e}}_i{\boldsymbol{e}}_i^\top,\quad {\boldsymbol{D}}_r({\boldsymbol{x}}_t) = {\boldsymbol{I}}- {\boldsymbol{e}}_j{\boldsymbol{e}}_j^\top,\quad {\boldsymbol{B}}({\boldsymbol{x}}_t) = {\boldsymbol{e}}_i {\boldsymbol{e}}_j^\top\] where \({\boldsymbol{e}}_i\) is a standard basis vector of \(\mathbb{R}^n\). The term \({\boldsymbol{D}}_{l} {\boldsymbol{H}}_t {\boldsymbol{D}}_{r}\) preserves probability mass for configurations compatible with the observation (outside the cross of row \(i\) and column \(j\)), while \({\boldsymbol{B}}({\boldsymbol{x}}_t)\) directly injects the certainty of the observed fact. After this update, \({\boldsymbol{H}}_t\) may leave the Birkhoff polytope. To recover a valid marginal distribution, we apply the Sinkhorn-Knopp algorithm [62] in the decoder (to keep the main recursion linear), which iteratively normalizes rows and columns to project back onto the Birkhoff polytope.

5.3 Stability Analysis: When Does Each Representation Fail?↩︎

We now characterize when linear RNNs can and cannot maintain stable belief tracking under each representation. The central issue is that the PFSA-SR belief update (Eq. 2 ) requires nonlinear renormalization (\(f(x) = x/\|x\|_1\)) after each reveal. A linear RNN must defer this normalization to the decoder, propagating the unnormalized state \(h_t\) through time. We analyze the consequences for each representation.

Joint representation: unstable under partial reveals. Consider the unnormalized linear recurrence that defers normalization: \[\label{eq:unnorm95message} h_{t+1} = T_{\sigma_t} Z_{\sigma_t} h_t, \qquad b_t = h_t/\|h_t\|_1\tag{3}\] The “survival probability” at step \(t\) is \(s_t := \|Z_{\sigma_t} b_t\|_1 \in (0,1]\), representing the fraction of belief mass consistent with the reveal. The unnormalized state carries the cumulative product \(\|h_{t}\|_1 = \prod_{k=0}^{t-1} s_k\) in its magnitude. Since \({\boldsymbol{B}}({\boldsymbol{x}}_t)\) cannot replenish mass in the joint case (as discussed in §4.2.1), whenever partial reveals repeatedly prune mass (i.e., \(s_k<1\) for many steps), \(\|h_t\|_1\) shrinks exponentially, and the state vanishes in finite precision. This is illustrated in Figure 9.

This is a known numerical issue in HMM/Bayesian filtering: unnormalized forward messages \(\alpha_t(i) = p(y_{1:t},x_t = i)\) encode the observation-prefix likelihood in their scale and quickly underflow for long sequences, motivating per-step scaling (“scaled forward-backward”) or log-domain implementations (e.g. log-sum-exp) for numerical stability [63], [64], both of which break the linearity of the recurrence.

5.3.0.1 Marginal representation: stable with sufficient reveals, unstable under adversarial sequences.

In contrast, the marginal representation can leverage \({\boldsymbol{B}}({\boldsymbol{x}}_t)\) to replenish decaying entries. When a reveal “position \(i\) contains element \(j\)” occurs, \({\boldsymbol{B}}({\boldsymbol{x}}_t) = {\boldsymbol{e}}_i {\boldsymbol{e}}_j^\top\) directly injects a 1 at entry \((i,j)\), regardless of its prior value. This acts as a “reset” that prevents accumulated decay.

Stable case: If each position/element pair is revealed periodically, every entry in the marginal matrix \({\boldsymbol{H}}_t\) receives periodic replenishment. The decay between reveals is bounded, and finite precision suffices.

Adversarial case: An adversarial reveal sequence can cause specific entries to decay indefinitely. Consider \(n=3\) elements starting at identity \({\boldsymbol{H}}_0=I\). The adversary repeatedly: (1) applies a probabilistic swap mixing elements 2 and 3, then (2) reveals “position 2 contains element 2”. The reveal replenishes entry \((2,2)\) to 1, but element 3’s mass (at positions 2 and 3) is never directly revealed, it can only be inferred by exclusion. Each cycle halves the unrevealed mass (\(1 \to 0.5 \to 0.25 \to \cdots\)) until it vanishes in finite precision. See Figure 10 for an illustration.

Summary. Under the representations analyzed here, the joint representation is numerically unstable for linear RNNs because reveals cannot replenish mass, and the marginal representation can be stable if reveals are distributed across all variables but admits adversarial sequences that force exponential decay. These constructions provide strong evidence that linear RNNs face an inherent difficulty with probabilistic state-tracking: without either nonlinear renormalization in the recurrence or constraints on the reveal distribution, stable belief maintenance appears infeasible.

One might note that finite precision makes the set of representable beliefs finite, so in principle the dynamics can be simulated by a DFA (and hence linearized via one-hot encoding). However, this is impractical: in the joint case, even a binary discretization of the \(n!\)-dimensional belief vector requires \(2^{n!}\) DFA states; in the marginal case, the belief lives in the Birkhoff polytope of dimension \((n{-}1)^2\), so discretizing each coordinate into \(k\) bins yields \(k^{(n-1)^2}\) states (e.g., \(10^{81}\) for \(n{=}10\), \(k{=}10\)).

Linear RNNs suffice for deterministic automata and full state reveals. Finally, we note two regimes where Linear RNNs suffice. First, if the system is a Deterministic Finite Automaton (DFA), the belief state remains a one-hot vector corresponding to the true state. Since no uncertainty is introduced, the normalization constant is always 1, and the update is linear. Second, if the system receives Full State Reveals (observing the exact current state), the belief is effectively reset to a one-hot vector. This "hard reset" can be implemented by the linear update \(h_{t+1} = 0 \cdot h_t + \mathbf{c}\), which zeroes out the history and injects unit mass at the revealed state. Consequently, in a probabilistic setting with partial and full reveals, if full state reveals occur frequently, they prevent the exponential decay of the state norm caused by partial reveals stabilizing the recurrence.

6 Conclusion↩︎

We study state-tracking in neural sequence models through the lens of code execution, bridging the gap between abstract automata benchmarks and the next-token prediction paradigm used to train language models.

Summary of contributions. First, we introduced Python REPL traces as a testbed for state-tracking under next-token prediction, showing that linear RNNs with extended eigenvalue spectra (DeltaNet\([-1,1]\)) can learn and generalize reliably with sparse state supervision, while Transformers fail even with dense reveals. Second, we provided evidence for a barrier facing linear RNNs in realistic code settings: when transitions are probabilistic or partially observable, exact belief tracking requires nonlinear renormalization. We formalized this through the PFSA-SR framework and exhibited adversarial reveal sequences that cause exponential decay in state norms under the natural joint and marginal representations, suggesting that stable belief maintenance is infeasible for linear RNNs under finite precision.

Limitations. Our experiments focus on synthetic permutation groups (\(S_n\)) rather than real-world code. While REPL traces are more realistic than prior sequence-to-sequence setups, they still simplify away parsing, control flow, and memory management.

Future work. Our findings suggest that probabilistic state-tracking is a promising benchmark for evaluating nonlinear RNNs [65], [66] and recent parallelization efforts [19]–[21]. Extending this work to real execution traces, e.g. from CRUXEval or system call logs, and investigating hybrid architectures that interleave linear recurrence with periodic nonlinear normalization similar to TTT [67] and Titans [68] is left as future work. Concurrent work demonstrates that Kalman-style filtering can be parallelized for language modeling [29].

Acknowledgments↩︎

We would like to thank Jan Tönshoff, Heiner Kremer, Fabian Falck, Alicia Curth, Teodora Pandeva, Hari Govind V K, and Andrey Rybalchenko for insightful discussions throughout the project.

Author Contributions↩︎

J.S. proposed to use code execution traces in the form of REPL traces and iterated on the format with R.G., K.K. and B.R. Identified failure cases under partial observability in code based on execution traces, conceived the possibility for failure of probabilistic state-tracking in linear RNNs for the joint representation in this setting. Developed the software implementation, ran the primary experiments, and led the manuscript writing. R.G. iterated on the REPL traces format, formalized the probabilistic finite-state example by J.S, proposed marginal state-representations and made connection to prior work on recurrences, and contributed heavily to the manuscript writing. K.P. did preliminary experiments on comparing linear and non-linear RNNs comparing their practical probabilistic state tracking capabilities. K.K., H.B., and B.R. provided project supervision, iterated on the format of the REPL traces, guided the overarching research direction, and contributed to manuscript revisions. B.R. provided additional conceptual framing based on parallel work in code world models.

7 From Permutation Groups to Variable Tracking in Code↩︎

7.1 Experimental Details↩︎

Model & Training. We train a DeltaNet and Transformer (\(\approx\)​265M parameters) using the implementation from flash-linear-attention [69]. The architecture consists of 18 layers, hidden dimension \(d=512\), 8 heads (\(d_{head}=128\)), MLP expansion factor 4, and SwiGLU activations. Optimization is performed using AdamW [70] with a peak learning rate of \(5\times 10^{-4}\), a cosine decay schedule [71] (minimum LR ratio 0.2), and 5% warmup. We use a per-device batch size of 3 with 12 gradient accumulation steps in BF16 mixed precision. We trained on single Nvidia A100s for each model.

Curriculum Learning. To ensure stability, we use a four-stage curriculum of 15,000 samples each. We use full permutations on 5 variables (\(S_5\)), progressively increasing the trace length (\(L\)) and reveal spacing (\(S\)): \[(L, S) \in \{(8, 1), (16, 2), (32, 4), (64, 8)\}\]

7.2 Interpretability↩︎

8 Probabilistic State-Tracking↩︎

8.1 Example: Linear RNN implementing probabilistic finite-state automaton tracking the join↩︎

This section details the explicit arithmetic of a stochastic update on the permutation group \(S_3\) and discusses the theoretical implications of norm decay in Linear RNNs.

Scenario Setup and Initialization (\(t=0\)). We model the system state \({\boldsymbol{H}}_t \in \mathbb{R}^6\) over the permutation group \(S_3\). The basis vectors (universes) correspond to the six possible permutations of \([1,2,3]\): Identity \([1,2,3]\) (Index 1), Swap 1-2 \([2,1,3]\) (Index 2), Swap 1-3 \([3,2,1]\) (Index 3), Swap 2-3 \([1,3,2]\) (Index 4), Cycle Left \([2,3,1]\) (Index 5), and Cycle Right \([3,1,2]\) (Index 6). We begin at Step 0 with perfect certainty at the Identity configuration: \[{\boldsymbol{H}}_0 = \begin{bmatrix} 1 & 0 & 0 & 0 & 0 & 0 \end{bmatrix}^T \quad (\text{Universe 1: } [1, 2, 3]).\]

Step 1: The Stochastic Action (\(t=1\)). The system receives the input "Try to Swap 1-2" with a noise profile defined as a 50% chance of the intended Swap 1-2 and a 50% chance of an accidental Swap 1-3. To represent this ambiguity, we construct a "fuzzy" transition matrix \({\boldsymbol{A}}_{fuzzy}\) by averaging the permutation matrices of the two outcomes: \({\boldsymbol{A}}_{fuzzy} = 0.5 \cdot {\boldsymbol{A}}_{\text{swap12}} + 0.5 \cdot {\boldsymbol{A}}_{\text{swap13}}\). Specifically, \({\boldsymbol{A}}_{\text{swap12}}\) maps Identity to Index 2, while \({\boldsymbol{A}}_{\text{swap13}}\) maps Identity to Index 3. The resulting update is: \[{\boldsymbol{H}}_1 = {\boldsymbol{A}}_{fuzzy} {\boldsymbol{H}}_0 = \begin{bmatrix} 0 & 0.5 & 0.5 & 0 & 0 & 0 \\ \mathbf{0.5} & 0 & 0 & 0 & 0.5 & 0 \\ \mathbf{0.5} & 0 & 0 & 0 & 0 & 0.5 \\ 0 & 0 & 0 & 0 & 0.5 & 0.5 \\ 0 & 0.5 & 0 & 0.5 & 0 & 0 \\ 0 & 0 & 0.5 & 0.5 & 0 & 0 \end{bmatrix} \begin{bmatrix} 1 \\ 0 \\ 0 \\ 0 \\ 0 \\ 0 \end{bmatrix} = \begin{bmatrix} 0 \\ \mathbf{0.5} \\ \mathbf{0.5} \\ 0 \\ 0 \\ 0 \end{bmatrix}.\] The state is now diffused; probability mass is split between Universe 2 (\([2,1,3]\)) and Universe 3 (\([3,2,1]\)).

Step 2: The Observation (\(t=2\)). We subsequently receive the observation "Position 1 contains Object 3". To process this, we construct a diagonal observation matrix \({\boldsymbol{A}}_{obs}\) that acts as a filter. We check every basis vector: Index 1 (\([1,2,3]\)) is false; Index 2 (\([2,1,3]\)) is false; Index 3 (\([3,2,1]\)) is true (keep); Index 6 (\([3,1,2]\)) is true (keep). All others are zeroed out. Applying this filter to the smeared state \({\boldsymbol{H}}_1\): \[{\boldsymbol{H}}_2 = {\boldsymbol{A}}_{obs} {\boldsymbol{H}}_1 = \begin{bmatrix} 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & \mathbf{1} & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & \mathbf{1} \end{bmatrix} \begin{bmatrix} 0 \\ 0.5 \\ 0.5 \\ 0 \\ 0 \\ 0 \end{bmatrix} = \begin{bmatrix} 0 \\ 0 \\ \mathbf{0.5} \\ 0 \\ 0 \\ 0 \end{bmatrix}.\] The model has correctly identified that we are in Universe 3 (\([3,2,1]\)). The ambiguity created in Step 1 was resolved because Universe 2 is inconsistent with the observation. The final vector magnitude is \(0.5\), representing the joint probability of the path: \(P(\text{Path}) = P(\text{Slip}) \times P(\text{Consistent}) = 0.5 \times 1.0 = 0.5\).

The Mechanics of Information Decay. A critical limitation of Linear RNNs when tracking probabilistic states is the phenomenon of norm decay. Unlike non-linear models (e.g., Transformers with Softmax) which re-normalize their internal state, a Linear RNN performs purely multiplicative updates (\({\boldsymbol{H}}_t = {\boldsymbol{A}}{\boldsymbol{H}}_{t-1}\)). In a stochastic setting, transition matrices often have eigenvalues \(|\lambda| < 1\) due to diffusion or filtering. Consequently, \(\|{\boldsymbol{H}}_t\| \approx \lambda^t \|{\boldsymbol{H}}_0\|\). For example, in a "Noisy Identity" transition where the system retains state with \(p=0.9\), the norm decays to \(\approx 0.00002\) after 100 steps. While floating-point standards allow for small numbers, the signal eventually vanishes relative to numerical noise.

The Role of B: The Gated Reset. The matrix \({\boldsymbol{B}}\), which is typically zero during standard tracking, serves as a solution to decay via a Gated Reset Mechanism. Let \({\boldsymbol{x}}_t\) be a binary indicator where \({\boldsymbol{x}}_t=1\) indicates a "Reset". We parameterize the weights as \({\boldsymbol{A}}({\boldsymbol{x}}_t) = (1 - {\boldsymbol{x}}_t) \cdot {\boldsymbol{A}}_{\text{step}}\) and \({\boldsymbol{B}}({\boldsymbol{x}}_t) = {\boldsymbol{x}}_t \cdot {\boldsymbol{H}}_{\text{prior}}\). When a reset is triggered (\({\boldsymbol{x}}_t=1\)), the history is annihilated (\({\boldsymbol{A}}({\boldsymbol{x}}_t) = \mathbf{0}\)) and the bias term injects the prior (\({\boldsymbol{B}}({\boldsymbol{x}}_t) = {\boldsymbol{H}}_{\text{prior}}\)). The update becomes: \[{\boldsymbol{H}}_t = \mathbf{0} \cdot {\boldsymbol{H}}_{t-1} + {\boldsymbol{H}}_{\text{prior}} = [1/6, \dots, 1/6]^T.\] This "re-inflates" the state vector to full magnitude, readying the system to track a new sequence.

References↩︎