Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization

We thank the reviewer for the positive comments and thoughtful questions. Below, we provide point-by-point response to the main comments.

Q1. Why is the assumption of simple transitivity?

A1. The assumption of simple transitivity can be thought of as a natural generalization of cyclic group actions. Our analysis only requires that for each pair of $y_1, y_2$, there exists a unique group operation $g$ such that $g⋅y_1=y_2$​. In our context, this assumption allows us to study group-based reasoning tasks while keeping the structure simple for analysis.

Q2. Can the symmetric group assumption in Theorem 4.2 be relaxed/generalized?

A2. This is a great question. In fact, our proof can be directly generalized to certain subsets of the symmetric group that satisfy the following condition: for any pair of elements $y, y^{'}$ in the domain, the number of group elements mapping $y$ to $y’$ is a fixed constant. This includes subsets of symmetric groups such as the groups of all k-swaps in $S_n$, for $k \leq n$. These subsets are not necessarily groups (they can be conjugacy classes) but they provide exactly the regularity we need for successful feature learning. If this structure is removed entirely, the distribution of $(g,y)$ pairs becomes imbalanced, which can harm the learning of both single-step reasoning and context-length generalization. While we agree it would be very interesting to extend the analysis to other groups, it is beyond the current scope of our work.

Q3. It would be better to just say "symmetric group" for clarity.

A3. Thank you for your helpful suggestions! We will revise the wording in the revision.

Q4. Is there a setting in the LEGO setup where length generalization provably fails?

A4. Our theory links length generalization to the model’s ability to perform in-context retrieval: as context grows, the attention mechanism must still assign enough mass to every relevant clause. Two key factors determine whether generalization can succeed:

1. The peak attention attained by the relevant clauses within training context length;
2. The rate at which attention is diluted by irrelevant clauses as context length increases.

In the case of a simply transitive group action, the learned features are sufficiently “pure,” enabling sharply focused attention and yielding a polynomial bound on generalization length. In contrast, in the symmetric group action setting, the learned features do not have a similarly strong guarantee, limiting the length generalization to only a constant multiple of the training length. While our exact theoretical predictions may not transfer directly to all LEGO settings due to differences in task structure and objective, we hope our analysis offers useful insight into when and why length generalization can succeed or fail, especially in realistic transformer-based reasoning scenarios.

We thank the reviewer again for your comments. We are more than happy to answer your further questions.

We thank the reviewer for the time and thoughtful comments on our work. Below we address your concerns.

Q1. 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.

A1. We would like to respectfully clarify that our title is not an overclaim in the context of existing theoretical work on reasoning with transformers.

1. State tracking has been widely adopted as a standard synthetic testbed for evaluating the reasoning abilities of language models. The seminal work [15] established foundational expressivity results for transformers in algorithmic reasoning and provided comprehensive empirical evidence across architectures, from RNNs to transformers, on tasks ranging from Dyck languages to non-solvable groups. It showed that sufficiently deep transformers can solve most group word problems within the training length. Subsequent studies [12–13] further demonstrated the expressivity gap: constant-depth transformers without CoT cannot solve problems beyond TC0, whereas adding a scratchpad (analogous to CoT) enables transformers to solve NC1-complete problems such as S5 [1,2,15,16]. Our work bridges the gap between theory and empirical observations. We go beyond prior results by providing a formal proof that transformers can learn a length-generalizable CoT solution, overcoming the expressivity limitations of the base architecture. This confirms that transformers can learn solutions that are provably more powerful than what the architecture can express without CoT.
2. To our knowledge, no prior theoretical paper on CoT has established guarantees for all reasoning tasks [1-7]. From the perspective of optimization analysis, existing theoretical studies on CoT [8-10] have exclusively addressed specific, simple tasks within TC0, such as parity or linear regression. In contrast, our work represents significant advances: 1). We provide the first rigorous analysis of transformer training dynamics for a broader and more challenging class of CoT-style reasoning tasks. Specifically, we study tasks beyond parity within TC0 (simply transitive setting) as well as tasks within NC1 (symmetry group setting); 2). We are the first to establish convergence guarantees that demonstrate not only the transformers' ability to solve these tasks but, crucially, their capacity for length generalization.

Q2. The evaluation section is weak

A2. We would like to emphasize that the empirical part primarily serves as numerical verification of our theoretical results, specifically illustrating length generalization in the cyclic case. Due to space constraints, we omitted detailed dataset descriptions; however, the synthetic data and loss function strictly follow the LEGO structure analyzed theoretically in our paper. Basic setup about network architectures and optimizers is briefly summarized in the paper. Other detailed training setups mentioned by the reviewer, which are less directly related to our theoretical focus, have been intentionally omitted to maintain clarity and conciseness.

