Research map / Large language models / Reasoning and chain of thought
Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that a one-layer NoPE transformer trained by gradient descent learns LEGO state-tracking CoT, generalizes to much longer chains for simply transitive actions, and a recursive self-training scheme extends solvable length for symmetry (NC^1-complete) actions.
Problem
Whether transformers trained with gradient descent can provably learn chain-of-thought state tracking beyond TC^0 (NC^1-complete symmetry-group tasks) and whether the learned reasoning length-generalizes to longer chains.
Approach
Analyzes a one-layer softmax-attention plus FFN transformer without positional encoding on LEGO state tracking; shows that the algebraic structure of the group action (simply transitive vs. symmetry) determines attention concentration and hence length generalization, and that recursive self-training on greedy self-labeled traces doubles solvable length per stage.
Key results
- For simply transitive actions, training on constant-length chains yields near-perfect accuracy up to length d^{c*} (Theorem: Acc_L ≥ 1 - 1/poly(d) for L ≤ O(d^{c*})).
- For symmetry actions, standard CoT training only generalizes to a constant factor of training length; recursive self-training doubles solvable length per stage, reaching length |X| after Θ(log d) stages.
- First optimization guarantee that a constant-depth transformer learns an NC^1-complete problem (S_n word problem) with linear-step CoT under GD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use simply transitive rather than symmetry group actions | One-layer NoPE transformer on LEGO · LEGO length generalization | accuracy | Helped |
| Use symmetry group actions rather than simply transitive actions | One-layer NoPE transformer on LEGO · LEGO length generalization | accuracy | Hurt |
| Add recursive self-training on the model's own CoT traces | Transformer on symmetry-action LEGO · LEGO reasoning length | accuracy | Helped |
| Train on longer self-labeled reasoning traces at successive stages | Transformer on symmetry-action LEGO · LEGO reasoning length | accuracy | Helped |
Limitations
- Analysis restricted to one-layer block with NoPE, block-sparse attention, and orthonormal embeddings.
- Embedding dimension not logarithmic in problem length; asymptotic regime d to infinity.
Experimental setup
Datasets: LEGO synthetic state-tracking data
Benchmarks: LEGO
Metrics: accuracy along chain
Baselines: standard CoT training
Builds on: LEGO Barrington's theorem tensor power method
Closest papers by meaning
- Transformers Provably Solve Parity Efficiently with Chain of Thought (ICLR 2025 · similarity 0.92)
- Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent (NeurIPS 2025 · similarity 0.92)
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization (NeurIPS 2024 · similarity 0.90)
- Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data? (ICLR 2025 · similarity 0.90)
- When can transformers reason with abstract symbols? (ICLR 2024 · similarity 0.90)
- A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.