Title: Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformer theory

training dynamics

multi-step reasoning
Abstract: Transformers have demonstrated remarkable capabilities in multi-step reasoning tasks. However, understandings of the underlying mechanisms by which they acquire these abilities through training remain limited, particularly from a theoretical standpoint. This work investigates how transformers learn to solve symbolic multi-step reasoning problems through chain-of-thought processes, focusing on path-finding in trees. We analyze two intertwined tasks: a backward reasoning task, where the model outputs a path from a goal node to the root, and a more complex forward reasoning task, where the model implements two-stage reasoning by first identifying the goal-to-root path and then reversing it to produce the root-to-goal path. Our theoretical analysis, grounded in the dynamics of gradient descent, shows that trained one-layer transformers can provably solve both tasks with generalization guarantees to unseen trees. In particular, our multi-phase training dynamics for forward reasoning elucidate how different attention heads learn to specialize and coordinate autonomously to solve the two subtasks in a single autoregressive path. These results provide a mechanistic explanation of how trained transformers can implement sequential algorithmic procedures. Moreover, they offer insights into the emergence of reasoning abilities, suggesting that when tasks are structured to take intermediate chain-of-thought steps, even shallow multi-head transformers can effectively solve problems that would otherwise require deeper architectures.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies multi-step reasoning with chain-of-thought (CoT) in a tree search setting. The tasks considered are the backward path finding task where a path from the goal leaf back to the root must be chained together, and the forward path finding task where the backward path must be found and reversed. A one-layer two-head transformer construction is shown to achieve both tasks by using appropriately long CoT. Moreover, it is shown that GD learns such solutions from zero initialization and the trained model can generalize to unseen tree combinations.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n* The considered CoT path-finding problem is interesting and very relevant to the current test-time scaling paradigm, which lacks a strong theoretical foundation.\n* The main results effectively show that CoT allows transformer models to implicitly think backwards from the goal to find the solution, which otherwise would require brute-force forward search over the tree to find.\n* The constructions are paired with a generalization guarantee and dynamical analyses which show well behaved solutions may be realistically found with GD.\n\n**Weaknesses**\n* It is not clear what or why the model needs to learn from a learning-theoretic perspective, since the desired algorithm and training distribution is fixed (and the model benefits from knowing this target mechanism/construction, e.g., as we only train $\\theta=(B_i,C_i)$). It seems that the token embeddings are arbitrary and the model can learn the appropriate weight matrices, but this is a rather artificial learning goal compared to learning (say) the tree structure, which is assumed to be fixed to uniform over perfect binary trees.\n* The GD analysis is quite complicated and would benefit from a section summarizing the proof ideas and overall strategy.\n\nQUESTIONS:\n* Are only the token embeddings being learned from a learning theoretic perspective? The paper mentions that \"our results can be easily extended to other types of trees\", how can the tree structure and/or distribution be made more general?\n* Can the authors provide a high level summary of how the CoT gradient dynamics is made tractable? What are the straightforward or nontrivial parts of the analysis? I am personally interested in seeing what insights can be extended to other CoT setups.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates how one-layer transformers handle the multi-step reasoning tasks. The authors design synthetic tasks and theoretical analyze the potential solutions of the model parameters. Besides, the authors discuss that the model parameters are guided to the potential solutions after gradient descent.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength:**\n\n1. This work theoretical analyze how single-layer transformers handle the multi-step tasks with attention heads. The analysis is inspiring for understanding how existing pretrained models perform multi-step tasks.\n2. The paper is well written, and the presentation logic is clear. The authors design synthetic multi-step tasks and present potential solutions of model parameters, then they verify with gradient descent the model will converge to these solutions.\n\n\n**Weakness:**\n\n1. Even if I understand that under the simplified setting, the authors can present the analysis in a more solid way, it seems a bit difficult to further extend the analysis for more practical transformers, i.e., pretrained models. First, some important modules (such as layer-normalization) are ignored. Besides, the training (i.e., optimizer and initial weights) has some assumptions. Thus, I’m wondering if such analysis can only be limited to the over-simplified setting. Do the authors have any idea if this conclusion is possible to extend to more practical transformers’ architecture or even pretrained models?\n2. This is relevant to the first point. The authors analyze the gradient descent manually. While there are many assumptions about the model weights. Even if we cannot deny the possibility that they might hold in practice, will the training results be like those with optimizer (e.g., Adam) and randomly initialized weights?\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to show that a multi-head transformer can learn symbolic multi-step reasoning during training. Specifically, the paper shows that a single-layer transformer can solve the two-stage path-finding in trees problem (finding the path from a goal node to the root, which requires the model to first find the path from the root to the goal, then reverse it) within a single autoregressive pass. The author attributes the capability of a single-layer transformer in this task to the multi-head mechanism, which has one head that identifies the next node and the other reverses the path. The authors prove mathematically that a Transformer can learn this behavior through training. They also run experiments to show that the trained model can solve new trees it hasn't seen before.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n1. The paper is generally well-written.\n2. The paper provides an interesting viewpoint on understanding the multi-head mechanism of the transformer.\n\n**Weakness**\n1. The paper claims that the multi-head mechanism allows the shallow transformer to perform complex reasoning tasks since different attention heads can specialize and coordinate autonomously. However, this raises a question.\n2. The experimental settings and the tasks are weak; they only show that a single-layer multi-head transformer can converge on the tree path-finding tasks, and there is no ablation study or any other evidence to support the paper’s claim that different attention heads learn to specialize and coordinate autonomously.\n3. The paper only discusses the situation of single-layer transformers; however, this may not apply to multi-layer transformers. For example, can we hypothesize that for a two-layer transformer, the two layers will also learn to specialize and coordinate autonomously to solve the two subtasks?\n\nQUESTIONS:\n1. To support Theorem 3, I would like to see experiments that compare the results of single-head and multi-head transformers (with comparable parameters) on the forward reasoning tree path-finding tasks. Can we observe that only the multi-head transformer can finish the task and generalize to unseen trees?\n2. Since the experiments are focused only on a single-layer transformer, is it possible that the capability of complex multi-step reasoning also comes from the multi-layer?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper claims to formally prove that shallow transformers are able to learn multi-step reasoning on trees via gradient descent. Specifically, the authors consider the task of learning backward and forward reasoning on a tree-shaped graph, i.e. given a tree and a goal node that is a leaf node of the tree find a path from the root of the tree to that goal node. They start by first claiming that a parametrisation of a shallow attention block exists that can perform both backwards and forwards reasoning under a certain type of node embedding when applied repeatedly in a chain-of-thought-like fashion. The authors then continue to claim that gradient descent actually converges to this constructed parametrisation when using a training set with specific properties.\n\nSTRENGTHS AND WEAKNESSES:\n#### Strengths\n1. The abstract and introduction are well-written and also motivate the problem well: we want to see if one can prove that transformers, i.e. attention blocks, are able to learn reasoning in practical circumstances.\n2. Figure 2 gives a nice visual abstraction of what behaviour you would like the attention block to have.\n#### Weaknesses\nI have severe concerns about this paper. In particular, I suspect a significant portion is produced by generative AI. Concretely, none of the theoretical results claimed in the paper are actually correctly proven as the so-called \"proofs\" in the appendix are filled with mistakes or seem to (try) to prove unrelated statements. In short, the paper really does not have anything of any substance. Please see a couple of precise examples below.\n\n1. The authors claim to prove the existence of a parametrisation that makes transformers exactly perform backward reasoning (Lines 190-191). However, the statement of Theorem 1 only says that a construction exists that converges to backward reasoning under a certain limit to infinity that is not attainable in practice. While such a result could still be interesting, there is no proof given, not even in the appendix, of this statement. The appendix contains a \"Lemma 1\" that seems to somewhat use the same notation, but is apart from that not showing anything as stated.\n2. Theorem 2 claims a similarly plagued statement for forward reasoning, which again is only a convergence claim for an impossible limit to infinity. While now the appendix does (seem) to contain a (trial) of a proof for Theorem 2, unfortunately the proof is filled with mistakes. First, Theorem 2 claims to show the existence of parametrisations that adhere to certain constraints (Equation 10 and 11), but the proof simply assumes these parametrisations exist in the first place. Even then, the proof claims a series of equalities (Equations 27-33) that should show that in the limit the softmax goes to a unit vector in the infinite limit. However, since Equation 29, 31 express equality to some positive constant multiplied with $\\alpha_1$, that means that the softmax would go to a uniform vector when $\\alpha_1 \\to +\\infty$. Hence, this proof is false.\n3. Given that Theorems 3 and 4 bases itself on the constructions of Theorem 1 (Lines 295-299) and Theorem 1 was not proven, it follows that Theorems 3 and 4 are also unproven. However, I still want to put attention to the provided \"proofs\" in the appendix. From Line 526 onwards, we get page after page of mathematical derivations that do not go anywhere. This sort of exposition is a strong indication that modern generative AI was used to construct those \"proofs\".\n#### Verdict\nIt is highly likely that this paper was mostly produced with generative AI. Even if we do not assume that it is, the claims of the paper are not proven, neither theoretically or practically. The paper tries to emulate formality by generating page after page of mathematical gibberish in the appendix, which does not hold up to close scrutiny. Clear reject.\n\nQUESTIONS:\nPlease find my questions intertwined with the weaknesses above.\n\nLIMITATIONS:\nNo.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers a symbolic multi-step reasoning problem: for a given reasoning tree, the model should first conduct backward reasoning for the root, and then forward reasoning for a complete reasoning path. The authors provide a mechanism on how a single-layer softmax transformer solves this problem, accompanied by training dynamics analysis. This work provides insights for understanding CoT reasoning in LLMs.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths*\n\n1. This work proposes a symbolic multi-step reasoning task, which mimics the reasoning procedure for real-world LLMs.\n2. This work provides a detailed mechanistic explanation on how a single-layer transformer can solve this task, which incorporates a detailed designed transition point between two attention heads, which is both interesting and insightful.\n3. This work provides training dynamics analysis, accompanied by experiments to validate their findings.\n4. The mechanistic interpretation, together with additional theoretical analysis for the forward process, takes an important step for theoretically understanding multi-step reasoning processes in transformers.\n\n*Weaknesses*\n- The backward pass process constructs a transformer adapting copying mechanism, which is well studied in many theoretical analysis papers [1][2].\n- In the construction for forward reasoning, each position has two nodes and a flag, which seems too artificial.\n\n[1] Huang, Yu, Yuan Cheng, and Yingbin Liang. \"In-context convergence of transformers.\" arXiv preprint arXiv:2310.05249 (2023).\n\n[2] Nichani, Eshaan, Alex Damian, and Jason D. Lee. \"How transformers learn causal structure with gradient descent.\" arXiv preprint arXiv:2402.14735 (2024).\n\nQUESTIONS:\nI think this is a good work. If any question, I think in the experiments in the appendix, the loss curve for backward reasoning seems not to converge in 1000 epochs, as the loss is still very high. Could you provide the full training loss curve?\n\nLIMITATIONS:\nPlease refer to the Weakness part in the previous section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}