Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization
Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi, Yuxin Chen
Abstract
The ability to reason lies at the core of artificial intelligence (AI), and challenging problems usually call for deeper and longer reasoning to tackle. A crucial question about AI reasoning is whether models can extrapolate learned reasoning patterns to solve harder tasks with a longer chain-of-thought (CoT). In this work, we present a theoretical analysis of transformers learning on synthetic state-tracking tasks with gradient descent. Specifically: 1). We prove how the algebraic structure of state-tracking problems governs the length generalization of learned reasoning in transformers. In doing so, we formulate the attention concentration mechanism, linking the retrieval robustness of the attention layer to the task structure of longcontext state tracking problems. 2). Moreover, we prove that a transformer can provably self-improve via a recursive self-training scheme that progressively extends the range of solvable problem lengths. We show that the model can achieve abilities outside the coverage of the base model in recursive training, different from prior theoretical works on self-improvement. To our knowledge, we provide the first optimization guarantee that constant-depth transformers provably learn NC 1 -complete problems with CoT, significantly going beyond prior art confined in TC 0 , unless the widely held conjecture TC 0 ̸ = NC 1 fails. Finally, we present a broad set of experiments supporting our theoretical results, confirming the length generalization behaviors and the mechanism of attention concentration.
- The authors contributed equally, and the author order was determined by a coin flip. Given the density of results, we recommend the full version for readability: arXiv:2511.07378 1 For background on circuit complexity and a detailed review of expressiveness results, see Section 2 and [32,33].
39th Conference on Neural Information Processing Systems (NeurIPS 2025).
Answer: [NA] Justification: Error bars are not applicable in our setting; instead, we report results averaged over a sufficiently large number of independent runs to reduce variance and provide a stable estimate of performance.
Guidelines:
• The answer NA means that the paper does not include experiments.
• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).
• It should be clear whether the error bar is the standard deviation or the standard error of the mean.
• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 0d24cd2d-9c44-4d1b-9b1a-d36b49f0ef6cCited by top-tier papers7
- Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous ThoughtHanlin Zhu, Shibo Hao, Zhiting Hu, Jiantao Jiao et al.ICLR 2026 · 21 citations
- VideoNSA: Native Sparse Attention Scales Video UnderstandingEnxin Song, Wenhao Chai, Shusheng Yang, Ethan Armand et al.ICLR 2026 · 11 citations
- Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient DescentTong Yang, Yu Huang, Yingbin Liang, Yuejie ChiNeurIPS 2025 · 8 citations
- On the Emergence of Implicit Curriculum in RLVR Learning DynamicsYu Huang, Zixin Wen, Yuejie Chi, Yuting Wei et al.ICML 2026 · 6 citations
- Breaking the Reversal Curse in Autoregressive Language Models via Identity BridgeXutao Ma, Yixiao Huang, Hanlin Zhu, Somayeh SojoudiICML 2026 · 2 citations
Builds on52
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Large Language Models are Zero-Shot ReasonersTakeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo et al.NeurIPS 2022 · 8,168 citations
- Faith and Fate: Limits of Transformers on CompositionalityNouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li et al.NeurIPS 2023 · 728 citations
- Self-Consistency Improves Chain of Thought Reasoning in Language ModelsXuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V. Le et al.ICLR 2023 · 681 citations
- YaRN: Efficient Context Window Extension of Large Language ModelsBowen Peng, Jeffrey Quesnelle, Honglu Fan, Enrico ShippoleICLR 2024 · 508 citations
Related papers
- Generalizing Reasoning Problems to Longer LengthsChangnan Xiao, Bing LiuICLR 2025
- Chain-of-Thought Gradient DescentHong-Yu Chen, Venkat Ganti, Hude Liu, Jerry Yao-Chieh Hu et al.ICML 2026 · 36 citations
- A Little Depth Goes a Long Way: The Expressive Power of Log-Depth TransformersWilliam Merrill, Ashish SabharwalNeurIPS 2025 · 62 citations
- Transformers Provably Solve Parity Efficiently with Chain of ThoughtJuno Kim, Taiji SuzukiICLR 2025
- Can Transformers Reason Logically? A Study in SAT SolvingLeyan Pan, Vijay Ganesh, Jacob D. Abernethy, Chris Esposo et al.ICML 2025