Q3. Sometimes assumptions are stated, but without enough explanation, For example, the authors stated that "we use no positional encoding (NoPE) in our attention mechanism"

A3. We thank the reviewer for raising the concern regarding assumptions stated without sufficient explanation. Specifically, for the use of "no positional encoding (NoPE)," we note that it has been widely observed empirically that standard positional embeddings inherently fail to achieve length generalization. This limitation arises because traditional positional embeddings implicitly introduce approximate biases that typically favor local positions and thus fundamentally hinder extrapolation beyond the training sequence length. In contrast, our theoretical analysis supports an alternative approach and suggests an alternate transformer design with no positional embedding that provably achieves length generalization, thereby offering insights and guidance for current empirical research. Furthermore, our theory implicitly reveals that the inability of existing networks to generalize to unseen lengths may stem from the presence of positional embeddings, thereby identifying a potential underlying reason why current settings fail in this regard. In the revision, we will carefully include additional discussion and remarks to explicitly clarify these assumptions and the rationale behind our design choices.

We thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

References:

[1] Feng et al. Towards revealing the mystery behind chain of thought: a theoretical perspective.

[2] Li et al. Chain of thought empowers transformers to solve inherently serial problems.

[3] Merrill et al. The expressive power of transformers with chain of thought.

[4] Abbe., et al. How far can transformers reason? the globality barrier and inductive scratchpad.

[5] Kim et al. Metastable dynamics of chain-of-thought reasoning: Provable benefits of search, rl and distillation.

[6] Huang et al. A formal framework for understanding length generalization in transformers.

[7] Golowich et al. The role of sparsity for length generalization in transformers.

[8] Huang et al. Transformers learn to implement multi-step gradient descent with chain of thought.

[9] Kim et al. Transformers provably solve parity efficiently with a chain of thought.

[10] Wen et al. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency.

[11] Kim et al. Entity tracking in language models.

[12] Merrill et al. A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers.

[13] Merrill et al. The Illusion of State in State-Space Models.

[14] Zhang et al. Unveiling Transformers with LEGO: a synthetic reasoning task.

[15] Liu et al. Transformers learn shortcuts to automata.

[16] Wen et al. RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval.

Dear Reviewer vE9B,

As the author-reviewer discussion period will end soon, we would like to check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards,

Authors

Dear Reviewer vE9B,

As the author-reviewer discussion period ends soon, we would greatly appreciate it if your could check our response promptly. If our response resolves your concerns, we kindly ask you to consider raising the rating of our work. We are also more than happy to answer your further questions. Thank you very much for your time and efforts!

Many thanks,

Authors

We thank the reviewer for providing the valuable feedback! We will enhance the readability of the figure and fix the typo in the revision. Below we address your concerns.

Q1. While the limitations and assumptions are acknowledged, their impact on generalizability is unclear

A1. We thank the reviewer for raising this general yet important point. Indeed, as in most theoretical works, certain simplifying assumptions are necessary to keep the analysis tractable and focused. Within the clearly stated scope of our theoretical model and assumptions, we have provided rigorous explanations of length generalization. While extending these results beyond our current setting would be valuable future work, the primary goal of this paper is to provide some novel insights from a theoretical perspective. We hope that our answers to your later questions can explain to you why our setting shares the structural similarities with real-world CoT performed by actual language models, and how to interpret our idealized results to shed light on practical applications.

Q2. It is unclear whether the results truly apply to CoT, given its different sequence structure and attention patterns.

A2. Through our technical construction, we aim to create a setting where the structure of step-by-step reasoning is preserved and is amenable to theoretical analysis. While it may appear different from natural CoT traces on the surface, we emphasize that our design abstracts two essential components of real-world CoT reasoning:

1. Each CoT step typically focuses on a localized subproblem, even when the overall task involves long-range dependencies. We represent this structure using a clause-based data format, which allows us to formalize how reasoning steps are represented, sampled, and generated during training.
2. Long-range dependencies in CoT are often handled by the retrieval function of attention. This is an important topic of study in both theoretical analyses [3,6,7] and practical transformer components such as induction heads [8]. In our framework, we build on the LEGO formulation and adopt the "variables" in LEGO as retrieval queries, thereby enabling transformers to attend to relevant context.

In summary, while our problems are abstract examples of chain-of-thought reasoning, they preserve key structural patterns observed in real-world reasoning traces. We hope these abstractions can offer insights into how attention mechanisms may function in practical CoT models.

