mlsanityData to 14 Oct 2026

Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought

arXiv 2502.21212 ↗First posted 28 Feb 2025cs.LG cs.AIcited by 32

On the map: LLM reasoning → Prompting and chain-of-thought reasoning

Abstract (arXiv)

Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs), particularly in arithmetic and reasoning tasks, by instructing the model to produce intermediate reasoning steps. Despite the remarkable empirical success of CoT and its theoretical advantages in enhancing expressivity, the mechanisms underlying CoT training remain largely unexplored. In this paper, we study the training dynamics of transformers over a CoT objective on an in-context weight prediction task for linear regression. We prove that while a one-layer linear transformer without CoT can only implement a single step of gradient descent (GD) and fails to recover the ground-truth weight vector, a transformer with CoT prompting can learn to perform multi-step GD autoregressively, achieving near-exact recovery. Furthermore, we show that the trained transformer effectively generalizes on the unseen data. With our technique, we also show that looped transformers significantly improve final performance compared to transformers without looping in the in-context learning of linear regression. Empirically, we demonstrate that CoT prompting yields substantial performance improvements.

Key facts. Cited by 32 papers in the map (data through 14 Oct 2026). Placed in Prompting and chain-of-thought reasoning, one of 17 directions in LLM reasoning. Builds on 38 of its 59 references that are in the map. First posted on arXiv 28 Feb 2025.

Builds on

38 of 59 references.

Closest ancestors · 10

Same direction, newest first

  1. 2024Unveiling the Statistical Foundations of Chain-of-Thought Prompting Methods25 Aug
  2. 2024Iteration Head: A Mechanistic Study of Chain-of-Thought4 Jun
  3. 2024Chain of Thought Empowers Transformers to Solve Inherently Serial Problems20 Feb
  4. 2023Why Can Large Language Models Generate Correct Chain-of-Thoughts?20 Oct
  5. 2023The Expressive Power of Transformers with Chain of Thought11 Oct
  6. 2023Analyzing Chain-of-Thought Prompting in Large Language Models via Gradient-based Feature Attributions25 Jul
  7. 2022Towards Understanding Chain-of-Thought Prompting: An Empirical Study of What Matters20 Dec
  8. 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them17 Oct
  9. 2022Least-to-Most Prompting Enables Complex Reasoning in Large Language Models21 May
  10. 2022Self-Consistency Improves Chain of Thought Reasoning in Language Models21 Mar

Other references · 28

  1. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  2. 2023GPT-4 Technical Reportcited by 18,396
  3. 2021Show Your Work: Scratchpads for Intermediate Computation with Language Modelscited by 886
  4. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 684
  5. 2022Language Models are Multilingual Chain-of-Thought Reasonerscited by 530
  6. 2022Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thoughtcited by 337
  7. 2023Trained Transformers Learn Linear Models In-Contextcited by 337
  8. 2022Selection-Inference: Exploiting Large Language Models for Interpretable Logical Reasoningcited by 334
  9. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 298
  10. 2022Transformers Learn Shortcuts to Automatacited by 268
  11. 2023Looped Transformers as Programmable Computerscited by 237
  12. 2023REFINER: Reasoning Feedback on Intermediate Representationscited by 205
16 more
  1. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 170
  2. 2024The Evolution of Statistical Induction Heads: In-Context Learning Markov Chainscited by 118
  3. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 118
  4. 2023In-Context Convergence of Transformerscited by 115
  5. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 102
  6. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 98
  7. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  8. 2023Transformers as Support Vector Machinescited by 74
  9. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 74
  10. 2023JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attentioncited by 52
  11. 2023Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Modelscited by 50
  12. 2024Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformerscited by 50
  13. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 49
  14. 2024Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscapecited by 46
  15. 2024Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chainscited by 40
  16. 2023CausalLM is not optimal for in-context learningcited by 26

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 18; 2026: 14.

32 papers.

Sort:

By area:

  1. 2025Scaling Latent Reasoning via Looped Language Modelscited by 165
  2. 2025Learning without training: The implicit dynamics of in-context learningcited by 43
  3. 2025Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalizationcited by 14
  4. 2026Solve the Loop: Attractor Models for Language and Reasoningcited by 13
  5. 2025Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoningcited by 9
  6. 2025Beyond Black-Box Obfuscation: Mechanistic Analysis and Defense of White-Box Monitorscited by 6
  7. 2026Outcome-Based RL Provably Leads Transformers to Reason, but Only With the Right Datacited by 6
  8. 2025Too Easily Fooled? Prompt Injection Breaks LLMs on Frustratingly Simple Multiple-Choice Questionscited by 4
  9. 2025Understanding the Role of Training Data in Test-Time Scalingcited by 4
  10. 2025Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Trainingcited by 4
  11. 2025Rectifying LLM Thought from Lens of Optimizationcited by 4
  12. 2025Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regressioncited by 3
20 more
  1. 2025Provably Learning from Modern Language Models via Low Logit Rankcited by 3
  2. 2026Theoretical Perspectives on Data Quality and Synergistic Effects in Pre- and Post-Training Reasoning Modelscited by 3
  3. 2026Continuous Latent Contexts Enable Efficient Online Learning in Transformerscited by 3
  4. 2025Deciphering Trajectory-Aided LLM Reasoning: An Optimization Perspectivecited by 2
  5. 2025Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learningcited by 2
  6. 2025Bounds of Chain-of-Thought Robustness: Reasoning Steps, Embed Norms, and Beyondcited by 2
  7. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  8. 2026Transformers Provably Learn to Internalize Chain-of-Thoughtcited by 2
  9. 2025Multi-Layer Attention is the Amplifier of Demonstration Effectivenesscited by 1
  10. 2025Convergence dynamics of Agent-to-Agent Interactions with Misaligned objectivescited by 1
  11. 2026Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from k-Paritycited by 1
  12. 2025Chain-of-Thought Enhanced Shallow Transformers for Wireless Symbol Detectioncited by 0
  13. 2025Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functionscited by 0
  14. 2026Enhancing Sentiment Classification and Irony Detection in Large Language Models through Advanced Prompt Engineering Techniquescited by 0
  15. 2026Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMscited by 0
  16. 2026Beyond What Seems Necessary: Hidden Gains from Scaling Training-Time Reasoning Length under Outcome Supervisioncited by 0
  17. 2026Feature Resemblance: Towards a Theoretical Understanding of Analogical Reasoning in Transformerscited by 0
  18. 2026How Transformers Learn to Plan via Multi-Token Predictioncited by 0
  19. 2026The Power of Power Law: Asymmetry Enables Compositional Reasoningcited by 0
  20. 2026An Asymptotic Theory of Chain-of-Thought in In-Context Learningcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Enhancing Auto-regressive Chain-of-Thought through Loop-Aligned Reasoning
  2. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis
  3. 2026Adaptive Loops and Memory in Transformers: Think Harder or Know More?
  4. 2025The Kinetics of Reasoning: How Chain-of-Thought Shapes Learning in Transformers?
  5. 2023Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective
  6. 2026Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers
  7. 2025Mind the Gap: Bridging Thought Leap for Improved Chain-of-Thought Tuning
  8. 2026Barriers to Universal Reasoning With Transformers (And How to Overcome Them)