Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent

Response to Reviewer WYm4

Thank you for your time reviewing our paper and your positive feedback! If our responses resolve your questions, we'd appreciate your consideration in raising the score.

W1/ Q1 (a): regarding what is learned by the model

Thank you for your insightful question.

What is learned: the transformer's internal algorithm, not token embeddings: First, to address the core question of what is being learned: the model learns the transformer's internal weights ($\theta=${$B_i$, $C_i$}), which collectively implement a multi-step reasoning algorithm, instead of memorizing the fixed token embeddings or the tree structure.

• Our training distribution is fixed to uniform over perfect binary trees, since the number and order of the nodes are random, this forces the model to learn the attention weights to perform abstract rules of:

  • finding the parent node: implement a "lookup" operation to find the parent of a given node in the edge list,
  • chain operations: autoregressively apply this lookup to traverse the tree,
  • identify termination/turning Points: recognize the root node to stop the backward chain and begin the forward one.

  Our convergence results show that the transformer achieves this by learning to specialize heads for different subtasks (path traversal vs. stage control) and coordinating their actions, which is highly non-trivial. In contrast, simply memorizing structure or token embeddings could not lead the training loss to 0, since any two embeddings could be a parent-child pair. This is also supported by our generalization results (Theorem 4 and Theorem 8), which demonstrate that a model trained only on perfect binary trees successfully generalizes to any tree at test time.

• Some parts of the weight matrices are fixed (in line 178-179 and 212-214), which is a standard practice in the literature for theoretical analysis of transformer training dynamics (see, e.g., [Huang et al., 2023],[Nichain et al., 2024]). This allows us to clearly trace how the attention patterns are formed without loss of optimality, i.e., a global minimum of the training loss can be achieved within this simplified parameter space, without cluttering the analysis with excess details that do not bring new insights.

Q1 (b): how can the tree structure and/or distribution be made more general?

The assumption of using a uniform distribution over perfect binary trees during training (Assumption 3) provides a balanced and symmetric data distribution, which simplifies the analysis of the training dynamics under gradient descent. However, this restricted setting is not a necessary condition for the transformer to learn the underlying path-finding algorithm. The crucial requirement for the model to learn the generalizable rule—rather than merely memorizing training instances—is that the training distribution has sufficient diversity. Specifically, the transformer can be forced to learn the algorithm as long as any pair of token embeddings has a strictly positive probability of appearing as a parent-child edge in a training sample. This ensures the model learns the rule of path traversal instead of overfitting to the specific paths seen during training.

It's possible to extend our analysis to handle more varied and less-structured trees with a proper handling of the imbalance of data. For example, our technique may combine with [Huang et al. 2023], which contains the convergence analysis of transformers when the feature distribution in in-context learning prompts is imbalanced. They show that transformers can successfully learn even with such imbalanced data, first learning from the dominant features and later from the under-represented ones.

W2/Q2 (a): regarding summary of our convergence proof

Thank you for your valuable feedback! We will include the following proof outline for Theorem 3:

• Our goal to prove that with gradient descent, the training loss converges to zero. This is achieved by showing that the learned matrix $H^{(t)} := A^\top B^{(t)} A$ converges to a diagonal matrix $\alpha I_S$. Specifically, we show there exists $T_1=\widetilde{O}\left(\frac{SN}{m\eta}\right)$ and $T_2=T_1+\widetilde{O}\left(\frac{mS}{\eta\epsilon}\right)$ that satisfy the following properties.

• At Phase I $(t = \{0, 1, ..., T_1\})$, our Lemma 5-7 guarantee that the diagonal of $H^{(t)}$, $H_{j,j}^{(t)}$ ($j\in[S]$) are strictly increasing, and the off-diagonal elements satisfy $|H_{i,j}^{(t)}\lesssim \log N/N|$ for all $i\in[S]\setminus\{j\}$.

• At Phase II $(t = \{T_1+1, ..., T_2\})$, we show that $H_{j,j}^{(t)}$ keeps on strictly increasing, and the off-diagonal elements $H_{i,j}$ are strictly decreasing and satisfy $-\frac{H_{j,j}^{(t)}}{S}\lesssim H_{i,j}^{(t)}\lesssim \frac{\log N}{N}.$

