A Fine-Grained Hierarchy and Benchmark for State Tracking in Language Models

Jackson Petty; William Merrill; Ashish Sabharwal

cs.AINov 2, 2026

A Fine-Grained Hierarchy and Benchmark for State Tracking in Language Models

Authors: Jackson Petty, William Merrill, Ashish Sabharwal

Organizations: Allen Institute for AI · New York University · Toyota Technological Institute at Chicago

Abstract

Theoretical limitations of transformers on state tracking---the ability to track information about entities or variables over long contexts---have motivated new recurrent language modeling (LM) architectures such as Gated DeltaNet (GDN). These architectures have, in particular, been benchmarked against the $S_5$ word problem, a theoretically hard state-tracking task that captures abilities such as code evaluation under arbitrary variable assignment. It, however, remains open whether other, simpler forms of state tracking might both be more relevant in practice and already explain observed differences between these architectures. We propose evaluating state tracking using a hierarchy of word problems with four increasing levels of difficulty: r-trivial, corresponding to immutable assignment to variables in programming; aperiodic, where assignments are mutable; soluble, where assignments are self-referential; and general finite state tracking. Theoretically, we show that the class of problems transformers are thought to length-generalize on does not even contain all aperiodic state tracking. On the other hand, we prove that GDN with positive eigenvalues expresses all aperiodic problems, and the regular languages recognized by certain transformers are precisely all soluble ones, thus providing a fine-grained picture of state tracking capabilities of various architectures. Taking these ideas to practice, we construct StateBench, a 4-tier long-horizon state-tracking benchmark. Our systematic direct training experiments across architectures largely align with theoretical expectation, albeit with interesting observations at very large scale length-generalization. Additionally, we also evaluate pretrained models---both open hybrid and closed frontier ones---on StateBench, finding that none of them robustly solves even soluble state tracking, suggesting a direction for future improvement.

Figures & tables

Panel of Figure 1 — Figure 1: Placement of various language model architectures within the hierarchy of state-tracking (monoid word problem) complexity, along with examples of the kinds of tasks that fall into these categories. Areas below the long horizontal line indicate circuit complexity classes that, when restricted to regular languages, are equal to the monoid classes above them. The two log-precision architectures extend into $\mathsf{TC}^{0}$ . Dashed lines represent conjectured strict boundaries between classes.

Figure 2: left: A visualization of an

8

-assignment general finite sequence with

4+1

variables. Each variable is assigned an initial value (e.g.,

\texttt{b}\leftarrow 10

\texttt{c}\leftarrow 9

). At step

t

, a variable may be overwritten by a new constant value (as in

\texttt{a}\leftarrow 37

t_{2}

) or it may be swapped with another (as b and d are swapped between

t_{6}

and

t_{8}

). right: Five possible final states for the sequence, with the correct one (

\delta

) annotated with a green star.

Figure 3: Direct Training: GDNs generalize perfectly on r-trivial, aperiodic, and soluble data, and partially on general finite data; Transformers fail to generalize beyond their training distribution on r-trivial, aperiodic, and soluble data and fail to even learn in-distribution general finite data. Error bars show Wilson confidence intervals.

Figure 4: Sensitivity Analysis: Histograms of sample sensitivities in (left) the four complexity distributions in StateBench and (right) code samples from CruxEval.

Figure 5: Pretrained models underperform theoretical expectations and direct-training results, showing consistent performance degradation on aperiodic data and completely failing on general finite data. Error bars show Wilson confidence intervals.

Figure 6: Frontier LLMs struggle on long-horizon soluble and general finite state-tracking tasks. Error bars show Wilson confidence intervals.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

	Monoid Complexity Class
	R-trivial	Aperiodic	Soluble	General finite
Max. dependency length (K)	$0$	$0$	$21.3$	$32.1$
Sensitivity	$0.08$	$0.08$	$0.76$	$0.75$
	[0.00–0.91]	[0.00–0.91]	[0.06–1.00]	[0.04–1.00]
Prompt tokens (K)	$1.16$	$1.16$	$1.34$	$1.25$
	[0.07–981]	[0.07–981]	[0.09–768]	[0.06–939]

Appendix

