Research map / Large language models / Reasoning and chain of thought
Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought
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.
Analyzes gradient-flow training dynamics of a one-layer linear transformer with chain-of-thought on an in-context weight prediction linear regression task, showing CoT learns multi-step GD with near-exact recovery while the no-CoT model is limited to one GD step.
Problem
Understanding how transformers are trained on chain-of-thought data to perform multi-step reasoning, and whether CoT training provably outperforms direct-answer training, in the setting of in-context linear regression.
Approach
Introduces in-context weight prediction (predicting the ground-truth weight vector rather than a query label) and a CoT objective summing next-step prediction losses over GD intermediate states; proves no-CoT one-layer linear attention is limited to one GD step with error Theta(d^2/n), whereas gradient flow on the CoT loss converges to a multi-step GD solution with polynomially small error.
Key results
- A one-layer linear transformer without CoT attains at best Theta(d^2/n) evaluation loss, equivalent to one GD step, which is Theta~(d) for n=Theta~(d).
- With Theta(log d) CoT steps, the constructed one-layer transformer achieves evaluation loss O(1/poly(d)), and gradient flow on the CoT loss converges to this solution in time O(log d + log(1/eps)).
- Experiments with d=10, n=20 show trained weights converge to the predicted block structure, and evaluation loss with CoT approaches near zero at k=10 while the no-CoT optimum stays much larger.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add autoregressive CoT with Θ(log d) gradient steps | One-layer linear transformer · In-context weight prediction | evaluation loss: from Θ(d/polylog d) to O(1/poly d) | Helped |
| Use one-layer transformer without CoT instead of CoT with Θ(log d) steps | One-layer linear transformer · In-context weight prediction | evaluation loss: Θ(d^2/n) -> O(1/poly d) | Hurt |
| Increase CoT inference steps from 10 to 40 | Trained one-layer transformer · Synthetic in-context weight prediction | evaluation loss | Helped |
| Use L = Θ(log d) loops | Looped transformer · Linear regression ICL | loss: loss lower bound \frac{d^{5/2}L\cdot 4^L}{\sqrt{n}} -> ξ | Helped |
Surprising findings
- The optimal no-CoT one-layer transformer is still a single GD step, so CoT is the only route to near-exact recovery in this setting.
- The looped transformer analysis in the paper's appendix claims a bound that is asymptotically better than the earlier multi-step looped result for all n, d, L.
Limitations
- Analysis uses a one-layer linear self-attention without softmax and a special reparameterization and initialization rather than standard softmax attention with Q,K,V.
- Population loss is used instead of finite-sample training, so no sample complexity guarantee is given.
- Focuses on iterative tasks (multi-step GD); compositional reasoning is not addressed.
Experimental setup
Datasets: synthetic Gaussian linear regression data
Benchmarks: synthetic in-context weight prediction (d=10, n=20)
Metrics: evaluation loss weight matrix structure
Baselines: one-layer linear transformer without CoT looped transformer multi-step GD
Builds on: linear self-attention ICL as gradient descent looped transformer multi-step GD (Gatmiry et al.) Zhang et al. 2024 training dynamics of linear transformers
Closest papers by meaning
- Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.91)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.90)
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought (NeurIPS 2026 · similarity 0.90)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.89)
- Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.88)
- Analyzing the Power of Chain of Thought through Memorization Capabilities (NeurIPS 2025 · similarity 0.88)
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.