Transformers trained via Reinforcement Learning (RL) with outcome-based supervision can spontaneously develop the ability to generate intermediate reasoning steps (Chain-of-Thought). Yet the mechanism by which sparse rewards drive policy gradient to discover such systematic reasoning remains poorly understood. We address this by analyzing the policy gradient dynamics of single-layer linear Transformers on a synthetic graph traversal task that cannot be solved without Chain-of-Thought but admits a simple iterative solution. We prove that despite training solely on final-answer correctness, policy gradient drives the Transformer to converge to a structured, interpretable algorithm that iteratively traverses the graph vertex-by-vertex. We characterize the distributional properties required for this emergence, identifying the critical role of "simple examples": instances requiring fewer reasoning steps. When the training distribution places sufficient mass on these simpler examples, the Transformer learns a generalizable traversal strategy that extrapolates to longer chains; when this mass vanishes, policy gradient learning becomes infeasible. We corroborate our theoretical results through experiments on synthetic data and with real-world language models on mathematical reasoning tasks, validating that our theoretical findings carry over to practical settings.
Figures & tables
Figure 1: Illustration of the chain identification task. The input contains edges forming two disjoint chains Ca=(a1,…,an) and Cb=(b1,…,bn) , and a starting vertex vstart (here ak ). The task is to identify the terminal vertex of the chain containing vstart (here an ).
Real-World Pre-Training
Synthetic Pre-Training
Training Dist.
Test Acc. ( Hard- 15 )
Training Dist.
Test Acc. ( Hard- 20 )
Hard- 15
52.9%
Hard- 20
43.15%
Mixed- 15
96.2%
Mixed- 20
84.77%
Mixed- 10
91.6%
Mixed- 17
78.68%
Table 1: Test accuracy on Hard- n under various post-training distributions, in two real-world post-training experiments on a mathematical reasoning task ( Section 6.1 ). Left: real-world pre-training using Qwen2.5-3B-Instruct and n=15 . Right: synthetic pre-training using a model pre-trained on chain-of-thought traces of simpler instances [ 73 ] and n=20 . Full details in Section I.2 .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Chain size n
Test Acc.
Chain Trav.
4
100%
100%
8
100%
99.3%
12
99.3%
94.2%
Appendix
Table 2: In line with our theory ( Section 5 ), on the chain identification task ( Section 3 ), a Transformer trained with simple examples achieves high accuracy through the efficient chain traversal (reasoning) algorithm, thereby demonstrating an implicit bias towards efficiency ( Section 4 ). In this table, “Test Acc.” and “Chain Trav.” refer to the fraction of test examples that are solved accurately and using the chain traversal algorithm, respectively. For further details on this experiment see Sections I.1.1 and I.1.4 .
Figure 2: In line with our theory ( Section 5 ), on the chain identification task ( Section 3 ), the accuracy of a Transformer on complex test examples is higher when its training data includes simple examples, despite such simple examples inducing a distribution shift. This figure shows the accuracy on complex test examples throughout training, for different training data distributions: ones entailing only complex examples (small k ), and ones entailing a mix of complex and simple examples (large k ). For further details on this experiment see Sections I.1.3 and I.1.4 .
Hyperparameter
Value
RL algorithm
GRPO
KL coefficient β
10−4
Generations per prompt
8
Sampling temperature (training)
0.9
Max completion length (training)
1,024 tokens
Max prompt length (training)
512 tokens
Appendix
Table 3: Hyperparameters used for fine-tuning Qwen2.5-3B-Instruct and Llama-3.2-3B-Instruct on the equations-reasoning task.
Figure 4: Training curves in the real-world pre-training setting. In both panels, the model is evaluated on 15 -Hard while being post-trained either on 15 -Mixed (blue) or 15 -Hard (red).
Training Distribution
Test Acc. ( Hard- 15 )
Hard- 15
50.2%
Mixed- 15
99.3%
Mixed- 10
98.6%
Appendix
Table 4: Same setting as the real-world pre-training experiment in Table 1 , but using Llama-3.2-3B-Instruct as the real-world pre-trained base model. “Test Acc.” is the fraction of 15 -Hard test examples solved correctly.
Real-World Pre-Training
Synthetic Pre-Training
Training Dist. Complexity Levels
Test Acc. ( Hard- 15 )
Training Dist. Complexity Levels
Test Acc. ( Hard- 20 )
{5,10,15}
96.80%
{10,13,17,20}
80.20%
{13,14,15}
52.90%
{17,18,19,20}
48.22%
Appendix
Table 5: Test accuracy on Hard- n under post-training distributions with the same number of distinct complexity levels. Within each setting, the first row includes simple examples and the second contains only harder examples.
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 poly(n) samples and L=log2k 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.
Reinforcement learning (RL) has become a standard tool for post-training language models on reasoning tasks, where the policy is updated by reward feedback while exploring the space of responses. Despite its empirical success, theoretical understanding of RL post-training remains limited, in particular of why on-policy exploration combined with a neural reward model is effective. In this paper, we address this question by modeling the reward as a hierarchical function on the response space: the reward consists of infinitely many local components, each of which becomes relevant only after the preceding ones have been resolved. We show that a natural Transformer-based actor--critic algorithm, which alternates between sampling from the current KL-regularized policy, fitting a Transformer critic to the observed rewards, and updating the policy, achieves the minimax optimal rates in the query budget and in the regularization strength up to logarithmic factors, and is minimax optimal for a fixed number of prompts. In contrast, we prove that sampling from the fixed reference distribution, as in offline reward modeling, can limit regret decay to a logarithmic rate. These results show that on-policy exploration progressively zooms in on the region where the reward is concentrated, and quantify its benefit for RL post-training.
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.
Yuxuan Cai, Yuan Cao
School of Computing & Data Science, The University of Hong Kong