Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization

AI summary

Overview

• Research area: Machine learning theory — specifically, the theory of transformer expressiveness, learnability, and chain-of-thought (CoT) reasoning, connecting training dynamics with circuit complexity.
• Technical level: Advanced. The paper is a theoretical analysis (proofs, theorems, asymptotic assumptions) with supporting synthetic experiments. Readers need familiarity with transformers, gradient descent, and basic circuit complexity (TC⁰, NC¹, P/poly).
• Scope: The paper proves, for a one-layer transformer with no positional encoding trained by gradient descent on synthetic LEGO state-tracking tasks, how the algebraic structure of the task (simply transitive vs. symmetry group actions) governs the length generalization of learned CoT, and that a recursive self-training scheme can extend the solvable reasoning length.

What This Paper Is About

Prior work showed that constant-depth transformers without CoT are limited to the circuit class TC⁰, while with O(L) CoT steps they can express NC¹. But it was unknown whether gradient-based training actually learns such reasoning, and whether the learned reasoning generalizes to longer problems than seen in training. This paper analyzes a minimal one-layer transformer trained with gradient descent on synthetic state-tracking (LEGO) tasks, and proves that it learns CoT reasoning beyond TC⁰, that the degree of length generalization depends on the group action structure, and that a recursive self-training curriculum can bootstrap the solvable problem length.

Key Contributions

1. Provable CoT learning for state tracking. The authors prove that one-layer transformers trained via gradient descent can provably learn to solve state-tracking problems for both simply transitive and symmetry group actions using CoT reasoning (Theorem 1.1, informal). Because state tracking for symmetry groups is NC¹-complete (Barrington), this is stated as the first optimization-based training guarantee for problems beyond TC⁰, matching the linear-CoT expressiveness result of [LLZM24].

2. A theory of length generalization governed by algebraic structure. The paper identifies an attention concentration mechanism explaining how far reasoning generalizes: simply transitive actions yield sharp attention concentration and generalization to problems of length d^{c*} (with 0 < c* < 1 a constant), whereas symmetry actions introduce distractors that dilute attention and limit generalization to only a constant factor of the training length (Theorem 1.2, informal).

3. Recursive self-training extends solvable length. For settings where direct length generalization is limited (symmetry actions), the authors prove that a recursive self-training curriculum, where the model trains on its own self-labeled CoT traces, bootstrap the solvable length: at each stage k > 1, learning on traces of length 2^k generalizes to length 2^{k+1}, so after Θ(log d) stages the model solves problems of length d, the maximal length in their setting (Theorem 1.3, informal).

4. Empirical validation. Experiments on synthetic LEGO tasks corroborate the predicted separation between the two algebraic structures, the effectiveness of recursive self-training, and the attention-concentration mechanism.

Main Findings

• Length generalization depends on group action type. For the simply transitive (cyclic) group, training on short problems yields nearly perfect accuracy at much longer lengths; for the symmetry group, training yields only constant-factor length generalization (Figure 1(a), over cyclic C₆ vs. symmetry S₅ tasks).

• Attention concentration is the proposed mechanism. At convergence, attention heatmaps for a LEGO task of length L = 5 show two distinct diagonal bands of concentrated mass: when the query is the answer clause Z_{ans,ℓ}, attention concentrates on the answer clause Z_{ans,ℓ} and the predicate clause Z_{pred,ℓ+1} (Figure 2). Simply transitive actions enable sharp concentration; symmetry actions introduce distractors that dilute focus.

• Recursive self-training shifts the length–accuracy curve. A double-and-self-labeled curriculum on the symmetry task consistently shifts the length–accuracy curve rightward, and after several stages performance matches the simply transitive case (Figure 1(b)).

• Beyond TC⁰ under a standard conjecture. Since state tracking with symmetry groups is NC¹-complete, the learning guarantees imply constant-depth transformers provably learn NC¹-complete problems with CoT, outside TC⁰ unless the widely held conjecture TC⁰ ≠ NC¹ fails.

• Contrast with prior learnability results. Earlier optimization analyses covered only TC⁰ tasks (e.g., parity and linear regression) that constant-depth transformers can already solve without CoT; prior length-generalization theory provided existence or statistical guarantees and non-gradient-based learnability, whereas this work shows gradient-based optimization actually finds length-generalizing solutions.

Methodology in Plain English

The authors build a deliberately minimal but analyzable setup:

• Model: A single decoder transformer block consisting of a softmax attention layer with no positional encoding (NoPE) followed by a position-wise feed-forward network, so that F = FFN ∘ Attention. The feed-forward network uses a smoothed, continuously differentiable variant of ReLU ("sReLU") with design parameters q (a large even integer) and ϱ = Θ(1/polylog(d)).