• After Phase II ($t\geq T_2+1$): $H_{j,j}^{(t)}$ (resp. $H_{i,j}^{(t)}$) keeps on increasing (resp. decreasing), and $H_{j,j}^{(t)}\gtrsim\log\left(\frac{m}{\epsilon}\right).$ Furthermore, the training loss $L_{\text{train}}(\theta^{(t)})\leq \epsilon$.

Similarly, we'll also add the following proof outline for Theorem 5 in our revised paper.

• We first define $U_l$, $V_l$ ($l=1,2,3$) as we do now in Eq.(179). Then we have for all $i,j$: $U_{l,i,j}^{(t+1)}=U_{l,i,j}^{(t)}-\eta\sum_{k=1}^{2m}\left(\delta x_{l,i,j}^{(t,k)}+\delta y_{l,i,j}^{(t,k)}\right),$ and $V_{l,i,j}^{(t+1)}=V_{l,i,j}^{(t)}-\eta\sum_{k=1}^{2m}\delta z_{l,i,j}^{(t,k)},$ where the gradients $\delta x, \delta y, \delta z$ are defined in (178) and are computed in Lemma 12. This allows us to analyze the dynamics of $U_l^{(t)}$ and $V_l^{(t)}$ separately. Moreover, the symmetry/special structure of the training distribution allows us to rewrite $U_1,U_2,V_1,V_2$ to the matrix forms given in (182) and (183) (note that there shouldn't be subscripts $i,j$ in (182) and (183), we'll fix this typo). Regarding $U_3^{(t)}$ and $V_3^{(t)}$, Lemmas 18 and 13 guarantee that $U_{3}$, $V_3$ can be written into the matrix forms given in (261) and (213), respectively.

• In Section C.3.2 we analyze the training dynamics of $C_1,C_2,C_3$ with three phases: Phase I.1-C ($t\in[T_1^C]$), Phase I.2-C ($t\in[T_1^C,T_2^C]$) and Phase II-C ($t\in[T_2^C,T_3^C]$), by tracking changes of the scalars defined in the matrix forms of $U_l, V_l$ ($l\in[3]$). The training dynamics of the three phases are respectively summarized in Lemmas 14, 15, and 16. At the end of this section, we give Lemma 17 to prove the convergence of $\sum_{k=1}^{2m}\Delta_z^{(k)}$, i.e., we show that after $t>T_3^C$, $\sum_{k=1}^{2m}\Delta_z^{(k)}\leq \epsilon/3$.

• In section C.3.3, we analyze the training dynamics of $B_1,B_2,B_3$ with three phases: Phase I.1-B ($t\in[T_1^B]$), Phase I.2-B ($t\in[T_1^B,T_2^B]$) and Phase II-B ($t\in[T_2^B,T_3^B]$), where the change of key scalars are concluded in Lemmas 19, 20 and 21, respectively. At the end of this secion, we give Lemma 22 that shows $\sum_{k=1}^{2m}\left(\Delta_x^{(t,k)}+\Delta_y^{(t,k)}\right)\leq \frac{2\epsilon}{3}.$

• Theorem 5 is obtained by combining Lemma 17 and Lemma 22.

Q2 (b): regarding nontrivial parts of the analysis and extension to other CoT setups.

• Nontrivial parts of the analysis: The primary novelty and non-triviality of our analysis lie in tackling the complexity of multi-stage reasoning and the coordination it requires, which has not been addressed in prior theoretical work on CoT training dynamics.

  • Two-stage reasoning in a single autoregressive pass: To our knowledge, this is the first theoretical analysis of how a transformer, using a single set of parameters, learns to execute two distinct, sequential subtasks within a single generation. This goes beyond existing analyses of single-algorithm emulation (e.g., parity checking or linear regression). The model must not only learn the algorithm for each subtask but also learn to autonomously identify the "turning point" (line 224) and switch its behavior accordingly.

  • Head specialization and coordination: our work provides a provable mechanism for how multi-head attention enables this complex reasoning. We show how two attention heads learn to specialize: one head becomes a path-traversal head, while the other becomes a stage-control head. Analyzing their coordination is highly non-trivial, as their parameters are trained simultaneously and their outputs are coupled.

  • Complex coupled training dynamics: the forward reasoning construction involves six trainable matrix blocks ($B_i$, $C_i$ for $i=1,2,3$). Analyzing their simultaneous convergence is a significant technical challenge. The gradient for each block is influenced by the state of all other blocks, creating a complex, coupled system. To manage this, we developed a careful multi-phase analysis (as outlined in our proof sketch for Theorem 5) to decouple the dynamics and track the intricate evolution of the key parameter matrices ($U_i$, $V_i$ defined in Eq. (16) and (17)) towards their target configurations.

