Learning Dynamics of Chain-of-Thought State Tracking in a Solvable Transformer Model

Abstract: Chain-of-thought generation can turn a multi-step computation into a sequence of locally checkable state updates, but the training dynamics by which transformers acquire such updates remain poorly understood. We study this question in a solvable setting: a simplified one-block transformer trained by supervised next-token prediction on state sequences generated by composing permutations. The architecture separates fixed-lag action retrieval, learned by RoPE attention, from a specialized MLP logic module that applies the retrieved permutation to the current state. Using a statistical-physics mean-field description, we derive dynamics for three order parameters measuring attention retrieval, teacher-matrix alignment, and off-target logic overlap. These equations quantitatively match simulations for the order parameters and, combined with a logit-distribution approximation, qualitatively predict the sharp transition in final rollout accuracy. The analysis reveals staged learning: the logic module first learns a mixed heuristic; attention then locks onto the relevant action, enabling efficient MLP alignment. Together, these results provide a controlled mechanistic account of how attention-based retrieval and MLP-based logic co-develop during chain-of-thought state tracking.

• The paper derives mean-field ODEs for attention retrieval, teacher alignment, and off-target overlap, showing that simulations across 100 seeds follow a consistent two-stage learning process.
• The solvable transformer first develops a mixed MLP-based logic heuristic, then attention learns to retrieve the correct fixed-lag action, enabling stronger alignment with the teacher’s permutation matrices.
• The paper explains sharp final-answer emergence as a compositional effect: modest gains in per-step accuracy compound across sequence length, while the specialized architecture limits direct transfer to standard language models.

Overview and motivation

This paper analyzes how a transformer acquires chain-of-thought state tracking by studying a deliberately simplified, analytically solvable architecture trained on a permutation-composition task. The central question is not whether transformers can solve such tasks—prior work has established that they can—but how the two core submodules of a transformer block, attention and the MLP, co-evolve during gradient-based training (2606.18164). The authors construct a student-teacher setup in which a one-block transformer must reproduce state sequences generated by sequentially composing permutations, with an exact teacher implementing the correct dynamics.

The design isolates functional roles: a single RoPE attention head learns fixed-lag positional retrieval of the relevant action token, while a specialized MLP-style "logic module" applies the retrieved permutation to the current one-hot state. Fixed unit embeddings, fixed values (one-hot for actions, zero for states), and a constrained GLU construction make the retrieval–logic coupling tractable while preserving the qualitative division of labor observed empirically in full transformers, where attention handles contextual mapping and feed-forward layers handle value mapping and knowledge storage (2606.18164).

Task and architecture

The task generates sequences of length $2L+1$: $L$ i.i.d.-drawn action indices $g_1,\dots,g_L$ from a vocabulary of $d_g$ permutations of $[N]$, followed by an initial state $s_0$ and $L$ recursively computed states $s_t = \sigma_{g_t}(s_{t-1})$. Under teacher forcing, the model predicts each state token; at inference it must autoregressively generate the entire suffix from the prefix $(g_1,\dots,g_L,s_0)$, so the final answer requires error-free execution of all intermediate steps.

The key architectural simplification is that the required retrieval lag is constant: to predict state $s_j$ at position $L$0, the model needs action $L$1 at relative offset $L$2. RoPE makes the attention logits depend only on this relative offset, so a single head can learn the retrieval. The logic module computes logits as $L$3 where $L$4 is the attention-output-weighted sum of learnable matrices. The authors show this module is exactly representable as a GLU MLP with fixed up-projection and gate projections and learnable down-projection blocks, establishing that the specialization is a constrained instance of a standard transformer component rather than an ad hoc device.

Training uses cross-entropy on the suffix tokens only, with SGD (learning rate 0.5), $L$5, $L$6, $L$7, and results averaged over 100 seeds.

Order parameters and mean-field dynamics

The analysis reduces training dynamics to three macroscopic order parameters:

• Attention retrieval $L$8: average attention mass on the correct action position across prediction steps ($L$9 is perfect retrieval, $g_1,\dots,g_L$0 is uniform).
• Teacher alignment $g_1,\dots,g_L$1: normalized trace overlap between each learned matrix $g_1,\dots,g_L$2 and its matching teacher permutation matrix.
• Off-target overlap $g_1,\dots,g_L$3: overlap between learned matrices and non-matching teacher permutation matrices.

Using a mean-field ansatz—replacing seed-specific logits by ensemble expectations over uniformly drawn actions and initial states—the loss closes exactly in $g_1,\dots,g_L$4 up to an initialization-dependent constant. The mean correct logit takes the form

$g_1,\dots,g_L$5