Q3. The model ties query/key projections and omits value projections. Why isn't this assumption explicitly stated, given it deviates from standard transformer architectures?

A3. Thanks for the question. We adopted a quite common simplification in existing theoretical works of transformer training dynamics due to technical tractability [1-5]. Our assumptions delegated all the functionalities of attention to a single attention weight Q, which leaves room to analyze the attention pattern between different tokens while drastically simplifying the analysis because we dropped many interplays between different weight updates. In fact, one can view our setting as we adopted a slightly stronger attention layer, since now the $Q^TK$ is full rank and the value matrix can be thought of as absorbed into the MLP. Thank you for highlighting this point, and we will explicitly include additional references and clarify this assumption to better justify our architectural choices in the revised manuscript.

Q4. Assumption 3.4 appears to be very strong

A4. As we pointed out in line 226 in the paper, setting certain parts of the weight matrix to zero aligns with common practices in existing theoretical analyses of transformer training dynamics [1-5]. Technically, for our task, only the blocks $Q_{4,3}$ and $Q_{4,4}$ play a meaningful role in learning: $Q_{4,3}$ is important for in-context retrieval, while $Q_{4,4}$ puts weights on the current state. The remaining components either receive no gradient or capture task-irrelevant biases that complicate the analysis. The delicate balance between these two blocks, and their interaction with irrelevant tokens in the context, is the major technical challenge in our attention dynamics analysis. In fact, some existing works [1-2] assume even stronger simplifications, such as being a scalar multiple of the identity matrix, with only a single learnable parameter, for similar reasons. Even though we have made some compromises in the setting to allow tractable analysis, we hope that our methods can offer some insights on the reasoning ability of general language models. In fact, our result gives a clear example of how transformers can learn to implement the induction head [9], showing that our setting can reproduce certain results on the inner workings of transformers studied in prior works.

Q5. How well do the results extend to RoPE or other positional encodings where clause boundaries are not explicitly separated

A5. As discussed in Q2, we believe that achieving length generalization beyond the training context requires the model to overcome certain proximity biases. However, we do not aim to discredit the importance of locality that is better captured by other PE such as RoPE, as our clause-based structure actually absorbed all the local correlations to make the analysis possible. Regarding the extension, we conjecture that currently used positional encodings like RoPE or ALiBi are inherently limited in their length generalizability, due to their position-dependent nature (either absolute or relative), such as shown in some empirical study of architectural design [9]. Prior theoretical works on CoT expressivity [6, 7] also rely on fixed positional encoding schemes and cannot generalize beyond a fixed problem length. We hope our work could motivate further exploration of architectural designs that can enable transformers to extrapolate beyond training context lengths, whether for CoT reasoning or to solve other problems.

Q6. The experiment could benefit from a non-CoT baseline to display the difference in length generalization empirically.

A6. Following the reviewer’s suggestion, we conducted additional experiments with a non-CoT baseline in both the simply transitive (cyclic) and symmetric cases to evaluate length generalization. Specifically, we trained the same one-layer transformer under the identical task setup described in the paper, but without intermediate answer clauses, i.e., directly predicting the final answer. For non-CoT models, we explored eight learning rates: $[5e-6, 1e-5, 2e-5, 5e-5, 1e-4, 2e-4, 5e-4]$, whereas our CoT model was trained with a single learning rate ($1e-4$). Additionally, the CoT model was trained exclusively on LEGO sequences of length 5, while the non-CoT baseline was trained on sequences of lengths 5 and 10.

Results show that the transformer with CoT achieves length generalization in our setting, while the non-CoT baseline does not. We also plan to include additional multi-layer transformer experiments to further investigate this setting.

We thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

[1] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

[2] Nichani et al. How Transformers Learn Causal Structure with Gradient Descent

[3] Kim et al. Transformers provably solve parity efficiently with a chain of thought

[4] Zhang et al. Trained transformers learn linear models in-context

[5] Kim et al. Transformers learn nonlinear features in context: nonconvex mean-field dynamics on the attention landscape

[6] Feng et al. Towards revealing the mystery behind chain of thought: a theoretical perspective

[7] Li et al. Chain of thought empowers transformers to solve inherently serial problems

[8] Olsson et al. In-context Learning and Induction Heads

[9] Ren et al. Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

I thank the authors for their thorough reply to my review, as well as to those of my fellow reviewers. In particular, the response to X31B's review highlighted some gritty details of this paper, such as that the task is indeed NC1-complete in the case of non-solvable groups. I think the changes and clarifications outlined by the authors will strengthen the paper substantially.

