Scaling Laws for State Dynamics in Large Language Models

TL;DR

This paper investigates how Transformer-based LLMs encode and update internal state to model deterministic state dynamics, and how performance scales with state-space size and transition sparsity. It formalizes the tasks as deterministic finite automata, spanning Box Tracking, Abstract DFA Sequences, and Complex Text Games, and measures next-state prediction accuracy as a function of $|Q|$ and $|T|$, while applying activation patching to locate responsible components. The main findings show that next-state accuracy degrades as state-space grows and transitions become sparser; GPT-2 XL reaches about 70% accuracy in low-complexity settings but falls below 30% when $|Q|>5$ or $|Q|>10$ depending on domain, and DFA tasks show limited joint state-action reasoning with action information not reliably routed to the final token. Mechanistic analysis reveals state tracking arises from distributed interactions among next-token heads rather than explicit symbolic computation, with induction heads and Name-Mover heads contributing but not forming a single symbolic state machine. These results have implications for evaluating and interpreting LLMs on stateful tasks and for designing architectures that can robustly track complex state dynamics.

Abstract

Large Language Models (LLMs) are increasingly used in tasks requiring internal state tracking, yet their ability to model state transition dynamics remains poorly understood. We evaluate how well LLMs capture deterministic state dynamics across 3 domains: Box Tracking, Abstract DFA Sequences, and Complex Text Games, each formalizable as a finite-state system. Across tasks, we find that next-state prediction accuracy degrades with increasing state-space size and sparse transitions. GPT-2 XL reaches about 70% accuracy in low-complexity settings but drops below 30% when the number of boxes or states exceeds 5 or 10, respectively. In DFA tasks, Pythia-1B fails to exceed 50% accuracy when the number of states is > 10 and transitions are < 30. Through activation patching, we identify attention heads responsible for propagating state information: GPT-2 XL Layer 22 Head 20, and Pythia-1B Heads at Layers 10, 11, 12, and 14. While these heads successfully move relevant state features, action information is not reliably routed to the final token, indicating weak joint state-action reasoning. Our results suggest that state tracking in LLMs emerges from distributed interactions of next-token heads rather than explicit symbolic computation.

Figures (23)

• Figure 1: GPT2 Model Performance (Small, Medium, Large, X-Large)
• Figure 2: Comparison of activation patching results for GPT-2 XL. Left: residual stream patching across token positions. Right: attention head patching.
• Figure 3: Aggregated attention pattern from the final token position for the top 5 heads identified by activation patching in GPT-2 XL, on an example prompt in Box Tracking domain. The queried object is "watch", and its correct final location is "Box A". The attention pattern shows strong focus on "A" especially in the phrase "Move the watch into Box A".
• Figure 4: TinyStories Model Performance (8M, 28M, 33M)
• Figure 5: GPT2 Model Performance (Small, Medium, Large, X-Large)