Title: Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformers

chain of thought

length generalization

deep learning theory

non-convex optimization
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 long-context 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 $\text{NC}^1$-complete problems with CoT, significantly going beyond prior art confined in $\text{TC}^0$, unless the widely held conjecture $\text{TC}^0 \neq \text{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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper theoretically studies length generalization in trained transformers for LEGO tasks (i.e., state tracking under $G$-actions). The studied transformer model consists of a constrained softmax attention and smoothed ReLU feedforward layer, without positional encoding. The model is trained on CoT loss via a curriculum of one-step and two-step reasoning. It is shown that for simply transitive actions, the CoT model trained for $poly(d)$ time can generalize up to length $poly(d)$ where $d$ is problem size. Moreover, when $G$ is the symmetric group on the state space, a bootstrap self-training curriculum of progressively doubling length can generalize to longer tasks. This also constitutes a novel optimization guarantee for constant depth transformers beyond the $TC^0$ class.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe paper rigorously studies a problem of crucial importance, namely the acquisition mechanism of length generalization in CoT transformers. The considered LEGO task is an empirically well-studied synthetic task, and the one-layer transformer model is also very natural (e.g., incorporates a smooth ReLU feedforward layer, compared to other works which fix or discard the feedforward layer). While there are some structural assumptions and a curriculum learning setup is required to separate the parameters dynamics, I think these are still quite weak for optimization guarantees. The paper is overall easy to read and obtains solid generalization guarantees which displays the potential of length generalization from a theoretical perspective.\n\n**Weaknesses/Questions**\n\nThe settings of the two results are restricted to simply transitive actions and action of the full symmetric group for Theorems 4.1 and 4.2, respectively. While the difficulty of analyzing even these settings is evident from the laborious treatment in the appendix, I am curious as to where the technical difficulty arises when attempting to study more general actions, and what we may achieve with future works.\n\n* Why is the assumption of simple transitivity necessary for Theorem 4.1?\n* The symmetric group assumption of Theorem 4.2 is rather restrictive in that it necessitates $|\\mathcal{Y}| < \\log\\log d$ so that $|G|$ is logarithmic. With Cayley's theorem in mind, can this be relaxed/generalized to (certain) subgroups of symmetric groups?\n* The heading of Section 4.2 mentions non-solvable groups. While it is true that large symmetric groups are non-solvable, unless the theory actually leverages the group-theoretic notion of solvability, it would be better to just say \"symmetric group\" for clarity.\n* Is there a setting in the LEGO setup where length generalization provably fails? (e.g., lower bounds without CoT which have been demonstrated in e.g. the parity papers)\n\nSome typos:\n\"action group action\" in line 242, \"atetntion head\" in line 301\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-19"}

comments: SUMMARY:
The paper proves the learning/generalization of CoT for a 1-layer transformer for LEGO, a reasoning task defined by the authors which involves performing a limited set of group operation on state spaces. The paper also proposes a recursive self-training scheme that enables CoT length extrapolation on the LEGO task.

STRENGTHS AND WEAKNESSES:
Strength:
1. The paper is in general written in a clean manner with clear math annotation usage.
2. There is some novelty in proofing CoT for a 1-layer transformer for LEGO, a reasoning task defined by the authors which involves performing a limited set of group operation on state spaces. The recursive self-training also seems an interesting idea that can be applied to some other reasoning problem.

Weakness:
1. The paper's title seems to suggest it proves for all reasoning tasks but in reality the paper only tackles a limited subset of the reasoning tasks which only include a limited set of group transformations on a state space. 
2. The evaluation section is weak. There is no detailed description of the datasets (how big is the training set? some example data), the training setup (how many GPUs? trained for how long). There is also little analysis of the results.
3. Sometimes assumptions are stated but without enough explanation. For example, the authors stated that "we use no positional encoding (NoPE) in our attention mechanism" on line 202, but never explained in details why this is the case. Modern LLMs all use positional embeddings which makes the proof not applicable to them.

QUESTIONS:
None

LIMITATIONS:
none
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a theoretical analysis of the ability of transformers to learn state-tracking problems belonging to $\\text{NC}^1$ by learning to produce the intermediate reasoning steps, as in the popularly used Chain of Thought setting. This means that sequential prediction methods such as CoT provably extend the capabilities of transformers, unless $\\text{NC}^1 = \\text{TC}^0$, which is an open problem but is popularly believed to be false. The proofs do make some assumptions on the dimensionality and structure of the transformer, which means they do not immediately apply to the popular transformer architectures as used by practitioners. The scope of the result nonetheless appears to be significant, since it brings seminal work on the expressivity of transformers closer to the CoT paradigm typically of use today.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper is fairly readable. There are a few minor typos, but nothing major.\n- As far as I can tell, the main results are quite exciting.\n- By basing their study off a task that has also been used to study similar aspects of transformers empirically in the past, the paper does provide some linkage to the broader literature.\n\nWeaknesses:\n- While the limitations and assumptions are acknowledged well in the paper, their downstream impact on the generalizability of the results is not clear to me.\n- Because of the clause-embedding mechanism, and lack of other positional encoding, it is unclear to me whether these results actually correspond as well to CoT as the authors claim, in which the structure of the sequence (and thus the attention mechanism) is substantially different.\n- The experiment could benefit from a non-CoT baseline to display the difference in length generalization empirically. (I realize empirical performance is not the subject being studied here, but it would give some sense of the practical implications of the very theoretical results.)\n\nQUESTIONS:\nSuggestions:\n- Please enhance the legibility of Figure 1 substantially. The text is impossible to read when printed out.\n- Please add a non-CoT baseline to your experiment to give some credence to the impact on downstream performance implied by your theory.\n- Some typos:\n  - P.3, Theorem 1.2, second bullet point: \"from its own outputs**, .**\"\n  - P.4, example of LEGO sentence: in answers, $x_L$ should likely be $x_{L'}$\n  - P.4, in Eqn (2), $L$ should be $L'$ (per the preceding sentence)\n  - P.9, \"Note that **in by**\"\n  - P.9, Experiments, \"transformer **with attention with two atetnion heads**\"\n\nQuestions:\n- If I understood your construction on P.5-6 correctly, your transformer in effect uses the same projection matrix $Q$  for the queries and keys, while not using a projection at all for the values. Is that correct? If so, why is this not explicitly listed as an assumption in the paper, since it does not correspond to the more general transformer architecture?\n- Assumption 3.4 appears to be very strong. Could you help me build some intuition for why this is necessary? In particular, why $(p, q) \\notin \\{(4, 3), (4, 4)\\}$?\n- My understanding is that your clause-wise embedding structure is what allows the model to disentangle the differing meanings of the tokens without positional encodings. Is that correct? If so, how well do you expect that your results would carry over to a RoPE (or other PE) setting in which the clauses are **not** hard-coded to be separated, as is the typical case in language modeling?\n\nLIMITATIONS:\nPlease see questions above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a theoretical analysis of a one-layer transformer's ability to learn Chain-of-Thought (CoT) reasoning via gradient descent. Focusing on the synthetic \"state tracking\" task, the authors prove that the model can learn to solve instances of this problem that belong to the complexity class NC1. The work further investigates length generalization, showing that for simple algebraic structures (cyclic groups), the learned solution generalizes to longer sequences. For more complex structures (symmetric groups), where direct generalization is limited, the authors propose a \"recursive self-training\" curriculum and prove that it allows the model to achieve similar polynomial-length generalization.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- **Addresses a Fundamental Question:** The paper tackles a central and challenging question in modern AI theory: whether the reasoning capabilities observed in large language models can be provably learned via optimization, moving beyond expressivity results alone.\n- **Detailed Technical Analysis:** The paper provides a detailed and technically involved proof for its specific setting, demonstrating a mechanism by which attention concentration can emerge from a staged training process to solve a sequential task.\n- **Formalization of Self-Training:** The theoretical formalization and proof for the \"recursive self-training\" curriculum is a novel contribution. It provides a potential mathematical explanation for the empirical success of self-improvement methods that have been observed in practice.\n\n### Weaknesses\n\n- While I completely acknowledge the necessity of studying why CoT is helpful from an optimization perspective, I think that the single special task is not representative enough for demonstrating the underlying mechanism why models can learn to leverage CoT. Specifically, the authors use the state tracking task in NC1 for demonstration, but they do not explicitly point out that whether state tracking is hard and representative enough to distinguish  TC0 and NC1. For instance, is it NC1-complete? In your citation [24], the paper prove their theorems on differnt tasks including a whole class of dynamic programming problems, making their results universal and convincing. Without showing that the theorem holds for a general class of tasks (rather than hold specifically for this task because of its special structure), I do not think that \"explaining why CoT can be learn from an optimization perspective\" can be claimed. \n- The focus on one-layer transformer makes the optimization perspective more or less similar to a construction (expressivity perspective), as Alg 2 constructes optimization in a manual and step-by-step way for different layers. Since the authors try to go beyond expressivity, and since, in theory, a two layer transformer is more expressive, I think the authors should also prove, at least more generally, that a two / general layers of transformers can also be optimized to learn CoT well. Otherwise, the whole proof for optimization seems to be not generalizable, making the theorems less convincing.\n- To the best of my knowledge, your citation [24] proposes \"transformers could simulate polynomial-size circuits using poly(L) steps of CoT\" earlier than [25]. What is the reason that you choose to base on a later work?\n\nQUESTIONS:\n1. Could you elaborate on the role of the smoothed ReLU in your theoretical analysis? Are there specific properties of this function (e.g., its behavior for negative inputs or its specific polynomial form) that are essential for the proof to go through, as opposed to other common smooth activations?\n2. In the recursive self-training curriculum (Algorithm 2), new data is generated using the model's greedy output (Definition 4.3). From a theoretical standpoint, what challenges would arise if stochastic sampling (e.g., with a non-zero temperature) were used instead of greedy decoding to generate the self-training data? Would the convergence guarantees for length generalization still hold?\n3. Theorem 4.2 provides an accuracy of 1−O(2^{−k/poly(d)}) for the bootstrapped model. Could you provide more intuition on how errors might propagate or accumulate as the reasoning length $L$ grows to be a large polynomial in $d$? Does the model become more susceptible to single-step failures as the chain lengthens, and does the theory provide a way to bound this cumulative error?\n\nLIMITATIONS:\nI do not think that equaling limitations with assumptions makes sense. I suggest the reader to position their work in a more objective way.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}