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

arXiv: https://arxiv.org/abs/2502.21212 · HTML: https://mlsanity.com/paper/2502.21212/

## 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.

## Placed in

- LLM reasoning / [Prompting and chain-of-thought reasoning](https://mlsanity.com/lr/prompting-and-chain-of-thought-reasoning/)

## Builds on (most cited first, 30 of 38)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Least-to-Most Prompting Enables Complex Reasoning in Large Language Models](https://mlsanity.com/paper/2205.10625/) (2022)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Language Models are Multilingual Chain-of-Thought Reasoners](https://mlsanity.com/paper/2210.03057/) (2022)
- [Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought](https://mlsanity.com/paper/2210.01240/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Selection-Inference: Exploiting Large Language Models for Interpretable Logical Reasoning](https://mlsanity.com/paper/2205.09712/) (2022)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Chain of Thought Empowers Transformers to Solve Inherently Serial Problems](https://mlsanity.com/paper/2402.12875/) (2024)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [Towards Understanding Chain-of-Thought Prompting: An Empirical Study of What Matters](https://mlsanity.com/paper/2212.10001/) (2022)
- [REFINER: Reasoning Feedback on Intermediate Representations](https://mlsanity.com/paper/2304.01904/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [The Expressive Power of Transformers with Chain of Thought](https://mlsanity.com/paper/2310.07923/) (2023)
- [Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models](https://mlsanity.com/paper/2310.14491/) (2023)
- [Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers](https://mlsanity.com/paper/2409.10559/) (2024)

## Built on by (most cited first, 30 of 32)

- [Scaling Latent Reasoning via Looped Language Models](https://mlsanity.com/paper/2510.25741/) (2025)
- [Learning without training: The implicit dynamics of in-context learning](https://mlsanity.com/paper/2507.16003/) (2025)
- [Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization](https://mlsanity.com/paper/2511.07378/) (2025)
- [Solve the Loop: Attractor Models for Language and Reasoning](https://mlsanity.com/paper/2605.12466/) (2026)
- [Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoning](https://mlsanity.com/paper/2506.04611/) (2025)
- [Beyond Black-Box Obfuscation: Mechanistic Analysis and Defense of White-Box Monitors](https://mlsanity.com/paper/2505.14300/) (2025)
- [Outcome-Based RL Provably Leads Transformers to Reason, but Only With the Right Data](https://mlsanity.com/paper/2601.15158/) (2026)
- [Too Easily Fooled? Prompt Injection Breaks LLMs on Frustratingly Simple Multiple-Choice Questions](https://mlsanity.com/paper/2508.13214/) (2025)
- [Understanding the Role of Training Data in Test-Time Scaling](https://mlsanity.com/paper/2510.03605/) (2025)
- [Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Training](https://mlsanity.com/paper/2511.07372/) (2025)
- [Rectifying LLM Thought from Lens of Optimization](https://mlsanity.com/paper/2512.01925/) (2025)
- [Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression](https://mlsanity.com/paper/2508.07571/) (2025)
- [Provably Learning from Modern Language Models via Low Logit Rank](https://mlsanity.com/paper/2512.09892/) (2025)
- [Theoretical Perspectives on Data Quality and Synergistic Effects in Pre- and Post-Training Reasoning Models](https://mlsanity.com/paper/2603.01293/) (2026)
- [Continuous Latent Contexts Enable Efficient Online Learning in Transformers](https://mlsanity.com/paper/2605.09867/) (2026)
- [Deciphering Trajectory-Aided LLM Reasoning: An Optimization Perspective](https://mlsanity.com/paper/2505.19815/) (2025)
- [Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2508.16027/) (2025)
- [Bounds of Chain-of-Thought Robustness: Reasoning Steps, Embed Norms, and Beyond](https://mlsanity.com/paper/2509.21284/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Transformers Provably Learn to Internalize Chain-of-Thought](https://mlsanity.com/paper/2605.28600/) (2026)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [Convergence dynamics of Agent-to-Agent Interactions with Misaligned objectives](https://mlsanity.com/paper/2511.08710/) (2025)
- [Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from k-Parity](https://mlsanity.com/paper/2601.22450/) (2026)
- [Chain-of-Thought Enhanced Shallow Transformers for Wireless Symbol Detection](https://mlsanity.com/paper/2506.21093/) (2025)
- [Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions](https://mlsanity.com/paper/2510.18638/) (2025)
- [Enhancing Sentiment Classification and Irony Detection in Large Language Models through Advanced Prompt Engineering Techniques](https://mlsanity.com/paper/2601.08302/) (2026)
- [Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs](https://mlsanity.com/paper/2602.00576/) (2026)
- [Beyond What Seems Necessary: Hidden Gains from Scaling Training-Time Reasoning Length under Outcome Supervision](https://mlsanity.com/paper/2602.00927/) (2026)
- [Feature Resemblance: Towards a Theoretical Understanding of Analogical Reasoning in Transformers](https://mlsanity.com/paper/2603.05143/) (2026)
- [How Transformers Learn to Plan via Multi-Token Prediction](https://mlsanity.com/paper/2604.11912/) (2026)

## 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.