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

Chain-of-thought generation can turn a multi-step computation into a sequence of locally checkable state updates, but the training dynamics of such state tracking 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 token retrieval 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 correct-token retrieval in attention as well as on- and off-target logic alignment. These solutions quantitatively match simulations for the order parameters and, combined with a Gaussian approximation to the logit distribution, qualitatively predict the sharp transition in final rollout accuracy. The analysis reveals staged learning: the logic module first learns a mixed heuristic, after which attention locks onto the relevant tokens and efficient MLP alignment is enabled. The same mechanism emerges in standard one-block transformers and is missing when either subsystem is frozen at initialization. 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.