Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (2511.07378v1)

Abstract: The ability to reason lies at the core of 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 longer chain-of-thought (CoT). In this work, we present a theoretical analysis of transformers learning on synthetic state-tracking tasks with gradient descent. We mathematically prove how the algebraic structure of state-tracking problems governs the degree of extrapolation of the learned CoT. Specifically, our theory characterizes the length generalization of transformers through the mechanism of attention concentration, linking the retrieval robustness of the attention layer to the state-tracking task structure of long-context reasoning. Moreover, for transformers with limited reasoning length, we prove that a recursive self-training scheme can progressively extend the range of solvable problem lengths. To our knowledge, we provide the first optimization guarantee that constant-depth transformers provably learn $\mathsf{NC}^1$-complete problems with CoT, significantly going beyond prior art confined in $\mathsf{TC}^0$, unless the widely held conjecture $\mathsf{TC}^0 \neq \mathsf{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 paper demonstrates that transformers trained with chain-of-thought can provably learn length generalization through attention concentration and FFN parameter shaping.
• It reveals that the algebraic structure of tasks, contrasting simply transitive and symmetry group actions, critically affects the model’s ability to generalize over increased sequence lengths.
• Empirical validations show that recursive self-training enables transformers to extend reasoning capabilities, confirming theoretical performance guarantees on synthetic state-tracking tasks.

Provable Transformer Learning and Length Generalization in Chain-of-Thought Reasoning

Introduction

This paper provides a rigorous optimization-theoretic analysis on the learning and generalization capabilities of transformers equipped with chain-of-thought (CoT) reasoning. Employing synthetic state-tracking tasks built on group-theoretic actions (especially the LEGO framework), the paper precisely characterizes how transformer architectures trained via gradient descent provably acquire length-generalizing reasoning skills and how these generalization properties depend on the algebraic structure (i.e., simply transitive versus symmetry group actions) of the underlying task. The theoretical contributions are complemented by empirical demonstrations, establishing concrete performance guarantees for both direct length generalization and recursive self-improvement via self-training.

Background: Circuit Complexity and Transformer Expressivity

A recurring theme in the analysis is the connection between reasoning depth and circuit complexity classes: $\mathsf{TC}^0$ (constant depth, threshold circuits), $\mathsf{NC}^1$ (log depth, bounded fan-in circuits), and beyond. Previous expressivity results established that constant-depth transformers, sans CoT, are fundamentally limited to $\mathsf{TC}^0$ [feng2023revealing, li2024chainthought]. By contrast, augmenting transformers with $O(L)$ intermediate CoT steps for input length $L$ allows simulation of sequential computation in $\mathsf{NC}^1$, tackling problems with inherently serial requirements (such as non-solvable group word problems), which are presumed to be outside $\mathsf{TC}^0$ unless the open conjecture $\mathsf{TC}^0 \neq \mathsf{NC}^1$ fails.

Problem Setting: State-Tracking via Group Actions

The studied reasoning tasks employ the LEGO synthetic language, which formalizes state-tracking over sequences of group actions. Given an initial state $y_0$ and a sequence of group elements $g_1, \dots, g_L$, the task is to compute $y_L = g_L \circ \cdots \circ g_1 (y_0)$, where the group operation is determined by either:

• Simply transitive action (e.g., cyclic group $C_n$): each pair $(y_i, y_j)$ admits a unique $g: g(y_i) = y_j$.
• Symmetry group action ($S_n$): for $(i, j)$, a large stabilizer induces many $g$ such that $g(i) = j$; more context ambiguity.

CoT training is performed without positional encoding (NoPE), using a 1-layer transformer with softmax attention and a feed-forward network (FFN). The sequence is tokenized such that each clause occupies fixed positions for variable, group action, and value tokens.

Main Theoretical Results

CoT Learning Beyond $\mathsf{TC}^0$

The analysis shows that gradient descent, under the specified architecture and initialization, can fit both types of LEGO tasks, including $\mathsf{NC}^1$-complete symmetry group word problems. This is achieved by factorizing the reasoning process into attention-based content retrieval (routing to relevant context) and FFN-based application of the group action.

Length Generalization: Algebraic Structure Matters

A key technical result is that length generalization is governed by attention concentration, and the degree of generalization fundamentally depends on the group action type:

• Simply transitive (cyclic) case: Training on short sequences directly yields generalization up to $O(d^{c^*})$ length (where $d$ is vocabulary size), with accuracy falling only as $O(1/\mathrm{poly}(d))$. Sharp attention focus on relevant tokens persists as input grows.
• Symmetry action (permutation group) case: Standard training only generalizes to a small multiple of the training length. Distractors (ambiguous actions under context expansion) destructively dilute attention, limiting extrapolation.

Figure 1: Length generalization results of cyclic ($C_6$) vs. symmetry ($S_5$) tasks.

Recursive Self-Training and Self-Improvement

For tasks where initial length generalization is sublinear, a recursive self-training curriculum—using previously learned CoT traces as bootstrapped data—provably extends the transformer’s solvable length. After $O(\log d)$ doubling stages, the model reaches maximal task length, sustaining near-perfect accuracy.

Figure 2: The illustration of how different components of the attention matrix Qb are used to route the attention to the appropriate locations. The query clause is $\Zb_{\ans,1}$.

Technical Mechanisms: Attention Concentration and Feature Shaping

At the core of these guarantees is the learning of content-based attention patterns. The block-sparse attention parameterization enables the query to retrieve either the relevant predicate clause or the immediate previous answer clause, depending on matching variable tokens (as opposed to positional indices). The FFN parameters are shaped via a curriculum in which neuron–feature pairs (for each token/class) are learned sequentially under the growth of dominating diagonal features and the suppression/interference of confounding ones. This process is formally analyzed using tensor power method dynamics.

Figure 3: $C_6$

Figure 4: Ground-truth permutation

Empirical Validation

Empirical experiments confirm the theoretical findings, including:

• Strong extrapolation in cyclic group tasks: Models trained at $L=5$ generalize successfully up to $L=160$.
• Limited extrapolation in symmetry group tasks: Accuracy plateaus after moderate length increase unless recursive self-training is used.
• Bootstrapping with self-training: Each curriculum doubling expands the range of accurate reasoning.
• Persistent attention concentration: Heatmaps show focused attention on relevant predicate and answer clauses at convergence, matching theoretical predictions.

Figure 5: Cyclic ($C_6$) task

Figure 6: $C_6$

Implications and Guidance for Model Design

Practical

• Transformer depth and CoT: One-layer transformers with CoT are sufficient to solve highly sequential tasks given appropriate curriculum and bootstrapping, matching the expressivity of multi-layer designs.
• NoPE for Length Generalization: The use of NoPE (no positional encoding) is critical; positional bias hinders length extrapolation due to locality preference.
• Attention Routing by Content, Not Position: Learning depends on variable identity matching, not token position.
• Recursive self-training should be incorporated for tasks with high context ambiguity to iteratively extend reasoning length.

Theoretical

• Optimization guarantees move beyond expressivity: This work closes the gap between circuit-theoretic expressive power and what optimization can achieve; prior optimization analyses were strictly limited to $\mathsf{TC}^0$.
• Attention concentration is the central mechanism for generalization: Algebraic task structure directly modulates how robustly attention can be focused and maintained over extended contexts.

Future Directions

• Extending to realistic, noisy, or non-synthetic tasks: While analysis is presently limited to synthetic tasks, the mechanisms should inform architecture, curriculum, and pretraining strategies for practical long-context models.
• Further research into mitigating context rot: Since attention dilution hampers generalization, strategies such as context engineering or more sophisticated attention variants may further improve performance and mitigate context rot.

Conclusion

The paper delivers formal optimization theory for transformer learning under CoT training, characterizing the necessary architectural and data-generation conditions for length generalization and self-improvement. By bridging circuit complexity and stochastic optimization, and by revealing the central role of attention concentration, this work sets a foundation for deeper mechanistic understanding and principled model development in CoT reasoning and length-extrapolating sequence modeling.