Transformers with CoT Discover Heavy-Ball Acceleration Through Training

Yuxuan Cai; Yuan Cao

cs.LGOct 22, 2026

Transformers with CoT Discover Heavy-Ball Acceleration Through Training

Authors: Yuxuan Cai, Yuan Cao

Organizations: School of Computing & Data Science, The University of Hong Kong

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.

Figures & tables

Figure 1: Comparison of gradient descent, Heavy-Ball, a transformer trained solely based on gradient-descent supervision, and a transformer trained first using the gradient-descent teacher and then using ground-truth regression weights.

Figure 2: In-context weight prediction performance of several models and baseline algorithms.

Figure 3: Examples of the heatmaps of learned weight matrices after the two-stage training for $d=16,n=32$ (left), together with the progress of second-stage-only training for $n=32$ under four different combinations of dimension $d$ and learning rate (right). For second-stage-only training, nonfinite gradient values are marked in red. When nonfinite gradients occur, we do not stop training immediately; instead, we allow the model to retry using a newly generated independent in-context dataset. Training is terminated only after nonfinite gradients are observed for 20 consecutive retries. Dashed lines indicate termination.

CoT Reasoning

Chain-of-thought (CoT) reasoning studies how language models use intermediate steps, visible or latent, to solve multi-step tasks.

See all 228 papers

Transformers Provably Learn to Internalize Chain-of-Thought
May 27, 2026Yixiao Huang, Hanlin Zhu, Zixuan Wang +4Transformer CoT Reasoning
UC Berkeley · Princeton University
Efficiently Representing Algorithms With Chain-of-Thought Transformers
Jun 18, 2026Yanhong Li, Anej Svete, Ashish Sabharwal +1Transformer CoT Reasoning
Allen Institute for AI · ETH Zürich
One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
Oct 12, 2026Lyumin Wu, Yuan CaoTransformer In-Context Learning
School of Computing & Data Science, The University of Hong Kong

May 27, 2026J/K move · Enter open · S save

Transformers Provably Learn to Internalize Chain-of-Thought

Yixiao Huang, Hanlin Zhu, Zixuan Wang, Jiantao Jiao+3

UC Berkeley · Princeton University

Transformer CoT Reasoning Latent CoT Reasoning

Open paper AskPDF arXiv

Chain-of-Thought (CoT) prompting substantially improves the sample efficiency of transformers, reducing the complexity of tasks like parity learning from exponential to polynomial in the input length. However, generating explicit reasoning steps at inference is computationally expensive. Implicit Chain-of-Thought (ICoT) has emerged as a promising empirical remedy that trains models to internalize intermediate steps within their hidden states, but its theoretical foundations remain poorly understood. We give the first theoretical analysis of ICoT, proving that an

L

-layer transformer trained under our proposed Log-ICoT curriculum learns

k

-parity with

\mathsf{poly}(n)

samples and

L = \log_2 k

training stages. This matches the sample efficiency of explicit CoT while eliminating its inference overhead, and extends prior one-layer parity guarantees to multi-layer architectures. Compared to standard ICoT, which removes thinking tokens one at a time, Log-ICoT removes them in geometric chunks, reducing the number of stages from linear in

k

to logarithmic. Experiments on multi-layer transformers confirm the theory and visualize how reasoning is progressively absorbed into deeper layers.

Abstract

CoT Reasoning

Explore similar work

Transformers Provably Learn to Internalize Chain-of-Thought

Efficiently Representing Algorithms With Chain-of-Thought Transformers

One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought