Chain-of-Thought Gradient Descent

We show that Chain-of-Thought (CoT) expands the expressiveness of Transformer in-context learning (ICL). Specifically, we show CoT enable efficient simulation of In-Context Gradient Descent (ICGD) for $N$-layer neural network. Different from CoT, a Transformer with fixed depth and hidden dimension has fixed ICL capacity in one forward pass. Simulating larger models or more optimization steps in-context requires deeper or wider Transformers. CoT removes this limitation by providing an expandable workspace via the sequence trajectory. This enables arbitrary-step and arbitrary-capacity ICGD within a constant-depth Transformer. Second, we provide a provable efficient guarantee unique to CoT through dynamical masking. The attention mechanism only process the relevant tokens for the current update step. This eliminates the redundant ``process everything'' cost of single-pass deep models. Specifically, we prove this CoT mechanism improves the computational cost of the prior best in-context result [Wu et al., ICML 2025] by $O(N)$. Numerical validations support our theory.

AI models often perform better when they write out intermediate reasoning steps before giving an answer, but it is not fully clear why these steps help. This paper gives a mathematical explanation: the written steps can act like extra workspace, allowing a fixed-size model to carry out learning procedures that would otherwise require a larger model. We also show how the model can save effort by looking only at the previous steps that matter for the next calculation. These results help explain why step-by-step reasoning can make AI models both more capable and more efficient.