• Generalizable insights for other CoT setups: The core principle our work reveals is that a shallow, multi-head transformer can learn to implement a multi-stage algorithms in a single autoregressive pass. This can be extended to other COT tasks that can be decomposed into simpler, sequential subtasksm such as dynamic programming, logical deduction or fact verification.

---

References:

[Huang et al., 2023] Y Huang, Y Cheng, Y Liang. In-context convergence of transformers.

[Nichain et al., 2024] E Nichani, A Damian, JD Lee. How transformers learn causal structure with gradient descent.

[Ahn et al., 2023] Transformers learn to implement preconditioned gradient descent for in-context learning.

[Yang et al., 2024] T Yang, Y Huang, Y Liang, Y Chi. In-context learning with representations: Contextual generalization of trained transformers.

Response to Reviewer sUWo(5)

Thank you very much for your valuable feedback and positive evaluation to our paper! The two comments of yours in the Weaknesses regarding how to extend our analysis to more practical scenarios are highly relevant and thus we address both points together in the following.

We believe our conclusions offer a strong foundation that can be extended to more practical scenarios, albeit with increasing analytical complexity:

• More complex components (e.g., layer normalization): Incorporating modules like LayerNorm would certainly make the gradient expressions more complex, as it introduces dependencies on the statistics of the activations. However, the fundamental analytical technique of tracking the evolution of key quantities (like $H:=A^\top B A$ (line 295) for our backward reasoning task and $U_i$, $V_i$ defined in Eq.(16), (17) for our forward reasoning task) through gradient dynamics would still apply. It is a tractable, though more involved, next step.

• Analysis with optimizers: Our analysis of vanilla gradient descent serves as the theoretical baseline. Extending this to optimizers like Adam is a well-defined research direction. It would require tracking not just the gradients but also their first and second moments. Theoretical analyses of Adam exist (e.g., [Chen et al., 2018]), and integrating them with our framework would be a valuable, though non-trivial, endeavor to understand how adaptive methods might accelerate or alter the learning phases we identified.

• Initialization:

  • Fixed Weights: Fixing parts of the weight matrices is a standard and principled approach in the literature on transformer training dynamics (e.g., [Huang et al., 2023], [Nichani et al., 2024]). This simplification is crucial for isolating and clearly tracing how specific mechanisms, like attention patterns, are learned. It allows us to prove that a global minimum can be reached within a well-defined parameter space. Relaxing this assumption would lead to a more complex, coupled system of matrix dynamics, which is a direction for future work.
  • Random/Pre-trained Start: Analyzing training from a random initialization is a common extension in theoretical studies (see, for example, [Yang et al., 2024], [Du et al., 2019]) and could be integrated with our work. Furthermore, starting from a well-chosen pre-trained state may actually simplify the analysis by allowing the model to bypass the initial learning phases, providing a more direct path to the specialized solution.

While a full, end-to-end analysis of a large, pre-trained model remains beyond the current reach of tractable theory for the entire field, our work pushes the frontier of what can be formally understood by being the first theoretical analysis to demonstrate how a multi-head transformer can learn to execute distinct, sequential subtasks within a single autoregressive pass.

---

[Yang et al., 2024] T Yang, Y Huang, Y Liang, Y Chi. In-context learning with representations: Contextual generalization of trained transformers.

[Chen et al., 2018] X Chen, S Liu, R Sun, M Hong. On the convergence of a class of adam-type algorithms for non-convex optimization.

[Du et al., 2019] S Du, J Lee, H Li, L Wang. Gradient Descent Finds Global Minima of Deep Neural Networks.

Response to Reviewer LK9v

Thank you for your time reviewing and your valuable feedback. If our responses resolve your questions, we'd appreciate your consideration in raising the score. Please don't hesitate to request any further clarification.

W1: 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.

Thank you for the comment. It seems the question regarding the multi-head mechanism is incomplete. We would be more than happy to respond once the full question is clarified.

Based on the partial question, we would like to clarify that the multi-head mechanism primarily facilitates the automation of task switching between backward and forward reasoning. However, it is the chain-of-thought mechanism that enables the shallow transformer to perform multi-step reasoning. Specifically, each step in the chain decodes one additional node in the path toward the root during backward reasoning, or one additional node in the reverse path toward the goal node during forward reasoning.

