Transformers with CoT Discover Heavy-Ball Acceleration Through Training
Abstract
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.
Related
- Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers
- Non-Acceleration of the Heavy-Ball Method in 1D
- Acceleration Survives Heavy Tails and Generalized Smoothness
- Acceleration of Diffusion Language Model through Discrete Average Generator
- Distributionally robust linear regression through the lens of adversarial training
BibTeX
@misc{cai2026transformers,
title = {{Transformers with CoT Discover Heavy-Ball Acceleration Through Training}},
author = {Yuxuan Cai and Yuan Cao},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.31089},
eprint = {2610.31089},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.31089},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Transformers with CoT Discover Heavy-Ball Acceleration Through Training
AU - Yuxuan Cai
AU - Yuan Cao
PY - 2026
DA - 2026/10/22
PB - arXiv
DO - 10.48550/arxiv.2610.31089
UR - https://arxiv.org/abs/2610.31089
AB - 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.
ER -