• Task: Synthetic LEGO state-tracking sentences. Each clause is encoded as a fixed-length 5-token tuple; predicate clauses specify an action linking two variables and answer clauses assign values to variables. Solving a sentence means composing actions along a path and propagating states, which maps naturally onto step-by-step CoT generation.

• Data: A LEGO distribution D^L where variables are sampled uniformly at random without replacement, the initial value is uniform, actions are sampled uniformly with replacement, and intermediate values are computed recursively.

• Theory: The authors analyze the gradient-descent training dynamics and track the evolution of attention patterns, deriving how the stabilizer size of the group action (unique for simply transitive actions, multiple for symmetry actions) controls attention concentration at convergence, and thereby the degree of length generalization.

• Experiments: They run experiments on their theoretical setup using cyclic (C₆) and symmetry (S₅) tasks, measuring accuracy as a function of problem length, running the recursive self-training curriculum, and visualizing attention heatmaps at convergence.

Assumptions include an asymptotic regime where the vocabulary size d and number of variables |X| tend to infinity, with |G| ≤ log^{C₀} d for some constant C₀ ∈ [1, 100).

Why This Matters

• Closes a gap between expressiveness and learnability. Prior theory showed transformers can express sequential algorithms with CoT; this paper argues they can actually learn them by gradient descent, which is a stronger and more practically relevant claim.

• Explains when and why length generalization works. The attention-concentration account gives a structural criterion — the algebraic structure of the task — for predicting whether a model trained on short reasoning chains will extrapolate to longer ones.

• Provides a theoretical basis for self-improvement. The recursive self-training guarantee is described as the first optimization guarantee of self-improvement for transformer networks, offering formal support for curricula that train models on their own longer traces.

• Real-world applications (as the paper frames the underlying state-tracking skills):

  • Narrative entity tracking in long documents.
  • Chess move analysis and game-state updates.
  • Code execution and evaluation.
  • Multi-step reasoning tasks where models must track how state evolves across many steps.

• Industry relevance. Frontier reasoning models (such as OpenAI-o1 and DeepSeek-R1, cited in the paper) produce long CoT traces at inference time, often elicited via reinforcement learning or supervised fine-tuning distilled from longer chains. Understanding when longer reasoning generalizes — and when it degrades due to context rot — matters directly for training and inference strategies in deployed reasoning systems.

Future Directions

• Relaxing the architectural and data assumptions. The analysis uses a one-layer NoPE transformer with orthonormal embeddings and specific scaling assumptions; whether the guarantees extend to deeper models, positional encodings, or attention variants remains open.

• Bridging to realistic tasks and distributions. The results are for synthetic LEGO state-tracking; extending the theory to richer reasoning tasks and natural data distributions is an open question.

• Understanding self-training more broadly. The recursive self-training guarantee is specific to the doubling curriculum (2^k to 2^{k+1}); whether other self-improvement schemes, or schemes with noisy or imperfect self-labels, admit similar guarantees is not established here.

• Practical mitigation of context rot and limited extrapolation. The paper notes that empirical evidence on length generalization is mixed and that performance can degrade as context grows; translating the attention-concentration diagnosis into concrete training or architectural interventions is a natural next step.

Target Audience

This paper is most valuable to machine learning theorists and graduate students working on transformer expressiveness, optimization guarantees, and learning theory; to researchers studying chain-of-thought reasoning and length generalization; and to practitioners in reasoning-model development who want a formal lens on when long-chain reasoning can be expected to extrapolate and how self-training curricula might extend it. Readers without a background in circuit complexity or optimization analysis will find the full proofs challenging, though the informal theorem statements and empirical figures are accessible.

Authors’ abstract

The ability to reason lies at the core of artificial intelligence (AI), and challenging problems usually call for deeper and longer reasoning to tackle. A crucial question about AI reasoning is whether models can extrapolate learned reasoning patterns to solve harder tasks with longer chain-of-thought (CoT). In this work, we present a theoretical analysis of transformers learning on synthetic state-tracking tasks with gradient descent. We mathematically prove how the algebraic structure of state-tracking problems governs the degree of extrapolation of the learned CoT. Specifically, our theory characterizes the length generalization of transformers through the mechanism of attention concentration, linking the retrieval robustness of the attention layer to the state-tracking task structure of long-context reasoning. Moreover, for transformers with limited reasoning length, we prove that a recursive self-training scheme can progressively extend the range of solvable problem lengths. To our knowledge, we provide the first optimization guarantee that constant-depth transformers provably learn $\mathsf{NC}^1$-complete problems with CoT, significantly going beyond prior art confined in $\mathsf{TC}^0$, unless the widely held conjecture $\mathsf{TC}^0 \neq \mathsf{NC}^1$ fails. Finally, we present a broad set of experiments supporting our theoretical results, confirming the length generalization behaviors and the mechanism of attention concentration.