Transformers with CoT Discover Heavy-Ball Acceleration Through Training
This paper asks whether transformers that generate intermediate reasoning steps can discover new algorithms beyond the teacher they are trained to imitate. It designs a two-stage training protocol for in-context linear regression—first supervised by gradient descent, then by ground-truth weights—and proves that a one-layer linear self-attention transformer learns the Heavy-Ball method. The result matters because it gives a rigorous case where chain-of-thought reasoning internalizes acceleration rather than simply memorizing the supervised update.
The paper delivers a crisp theoretical result: a linear self-attention transformer trained with CoT discovers Heavy-Ball acceleration after a two-stage curriculum, even though the Heavy-Ball algorithm is never provided as supervision. The authors back this claim with a convergence proof for both stages, an explicit characterization of the emergent iterates, and small-scale experiments that align closely with the theory. The analysis is technically intricate and leans on strong assumptions, but the core phenomenon is genuine and well-supported.
The strongest aspect is the careful proof that the learned algorithm is not just optimized, but specifically approximates the Heavy-Ball method. The authors introduce auxiliary parameters and recursively defined matrix polynomials to track the autoregressive CoT dynamics without teacher forcing, which is a meaningful technical advance. Empirically, the transformer trained with the two-stage procedure closely matches the Heavy-Ball method, while ablations recover plain gradient descent, confirming the theoretical prediction that both components are necessary.
The theoretical guarantees rely on a stylized setting that is distant from practice. Assumption 3.1 requires freezing specific blocks and a structured initialization, while Assumption 3.2 forces $n$ and $d$ to be proportional and demands the super-polynomial scaling $n, d = \Omega(\exp(\log^4(1/\epsilon)))$, which the authors acknowledge as relatively strong.
Moreover, the training is analyzed as continuous population gradient flow rather than discrete stochastic optimization. The experiments explicitly depart from theory by using unfrozen Gaussian initialization, removing teacher forcing in Stage 1, and choosing dimensions where $d/n \not< 0.1$; while the authors argue this shows robustness, it also means the experiments are not a direct validation of the theorem.
The evidence supports the narrow claim: the paper shows convergence to a specific polynomial update, connects it to the optimal polynomial associated with Heavy-Ball under the MP law, and verifies the trajectory shape empirically. The comparison to prior work is fair and well-motivated—the authors clearly explain why prior settings with concentrated sample covariances cannot distinguish Heavy-Ball from gradient descent. However, the paper does not rule out that other accelerated schemes could emerge or provide a general principle for which algorithms are discoverable; the result is tightly coupled to the MP spectral structure and the two-stage curriculum.
Reproducing the exact theorem would require implementing constrained gradient flow with frozen weights and exact population expectations over Gaussian data, which is feasible but delicate. The experiments are more standard—online Adam on small synthetic datasets—but intentionally deviate from the theoretical protocol to test robustness. No code or data repository is mentioned in the provided text, so independent reproduction of the experimental curves would require reimplementation from the stated hyperparameters.
Modern transformers often solve complex tasks by generating and refining intermediate reasoning steps. To develop a theoretical understanding of chain-of-thought (CoT) reasoning, a notable recent study by Huang et al. (2025) viewed transformers as algorithm executors, and showed that when transformers' CoT steps are supervised by a gradient-descent teacher, they can learn to implement gradient descent for in-context weight prediction in linear regression. In this paper, we study the same in-context weight prediction task but ask whether CoT reasoning can go beyond merely imitating the teacher algorithm. We show that, after initial supervision from a gradient-descent teacher, further trajectory-level supervision using the true regression weights leads the model to implement the Heavy-Ball method, thereby achieving acceleration over the original teacher. Notably, this accelerated algorithm emerges without supervision from any Heavy-Ball teacher. Our results provide a concrete theoretical example in which transformer reasoning can discover new algorithms through training, rather than merely imitate the algorithm used for supervision.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.