TLDR: A new research paper theoretically proves how transformers learn Chain-of-Thought (CoT) reasoning and achieve length generalization. It shows that the algebraic structure of tasks dictates generalization ability, with ‘attention concentration’ being a key mechanism. For complex tasks, a ‘recursive self-training’ scheme is proven to extend reasoning length, enabling constant-depth transformers to solve problems beyond simpler computational limits. Empirical results on synthetic tasks support these theoretical findings, offering insights into model architecture and context handling.
Artificial intelligence is constantly striving to achieve more human-like reasoning capabilities. A significant advancement in this area has been Chain-of-Thought (CoT) reasoning, where large language models (LLMs) break down complex problems into intermediate steps before arriving at a final answer. This method has shown impressive results on challenging tasks, but a fundamental question remains: can these models extrapolate their learned reasoning patterns to solve even harder tasks that require longer chains of thought?
A recent research paper, titled ‘Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization,’ by Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi, and Yuxin Chen, delves into this crucial question. The researchers provide a theoretical analysis of how transformers, a core architecture in LLMs, learn CoT reasoning and whether this ability can generalize to longer problem sequences.
The study addresses two main research questions:
- Can transformers, trained with a common optimization method called gradient descent, learn CoT reasoning for problems that inherently require sequential thinking, going beyond simpler computational tasks?
- Can this learned reasoning ability extend to problems that demand longer chains of thought than those seen during training?
To tackle these questions, the researchers analyzed a simplified transformer model, consisting of a single attention layer and a feed-forward network, trained without positional encoding (NoPE). They tested this model on synthetic ‘state-tracking’ tasks, specifically using a framework called LEGO (Learning Equality and Group Operations). These tasks mimic core LLM skills like tracking entities and updating game states, providing a controlled environment for theoretical analysis.
The findings reveal fascinating insights into how transformers learn and generalize. The researchers mathematically proved that the algebraic structure of the state-tracking problems significantly influences how well the learned CoT can extrapolate. For tasks with ‘simply transitive’ group actions (where there’s a unique way to get from one state to another), transformers trained on short reasoning chains could generalize to much longer problems. This impressive generalization is attributed to a mechanism called ‘attention concentration,’ where the attention layer effectively focuses on relevant information even in longer contexts.
However, for tasks with ‘symmetry group’ actions, which are inherently more complex and require more nuanced reasoning, the models showed limited length generalization. This means they could only extrapolate to problems that were a constant factor longer than their training data. The challenge here lies in the presence of ‘distractor’ clauses, which dilute the model’s attention and make robust retrieval harder.
To overcome this limitation, the researchers introduced a ‘recursive self-training’ scheme. This method involves training the model on its own generated CoT traces, progressively extending the range of solvable problem lengths. They proved that this self-training approach could bootstrap the model’s reasoning capabilities, allowing it to solve problems up to the maximal possible length in their setting. This offers a theoretical guarantee for the self-improvement observed in many advanced AI models.
Crucially, this work provides the first optimization guarantee that constant-depth transformers can learn problems beyond a complexity class known as TC0, reaching NC1-complete problems with CoT. This is a significant theoretical leap, as NC1 problems are conjectured to require inherently serial computation, highlighting the power of CoT reasoning in transformers. You can read the full paper for more technical details at arXiv:2511.07378.
The theoretical predictions were strongly supported by a series of experiments on the synthetic LEGO tasks. These experiments confirmed the distinct length generalization behaviors for different group actions, the effectiveness of recursive self-training in extending reasoning length, and the predicted patterns of attention concentration within the models.
Also Read:
- Enhancing Language Models for Deeper Reasoning Through Recurrence
- CoT-X: Bridging Advanced AI Reasoning with Practical Efficiency
This research offers valuable insights into the fundamental mechanisms of CoT reasoning in transformers, shedding light on both their capabilities and limitations. It also provides architectural guidance, suggesting that models without positional encoding (NoPE) can be advantageous for length generalization by relying on content-based retrieval rather than absolute position. The findings also contribute to understanding ‘context rot,’ a phenomenon where LLM performance degrades with longer inputs, by characterizing how attention dilution affects retrieval in extended contexts.