W2: the experimental settings and the tasks are weak...

• On the "weakness" of the task and setting:

  • Our goal is not to achieve state-of-the-art performance on a complex benchmark, but to create a controlled and clean setting where we can provably analyze the learning dynamics of a fundamental reasoning pattern. As we state in our introduction (lines 25), our work is motivated by the "emergence of reasoning ability" where extending intermediate reasoning steps (i.e., CoT length) improves performance.

  • Our key research question, inspired by empirical findings in works like [Brinkmann et al. 2024], is whether a shallow model can solve a complex, multi-stage task if provided with sufficient CoT "scratchpad" space. The forward reasoning task, which requires two distinct sub-steps (backward path-finding then reversal), is an ideal testbed for this. Showing that a one-layer transformer can solve this task—a task that might otherwise require deeper architectures (see [Brinkmann et al. 2024] for details)—is precisely the central, and we believe surprising, result of our paper (lines 92-102, 251-254).

• On the lack of evidence for head specialization and coordination:

  • Empirical validation: While we did not conduct an explicit ablation study, our current experimental results presented in Appendix E (Figures 8 and 9) have already provided strong evidence supporting the claim that different attention heads learn to specialize and coordinate autonomously. This is illustrated particularly in Figure 9 (page 90), where we plot the training dynamics of the projection matrices $U_i$ and $V_i$ for $i = 1, 2$, corresponding to attention heads 1 and 2. As shown in the figure, the entries of these matrices converge toward values that align with those theoretically constructed and proved for solving the backward and forward reasoning steps (as explained further below). This demonstrates that the two heads learn to specialize in distinct sub-tasks and coordinate effectively, where head 1 handles the prediction of next node on the path, head 2 identifies the task switch from backward to forward. Importantly, this behavior emerges naturally from training via standard gradient descent, without any hard-coding enforcing such specialization. This provides direct, quantitative evidence that the heads learn their distinct, coordinated roles autonomously during training.

  • Construction (existence): In Section 3.2, we provide an explicit construction for a two-head transformer that solves the forward reasoning task. As detailed in Theorem 2 and the subsequent discussion (lines 237-250), we demonstrate precisely how the heads must specialize: head 1 (path traversal) is responsible for recursively finding the next node in the path, and head 2 (stage controller) manages the reasoning stage, tracking whether the model is in the initial "goal-to-root" phase or the subsequent "root-to-goal" phase, and triggers the transition.

  • Optimization (learnability): In Section 4.2, we further prove that these specialized roles are not just hypothetically possible, but are learnable via gradient descent. As mentioned in line 317, our convergence analysis and Theorem 5 show that the model parameters provably converge to the specialized matrices specified in our construction Eq.(10), (11). This is a highly non-trivial result that forms the technical core of our paper. As far as we know, we are the first to theoretically analyze how a transformer use the same parameters to do two different subtasks.

W3: no discussion on multi-layer transformers

Thank you for your comment.

• As we articulate in our introduction, instead of focusing on the scaling of model depth, we focus on the scaling of CoT length. The reviewer's question about a two-layer model is insightful, but it implicitly favors the "scaling depth" paradigm. Our paper's contribution is to show, provably, that the "scaling CoT" paradigm is a viable alternative for complex reasoning, even in shallow models. Therefore, extending our analysis to deep models would address a different research aspect and depart from our current central message. Thus, we will leave this direction for future exploration.

• A rigorous theoretical analysis of the training dynamics of deep, non-linear transformers is a notoriously difficult, and largely open, problem in the field. Our work pushes the frontier of what is theoretically tractable. By analyzing a one-layer, multi-head model, we are able to provide the first theoretical analysis of how a transformer can learn to execute two distinct subtasks sequentially using the same set of parameters within a single autoregressive pass. This is a significant and non-trivial contribution.

Q1: regarding experiments that compare the results of single-head and multi-head transformers

Thank you for this question. We speculate that you mean Theorem 5, not theorem 3.

• We did not include experiments comparing single-head and multi-head transformer performance because it is not straightforward to implement a single-head transformer that can autonomously switch between backward and forward reasoning. In our attempts, such experiments consistently failed to converge. This is not surprising. At a high level, it is highly non-trivial for a single attention head with fixed parameters to reliably perform the two distinct and concurrent subtasks of path traversal and stage control. Without any structural guidance, it is unclear how a single head could coordinate these subtasks effectively. In contrast, our results demonstrate that a two-head architecture provides a natural and learnable decomposition of the problem, with each head specializing in one of the subtasks. This specialization emerges purely from training dynamics, and highlights the power of multi-head attention in enabling implicit task coordination.