There are two things that are still a bit unclear to me:

1. While the discussion here has convinced me that the assumptions are justified in the sense that they represent the frontier of what we are currently able to say about the theoretical properties of transformer training dynamics, it is still not clear to me why it is $Q_{3, 4}$ and $Q_{4, 4}$ in particular that allow the model to solve this task. Can the authors expand on this a bit? Is it because state-tracking can be implemented with simple one look-behind operator + one update to the current state per timestep?
2. The extended empirical results are very confusing to me. They do not appear to suggest that the non-CoT model fails to length-generalize, but rather that it fails to learn the task at all (since performance is roughly flat, and very low even on in-distribution lengths 5-10). Do the authors have any thoughts on why this might be? The task itself does not appear so difficult that it cannot be solved in-distribution without CoT, but perhaps my intuition is wrong.

Other than these two points I do not have any other requests to make of the authors at this point. I will await the replies of my fellow reviewers before officially updating my score in case new evidence is brought to light, but would at this point consider updating my overall score to 5 - Accept as I believe the discussion has improved the paper both in Quality and in Clarity. (I would encourage the authors to spend more time on producing a set of relevant and high-confidence experiments, but since the primary focus of this submission is theoretical I do not believe it should be a reason for the paper to not be accepted.)

We thank the reviewer for the time and thoughtful comments on our work. Below we address your concerns.

Q1: The single special task is not representative enough

A1:

1. We first emphasize that state tracking over group operations has been widely used as a benchmark for separating architectural expressivity in seminal works [1–4]. More importantly, state tracking problems that distinguish TC0 from NC1 or beyond have been used to provably separate not only different neural architectures, but also the expressivity of transformers with v.s. without CoT reasoning [5]. This makes state tracking a meaningful and theoretically grounded testbed for studying the capabilities of transformers, either using CoT or not. The task is also widely adopted in empirical studies of neural reasoning [6–7].
2. Importantly, it is a well-established theoretical fact that state tracking over non-solvable groups (e.g., $S_5$) is NC1-complete [1]. We highlight this in direct response to the reviewer’s question, as it is a standard and widely known result in circuit complexity.
3. Prior theoretical studies on CoT [8–10] necessarily focus on simple, hand-crafted tasks within TC0 (e.g., parity, linear regression), due to the inherent difficulty of analyzing optimization dynamics. In contrast, our work advances this line of research by studying both TC0- and NC1-level CoT tasks, and by providing the first convergence guarantees in these more expressive settings without overclaiming generality beyond this well-motivated domain.

Q2. The focus on one-layer transformers makes the optimization perspective more or less similar to a construction

A2. We believe there may be a misunderstanding regarding our optimization analysis. Our work is fundamentally different from and significantly more challenging than the expressivity argument the reviewer suggests.

1. This confusion likely stems from prior optimization-based CoT theory [8–10], which focuses on simple tasks (mostly the parity problem in TC0) where a single gradient step closely mimics a hand-crafted construction. In contrast, our setting requires studying the multi-phase training dynamics of both MLP and attention weights across many gradient steps, and there is no straight path from initialized weights to the optimal weights. The features learned in the network go through sophisticated trajectories that we cannot control and characterize completely but can only guarantee its convergence on a high level, which is very different from construction-style arguments. Thus this highly nontrivial learning process is nothing like expressivity type construction.
2. The reason we introduce a curriculum-style schedule (Alg 2) that separates the learning of MLP (focused on one-step group actions) from the learning of attention (enabling length generalization) is to manage the complexity of the intertwined dynamics between different parts of the network. This separation does not simplify the learning process into a staged construction. Both stages still undergo complex multi-phase optimization, and our analysis must carefully track how intermediate solutions influence subsequent stages. This level of coupling introduces deep technical challenges far beyond what construction-based proofs address.
3. The implication that one-layer transformers are “simple” to analyze is also misleading. The optimization landscape remains highly nontrivial even in this setting [11–13]. In fact, one-layer models likely represent the current theoretical frontier for understanding transformer training dynamics. Extending such results to deeper models, without strong layer-wise assumptions, is out of reach for current research.

Prior CoT learning theory has only addressed very simple tasks such as parity with no results on even cyclic group reasoning or length generalization. To our knowledge, we are the first to bridge the gap between the theoretical expressivity and learnability for transformers with CoT, and to provide a length generalization guarantees that go significantly beyond what previous work achieves. Especially since most prior expressivity results rely on absolute positional encoding and thus cannot length generalize.

Q3. What is the reason that you choose to base on a later work [25] rather than [24]?