Table 1: Descriptive statistics for StateBench. Cells give medians, with [min, max] ranges in grey below. Token counts and dependency lengths are in thousands.

Figure 7: Effect of Sensitivity: On general finite state tracking data, direct trained models perform worse on higher sensitivity data. Error bars show Wilson confidence intervals.

Figure 8: Curriculum Training: Sorting general finite sequences by sensitivity during training (right) boosts direct training performance of Transformer-RoPE models on in-distribution data compared to training on the same randomly-shuffled sequences (left, repeated) ; Transformer-NoPE, GDN, and Hybrid models are unaffected. Error bars show Wilson confidence intervals.

Figure 9: Pretrained base models perform better on RULER-formatted StateBench sequences (top) than on the same sequences formatted as spreadsheet manipulations (bottom) . Error bars show Wilson confidence intervals.

Figure 10: Pretrained models evaluated on StateBench problems by complexity class and formatting. Error bars show Wilson confidence intervals.

State Tracking in Language Models

See all 8 papers

Feb 16, 2026

Learning State-Tracking from Code Using Linear RNNs

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.

Julien Siems, Riccardo Grazzi, Korbinian Pöppel +3

University of Freiburg · Prior Labs · Microsoft Research +2

Recurrent Neural Networks State Tracking Linear RNNs Sequence Modeling

Ask PDF

Oct 20, 2026

PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking

Looped language models can run more iterations at test time, but additional iterations need not help them solve longer problems. For sequential state tracking, final-answer supervision does not specify how far each iteration should advance, and repeated updates can change the scale of the recurrent state. We introduce Progress-Aligned Recurrence (PAR), which trains the prediction at iteration

k

to match the state after the first

k

operations and rescales recurrent states to their first-iteration norms. On swap-

S_5

, a Qwen2.5-0.5B model trained on at most 32 moves achieves 95.5% mean full-state accuracy at 512 moves, or

16\times

the training length, compared with 55.7% reported for RLT and 0.8% for explicit chain of thought on the same backbone. Controlled

S_5

ablations show that removing norm anchoring reduces the extrapolation factor from

16\times

to below

2.2\times

, and replacing progress-aligned targets with final-answer supervision reduces it to at most

8\times

. A learned completion readout selects when to answer without a supplied stopping depth. When intermediate labels are unavailable, prefix self-distillation induces the correspondence from shorter problem prefixes using only final answers. The gains apply across sequential state-tracking tasks, establishing the task structure required for depth extrapolation. These findings support training intermediate predictions and state scale together to make additional recurrent depth useful for longer problems.

Pan Wang, Yihao Hu, Guangyuan Dong +8

Ant International · National University of Singapore

Length Generalization Recurrent Language Models State Tracking Process Supervision

Ask PDF

Feb 20, 2026

On the "Induction Bias" in Sequence Models

Despite the remarkable practical success of transformer-based language models, recent work has raised concerns about their ability to perform state tracking. In particular, a growing body of literature has shown this limitation primarily through failures in out-of-distribution (OOD) generalization, such as length extrapolation. In this work, we shift attention to the in-distribution implications of these limitations. We conduct a large-scale experimental study of the data efficiency of transformers and recurrent neural networks (RNNs) across multiple supervision regimes. We find that the amount of training data required by transformers grows much more rapidly with state-space size and sequence length than for RNNs. Furthermore, we analyze the extent to which learned state-tracking mechanisms are shared across different sequence lengths. We show that transformers exhibit negligible or even detrimental weight sharing across lengths, indicating that they learn length-specific solutions in isolation. In contrast, recurrent models exhibit effective amortized learning by sharing weights across lengths, allowing data from one sequence length to improve performance on others. Together, these results demonstrate that state tracking remains a fundamental challenge for transformers, even when training and evaluation distributions match.

M. Reza Ebrahimi, Michaël Defferrard, Sunny Panchal +1

Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc.

Transformer State Tracking Inductive Bias Data-Efficient Learning Recurrent Neural Networks Length Generalization

Ask PDF

Abstract

State Tracking in Language Models

Explore similar work

Learning State-Tracking from Code Using Linear RNNs

PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking

On the "Induction Bias" in Sequence Models