• As mentioned in our response to Weakness 1, our experiments in App. E directly supports our Theorem 5: Figure 8 confirms that the training loss for the forward reasoning task converges to zero, showing the two-head model successfully learns to solve the task; Figure 9 provides the crucial mechanistic evidence. By plotting the learned parameter entries over time, it shows that the attention heads' parameters indeed converge to the specialized matrix structures we constructed and analyzed. This empirically validates our theory that the heads learn their distinct, coordinated roles during training.

Q2: 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?

Thank you for your question. You are absolutely right to point out that architectural depth is a key mechanism for enabling complex reasoning in transformers. Your question highlights a crucial dichotomy in how transformers can achieve such capabilities: either by leveraging architectural depth or by utilizing longer intermediate reasoning steps (i.e., Chain-of-Thought length). Our work is specifically designed to investigate the latter paradigm, complementary to the empirical study [Brinkmann et al. 2024], which show that deeper models can indeed learn to perform multi-step path-finding implicitly across their layers.

---

Reference:

[Brinkmann et al. 2024] J Brinkmann et al. A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task.

Thank you very much for your follow-up question.

Our central argument is that the benefit of the multi-head mechanism in this specific, complex reasoning task stems from its ability to support architectural specialization, rather than simply an increase in parameter count. The transition from one head to two is not a matter of the model "suddenly" working due to more parameters, but rather a qualitative architectural change that enables different heads to learn distinct, complementary roles—in this case, path traversal and stage control. A single attention mechanism is tasked with the highly non-trivial challenge of learning both functions simultaneously, which appears to be significantly more difficult for gradient-based optimization.

To provide concrete evidence for this claim, we conducted a comparison experiment as you suggested. To ensure this model had sufficient capacity and was not unfairly constrained, we made its entire $W^{KQ}$ matrix and its $W^V$ matrix (they are both $(2d_1+d_2)$-by-$(2d_1+d_2)$ matrices) fully learnable, giving it more trainable parameters and flexibility than our structured two-head model. All other experimental settings were kept identical to those described in Appendix E.

The training loss progression across different training steps is as follows:

As the results clearly demonstrate, our training loss converges to 0 while the single-head transformer's loss stops decreasing around a high value 10.

Thank you again for your response! If our response has satisfactorily addressed your concerns, we would be sincerely grateful if you could consider updating your rating to reflect that. Of course, we are more than happy to address any further questions you may have.

Response to Reviewer WaDb

Thank you for your detailed review. We want to state firmly that all theorems and proofs in this work are our own original contributions. We hope our detailed responses below, including the one-paragraph missing proof of Theorem 1 (which was regrettably omitted) and clarify the other points, will resolve these concerns. If our response rectifies your primary concerns, we would be grateful if you would consider re-evaluating the paper and raising your score. We're happy to answer your further questions if any.

1.No proof is given for Theorem 1

We sincerely apologize for our accidental omission of the proof of Theorem 1 in our submitted version. Because the proof is short, we originally included it in the main text but then, due to page limit, decided to move it to the appendix, but apprarently neglected doing so. The proof follows directly by applying Lemma 1 and a brief further argument. For the reviewer’s reference, we present the proof below and will include it in the revised version. We appreciate your careful review.

When Assumption 1 holds, the existence of B satisfying (6) follows from Lemma 1. Letting $\alpha\rightarrow +\infty$, it follows that $softmax(A^\top B A)=I_S$ converges to an identity matrix. By (4), we know that at the $k$-th reasoning step, softmax($Y^{(k-1)\top}Bx_{-1}^{(k-1)}$) becomes a one-hot vector with the $j$-th entry being 1 if $y_j^{(k-1)}=x_{-1}^{(k-1)}$ and 0 otherwise (note that $x^{(k-1)}_{-1}$ will appear exactly once in $Y^{(k-1)}$).

Therefore, the model's output $\hat{o}^{(k)}$ will be the embedding of the edge that connects the current node $x_{-1}^{(k-1)}$ with its parent $x_j^{(k-1)\top}$. Thus by induction, $\widehat{o}^{(k)} = o^{(k)}$ for all $k\in[m]$.