A3. We clarify that prior works on expressivity mainly provide background context rather than form the main foundation of our study. While [24] was the first to show that transformers with CoT can go beyond TC0 in expressivity, their results do not establish a separation beyond NC1, as the CoT expressivity lower bound is only NC1-hard, even though the tasks may include P-complete problems. We refer more often to [25] because it offers a more comprehensive and aligned characterization of the complexity classes for transformers with and without CoT. That said, we note [25] is not the only work that proves such results, for example, [1] also establishes similar findings, albeit with a slightly weaker expressivity upper bound.

Q4. The role of the smoothed ReLU in your theoretical analysis

A4. The smoothed ReLU is chosen due to two key properties that are essential for our optimization analysis:

1. Its polynomial smoothness around zero is necessary for applying the tensor power method, which plays a central role in proving MLP feature emergence, as shown in prior works [14–16].
2. The symmetry around zero serves as a soft barrier that prevents some pre-activations from drifting too far into the negative region. This helps ensure convergence to a clean and interpretable set of features. We adopt this design specifically to enable our length-generalization result.

In summary, the use of smoothed ReLU makes the optimization analysis more tractable, allowing us to show that the transformer learns clean features for one-step state tracking. This enables us to focus on analyzing length extension and extrapolation, which is the core contribution of this work.

Q5. Challenges of using stochastic sampling (e.g., non-zero temperature) instead of greedy decoding for generating self-training data?

A5. The main challenge lies more in the clarity and tractability of the proof rather than in any fundamental technical obstacle. In fact, even with non-zero temperature, at the end of each training length $2^k$ for $k \in N^+$, the probability of generating an incorrect prediction has already fallen below 1/d with high probability over the dataset. A more involved analysis would show that the noise introduced by occasional incorrect predictions due to stochastic sampling does not significantly affect the learning of the attention weights. Sequences with slightly higher error rates occur with exponentially small probability, and thus have a negligible impact on the overall attention gradients. However, to analyze such gradients in detail would require extensive case-by-case discussion, which would make the proof significantly more cumbersome without offering fundamentally new insights.

Q6. Intuition on how errors might propagate or accumulate as the reasoning length grows to be a large polynomial in d

A6. This is a great question, and we agree that a clearer explanation of the limits of length extrapolation should be included in our paper. At a very high level, as context length grows beyond the training length, attention to important clauses becomes diluted, and the influence of incorrect group operations can accumulate for CoT prediction. If the total attention to misleading tokens becomes significant, the model’s prediction can deteriorate. In the simply transitive action case, the correct operation typically receives $1 - 1/d^c$ attention, while all others get $1/d^c$ in total. Thus, one would need roughly $L \cdot d^c$ incorrect operations in the context to confuse the model, allowing extrapolation to polynomially longer lengths. Beyond length $\gg |Y| \cdot L \cdot d^c$, predictions approach uniform. In the symmetric group action case, the learned representations are less balanced, and error arises for length $> C \cdot L$ with a small constant $C < 10$. Beyond that limit, the prediction of transformers for values $y \in \mathcal{Y}$ can quickly become uniformly random.

We thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

References:

[1] Merrill et al. The expressive power of transformers with chain of thought

[2] Merrill et al., A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers

[3] Merrill et al. The Illusion of State in State-Space Models

[4] Liu et al. Transformers learn shortcuts to automata

[5] Li et al. Chain of thought empowers transformers to solve inherently serial problems

[6] Kim et al. Entity tracking in language models

[7] Toshniwal et al. Chess as a Testbed for Language Model State Tracking

[8] Huang et al. Transformers learn to implement multi-step gradient descent with chain of thought

[9] Kim et al. Transformers provably solve parity efficiently with a chain of thought

[10] Wen et al. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency

[11] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

[12] Zhang et al. Trained transformers learn linear models in-context

[13]Li et al. A Theoretical Understanding of shallow Vision Transformers: Learning, Generalization, and Sample Complexity

[14] Cao et al. Benign Overfitting in Two-layer Convolutional Neural Networks

[15] Allen-Zhu et al., Towards understanding ensemble, knowledge distillation and self-distillation in deep learning

[16] Chen et al. Towards understanding the mixture-of-experts layer in deep learning

Dear Reviewer X31B,

As the author-reviewer discussion period will end soon, we would like to check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards,

Authors

Dear Reviewer X31B,

As the author-reviewer discussion period ends soon, we would greatly appreciate it if your could check our response promptly. If our response resolves your concerns, we kindly ask you to consider raising the rating of our work. We are also more than happy to answer your further questions. Thank you very much for your time and efforts!

Many thanks,

Authors