and the mean incorrect logit follows from conservation of the total logit sum under cross-entropy without weight decay. Gradient-flow updates then yield closed ODEs for the three order parameters in the thermodynamic limit ($g_1,\dots,g_L$6 at fixed ratios). The attention equation requires a closure for the gradient norm factor, derived as $g_1,\dots,g_L$7 with $g_1,\dots,g_L$8 determined by initialization variance and RoPE frequencies; notably, the diagonal approximation to the attention matrix overestimates the slope mid-training because off-diagonal entries behave unevenly, so the authors calibrate the prefactor at initialization instead.

The resulting ODE solutions match simulations quantitatively for all three order parameters over 100 seeds. The dynamics exhibit a clear staged structure:

Phase	Time	Behavior
Early ($g_1,\dots,g_L$9)	$d_g$0 plateau	Logic module learns a mixed heuristic; both $d_g$1 and $d_g$2 grow linearly
Transition	onset set by $d_g$3	$d_g$4 rises rapidly toward saturation
Late	$d_g$5	Matrices align strongly with teacher; $d_g$6 saturates

In the early regime the ODEs are solvable exactly: $d_g$7 and $d_g$8 grow linearly with slopes proportional to $d_g$9 and $[N]$0 respectively, while $[N]$1 stays flat. Near the transition, combining the linearized solution $[N]$2 with the attention equation gives Gaussian-like growth $[N]$3.

The mechanistic interpretation is a two-stage division of labor: because attention initially weights all actions nearly uniformly, the logic module can only learn a mixture of actions—a heuristic that raises both correct and off-target overlaps. Once attention begins discriminating the correct action, efficient teacher alignment becomes possible, which in turn reinforces correct retrieval. This MLP-first ordering is consistent with gradient-flow analyses of word co-occurrence recognition [yang-2024-words], condensation-to-rank-collapse analyses [chen-2026-condensation], induction-head formation studies [wang-2025-rich], and empirical findings that feed-forward layers learn distributional associations before attention mediates in-context reasoning [chen-2025-distributional]. The present work extends these observations to a setting where the staged structure is derived rather than inferred, and tied specifically to selecting the correct state-update rule.

Sharp emergence of final rollout accuracy

Final rollout accuracy—the probability that the last generated token $[N]$4 equals the true $[N]$5—exhibits a sharp transition from the chance level $[N]$6 to near unity. The paper models this via a recursion on per-step correctness, assuming (i) approximately Gaussian logit distributions with means obtained from the order-parameter solutions and variances held at their initialization values, and (ii) symmetric incorrect logits, giving per-step accuracy $[N]$7 as a one-dimensional Gaussian integral involving the standard normal CDF. Under these assumptions,

$[N]$8

This formula explains the emergent jump directly: next-token accuracy enters to the $[N]$9th power, so a gradual improvement in single-step accuracy compounds into an abrupt rise in end-to-end rollout success. The theoretical curve reproduces the empirical transition qualitatively but peaks slightly early, because empirical logit variances grow during training whereas the derivation holds them fixed; using empirical variances restores close agreement. Two approximations are acknowledged explicitly: training induces correlations between correct and incorrect logits that the independence assumption neglects, and logit distributions deviate from Gaussianity later in training.

An implication worth noting is that the sharpness of the transition is a compositional property of the task itself—not a signature of any sudden internal reorganization beyond the smooth order-parameter trajectories already described. The apparent "emergence" in rollout accuracy is thus attributable to multiplicative error accumulation over $s_0$0 steps.

Limitations and open questions

The paper is explicit about scope. The architecture is deliberately constrained—fixed embeddings, fixed values, a single head, a specialized logic module—and therefore does not constitute a standard transformer block, let alone a full-scale LLM; the authors state plainly that the goal is a controlled mechanistic model of algorithmic chain-of-thought, not a model of natural-language reasoning. The analytical results rest on a mean-field description whose approximations may behave differently in other architectures or tasks. Teacher forcing supplies the current state locally, so no second transport head is needed; whether the staged dynamics persist when the state itself must be retrieved is unaddressed. The robustness of the staged learning under reinforcement learning after pretraining, and the extent to which the picture transfers to standard architectures, are left as open questions. Hyperparameter robustness is supported empirically (learning succeeds for $s_0$1, $s_0$2, $s_0$3), though fewer RoPE frequencies limit learnability.

Conclusion

This work provides a closed, quantitative account of learning dynamics for chain-of-thought state tracking in a solvable transformer. Three order parameters—attention retrieval, teacher-matrix alignment, and off-target overlap—obey derived ODEs that match simulation, revealing a consistent sequence: the MLP-style logic module first acquires a mixed heuristic, attention then locks onto the relevant action at fixed lag, and efficient teacher alignment follows. A logit-distribution approximation further shows that the sharp rise in final rollout accuracy is a compositional consequence of per-step accuracy entering to the $s_0$4th power. The main limitation is architectural specialization, and the principal open question is whether the same staged attention–MLP co-development arises in less constrained models and under post-pretraining fine-tuning.