---

2. regarding the existence of parametrisations and limit of the softmax in Theorem 2

Thanks for the question. The proof of Theorem 2 is correct. We clarify the reviewer's concerns as follows.

Regarding the first point, the existence of the parameter matrices {$B_i$, $C_i$} that satisfy the constraints in (10) and (11) is guaranteed by Lemma 1, given the linear independence of embeddings from Assumption 2. We will add a sentence at the beginning of the proof to explicitly reference Lemma 1. Thank you for raising this.

Regarding the softmax under the limit, the reviewer seems to have a misunderstanding about the behavior of the softmax function in the limit. The core of our proof relies on the relative growth rates of the logits, not just the fact that they go to infinity. Specifically:

• softmax principle: given any vector $x=(x_1,...x_n)^\top\in R^n$, assuming $x_1>x_i$ for any $i\in\{2,...,n\}$, then softmax($\alpha x$)$\to (1,0,...0)^\top:=e_1$ when $\alpha\to+\infty$.

• Application to our proof: In our construction (Eq. 27-33), all logits for $\mu^{(0)}$ are scaled by $\alpha_1$. The coefficients of $\alpha_1$ are designed to have a unique maximum. Specifically, the coefficient for the correct parent edge is $1 + b_2$ (from Eq. (27)), which, under our stated constraints, is strictly greater than the coefficients for any other token (e.g., $b_2$ from Eq. (29)). Because the logit for the correct edge grows strictly faster than all others, the softmax output correctly converges to a one-hot vector. The same principle applies to $\nu^{(0)}$ converging to $e_{l+2}$. Thus, the proof holds.

3. mathematical derivations of Theorems 3 and 4 do not go anywhere...strong indication that modern generative AI was used to construct those "proofs"

We respectfully disagree with the reviewer’s assessment of the proofs of Theorems 3 and 4. These are rigorous, original proofs derived by the authors without AI assistance.

First, as clarified in our responses above, Theorems 1 and 2 are both valid and can be rigorously proven. Therefore, Theorems 3 and 4 (among other) are built on a solid theoretical foundation.

The proofs beginning at Line 526 are correct and lead to the conclusions stated in Theorems 3 and 5. To help clarify any confusion, we provide below an outline of the analytical roadmap that underpins our approach.

• On the Structure of the Proofs for Theorems 3 and 5: We acknowledge that the detailed derivations for the training dynamics in the appendix are dense. The lengthy details are necessary to ensure full mathematical rigor. The derivations follow a carefully developed multi-phase analysis to track the evolution of the model's parameters from a zero initialization to a configuration that solves the task. Such multi-phase analysis is necessary due to the complex, nonconvex dynamic of the optimization problem.

• To make the logic transparent, we'll include the following proof outlines in Appendix B to summarize the key steps for Theorem 3:

  • Goal: to prove that with gradient descent, the training loss converges to zero. This is achieved by showing that the learned matrix $H^{(t)} := A^\top B^{(t)} A$ converges to a diagonal matrix $\alpha I_S$. Specifically, we show there exists $T_1=\widetilde{O}\left(\frac{SN}{m\eta}\right)$ and $T_2=T_1+\widetilde{O}\left(\frac{mS}{\eta\epsilon}\right)$ that satisfy the following properties.

  • At Phase I $(t = \{0, 1, ..., T_1\})$, our Lemma 5-7 guarantee that the diagonal of $H^{(t)}$, $H_{j,j}^{(t)}$ ($j\in[S]$) are strictly increasing, and the off-diagonal elements satisfy $|H_{i,j}^{(t)}\lesssim \log N/N|$ for all $i\in[S]\setminus\{j\}$.

  • At Phase II $(t = \{T_1+1, ..., T_2\})$, our Lemma 8,9 guarantee that $H_{j,j}^{(t)}$ keeps on strictly increasing, and the off-diagonal elements $H_{i,j}$ are strictly decreasing and satisfy $-\frac{H_{j,j}^{(t)}}{S}\lesssim H_{i,j}^{(t)}\lesssim \frac{\log N}{N}.$

  • After Phase II ($t\geq T_2+1$): $H_{j,j}^{(t)}$ (resp. $H_{i,j}^{(t)}$) keeps on increasing (resp. decreasing), and $H_{j,j}^{(t)}\gtrsim\log\left(\frac{m}{\epsilon}\right).$ Furthermore, the training loss $L_{\text{train}}(\theta^{(t)})\leq \epsilon$.

• Similarly, we'll also add the following proof outline for Theorem 5 in our revised paper.

  • We first define $U_l$, $V_l$ ($l=1,2,3$) as we do now in Eq.(179). Then we have for all $i,j$: $U_{l,i,j}^{(t+1)}=U_{l,i,j}^{(t)}-\eta\sum_{k=1}^{2m}\left(\delta x_{l,i,j}^{(t,k)}+\delta y_{l,i,j}^{(t,k)}\right),$ and $V_{l,i,j}^{(t+1)}=V_{l,i,j}^{(t)}-\eta\sum_{k=1}^{2m}\delta z_{l,i,j}^{(t,k)},$ where the gradients $\delta x, \delta y, \delta z$ are defined in (178) and are computed in Lemma 12. This allows us to analyze the dynamics of $U_l^{(t)}$ and $V_l^{(t)}$ separately. Moreover, the symmetry/special structure of the training distribution allows us to rewrite $U_1,U_2,V_1,V_2$ to the matrix forms given in (182) and (183) (note that there shouldn't be subscripts $i,j$ in (182) and (183), we'll fix this typo). Regarding $U_3^{(t)}$ and $V_3^{(t)}$, Lemmas 18 and 13 guarantee that $U_3$, $V_3$ can be written into the matrix forms given in (261) and (213), respectively.

  • In Section C.3.2 we analyze the training dynamics of $C_1,C_2,C_3$ with three phases: Phase I.1-C ($t\in[T_1^C]$), Phase I.2-C ($t\in[T_1^C,T_2^C]$) and Phase II-C ($t\in[T_2^C,T_3^C]$), by tracking changes of the scalars defined in the matrix forms of $U_l, V_l$ ($l\in[3]$). The training dynamics of the three phases are respectively summarized in Lemmas 14, 15, and 16. At the end of this section, we give Lemma 17 to prove the convergence of $\sum_{k=1}^{2m}\Delta_z^{(k)}$, i.e., we show that after $t>T_3^C$, $\sum_{k=1}^{2m}\Delta_z^{(k)}\leq \epsilon/3$.

  • In section C.3.3, we analyze the training dynamics of $B_1,B_2,B_3$ with three phases: Phase I.1-B ($t\in[T_1^B]$), Phase I.2-B ($t\in[T_1^B,T_2^B]$) and Phase II-B ($t\in[T_2^B,T_3^B]$), where the change of key scalars are concluded in Lemmas 19, 20 and 21, respectively. At the end of this secion, we give Lemma 22 that shows $\sum_{k=1}^{2m}\left(\Delta_x^{(t,k)}+\Delta_y^{(t,k)}\right)\leq \frac{2\epsilon}{3}.$

  • Theorem 5 is obtained by combining Lemma 17 and Lemma 22.

We sincerely thank the reviewer once again for your time and effort. We hope that our responses above have assured you the originality and rigor of our proofs. If so, we would greatly appreciate it if you would consider re-evaluating the paper and possibly raising your score. Of course, we would be more than happy to address any further questions or concerns you may have.

Response to Reviewer Avae

Thank you very much for your insightful comments and your appreciation to our work!

W1: The backward pass process constructs a transformer adapting copying mechanism, which is well studied in many theoretical analysis papers [1][2].

Thank you for your insightful feedback. We agree that the core mechanism in the backward reasoning task—sequentially copying parent nodes—is conceptually related to well-studied copying or induction head mechanisms [1, 2]. This analysis provides a clear, simplified setting to introduce our problem formulation and the basic mechanics of single-head reasoning on trees, making the readers easier to understand our paper.

W2: In the construction for forward reasoning, each position has two nodes and a flag, which seems too artificial.

Thank you for this insightful comment. Our structured output, which includes a stage flag, is a deliberate design choice. It allows us to mechanistically prove how even a one-layer transformer can learn to autonomously switch between subtasks. While a deeper model like the one in [Brinkmann et al., 2024] might manage this implicitly across layers, our work shows how this control can be explicitly learned as a specialized head function within a shallow architecture, providing a clear illustration of multi-stage reasoning.

Q1: I 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?

Thank you for this excellent point. You are correct that the current plot only shows the initial training phase. We will update this figure in our revision to show the complete loss curve until convergence to 0.

---

Reference:

[Brinkmann et al. 2024] J Brinkmann et al. A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task.