Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
Authors: Tong Yang
Yu Huang
Yingbin Liang
Yuejie Chi
Venue: NeurIPS 2025 poster ()
Keywords: transformer theory
training dynamics
multi-step reasoning
Primary area: theory

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.

Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
Tong Yang* CMU
Huang Yu' UPenn
Yingbin Liang OSU
Yuejie Chili Yale
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 theoreti- cal standpoint. This work investigates how transformers learn to solve symbolic multi-step reasoning problems through chain-of-thought processes, focusing on path-finding trees. We analyze two intertwined tasks: 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 goal-to-root path and then reversing it to produce the root-to- 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 generaliza- tion guarantees to unseen trees. 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 single autoregressive path. These results provide mechanistic explanation of how trained transformers can implement sequential a algorithmic procedures. Moreover, they offer insights into the emergence of reasoning abilities, suggesting that when tasks are structured to take intermediate chain-of-the of-touht- steps, even shallow multi-head transformers can effectively solve problems that would otherwise require deeper architectures.
Introduction
Transformers [30]-the building blocks of large language models (LLes)-have demonstrated impressive capabilities tasks that require multi-step reasoning [32], where LLMs start to excel in solving hard problems via taking simpler step-by-step actions, and executing different subtasks within a single response. Many of these capabilities are further fueled by the remarkable phenomenon called emergence of reasoning ability, where simply extending the length of intermediate reasoning steps can dramatically boost accuracy [10].
Spurred by their remarkable success, understanding how transformers enable Chain-of oiee-off--Thhtttt (CoT) and multi-step reasoning has attracted considerable research attention recently. Most the theoretical studies the multi-step reasoning mechanism focus on the expressive power [9, 23, 25, 5], statistical properties [12, 28, 22] or learnability [1, 11, 34, 29, 19, 2] of the CoT mechanism. Notwithstanding, the training theory of transformers, which examines the optimization and generalization properties of
transformers trained via gradient-based methods, another important direction for understanding
DDepartment Electrical and Computer Engineering.
Carnegie Mellon
University; email:
tongyang@@ndrewwcom... andrew.emu.dde
Department of Statistics and Data Science, Wharton School, University of Pennsylvania: email:
yuh42 @wwrttnn..nn. wetharron.upen.ede.. upenn.edu
Department of Electrical and Computer Engineering, The Ohio State University; email: liang @ odu Departmeet Statisties Data Science, Yale University; email: yuejie.chi@yall yalleed
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Prompt:
Path-finding task: Goal-to-root (backward): + 1 - 2 +
Edge set: - 8, 4
7,3 5, > 1,
-2
Goal:
Root-to- tooaooooollloooo (forward))
Root:
4 2
Figure 1: An illustration of the path- auth-finding reasoning task in tree. Solving the forward task (finding the root-to-gost path) requires solving the backward task (finding the goal-to-root path) first.
the CoT mechanism. This line of work is more closely aligned with practical implementations, as it investigates how training enables transformers develop multi-step reasoning capabilities through COT. A small but growing number of studies have recently explored this direction, focusing on tasks such in-context learning, algorithmic emulation, and parity checking (e.g., [21, 33, 19, 13, 14]).
However, existing theoretical studies do not fully capture the complexity present in real-world appli- cations that transformers are expected to handle through CoT reasoning: (i) Graph-based structural complexity: Reasoning over structured relational dependencies introduces a richer level of complexity. It requires transformers to perform multi-hop traversal and extract rule-based generalizations, making it more representative of realistic symbolic reasoning scenarios. (ii) Mulli-stage reasoning and aulonomoux stage transition: In real-world CoT reasoning, achieving complex reasoning objective often involves performing multiple consecutive reasoning tasks. Crucially, the model must learn to autonomously determine when to switch between stages, reflecting hierarchical and self-regulated reasoning process.
1.1 Symbolic multi-step reasoning: path finding on trees
In response to the above gap, this paper resorts to a canonical symbolic multi-step reasoning problem of path-finding on frees, motivated by the empirical investigation in [4]; Figure 1 for illustration. The input prompt is the list of edges in the tree, the root node, and the goal node (which is a leaf node). We are interested in two fundamental reasoning tasks: (i) backward reasoning, which involves finding the path from the goal node to the root node; and (ii) forward reasoning, which aims to find the path from the root node to the goal node. While the backward reasoning task can be solved by sequentially chaining edges connected to the child node recursively, tackling the forward reasoning task requires solving the backward reasoning task (i.e., finding the goal-to-roof path) first and then reversing the path, as there can be multiple nodes connected parent node. Studying these two intertwined tasks has the potential to elucidate deeper understandings of CoT on the following front.
How CoT implements chaining over graphs. Both tasks require step-by-step reasoning, as one needs trace the edges connecting the goal node and the root node in recursive manner.
How CoT implements subtask control. The harder, forward reasoning task requires solving two subtasks in a single CoT path, where the transformer needs to automatically identify the turning point between the subtasks, when it traverses the root node.
The empirical study in [4] revealed that deep transformers, when tasked with forward reasoning (i.e.. outputting the poot-to-goal path), internally implement a backward chaining algorithm across their layers that effectively from the goal node, "climbs" the tree layer by layer towards the root node. The discrepancy between the internal reasoning direction (i.e., backward) and the required output format (i.e., forward) suggests that architectural depth is leveraged by these models to manage intermediate representations before producing the final answer. While inspiring, it is intriguing to ask if one can provably train shallow, one-layer, transformers to solving the path-finding task in trees via CoT, and if extending the reasoning length-rather than using deeper architectures-can be leveraged to solve the harder forward reasoning task. Answering these questions will lead to new insights into the multi-head attention mechanism that are unavailable in today's literature.
1.2 Our contributions
This paper investigates how one-layer transformer can learn symbolic multi-step reasoning- finding a tree-by employing the CoT mechanism to make sufficient intermediate reasoning steps. We provide comprehensive analyses covering transformer constructions. training dynamics, and generalization guarantees, demonstrating that when equipped with CoT, even one-layer multi-head transformers capable to solving complex tasks that require multi-stage and multi-step reasoning. Specifically, our contributions are as follows.
2

---

Construction, We provide explicit constructions detailing how one-layer transformer can perform both tasks using COT. While single attention head is sufficient for backward reasoning. it takes two to implement forward reasoning. Our construction elucidates how different attention heads, through coordination, can specialize for distinct reasoning subtasks.
Optimization. We prove that gradient descent successfully trains the transformers to acquire multi- step reasoning capability for solving both tasks via CoT.. Our multi-phase training dynamics analysis demonstrates that the model parameters indeed converge those specified in our construction, explaining how attention heads learn their respective roles in a non-asymptotic manner.
Generalization. We demonstrate that the learned reasoning mechanisms, including the specialized head functions, generalize effectively to correctly solve path-finding problems unseen tree structures. This demonstrates that transformers learn the generalizable rule for path- futd-finding rather than memorizing paths in training graphs. We also conduct experiments that empirically validate our theoretical predictions.
Our finding offers a mechanistic explanation for how CoT can empower even shallow models to tackle tasks that seemingly require the capabilities of deeper architectures, which is particularly illuminating in the study of forward reasoning. In this setup, the one-layer transformer, utilizing two attention heads, is trained to perform two sequential subtasks: first, it must generate the path in reverse an explicit CoT sequence; then, upon identifying the root, it must sequentially
order ggoal-to-rood output the path in the correct forward order (root-to-gode -oog.. Our optimization based analysis reveals how these two heads learn to coordinate: one head is responsible for identifying the next node in the current path-finding stage, while the other head acts as a stage controller, monitoring the current reasoning phase (e-g., backward chaining) and triggering the switch to forward chaining once the root node is detected. This learned specialization allows the one-layer, multi-head transformer to perform this entire two-stage process within single autoregressive pass.
Importantly, by explicitly generating the backward path as a "scratchpad" in CoT outputs, the one- layer model can then leverage this intermediate information to produce the forward path, rather than relying on depth like in [4]. This resonates with the emergent reasoning abilities in LLMs through test-time scaling [10], where generating longer intermediate steps via CoT often unlocks more sophisticated groblem-solving capabilities. Our two-stage setup illustrates how extending the reasoning trace allows one-layer model effectively perform complex task that, without such explicit intermediate steps, might necessitate a deeper model.
1.3 Related work
Theoretical understanding of transformers with COT. A growing body of research seeks to theoretically demystify the success of CoT reasoning in transformer models. Most existing studies examine how CoT enhances transformer expressiveness [9, 23, 25, 5], showing that polynomial- length CoT enables transformers to transcend TC", a class efficiently handled by constant-depte transformers without CoT [24]. Other works identify intrinsic limitations, establishing lower bounds on necessary intermediate steps for certain problems [27, 3, 2]. Another line of research addresses the learnabllity [1, 11, 34, 29, 18] and statistical properties [12, 28, 22] of CoT-enabled transformers.
Optimization theory for COT. Recent analyses explore optimization aspects of transformers trained with CoT [21, 13, 33, 18, 13, 17J. The most closely related works are [33, 18, 14], which analyze the optimization dynamics of one-layer transformers trained with CoT to solve parity-check tasks. However, these studies are limited to the single-head setting. In contrast, our work considers multi-head transformers solving a more challenging task involving reasoning over graphs as well as auto-transition over two subtasks, involving specializations and coordination of different heads.
Training dynamics of transformers. Many existing works have investigated the training dynamics of transformers across variety of tasks, including in-context learning [15, 35, 36, 20, 6], induction heads [26, 7], sparse token selection [31], suff-sueereieed learning [16]. Due the rapid growth of this area, we are unable to reference all related studies. Particularly relevant to our work are recent analyses of CoT reasoning [21, 13, 33, 18, 14]. However, the majority of these studies are limited to the single-head setting. While a few works have considered the multi-head case, they typically address only simple tasks such as linear regression, leaving open the question of how multi-head architectures can be trained to solve more complex reasoning problems.
3

---

Notation. Let [N] {I...., Nj. For matrix A, denote A. or A(:, - aa)s its last column. In denotes the identity matrix of size n x n. Let Il-1, represent the ep norm vector.
2
Formulation
We consider path-finding task in randomly generated trees [4], defined formally as below.
Tree. We consider the trees with distinct nodes, i.e., the nodes of the trees are generated by random sampling from set [S] without replacement. Let T (V(T), (7)) g(T) denote a tree with node set V(T), edge set ( and a goal node g(T), which a leaf node in the tree. Denote the root node as rT)). For each non-root node i, i.e., ((T),, let p(i) p(il7) denote the parent node of node i. Then the edge set given by E(T) (p((), (e(CT)VT)) For k > 2, we recursively define p² (1) = p()*=)(()) i.., (i) is the k-th ancestor if i.
Path finding. Given tree T, the backward path from the goal node to the root node is given by (T) = (T) pT)( p²(((T)) - -
r(7) pmm(T) (gT)), where m(T) >/ is the length the path. Accordingly, the forward path from the root node to the goal node can be obtained by reversing the backward path, yielding, (7) r(T)
p(g(T)) - g(T). We aim understand whether and how transformers can be trained to discover the underlying rule for identifying both backward and forward paths in tree T through multi-step reasoning. Since parent node may have multiple children, determining the forward path requires the transformer to autonomously perform two-stage reasoning: first, identifying the backward path, and then reversing it. Therefore, we consider two path-finding tasks:
p²(
Backward reasoning: find the goal-to-rool path.
Forward reasoning: find the goal-to-root root path and then reverse to produce the root-to-goal path.
Input embedding. For each node i [S], we let a, Rd denote its token embedding. We embed (a a. I2dd: where ap(i) and y a;
each edge (p(i).) i) in E(T) (x 1 denote the parent and child node embeddings, respectively. We I(T) = [T)) denote the number of edges in T, and let ".....,,([[) denote the edges E(T) in random order, and let (z,,, .y, be the embedding of edge Ej.
Transformer architecture and multi-step reasoning. For both reasoning tasks, consider one-layer transformer f with H > attention heads. Given an input matrix E, the transformer with parameter 8 computes the output
head, (E)
f(E:0) 0) = wo
: 1ead (E)
head, W E softmax (ET w W&E). vh [H], (I)
where softmax(-) is the coiumn-wise softmax function and 0 denotes the trainable parameters of the transformer. For simplicity, we follow standard practice [15, 13, 26] to combine w and Wh into a single matrix W.Q We consider step-by-step reasoning in an autoregressive manner. At each reasoning step k > O, given the input embedding (T), the output vector is taken as the last column of the output matrix, which is then concatenated with the input embedding to form the new input to the next reasoning step, i.e.,
k+ +1 ( (T:0) =
(T):)).-1. (T) (k++1) ((++)((:::)) = (T:;0), ((++1) (T:0)).
(2)
Let the number of reasoning steps be (T), and then Ô(T;0) (T ((K(() denote the output matrix of the model given T.
3 Construction of Transformers
In this section, we show that there exists parameter setting 0 such that the one-layer transformer can solve both forward and backward reasoning problems through multi-step chaining and reasoning. Whenever it is clear from the context, drop the dependency T and 0 in the notation.

---

3.1
Construction for backward reasoning
For the backward reasoning task, we are interested in outputting the path from the goal node the root node (g2r). Let the input embedding of the transformer be
E = (T) =
( Y(T)
=
Φ1
2((7)
ao
ag(T)
E (
x((
(3)
VieT)
ar(T)
ag
where dg Rd is used fill the empty slots. We set the ground label of T as the embedding of the reversed path:
(gar(T) (pla(T)) (T(T)) (p(a(T))
(pP(a(T))
(r(T) apm(T) ((((T))
E IR 2d xm(
We consider one-layer single-head transformer, where H 1. We impose the following parameters: let 0o WY I2d,- and set
WKQ (Od, xd:
Od, xds Od, xds
where B e Rddidd is trainable. In this case, 0 {B).
Multi-step reasoning. We set the first step g (T), and let ((+ g ((l--)(1 d1.:)
p4x((+++b) and (k-1) == g2r E (d1 + 1 2d1,;) E Rd x(++++k) obtained recursively following (2) for 1. Then according to the transformer architecture described above, we have the output ((k) at reasoning step A: as
(((-1)) g(k) = ((--1) softmax (Yik-1)T Bx(T-"))
(5)
(k-1) is the last column ( ()), Reasoning recursively for m(T) steps according to (2), where x x=1
obtain the output (T (6(1) ..... ((m(T))
We make the following assumption on the node embeddings for the backward reasoning case.
Assumption 1 (linear independent embeddings (backward reasoning)). Suppose ag Odl, and {as hets are linearly independent.
Let A := (a1...., as) Rd xS be the embedding matrix. construction the transformer that solves backward reasoning. Theorem 1 (Construction for backward reasoning). Under Assumption for any a R, there exists B = B, E Rd xd such that
The following theorem provides
A = alg-
(6)
Let 0 (Ba), then for any tree we have Öger(T (T)r(T)) as a +00.
Theorem suggests that it is possible to learn a sharp attention pattern- -when applying softmax to A BA-between the node embeddings, where the query node attends to only itself among nodes. Figure (b) provides an illustration of the chain of input-output pairs for each backward reasoning step. This also highlights the path-finding strategy that the transformer learns during training. At each step, the transformer focuses the current node the path (starting from the goal node) and uses its attention mechanism to assign high weight to the next node (i.c., its parent) in the tree. The identified node is then selected the next step in the path and subsequently encoded into the input matrix for the following iteration, enabling the model to reason sequentially through the structure.
3.2 Construction for forward reasoning
The forward reasoning task is more challenging than backward reasoning, it requires first identify- ing the path from goal to root, and then reversing to obtain the root-to-goal (r2g) path. Crucially, the transformer must learn to autonomously determine when to switch to the second stage of reasoning.
5

---

[xtan
K3
X0
(a) task illustration
(b) multi-step reasoning for finding the goal-to-root path
edges
output
bac backward reasoning
2:31
20
aawardd Ressoni
.611)
(orward
X
15
y000 200
2:31
STAGE
Turning point
STAGE
(c) two-stage multi-step reasoning for finding the root-to-goal path
Figure 2: The multi-step reasoning process of the constructed transformers for the backward and forward reasoning tasks. Color indicates the attention association and output in each step.
To this end, we introduce two stage token embeddings 8,, 8, E Rdda used to distinguish the two reasoning stages. We include the stage token embeddings in the input as follows:
((T)) Ergg( = Y(T) ((T)
F1
2((T)
dg
au
0g(T) E (2d) (2(2d++
91
3((T)
(r(T) $f
81
8f
Sp
where 8, marks the start of the first reasoning stage. We require the model to output the path first in the backward order, then in the forward order, along with stage indicator at each step. This results in total of K(T) 2m(() steps. We set the ground truth label (T) E [Rd +dg) x2m(() of T as the embedding of both the reverse path and the forward path along with their stage indicators:
Og(T)
(p(g(T))
ar(T) a (pm(T)-(o(T))
a apm
ap(o(T)) ag(T)
Oreg(T)
ap(a(T))
(pP()(T))
a (p
(8)
Sb
86
8f
8f
8f
We consider one-layer two-head transformer, where H 2. We impose the following parameters: let wo
-(4xx
(Od, xdi Ia
ld.
Odz
Izd, +d2. WY =
W. (0d, Id,), and set
Odi
Odz
Od, xdi Odixds xdi
Bi
Od, xds
Od, xd, Od Odzxd,
C1 C2
Od, xda 0d xd2 Css
w
B2
Od xd2
and
w.Q
Odzxdi
0d
B3
Od
with trainable matrices
B1, B2, Bg
and C1. C2 Cs. respectively.
In this case,
(Bi, B2, B3, C1, C2. C3).
Multi-step reasoning. We set the first step E = (T)((T) At the reasoning step > 1, let (1 (k-1))
(kk-1) (di + 1 2dj.:) E Rd (1+1++), and (kk-1) (k-k1) (2d, 1: 2d, d2,;) :) (1+1+)+ where E (k-1) is obtained recursively
((--1) = E//0
d1.:) E (++((++)k)) (k- = Eiko

---

from (2). Then we have
( softmax (xx(t1)T -1 + y(k-1)T B2y BBy + z(k-1)T 1)T (11
(k-1)
(x(--1)
g0k)
Z(k-1). -softmax (x(x-1)TC111 (k-1) y(k-1)TC2yl (k-1) + Z(k C3=-1)CC221111
:) is the output of the first attention head, predicting the next node based on the learned path encoded in (k-1) and ( and the lower portion (2d, + 1: 2d + d2.:) is the output the second attention head, indicating whether the second stage begins.
Here, the upper portion (1: 2d1.:)
The process of predicting { is called STAGE 1, where the model outputs the reverse path from goal to root similarly to backward reasoning. The prediction of g(m(T)+1) is the turning point, because in this step the model detects the root and switches to STAGE (indicated by 8y), where the model outputs the path in forward order from root to goal. The entire process includes reasoning recursively for (T) steps, and the final output is given by Örag (T; (o(")... ...... (
We make the following mild assumption on the embeddings for forward reasoning.
Assumption 2 (linear independent embeddings (forward reasoning)). Suppose (i) ao, 00,00,,-0 01, 0g linearly independent; (ii) 85, Sb are linearly independens.
Denote À = (ag. .i..... as) E Rd ((++) and Š = (sy. 8)). The following theorem provides construction one-layer two-head transformer that solves forward reasoning.
Theorem 2 (Construction for forward reasoning). Under Assumption 2, there exist 0 {Bs Ci that satisfies
1
AT B,A aan
0
A B.A anls, 3 B3 an (-hi by
b2
-b2)
(10a)
0
ATAA=
0
0
A ATCAA asts, ST 02
(
C1
-02
(10b)
-C1
When a E [0,1], > 0, by a/2) and c1 > 0, c2 E (0,1/2), we have (T:g(T:0) - (T) as
a1 02 +00.
Specifically, let m = mT) denote the path length and (+)1)) e [S] denote the node with embedding y w-1 (k-1) for any [2m]. Then we can show that when ai 02 +00, our construction ensures that at each reasoning step k; the model performs the following (see Figure 2 (c) for an illustration):
1. when & & m, i.e., the model is STAGE 1. Our construction ensures the first attention head attends to the embedding of (p(n(--12), n (k-1)), switches the parent-child order, and outputs the embedding of (n (n()-1) p(pp(( ((n((-1))), and the second head attends to and outputs the stage signal 86;
2. when t- = m + 1, i.e., the model is at the turning point. Our construction ensures that the first attention head attends to (y (k-1).T _x_ (k=-)).T] and outputs (k-)),T ",,,,,((())))) (k-1))TTT and the second attention head attends to and outputs the stage signal 8y, changing the stage signal.
3. when & > m + 1, i.e., the model is at STAGE 2. The first attention head attends to the embedding of (n2 (k-1) ((n)(-)))) generated in STAGE 1, switches the parent-child order, and outputs the embedding (p(n (k-)),, (kk-)), and the second attention head attends to and outputs s..
This process demonstrates that the first attention head is responsible for traversing the reasoning path, while the second attention head manages the stage signal, which dictates the direction of ruth-finding.
7

---

7
10
Figure 3: An example of a perfect binary tree of depth m 2 and distinct nodes.
The importance of Theorems and 2 lies in demonstrating that even single-layer transformer suffices to carry out the above logical steps with enough CoT steps. This highlights the surprising expressive power of even shallow transformer architectures for multi-step step reasoning, in contrast to leveraging multi-layer transformers for such a multi-step reasoning task in (author?) [4].
Training Dynamics and Generalization of Transformers
In Section showed the existence of one-layer transformer capable of solving both backward and forward reasoning tasks. However, those results do not address whether such transformers can be learned in practice. In this section, we analyze the training dynamics of the transformer both settings and demonstrate that gradient descent (GD) can successfully minimize the loss to O. Upon convergence, the transformer performs the desired backward and forward reasoning. This establishes that transformers can, in fact, be trained via GD to acquire multi-step reasoning capabilities. We further show that the trained model generalizes well unseen trees.
Training, We train the transformer using standard autoregressive supervised learning. We define (T) as the k-th column of O(7) := 0g2r(7) or (T)) the ground truth label for either case. For each tree T Prrain. where Priiin is the training distribution to be described later, at each reasoning step > 2, the input (T) = (E(A=2)(7) (T), (k--)) (T)) is the concatenation of the input at step k - 1 and the (A: 1 column of the ground truth output matrix O(T). We then obtain the output
Train (T: 0) = f(E f(Eta ((EE(--)) (T): 0).-1 at step Let (T;ai = "(,
denote the output matrix of the model given T.
We train the model by minimizing the squared error loss between the output and the label. The training loss is given by
B | 1
Oraain [7:0)
train (0)
(11)
For both cases, we initialize all parameters to be zero, and we train the model using gradient descent with learning rate 1j > 0, i.e.,
g(++) = (10) - V
(0(9), (0
Vt > O.
(12)
Here, ((1) is the parameter at the f-th iteration.
Training distribution. We make the following assumption the training set Pirain-
Assumption 3 (training distribution). Let m 2 3 and S 3/ 2m++ -1 1= N. The training distribution Ptrain is the uniform distribution over all perfect binary trees T (V(T), ((T)) of depth m with distinct nodes chosen from [S] and goal node g(T) being of the leaf nodes.
We note that we fix the tree to be a perfect binary tree of depth m during training only for simplicity of analysis and presentation, and our results can be easily extended to other types of trees, Also note that under Assumption 3, the number of nodes in any tree sampled from Ptrain is N = 2 1. An example of perfect binary tree generated from Ptrain is shown in Figure 3, where set 2.
Testing. At test time, given any tree T (that may not be in the training set), we recursively input
(k- E(k-l) (T: (E(T), (k (()))) (T;0))
'The depth (or height) of tree is the length of the longest path from the root down to any leaf, measured in the number of edges.

---

into the model at the k-th reasoning step and obtain the output (T:0) = f(E() (T:0): 0).
The test loss on is given by
Lisst (To (T:0) 2 - 1
(13)
4.1 Optimization for Backward Reasoning
For case of analysis, we make the following assumption on the node embeddings for backward reasoning, which is slightly stronger version of Assumption 1.
Assumption 4 (Orthonormal embeddings (backward reasoning)). Suppose (i) do Odg and (ii) [ashee[s] are orthonormal vectors in Rd
Note that commonly used one-hot embeddings are orthonormal [8], and orthogonality of token embeddings is also assumed in (author?) [15, 26, 6, 21] for analyzing the training dynamics of transformers.
Training dynamics. We have the following theorem regarding the convergence of the training error for backward reasoning.
Theorem 3 (Convergence of backward reasoning). We initialize ((0) = ((O) = Os, xdy Un- der Assumptions 3 and 4, for any learning rale V E (0, 1] and any (0, 1]. there exists T mS log (mm)) such that for any t 3 T, we have traiin (((t)) E c.
Theorem 3 establishes that to achieve e-aceuracy, takes at most Ô (1/e) iterations, omitting the dependency with other salient parameters. Recall in Theorem (cf. Section 3.1), we construct B such that H A BA- alg, ensuring the test loss approaches zero as a +0C. our convergence analysis of GD, we track the dynamic of (+) A ((*) A. Specifically, during training, the diagonal entries of (1()) continue grow while the off-diagonal entries remain small, indicating (C)) indeed converge toward our constructed solution.
Generalization. We further provide the following generalization result, which shows that the model can perform well on unseen trees.
Theorem 4 (Generalization of backward reasoning). Under the same setting as in Theorem 3, given any tree T with N distinct nodes chosen from [S] and a path length mn there exists small enough co (m, N, m) such that for any (0, ca). there exists T Ö mS such that after training the model for t 2 T steps on the training loss, have
N N + m -
Liear (T; (())) & 4 max (1,
E. m
Theorem 4 implies that for any tree T (including those not seen during training), the test loss converges to zero, provided the model sufficiently trained on Prrin This reveals an appealing insight: the transformer learns rule for path finding, rather than memorizing training examples.
4.2 Optimization for Forward Reasoning
Similar to Assumption we make the following assumption for analyzing the training dynamics of forward reasoning.
Assumption 5 (Orthonormal embeddings (forward reasoning)). Suppose (i) ag. a....., ag are orthonormal vectors in R:d;; and (ii) 8f; 86 are orthonormal vectors in Rdd
Training dynamics. The following theorem guarantees the convergence of GD on the forward reasoning task.
Theorem (Training dynamics of forward reasoning). We initialize g(0) {B(", " ,,,, all by zero. Suppose Assumptions and 5 hold, and N is sufficiently large. and learning rate nj & mN. Then for any n (0, mo] and Eo). where

---

No O and eo 1/poly(N) are sufficiently small, there exists T O such that for any 3 T, we have Ltrain &
( (m)) 3/2 log (
(m)
Theorem establishes that achieve e-aceuracy, it takes most Ô (1/63/2) iterations, omitting the dependency with other salient parameters. Define the following trainable matrices, which can be viewed counterparts to those constructed in (10a) and (10b):
U = AT (TT A RR( U.D) A B² A R
Ug ;- S BBg
B 5 R2222 (14a)
v = A ct ( Ae (($+1)x2 V A (t) Ae R Vi" ST c S R (14b)
Our convergence analysis reveals that t7(h and 1/ 1 3, converge to the matrices given in the construction (10a) and (10b), respectively. In particular, show that
U (resp. V the first row decreases (resp. increases), and the other entries stay small. For
For both U and V2( the diagonal entries are strictly increasing, and the other entries stay small.
For U and () we have
3,0,0 -
3,1,0
3,0,1 =
(1) 3,1,1
(+)) =
3,0,0
-v(()
3,1,0
V((1)
3,0,1
=
-v=(t)
3,1,1
(15)
addition, U,0,1 first decreases and then increases, U20000 (resp. 13,0,) is strictly decreasing
(resp. increasing). and V3,1,1 also grows large, but not monotonically.
At convergence, (tt)) (z(() are close to the matrices in (10a) and (10b) with large enough ar Or2, close to 1, bj > by 2 1/4, and 01 > > 1/4.
Generalization. We further provide the generalization guarantee for forward reasoning. We show that the model performs well unseen trees at test time, similarly to the generalization result in Theorem established for backward reasoning.
Theorem 6 (Generalization of forward reasoning). Under the same setting as in Theorem 5, given any tree T at test time with N distinct nodes chosen from [S] and path length m. there exists small
enough = =o(m, N, m) > O such that for any E (0, eol, there exists T = 5 m3323) such that after training the model for t > T steps on the training loss, we have
Ñ + 2m
2
N N
Liest TT;0) & 4 4mxx
1 C.
(16)
N + 2m
We also conduct experiments to validate the theoretical results. Due to the space limit, we refer the reader Appendix A.
Conclusion
This work provides a theoretical investigation into how one-layer transformers learn symbolic multi- step reasoning via gradient descent, focusing on path-finding trees. We demonstrate through explicit constructions that transformers can provably solve both backward (goal-to-root) and the more complex, two-stage forward (root-to-goal) reasoning tasks by employing CoT mechanisms. Our optimization analysis reveals that gradient descent successfully guides the model parameters these effective configurations, with multi-head attention mechanisms learning to specialize and autonomously coordinate for distinct subtasks- -uch as path traversal and stage control-within a single autoregressive pass. We also show that these learned abilities generalize to unseen tree structures, indicating the transformer acquires underlying algorithmic rules, rather than memorizing instances. These findings offer mechanistic interpretation of how shallow transformers leverage CoT for sequential, stati-stage reasoning procedures.
10

---

Acknowledgments and Disclosure of Funding
The work T. Yang and Y. Chi is supported part by the grants NSF CCF-2007911, DMS-2134080 and ONR 100014-1911-2444. T. Yang is also supported by the Wei Shen and Xuehong Zhang Presidential Fellowship at Carnegie Mellon University. The work of Y. Liang was supported in part by the U.S. National Science Foundation under the grants ECCS-251540 and DMS-2134145.
References
[1] Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can transformers reason? the globality barrier and inductive scratchpad. Advances in Neural Information Processing Systems, 37:27850-27895, 2024.
[2] Alireza Amiri, Xinting Huang, Mark Rofin, and Michael Hahn. Lower bounds for chain-of- thought reasoning hard-attention transformers. arXiv preprint arXiv:2502.02393, 2025.
[3] Pablo Barceló, Alexander Kozachinskiy, and Tomasz Steifer. Fhrenfeucht-haussle rank and chain thought. arXiv preprint arXiv:2501.1299) 2025.
[4] Jannik Brinkmann, Abhay Sheshadri, Victor Levoso, Paul Swoboda. and Christian Bartelt. A mechanistic analysis of transformer trained on a symbolic multi-step reasoning task. arXiv preprint 21Xv:22422.11917, 2024.
[5] Lijie Chen, Binghui Peng, and Hongxun Wu. Theoretical limitations of multi-layer transformer. arXiv preprint arXiv:24/2.02975. 2024.
[6] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality, arXiv preprint 2XXvv:2402.19442, 2024.
[7] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. arXiv preprint 2xXi::2409.10559) 2024.
[8] Krzysztof M Choromanski, Mark Rowland, and Adrian Weller. The unreasonable effectiveness of structured random orthogonal embeddings. Advances neural information processing systems, 30, 2017.
[9] Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36770757-70798, 2023.
[10] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12945 2025.
[11] Michael Hahn and Navin Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.00 23077777777 2023.
[12] Xinyang Hu, Fengzhuo Zhang, Siyu Chen, and Zhuoran Yang. Unveiling the statistical founda- tions of mfain-of-thoughh prompting methods. arXiv preprint arXiv:2408.145// 2024.
[13] Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gradient descent with chain of thought. arXiv preprint 2rX1v:250..12122.2000
[14] Ruiquan Huang, Yingbin Liang, and Jing Yang. How transformers learn regular language recognition: theoretical study on training dynamics and implicit bias. arXiv preprint 25K5:2505000920 2025.
[15] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.

---

[16] Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. A theoretical amalysis of self self-supervised learning for vision transformers. In The Thirteenth International Conference on Learning Representations, 2025.
[17] Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi, and Yuxin Chen. Transformers provably learn of-nn- reasoning with length generalization. In Advances in Neural Information Processing Systems, 2025.
[18] Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. arXiv preprint arXiv:2410.08633, 2024.
[19] Juno Kim, Denny Wu, Jason Lee, and Taiji Suzuki. Metastable dynamics of chain-of-thought reasoning: Provable benefits of search, RL. and distillation. arXiv preprint arXiv:2502.01694, 2025.
[20] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yt Chen. How do nonlinear transformers learn and generalize in in-context learning? arXiv preprint arXiv:2402. 15607, 2024.
[21] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. Training nonlinear transformers for chain-of-thought inference: A theoretical generalization analysis. arXiv preprint 2rXiv:2410.02167, 2024.
[22] Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of thought Compositionality through in-context filtering and learning. Advances in Neural Information Processing Systems, 36:22021-223 1220466 2023.
[23] Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint arXiv:2402.12875. 1, 2024.
[24] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint urXiv:2210. 10749, 2022.
[25] William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07992, 2023.
[26] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint 14X:::2402.14735, 2024.
[27] Binghui Peng, Srini Narayanan, and Christos Papadimitriou. On limitations of the transformer architecture. In First Conference on Language Modeling, 2024.
(28) Ben Prystawski, Michael Li, and Noah Goodman. Why think step by step? reasoning emerges from the locality of experience. Advances in Neural Information Processing Systems, 36:70926-
70947, 2023.
[29] Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halcrow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. Advances in Neural Information Processing Systems, 37:78320-78370, 2024.
[30] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[31] Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while callycconnected nets cannot. arXiv preprint arXiv:2406.. 06893, 2024.
[32] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-though prompting clicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
[33] Kaiyue Wen, Huaqing Zhang, Hongzhou Lin, and Jingzhao Zhang. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency. arXiv preprint 4rX1v:2410.05459, 2024.
12

---

[34] Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in sequence sequence tasks. arXiv preprint arXiv:2204.02892 2022.
[35] Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408. 10147, 2024.
[36] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can expected to generalize other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of the work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken low lighting. Or speech-to-ter -exxettto-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess..
13

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
We provide the full set of assumptions and a complete (and correct) proof.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear the supplemental material, the authors encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We provide the full set of information needed reproduce the main experi- mental results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
If the contribution a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reprodueibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (c.g., in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
14

---

(c) If the contribution is a new model (e.g.. large language model), then there should either be a way to access this model for reproducing the results or way reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case closed-source models, may be that access to the model is limited some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We do not provide open access to the data and code since the experiments are simple and are not central the contribution.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/ public/guidess /twgdd///oodeuubnnn ottp:/ooeeeniiiiiooooooooooooooe
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips. https/ccccppcc///eeeeeeeeeecccce
Stw.iddd//ooeeSumisssiiooo. oolooicoooooooooooooooooooioople for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and basclines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized ver- sions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide all the details necessary to understand the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
15

---

Justification: statistical significance is not reported because it is not central to the contribu- tion.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
to report I-sigma error bars, one should state The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [NA]
Justification:
The experiments are simple and be on single CPU.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't ' make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification:
We conform the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
16

---

Answer:
[NA]
Justification:
This is foundational research paper and does not have any societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses addition to attacks, mecha- nisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This is foundational research paper and does not have any societal impact.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We properly credit and mention the license and terms of use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
17

---

The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from partieular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package
should provided. For popular datasets, paperswi thcode. cwm/coo///////o datasets has curated dataset.
licenses for some datasets. Their licensing guide can help determine the license of
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We have no crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We have no crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
18

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We do not use LLMS the core methods.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: ://ner/is...../ cc/Conferences// (025/LLM) for what should or should not described.
19

---

A Numerical Experiments
In this section, we provide experimental validation of our theoretical results.
Setup. We construct the transformer model as in (5) and (9) for the backward and forward reasoning task, respectively. We use one-hot embedding: we embed d; one-hot vector with the i-th entry being 1. For training, since the expected loss not available, for both backward and forward reasoning,
for all i E [S], where e; is the
use stochastic gradient descent with batch size 256 on randomly generated perfect binary trees. The tree generation process is the same as described in Section 4 with 4, S 31 for backward reasoning, and 3, S 25 for forward reasoning. We use learning rates 1 and 0.2 for backward and forward reasoning. respectively. To validate the generalization performance, nodes as the test set. Different nodes of every tree in the test set are assigned with a unique number chosen from [S]. and each node have different numbers of children (range from 0 to 3).
randomly generate 1024 trees with various depths and number of
Paramater
200
600
500 1000
Epoch
Figure 4: Training dynamics of selected entries of H.
Experimental results for backward reasoning. Figure shows the training and test loss curves for backward reasoning, which (i) validates the convergence of the training process, and (ii) shows the generalization error also converges to O. Moreover, same as our construction and training dynamics
2,00
1.75
Tuxt Lont
1.50
1.25
1.00
0.75
0.50
0.75
0.00
2001
1401
5008
Epoch
Figure 5: Training and test loss curves for backward reasoning.
analysis, during the training process, the diagonal entries of ((*) A B(t) A (cf. (6) increases while the off-diagonal entries stay small. See Figure 4, where we plot the training dynamics of H1,1 and H1.2.
Experimental results for forward reasoning. Figure 6 shows the training and test loss curves for forward reasoning, which (i) validates the convergence of the training process, and (ii) shows the
generalization eтror also converges to 0. Moreover, by tracking 0,(0) v,()) (1 (1 2, 3) defined in (14), we can see that they converge to our construction, and the true training dynamics also matches our theoretical analysis, see Figure 7. For example, in the right top figure, we plot the curves of the four entries of Uz. Identical to our analysis (c.f. Appendix C.2.3), U3 -U U3.011 -U,,111
13 keeps increasing while U3.0.1 first decreases and then increases.
20

---

17.5
10.0
3.0
0.0
Epoch
Figure 6: Training and test loss curves for forward reasoning.
Tralies
Tonlle
-1.
-1.
1poch
Figure 7: Training dynamics of selected entries of U, V/2 for 1 = 1,2,3.
B
Construction Analysis
We first present the following lemma, which presents the existence of parameters for the desired attention patterns.
Lemma 1. Given [uN)) = Thm (uz)je(Na) C Rmm where u1..... UN,, are linearly indepen- dent, and v..... UN, linearly independent. Then for any fixed matrix A NNINN there exists a matrix E Rmixnn such that
u Bu, Aij, Vie [N1].j [N2].
Proof. We let
U == (ui...., uN.) E RMIXNi V = (v....., UN.) m2
(17)
and use & denote the Kronecker product. We have for any B R m1 xm2
UTB BV = A <> (V (UU) vec B) vec(A).
(18)
Note that since $1,.... uN, linearly independent, 01,-.., UN, UN, rank(U) N₁ and rank = N2, which implies
linearly independent, have
rank(V VT @ UT) rank(UT (AAUUT))ank(TT)) N2
Therefore, the linear system
(VT & UT b vec(A)
(19)
21

---

has solution E Rm (18) implies reshaping b to matrix B E R gives the desired
B.
B.1 Proof of Theorem 1
When Assumption holds, the existence of satisfying (6) follows from Lemma 1. Letting a - + 400 it follows that
softmax(A BA) = Is
converges to an identity matrix. softmax (Y(k-1)T Ba(k-1)) becomes a one-hot vector with the j-th entry being if (k-1) x=1 (k-1)
By (4), we know that at the k-th reasoning step.
and otherwise (note that x x11
(k-1) will appear exactly once in ((k-1),, and thus the model will output which is the embedding of the edge that connects the current node
@(k) (xr (R-1)T (k-
1)T
(k-1) with its parent or (k-1)T Thus by induction kk (kk) for all k [m].
F-1
B.2 Proof of Theorem 2
By Assumption 2, the existence satisfying (10) follows from Lemma 1.
For the forward reasoning task, there are 2m reasoning steps in total, where m represents the path length. For any & [2m]]. we let
11 = x(k-1)T x (k-1)
B1y=1 + ((--1) 1)T (k-1)
B2y + 1(T B3==1
T)T (k-1)
(20a)
(lk-1) = xx( (k-1) + ((k-1)T
C2 (1-1) (12-1)
+ ()T
1)T C3 (k-1)
(20b)
Let ui (1-1) (resp. (k-1)) denote the i-th entry of ((l--1) (resp. (k-1)).. We also denote (="),, /
(r = (I, T.T 2 for i E [2 + 2). where xj E Rdd, Rdd, and 2, E Rdd, and stands for the number of edges in the tree T.
In the following, we will prove by induction that under the assumptions in Theorem 2, for any k e [2m], we have
gkk) (8) 01, 02 +CC
(21)
Base case (k 1). Without loss of generality, we can assume
( (1:21,, 1) = (z) (TT (obo)) "
where ag is the embedding of the goal node. We compute the attention pattern by checking the entries in (20).
i = 1. Since y1 = 0g = y_1, 21 8y, and 200 = sb. by our construction (10), we have
Vk [m]
P1 (0) x. Buy (0)
y1 L + 21 B3= (0)
(1 + b2)
(22a)
(0) x Civ (0) wi Cau (0) + z Ca (0) (1 11-22aa 02)
v1
(22b)
i= ........l. We have
33 * ag:
y/ * y-1
(0)
and 2; 8f;
where the second relation holds because each node embedding appears exactly once in (0)),,
I + 1). Thus we have
yk [m]
(0) p x
(0)
B-y--1 + // Bay (0)
2, By= (0) b201,
(23a)
(0) 2.7 Ci C2y 2.7 Cz2 -C202, =i= 2 2....... I. (23b)
(0) (0)
22

---

ii= I ++ 1. Since we have
(1, 1) (ap 0, 87)
(0)
where a, is the embedding of the root node, and a, * V-i- We deduce
M1+1 = ag B1y _ + b (0) V/41 = ag Ciy + +w = (1 - -22))22
= (-a + b1)a1.
(24a)
(24b)
i = 1-+ 2. Since
E (:,++ (1_ (O)T (O)T 2 , a ag ; sb))
we have
(0) P142 (-a + 1 - b 14 = (2 + 02 02)
(25)
Combining the above relation, we see that
softmax(u(0) cl, softmax(v(0) + e142, as a8,100 + +00.
(26)
Therefore, by (9) we have
(X00)
softmax (1el01)
ag ap(g) 86
2(1) ( (0) softmax (p())
(8) (1)
as 01, 02 +.00
(27)
i... (21) holds for k =
Induction hypothesis. Suppose (21) is true for k p. 10-1(22 12 p 2m). Below will prove that (21) is also true for k p. By the induction hypothesis, the input at the p-th reasoning step satisfies
x(0-1) 1)
F(e-1) - (F,, - (F, 10,00-11) : F(o-1)
F(o-1)
as 01,02 - t00, (28)
z(o-1)
where x(1-1) = E(o-1)(1 d),;) E Bdx(U-+11) Fle-1) Elp-1) : 2d1,:)
Rdi x(+p+1),, and z 10=11) ((d- +1 2d, d2,;) Rdxx((+++)), For notation sim- plicity, drop the superscript (p 1) and let au. yi, and 2; denote the i-th column of (x-11) 1 and Z(r-1), respectively, and let T., Fy, and 2, denote the i-th column of X ((-11) (1-11) 1) and z ((p-1) respectively for all i E [2 + p 1]. We have ( W. x T =xx,, vr TT),, (x] Ui Ri - - V. )
Vi [ [1+22 2], as 0, 02 + +00, Vi= 3......... +++p 1.
(29a)
(29b)
We also let y - (1-1) = Vi+++ ا 2 and define
p 1) = (11-1) Bud
T)T Bi
(p-1) + f(P-1)T 1)T Bug (p-1) + 1(T-1)T BB
(30a)
1)--) = x ( 1) CirlT-1 1) y ( 1)T C2f( z/D-111 1)T C3
(30b)
We let n E [S] denote the node with embedding " (1-1) i.e. y (p-1) = an.
Case 1: p E (2...., m) (at Stage 1). In this case, we have x: x- (p-1) ac(n)= 1 ( and and
=(--1) Sb. where we define c(n) the child of node that's on the path. Same as the backward reasoning, without loss of generality, we can assume
(x (TI (
23

---

By similar argument as the backward reasoning, we can compute that for any E {2...., m,) we have
(1+ b2)
if 1,
(1 02 02),
i = I, if = 2, 2,...,,, ,
b201,
if = 2....,
C202
1p-1) 1)
(-a+ b2)
if = 1+1,
and D v(p-1)
(1 02 0g2
if = 1+1, 1,
(-a b2)
if I+ 2,
(1 + 02 022,
if 1+
byan.
if 1+ 3....,I+ p.
C202,
if 1 3, ...,I.. Ppp
(1 b2)a
if = + p+1
(1 + 02 (2),
if = = (31)
which suggests
F
softmax (zl--1)) -
((grr)) an Ap(n))
z(o-1) softmax (Dlo-1)) - Sb,
F(o-1)
as 01, 0/2 +000
(32)
Thus by (29) and the above relations, we have
(x
(x()-1) ( softmax (n(e-1))
an
olm)
ap(n) 8b
(8) 00))
as 011 02
+00.
(33)
Z(p-1) 1) softmax (vip-1))
Case 2: p= m + 1 (at the turning point). In this case n ny where n, the root node, and we have - = Gc(n-).
((p-1) on,,, and z(p-1) 8b. where we define c(n,) as the child of node n, that's on the path. Then analogous to the above argument, we can compute that
b2a1,
if e [4]
C202,
if E [4]
( (-a a 1+ b2)a1,,
i = 1 + 1, if = [ + 2,
(2 (2-02)a2
if /+ 1,
T1 (P-1)
aaabae bb)aa,,
and 2(p-1)
(1+ 02) a2,
ifi = |+ 2,
if = | + 3....,1+p2
C202,
if = | ...... +
(1-b2 b2)
if = = p+
+ (1+02)022
if = 1+ p 1. (34)
Recall that we require E (0, 1], b2 < a/2 and 02 1/2, which guarantees that
+1+, Pi+p+1 > Ps Vi * 1 +, and V4+1 > Di, vi *
(35)
and thus (
softmax (zl0-1))
One ac(n.))
Z softmax (z(P-1) -85,
((0-1)
88,
as 01, 012 - +00.
(36)
Thus by (29) and the above relations, we have
((=-1)
((--1) softmax (1(p-11)
anr ac(n.) 85
g(v)=
(8) (p))
a1, +00.
(37)
Z(p-1) softmax (u(e-1))
Case 3: E {m+22 2,... 2m) (at Stage 2). In this case, by (8), we have x (p-1) ap(n): V-1
(p-1) 1) an-
and z(p-1) s,, and ক T. -1+3+2m-p = a am() T1+3+2m-p a,,, Z7+3+2m-p = 5 Then by a similar argument as the backward reasoning, we can compute that for any P {m ...... 2m), we have
(1 - (1-b1)an bi)an, bio
if = 1, if = ,, 1
(-a bi)on,
if l ++ 1, if 1+ 2,
7 (1-
(-a+ b1)on:
(38)
bla1,
if = 1 +3....,1 m 2, 1+ 3 2m - p.
(1 b1)
if = 1 1+3 2m p.
bia1,
if [1=1+ 1 4 4+ 2m P. ...... P.
b1)a1,
di==+pp+
24

---

T28
129
Figure 8: An illustration of the node ordering of a perfect binary tree with depth m 3.
(1 c1)
if 1, if 2....:
0102
(1 + C1 022
if = | ++] 1, if = 1+ 2,
(1-1))
(1 01)
(39)
C1
if = |+ 3,...,
+ 2,1 1+ 3 2m
(1-4) 02,
if 1 3+ 2m - p.
-102
if + 4+ 2m ..... l+p.
(1+a))22
i =
This indicates
F(-11)
x(1-11) softmax (ulo-1))
(
On ac(m))
)
7(1-1)
softmax (uO-1)) 88 01, 02
+000
(40)
Thus by (29) and the above relations, we have
Y ((p-1) x((-1) softmax (1(o-11) 7(p-1) softmax (ulo-1))
an
o(p) =
Oc(n) $y
(p)) as a1 02 - +00.
(41)
C Convergence Analysis
We multi-phase analysis track the training dynamics of both backward reasoning and forward reasoning tasks, where the details can be found below.
Tree node ordering. We let N == 2 1 denote the number of nodes in the tree T Pirain- and let nh, 722. ... ny denote the nodes in the tree T. In particular, we let ni denote the root node, and for each E [2m 1], we let riz; and n22+1 denote the left child and right child of node n, respectively. For E [N]. we let s(n,) to denote the sibling node of n4. Without loss of generality, assume the goal node is nggm, and the sibling of node n2; is n22+ for i [2m 1). and thus the path from the root node to the goal node is
n1 - n2
n4
72mm.
(42)
See Figure 8 for an illustration of our ordering.
Let (8) (A) (1 di), (h)) = (8) (d1 + 2d1), 2(h) = o()) :1):::::))))))))))))))) g(k) (11 + +1 2d,). Under this ordering, we have
(K) an.mmkk'
(43)
Throughout, let n
(k) (+) denote the node index of the embedding x _11
(k-1)
C.1 Proof of Theorem
We assume Assumption 3 and Assumption 4 hold throughout the proof.
25

---

Loss simplification and gradient computation. We first simplify the training loss and compute the gradients needed for the proof, where we drop the dependency with the parameter 0 {B) whenever it is clear from the context. Define
(k)
2 Pran (k))
2(()
A = 1 2 ET
Pran Mran
-gia
(44)
2
It follows that
Drain ==
Σ k=1
(k) + (
(45)
Define Ho.j = 0 for all j E {0, 1, ,),
S}.
11 A BA,
(46)
and (z) (y((t) 12(++)x(5++) be as follows:
a (1)
a (k)
5x = HHis
8yi. = HHis
(k)
Vi,j = 0, 0,1,,..,,
(47)
Given tree T, we let N/(A) A!Aa (T) denote the number of times embedding a; appears in
the columns of (k), for all i 0, 1, ... S. We define the attention matrix (k)) a (T)T) 1(5+1) as follows:
N(R-11) expla Baj)
W [m]
= 0,1,
(48)
N1 (1 exp(a, a, Baj)
For any i,j E [S], if s(j) does not exist, we let a dz)), = O. Assumption 3 and Assummption 4 allow us to simplify the training loss, as in Lemma 2, whose proof is deferred to Appendix C.1.2. Lemma 2 (Loss simplification). Under Assumption 3, Assumption 4, we have for all E [m]]
(k)
1
Σ
(K) n- n
1
=}} :){{
(b
Σ
(o (k) + a200)) (1-a a (k) 2 (k)) 2
) (o
2
JE[S]
(49a)
A
(k)
1
Σ
E
1n
(k)
j}
Σ
(1) (is))
+ (1- (1- (1)) )2)
(49b)
2
JE[S]
where n denotes the node index of the embedding x
The following lemma computes the gradients (x((k) and (y)*), whose proof is also deferred to Appendix C.1.3.
Lemma 3 (Gradient computation). Under Assumption 3, Assumption we have for all ls [m].
50 0, 5y6g = for all j = {0, 1, S...,,,,, and
(k)
Vj [$]
(k)
1 E s
(003)
(K)
(k)
n
(k)
=
S
S-3
E
Σ
(k) + (1))
a (ppp)s)) )²
2 (1) (k)
= 5
25(S
1)
pe(S))
1
1
(k)
- -a0
aj n- =
(k)
(50a)
S
W e [S]b): : 60
E
(
(k))
(1))
aij n =j j
(10)
s
1
(1
(N)
a (0)
(1) (a
(M)
a (8)
) |n.00 =
a
26

---

S-3
Σ
(ps + as(p).s)
(k)
(k) 2 (k)
oi n- =
2S(S - 1)
[6[3]
(2)
1
E
(1
(R)
(1)
(8)
m = sss
-(S-1)
S(S - 1)
S-3
E
{a +
a
=1)
(K)
(10
j.id (s())
(50b)
and
Vj E [S]:
1 E 5
a JJ
Σ
(ps)
(1))))
+ (1- (1- (A)
0331
(1)))2
n =
(k)
(51a)
pe[S][(
Vie [S]M
806. = S
1 -
04
(k)) pp3
- -11-
-1-0
(k)
Oij
(8)
n-1 = j
(K)
pE|SRR
(51b)
(K) Here, n-1 denotes the node index of the embedding (k-1)
We'll repeatedly use the gradient inequality given in the following lemma, whose proof is provided in Appendix C.1.4.
Lemma 4 (Key gradient inequality). Under Assumption 3, Assumption 4, we have for all k E [m]].
(K) 531,3
8313
(k)
x
(1
(8)
(k)
(K)
N
a
(A) jJ
max pE(S)N(A
(R)
00.3
(k)
n (k)
5
N
+ s
(1- و.م Oijj.
(k)
(k)
(k)
OSESS pPo
max
p.j
(A)
n
(k)
(52)
where n (k) denotes the node index of the embedding (k-1)
C.1.1 Proof outline
Letting ((0 = A (()))) by the update rule (12), we have
(((++))
H b. 17 Σ
(1,8) + oyll)t)) Vi,j = 0, 0,1,...,.
(53)
where g((t.)), ( s(1.)) g(c.k) denote the value of (A) ((),, (R,, g(t), respectively. Further, let (e.k) be the value of a at the t-th iteration when replacing B by B(t,, which can be written in ai
(l)
terms of HI!!) as
NORK exp(HI
W [
(L.K)
21 N (k- exp (x)(41
(1,k) 1. Note that by symmetry, we have for any E [S], H7!)) are equal for all satisfying E, E.ali a i E [S]M and H(:) are equal for all j [S].
In what follows, we divide the training dynamics of (1)) into two phases to prove the convergence of
the backward reasoning task. We show there exists Ti Ō that
SN my
and T2 Ti Ô m
such
TJc
At Phase I (t =t=(0 1.... D: the diagonal elements H ( strictly increasing for all j [S]. and the off-diagonal elements
(t)
log N
N
for all i [S]M]].
27

---

At Phase II (t (T1 +, T2)): the diagonal elements H3 off-diagonal elements H are strictly decreasing, and satisfy
keep strictly increasing, and the
J.
H
10
log N N
S
After Phase II (t> T2 + 1): H (resp. H keeps increasing (resp. decreasing), and
1 (0) 33
2 log )
and the training loss Craain (O(t)) e.
Training dynamics of Phase I. Let 7, E Nu {+00] be the number of iterations of Phase I, defined
as
1
Ti max NEN: tEN:
(54)
25
where
exp(H
@00 =
(55)
(N 2x + exp( (1
which satisfies
(00 200)
Vj E [S].
(56)
In the following, we show by induction the following claim holds for all t [T] - 1]:
Induction Hypothesis Vie (0) J [S], e [8], i #j: H > 0, H N
3.
(57)
We start with the base case when t 0. Since H,, = 0 for all i.j. (57) holds true trivially. We next assume (57) holds up time t. We'll make use of the following lemma, whose proof is postponed Appendix C.1.5.
Lemma 5 (Gradient proportions at Phase I). Under Assumption 3, Assumption 4, Induction Hypothe- sis (57), we have for all E [m]]
-dx(thh)
(t.k)
(1.k)
(-ou(tah) 89 (1))))))
(t.k)
(58)
N
(t.k) We now establish (57) holds for + 1. From the expression of da and oyحJ
(1.K) (c.f. (50a) and
(5la)) we can see that H (t.kk is increasing, and thus H ((+1A)) > 0. Moreover, By Lemma 5 and
(53), we have
H
H!!
(n
(t.k) + 886.5 (1.)))
-7
<
N
9
H
33
(1)
N
9
17
m
(on(t.
(2.A0) +8y
))
9
H 53
(+11) (59)
where the first inequality follows from the induction hypothesis (57). Therefore, (57) holds for 1, and the induction is complete.
The following lemma guarantees that Phase terminates within T; 3( mm
(SNN
iterations, see
Appendix C.1.6 for the proof.
28

---

Lemma 6 (Upper bound of Ti). Under Assumption 3, Assumption 4, we have
6S(N + m)
NN+1 1
Ti &
log
(60)
mm
In addition, we have
H
3 log
N +
(61)
The above lemma, together with (57), implies that for all i [S]M[]]
log( N)
e (0, 1..... Ti) H x
(62)
N
Finally, we give a lower bound on H in the following lemma, whose proof is postponed to Appendix C.1.7.
Lemma 7 (Lower bound on H. Under Assumplion 3, Assumption 4, we have
19
e (0, 1,.1........ . Ti}
H i.j
H
(63)
S-
Combining (61) and Lemma 7, we have
log(N)
Vt E {0, 1, I.......TT. Ti} - HH x
(64)
S
Training dynamics of Phase II. Let T3 ENNu o {+00} be the number of iterations of Phase II, defined
T2 == max
LE N: x(0) < E1
(65)
2m
where
exp( (1)
[(()
(66)
(N m 2) exp( (1)) exp( ex +
Note that by definition (48) and symmetry, we have
(x)
aij
(L.k)
(N 2) exp (1))) exp(H (1) +
> x(t)
(67)
j E [S] and e [m]] Since E (0, 1]. we have T3 > T1.
We'll show by induction the following claim holds for all E {Ti + 1, T2):
m
Induction Hypothesis II:
Vie (0} o [S].j [S], #
(1.) jJj
(t,b)
+ +8y (t.k) ())
4ms
(68a)
5z j,j
+ +8y (2.b)
< 0.
(68b)
(11+1,1) We start with the base case when /. Ti + 1. By the definition of Ti, we have ajj. j [S]. For any j [S], when n (k) j. we have
2/2 for all
a jj
(T) +111 a
which implies Vi [STM))
o
(T++1k)
1
N
29

---

for all & [m], where use the fact that N²(
< 2 for any i, and
a jj
(Ti+1,k)
N --1
1 1
(69)
+ m 2
2
where we use the fact that
(t.k)
jj
(N exp( H + exp( +1
N -
VIEN: ::
(70)
(t.1)
(N + m - 2) expl H2 + exp( (1 +
N 2
ja
Since H 2 H we also have
+1,))
0,1
& 1 e
+1.4)
exp
11 (11-11)
j.
Thus we have N-
N
33
(T++1k)
max
(1,8)
N 005
(T++1.))
N - 3
2
N
1
N-1
2 > 0,
(T++1M)
max 0
(T++1k)
0.
jj
P.
(71)
By Lemma and the above relations, we have (68b) holds for = T1 +
Moreover, since وx +1.)
< O, we have
m
80 8yj.
(1)+1,))
(T) (
-8y5o
(T) (1+)
(51a) 1
s ajj
(1)+1.m)
(1 و
(Ti +1,mm
2 (K)
n- =
(65)
4mS
(72)
Thus we have (68a) holds for Tı 1.
We now continue with the induction, by assuming (68) holds for all iterations up to f >> T1 + 1 and
S), all j e [s], e (0, S), * j, and aim to establish it continues to hold at iteration + 1. By
the induction hypothesis, we have H (i ( j) decreases after iteration Ti + 1 while H(!
j.j keeps
(t.k)
increasing. Thus we have j.j monotonically increases after iteration T₃ + 1, and i.j (i + j) monotonically decreases after iteration Ti + 1. Therefore, by (71), we have
N= (t+1,))
(t+1,6)
(++1,8)
N 1 α
30
max
P.
N - 1 0.
N
N
a
jj
(Ti+1,k)
max O P.j
(1)+
- (o.
(Ti+1.b)
0,
(1+1,K) O.
max (++1,)
(T++1,k)
(Ti+1,A)
P.j
jj
max @ PJ
0
for all k e [m]] This suggests that (68b) hold at iteration / + 1.
Moreover, similar to (72), have
da (++1.)) - 65 x1) > 50
((+1A)
(+1mm) a jj
(1-a
(++mm)) 2
(65)
n =
2
4mS
(73)
Thus we have (68a) holds at iteration t + 1. The induction is complete.
Similar as Lemma 7, we give lower bound on H. after iteration Ti, see Appendix C. for its proof.
Lemma 8. For any i,j [S]. i + j. we have
Vt Ti +1: - -(H H H()++))
4 (H(0)
H (1)+))
& s
30

---

Combining (64) and Lemma 8, we have
Vt
- -H
3 + C.
4
log(NN
(74)
5-1
S
where C' is an absolute constant.
The next lemma states that the duration of Phase II, T2 Ti, is bounded by ① deferred to Appendix C.1.9, Lemma 9. We have
ms
The proof is
12mS
( 2m
)
T2 T1 &
log
(N 1M++m11)
(75)
1/c
After Phase II. By (50a) and (51a) we know that 5 is strictly increasing after iteration T2- Also, we can easily see that (68b) holds after iteration 72, which guarantees H (i [STM)] strictly decreasing after iteration T2. The following lemma guarantees that the training loss is less than after
iteration 72, and gives lower bound on (!) after iteration T2; see Appendix C.1.10 for its proof. Lemma 10. After
j.j T2, the training loss L(rain(O(t)) & Moreover, for any E [S], have
2m 1-m-11
Vt > T2:
H jJ
log (V (V--1)
(76)
C.1.2 Proof of Lemma 2
For notation simplicity, we let n- n- For any k E [m]] let (note that n2m is the goal node)
(k)
+ an ang 21 if i E [2 - 1]. i 2m.
(k) W
(77)
(k)
then have
(k)
(k)
On.
an
1
(w (k))2 2 (1
Wmm
(k)
)
2
1 + #
(K)
)2 (a)
-nn-mm-1 12 (ao.n..) )
(k) (b)
)
(k)
2
1-1
1
Σ [E[S]
E
1 (m
= 5)
Σ Te[s]
(a
(K)
a (k)
2 (1 a ((),,))
-1 - a -0)) (a()))
(1)
2
which gives (49a), where the second equality follows from our assumption that (0:)ie[s) are orthonor- mal, and the third equality we use the fact that 22m- is the parent of n Similarly.
(K)
(k) no On; - On
2
E
(b)
)2 + (1
a (b)
p )
2
Σ E
1n j}
Σ
(a (k) 2 (1- (k)g3
2
[c[S]
(e(s)
31

---

C.1.3 Proof of Lemma 3
For notation simplicity, we let n-
(k)
Computation (*), The fact that Ó = follows from our assumption that ag = Below we consider the case when j E [S]. By the chain rule, we have
(k)
P (k)
da p.j (k) HH
(K)
E
(78)
HH
do P.J
We compute the terms in the summand separately.
When p = j. by (49a), we have
a
1
12
aj
(K)
(k)
R-1 =j
(79)
δa (R)
pd 5
When p [S]M)] by (49a), we have
A (k)
E
1 1 - j.8(f) = pl 21
Σ
(ai + 2()))) 2
(R)
(k)
+ (1 - aj.
(1-aj,3 (2) + (2)6))))
2
+
2
E
1 {n-1 = j,,(1) # p, # pp 2
(ai a
(K) (k)
(1)
)² (1 - a3 - a ()) )2 + (a)b))))
(2)
+
E
1 (m-1 j. j pl (
(a + a (R)
a(1)
)2 + (1
jn a 20 +3
from which compute that when P [S]]i),
a (k)
1
(k)
(k)
n- 1- = p]
do (k)
S(S-1)
0
-
p.j
p.j
S-3
5(S
E
P.
a
p()).)
n-1 = j,8(f) * nn + #p
(80)
When p = 0, we have
24 (k)
1
(k)
aa a
(k)
S
E
0
n_ n-1=1 =
(81)
In addition, by (48), we have for all i, {0, (0.1,...,$)) S}.
dcx (k)
(2) T
(k) T
0B aij
a
ai aja
(82)
1-0
which gives
(k) 1 (111
00:
aij
when p = i j.
(k)
when p [S]Vi)] = J.
(83)
H p.q
p.j.
0,
otherwise.
Plugging (83) into (78), we have
S (k)
(k)
(R)
(1 (1- (1-o (K)
8
E
8ox
(K)
a
ap.
+
804.3
(K) ai
(84)
p.j
32

---

By (84), (79), (80) and (81), we can compute that
8n b0g =
-
E
(0.3
2
(k)
m===
n-1
j
s
E [1-aa]]
(k)
(k) (8)
==== = =pl
5(S 1)
pe(s]
S
E (a + (
(N)
(K)
= = * pm #D|
4/ 9p.
S(S
- 1)
pe(s]
B (1-
(k)
)
(2)
(2) (1 (8) |n
(85)
5
Note that the second term can be simplified follows:
B
11 ap. a3 pa m- = = pl
(k) (k) (8)
S(S 1)
p=[3]
(1110
(k)
a
(k)
33 *(),,]] 7-1 = }]
(k) (k)
(86)
and we can rewrite the third term in (85) as
S(S 1) 1)
S-3
Σ E
(t)
+ a
j, 0 p.j
(k) (k)
n- أ,,( # p. * p
(C(s)
5 5-3
3
E
Ops
(k)
+ (((r)
j 4p. R-1 = j.)(5) $ (p. s(p)). mi $ sp))
(k)
S(S - 1)
pe[s]
S-3
EF
(k)
+ aa
(K)
jj D. n-1 = j,) & sp)), nı $ sp))
(k) (8)
25(8- - 1)
pjj
ap)))
pE[S]
S-3
28(S - 1)
E
p.j
+ (k)
as(p)s)
ajs (p))))
(k)
n_1 = s() & tp. ((.)( m $ sp))
pe[s]
25(8- - 1)
S-3
E
a p.j
(R)
+ a
<()))))
(k)
2 (1)
aj
n- = j,ز( & s()).. nı & tp, s(D)
pE[S]
25(8- - 1)
S-
E
(k)
+ a
((p)
(k)
2 (k)
9j
n-1 j.p mm())m
pe[S]
S-3
E
Σ PE[S] (
(a
(M)
a (k)
(k)
n-1 = |]
m==
(87)
25(8 11)
s(p).))
where the second equality follows from the symmetry of the tree structure (a node and its sibling s(p) are symmetric). Thus we can rewrite (85) as
8n (K)
1
(0.)
(K)
2
(k)
m-==j
n-
jj j
5
25(S - 1)
S-3
E
Σ
p.j + a
(k)
((pp)ss)
(K)
2 (k)
aj
n- =
p
1
E
a j.j
(k)
a (k)
2 (k)
a i5 n-1 = j
(88)
which gives (50a). Similarly, for i E [ST\(í}, we have
da (k) = -E 1 (0.5
E
(8)
2
a
(k)
n-1 = j
+
1
(1-aa3
(1)
a (k)
(1),.))
a ajs n-1=1 =
(k) (k)
33

---

1
E (1-a (k) - (8)
ai, p. n_ -- =p]
(8) (1)
S(S 1)
S(S -1) 1)
S-3
Σ
E
[a/t) p.j
(k)
+ a ((p))s
(k)
ai. p...
(k) (k)
n-1 = j.s(f) + p. # pl
1
E
(1
(k)
(k)
(N)
a1-
(2)
7-1 = ==1= = =|
S(S -
E
[(!!]
(k)
a (()..)
(K)
(1)
(K)
1- j. * i, * 1]
S(S - 1)
1
(003)
(8)
(8)
(k) (K) 03 a ( 0i5 035 n-1 5
(2)
(K)
ai,j
n- = j 5
1
Σ
E (1- (M) -
0: 9.. 111 = أ,s(5) = p]
(8)
S(S-1)
S-3
Σ E
(k)
+ + (1) (k) (k)
n-1 = -=-==, + pmm # v]
S(S 1)
pJ
(p))))))
0 pJj
1
E [1
(k)
(1)
(k)
n-1 = j, = s]
0 jj
0i
S(S -)
S(S 1)
S-3 E
(o! a
(k) (1)
(0)).))
(1)
n-1
J. s() * i, ni * 1)
(89)
Similar as in (86), we can simplify the third term in (89) as
Σ E [1
( -
(k) (k) ajj 9 p.j
n-1 أ, r]
S(S 5(8-1)
5
E
[(1
o (k) - .1m).)
(4)
5, a
(k) (k)
.1)),
n - 5
(90)
and similar to (87), we can rewrite the fourth term in (89) as
S(S - 1)
5-3
Σ E
ps +
a (pp)s
Jj p.5
(k) (k)
n-1 n-1 = },) * p. m # pl
pE[s] pid
S-3
E
Σ pe[s] p
(0/))
(K) p.j
+ a afpps)
(K)
2 (8)
aij n_ = }
(91)
2 - 1)
a(p),)
Substituting (90) and (91) into (89) and reorganizing the terms, we have
Sx
1 - S
E
(0.))
(k)
m-1 =
5 + 1 1 (1 [
a (1-0, -
al
2is (0/ + =0th))) m_1 = 51
(1)
5-3
E
Σ
(o00 + a
)²
2 (x)
0i5
n- j
25(S -1) 1)
1
E
[1-1 ajj
(k)
- dijj
(k)
aij
(k)
jزز = :| n- j (s()),
- --
n-
S(S - 1)
5-3
E.
(k)
(k)
(k)
(92)
S(S
This gives (50b).
Computation of (k), The fact that 6y6 0 follows from our assumption that ag 0. below we consider the case when j E [S]. By the chain rule, we have
2 (k)
a (k)
dap
(R) Ps
8913
(8) HH.
E
pap
(K)
HH.
34

---

(83)
E
7 A
y
(k)
aij (-ais
(8)
aps)
(K)
2
(k) aij
(k)
(1)
(1- (3))
(93)
da
[k]
dais
P priti
p.j
By (49b), we have
E
(k)
11 =
if = j,
04
(k)
dap.)
(R)
E
1 5|
p.
(1)
if [S]]]]] p 0.
(94)
0,
Plugging (94) into (93), we have
(b)
1
Σ
E
(pp.s)
(1)
ajن n- j
(k)
E (1 a1- (k))
035 m1
S
8
pc
(S]
(k)
(k)
+ (1- o ( m- )
pp.
(95)
S
pe[S]]Li
which gives (51a).
For iorie [S]][]}, we have
Σ
E
p..
(K)
2i n- m
E - (0 (2)
(1- (1- + (a (11))) m_ |_-=1)
S
p
1
E
(k)
(1)
(1)
11-1 = b
5
E
aij
(k)
Σ
(ops)
(k)) 2
a a0
(1-
(k)
(k) (k)
azj
m-= =
(96)
which gives (51b).
C.1.4 Proof of Lemma 4
For notation simplicity, we let n- =n -- By (50b), for all j [S] and i [S]M
dx.
+
E
(o.5
(k)
(16
n- =j
5
1
(1
(k)
(k)
(k)
{a + +a (()))) 7- a-1 = 5]
(k)
a
((1)
S-3
2 (K)
Σ
+ a s(p).s)
n- j
25(8
-1] 1)
PC[S]
1
E
[1
(k)
a
(1)
a
(k)
n- = = -,s =i]
S(S
S-3
E
[a
(k)
(K)
(1)
Th jii (s(1).
S(S 1)
((()).))
1
E
[1-1 a
(k) - a
(k)
(1)
(8)
+ (k)
n j, =
S(S - 1)
'
aj
S(S 1)
2
(1- (k)
a
(k)
aij
(k) (ki
jo
n-1 = j,)( #
S(S E
S-2
(1-
(b)
(k)
03
(K) (k)
n-1 J. 8() x
(6)
35

---

S-3 E
Σ
(a + as(p)ss) ) 2 (k)
aig n2 = j
2 - 1)
[6[s]
(3)
1
B
[11
33
(b)
- ais 0/.5 n-1 = j, = = -
(k) (1))
(97)
S(S
1)
(c)
Note that (a) and (c) add up to the following:
+
S(S
1)
1
E
j.
(R)
- 0 .
(k)
01 3 + (1)
(k)
(k)
(s(1)s))
| - s.)(1) - $)
1
E [ (1- 33 a3 0i
(K)
(b)
(k)
ni-1 = j, {s] =]
S(S 1)
E
(1- 23 .5).3)
(k)
(K)
(k)
n_1 5.s(5) =
-(S
-(5-1) R
1) (1 j3
(l)
a (5)s))
(8)
(K)
(98)
We could also rewrite (b) as follows:
(b) S(S
S 2
E
[1
(k) (k)
a a()r.s))
Rijj
(k) (k)
a ax(s>s]
n1-1 j. * s(i)
- 1)
1
E
(1-0/% (k)
Σ
(k) (k)
S(S 1)
1
E
(1- 2
(k)
(4)
(b)
n2 J
S(S - 1)
S(S 11 1
1
E
[11
03
(k) (K)
00.
s(1)))
(K)
11
(1 (1-0 033
(9) (2)
S(S 1) E
n-
1
N W1 1 (1 j.
8 8- 1
(k)
(k)
(k)
> 13 20)s
n-1 n- =]]
S(S
1]
(()))))
ab.j
1
(1 (1-a jj - a .5).s)
(k)
(k)
2 (8)
S(S - 1)
E
(o())s
n-1=j =
-- (a) (c) by (98)
E [1- -
(x) دز (()))
(k)
(k)
(1) n_
(99)
S(N- -1) 1)
where the second line follows from the symmetry of the training distribution, and the last inequality
uses the fact that [. > 1(N > 13a-(h.
(k)
(k)
This is because by symmetry we have Ha
Ha,j, and N (=(11) = 1 because s(j) is not on the path.
alf)
Plugging (98) and (99) into (97), we have
8x + s
E
(0.3)
2 n- 5
(M)
>
S-2 2
TE. [111
(K) (k) (k) (2)
oi 2j n_1 = j.)(j) t #
1) S(S
1
E
0 j.j
(k)
(k)
(K)
5 n-
S(N - 1) S-
25(S- - 1)
E
Σ
(a + (1)
(k)
(((p)))))
ai
(K)
pets]
pe(3)ats).m
5
B
1 5
(k)
a
(k) (k)
aز n- = j
(d)
36

---

E
a
(K)
ai
(k) (1)
1 1-1 í. -
S(S - 1)
(e)
1
E (
(k)
(k) (k)
n_ j
aزs - am())s)
و
S(N
1)
(
1
E
Σ
(ops + a (k)
(8)
(6)) )²
2 (0)
a
n-1
25
pe[s] (
(g)
(100)
Following the same logic as in the proof of (99), we have that (e) can be expressed as
(e) 1
E [1
(k)
-a (())s)
(k) (k)
a ". 3 nh_ =
(1))
1) S(N
E (1 (1-aj, ajj 20)s
(k)
(())) 1(N > 13 (k)
m m-1=1 = 5|
(1- (1)
(N)
(h) (k)
n_1 = 5|
(101)
S(N - 1)
a())
Thus we have
(d) + (e)
S(N 1) [(11 jj x()),)
2
0i دز n1
(k)
(102)
Furthermore, (g) can be bounded as follows:
(g) E
max Pnd (1 (1-0 o a
(k)
(8) (k)
aij n- = 3
(k)
(103)
which gives
(f) + (g)
pE((]] Pp
max
(k)
1
00.3
(k)
(1 |
(k)
a a
(t) (k)
a n = 5 (104)
N -
Substituting (102) and (104) into (100), we have
(k)
+
E
[0.3
(k)
=
S
N -
E [(1
(k)
(k) (1)
n-1 = 3
S(N- 1)
jj
0ij j.j
E
[S](tt)
max
0 p.j
N-
1
-
U.
(k)
Ojj (1-0
(1
(k)
0i n-
(k)
= B (1
(k)
N 2 (k)
(k)
(k)
aij
N
max
P.
N - 1
00.5
1 = j
(105)
By (51b) we have that
5y
E
(b..)
2. m= m- = j
s
-E
agj
(k)
Σ
(pps)) - a ajj
(k)
-000 a
(k)
(8)
n- =
-5 E
(abb)
a
(8)
m
S
(k)
Σ
pp..
(K)
- -(1
(k)
(K)
n- n-
S
37

---

/ (1
(1-a53
(k)
aij
(R)
aj
(k)
max
(8) Ppd
n1 =
(106)
We have (52) follows from adding (105) and (106).
C.1.5 Proof of Lemma 5
Recall in Assumption 3 we require m >> 3, which implies N > 15 and (57) in turn implies H J.j
H In addition, when n = j. we have ((kk-)) 2, and (k-1-1 = 1. Thus the above facts together with the definition of a (c.f. (48)) lead to:
(1.h)
V S j. (1,M)
(107)
2
and
2 2exp(H (
Vt p. [SM
(t,bt
p.j
2 [s] N₂ (((.k-1) 1) exp (H())
95
2 eexp( H
exp(11 13 + (N - 1) exp(
p.j
2
(57)
2
(108)
exp( j.
/ + N
N
where the second line follows from the symmetry of our training distribution (i.e., Hi, 's are equal (t.1) for any
for any (8].))) Thus by our choice of Ti (see (54)) and the fact that ajj ajj l [m]. we have when n_1 j:
(t,k)
(t.k) V [T] 1 ajj 2
and 1 1
(l,k) ajj
(()).)
(t.K)
2
1
N
2
3
(109)
and
(1.k)
(t.k)
(107)
(f.k)
(108)
1
j.j
aspsS
max
PJ
0SpS5 3 p.j
max
2
N
(110)
pdb
pdf
Similarly, we have
N
(1.k)
(t.k)
1
(1.A) (107)
N 2
max (L.A)
1
N
max PE[STRU]
C..
2(N 1)
pE(S)
PPj
2(N - 1)
(108)
N 2
2
1
-3
2(N - 1)
N
2(N - 1)
2(N - 1)" (111)
where in the first inequality we also use the fact that Therefore, have
(t.k) 00
$
1 2.
since by (57) we have d
(1,b)
(t.k)
E
j.
(1,8)
aa
(t.k) a(1).)
(t.k)
(t.k)
N
N + 1
06p45 xF(3)
a (t.k) P.j
n- =
(111)
3
(t.k)
(t.k)
(t.k)
E
n_
2S(N
1)
(109)
28(N
9
1)
E
[
(t.R)
a
(()i)))
(1.A)
a
(t.
n2 j
(107)
S(N
9
1)
E
jj
(1,6)
(1:A)
a
(t.k)
n1 =
(50s)
N
9
8u (t.k)
jJ
(112)
Analogously, we have
(1.k)
(k)
(t.k)
E
a
j.j
aij
max
n-1 =j
osess
38

---

(10)
x
1
E
( 1 a1s (1.kk) (1:))
aij m_1 j
SN
(109)
2
(1 (1-ay رزز (t.kk (1.6)
E
dij n- =
5N
(00)
4
(1 (1-a وز jj
(1.k) 2
(1.ki
(51a)
9
(t.k)
E
n-
(113)
SN
N
Combining the above two inequalities with (52), we have the desired bound
80 (t.b) + 8yin
(t.k)
9
(ougfaa)) (t.b) + 55y
N-1
C.1.6 Proof of Lemma 6
By (50a) and (5la), we have
Vj E [S]
-dx (2,0)
رزحة > -80
(t.h)
1 E
ajs (1-0 (1:)))))) 2 n-1 = 5 (114)
(2.b)
Besides, when n (k)
j. we have ((---) and
exp (H!0) 1.j
(57) exp (1 x1-
(c.R) O3
W Tr,j
(115)
E Eisies N₂
(k-1) exp His
Lesies (k-1)
which indicates
V1 T::
((,k)
2' Vj [S],
N+
jj
where we use the fact that (t,R) distribution. By (114) we have
are equal for any j [S] due to the symmetry of our training
(1.b)
(1.b)
1
(N + + m
(116)
S(N + m)
2S(N + m)
(115) also implies that
Vt < Ti
H
1
9 log (IN 11)+1a9911 (11)
3log
N + 1
(117)
N-I
The above two expressions together with (53) imply that
N + 1
Ti & 3 log
2S(N m)
which gives the desired result.
C.1.7 Proof of Lemma 7
(k)
For notation simplicity, we let n_1 n= By (50b), we know that
(1.b)
E (1
(1,k)
(1.b)
((.k)
+.0)))
jj
o
Oi
ajs +
1
(a)
+
S(S - 1)
S 3
TP.
[
(1,b)
a
*()))
(L.)
(1,k)
m-
j
(
(118)
(6)
By symmetry and the fact that N²( (x(--1)
& 2 for all i [S], we have
2 (t.k)
(t.k)
jj - -60)))
a
(G:) & 1 ( > =]
and
a (1.b)
$ (t.1)
-
39

---

since (((-1) s(5)
1. Thus have
(a) < S(N 2)
4
E5 1
2 1)(1
(1,8)
(3(55s)
)2
Oj n = j
(1.2))
S(S - 1) - 2)
1(N 1)
E
(111- jos
o0)).))
(t.k)
)²
(t.k)
n2- = 5
(119)
Similarly, we have
3 3 223 a (t.10
(t.h)
(C.R) (t.K)
{ ( =1}
(t.b)
(t.8)
0i + a (i).
and
0i & 2033
N 2
Further, by (109) we have when & Ti.
(t,b)
(t.k)
1 2
1 i3 a a)()s
2 N
Thus we have
(b)
12(S - 3
E
1 N 11-1) 1)
3 13 (1 (1-
jj 0' (2113)) )²
(t.k) (2.b)
2 (k)
D3 nh11 3 m13
S(S 1)(N 2)²
12N 3)
E
(1
35 as()ss)
(1,b) ((.k)
aj n- = - m3
((.k)
S(S 1)(N 2)²
12N
E
(1 3j a
(t,b) (e.A)
a(iis))
2 ((.k)
Oid n2- 5
(120)
5(S- (I(N-)(N- -2)
Plugging (119) and (120) into (118), we have
(t.k) & Sr
16N
E (1
(c.k)
(t.k) 2 (1,b)
(50m)
16N
(2.b)
a ((t)))) 9s n-1 =j
[S 1)(N - 2) 6,33,
5(8 (N 2)
(121)
Analogously, by (5lb) we have
(t,
1
E
aii
(t:t)
((1-
(tah))
)
ajs + (t:1)
(e.b)
n21 = }]
(122)
Note that
(t.b) (1 (1- (1-a( (1.2))
(t.k) & 1 (N > 1] N -- (1 (1-a a )
2
((.k)) 2 (t.b)
ai
(123)
and when & & Ti.
1 - 1-a
(e,k)
(alxe)) 2
< 1 1IAAASt ( > >1} 4a
(t.k)
(109)
E 1
> 1 N-1
8
(1- (1-af
(t,h)
(1.))
ajs
N -1
(124)
Plugging the above two expressions into (122), we have
(1.A)
1
1
N (k-1)
} N (1-
10
(1
--aa joj
(1,b)
a j.j
(1.A)
1- = J
10
E [11111 (11-0(11)
(1,k)) ((.k)
(51a)
10
(1,))
ajs
n- =
8y13
(125)
S(S - 1)
S
Combining (121) and (125), we have
8x((gh) (t.k) + +5y
(t)k)
16N
(0x)))) + +5y((a))
((,))
19
(Sx((te)) + +80/0/0) (126)
$ S 1
(S 1)(N - 2)
This indicates that for all i.j E [S]. i # j.
19
Vt & Ti
H
H(0
S
40

---

C.1.8 Proof of Lemma 8
For notation simplicity, we let n- n 1 The proof is similar to that of Lemma 7, except that after iteration Ti, we have the following tighter bound:
3a (t.A)
3a (t.k)
(C.R) "0)J
&
1
(69)
and
(t.b)
jjj
(127)
2(N - 1)
2(N -
N
and thus
(a) & S(S - 1)(N 2) E (1
2N +1
(t,R)
2 (t.k)
- d((()is))
j n- =
(128)
where (a) is defined in (118). When & < [i + 1, we also have
(1,b) (f,
jaj (1 (1- U3 ((0(5))))
(1,k)
(1.h) (3,8)
(L.K)
2
jjj a8())
(69)
(129)
N
where the second inequality also uses the fact that (l.k) ajj
increases after t > Ti. Therefore, we have
(b) $ S(S 1) 2)
18(5 - 3)
E
1
Gj
(t.k)
as())s)
(t.k)
Oj n1 j, {s(s), mi
(t.k)
S(S ()
18
E
j.
(c.N)
a())))))
(t.k)
js
(t,k)
n_ = j, {s()) m1
(130)
where (b) is also defined in (118).
Plugging (128) and (130) into (121), we have
(t.k)
2.N + 19
(1-1
(t.h)
(t.k)
2 (t,b)
a;
a s(fiss) )
a a"j
n- =
S(S - 1)(N 2)
4
(1-ag)
(t,k) (1.k)
2 (1.kk
S(S 1)
j.j
a
a(j).)
jjj
R2 =
(50g)
4
(t.k)
8
(131)
S
Moreover, when / > Ti + 1, (123) still holds, and we have
(1-- (e.))
(t,k) 2
(t,k))
2
12 1 - 2.
(a(t)))² < 1 ( > 1]
(N- 1)2 < 1 (( > >1]
(69)
(t,k)
(N 1)2
j.
(132)
where the second inequality also uses the fact that jj (132) into (122), we have
(t.k)
increases after t > Ty. Plugging (123) and
(1.R)
2N +
10
E [111--
(1,6)
(1.k)
8y
S(S 1)(N -1)
ajj
a j.j
71 =
S(S - 1)
4
E
(1-ajhh
jjj
((,k)
aj
(1.A)
n_ = j
(5la)
$
S -
8y13
(t,k)
(133)
Combining (131) and (133), we have
(t.A) + 8y1
(t.k)
4
+ 5y (1:))
&
5 2---
This combining with (63) indicates that for all i, [S], # j.
Vt +1: -(H0-H
4
(H())
H (T (T+
& s 1
41

---

C.1.9 Proof of Lemma 9
By our choice of T2 (see (65)), we have
exp )
á (Ta)
(Ta) (N + m 2) exp( H 2 + exp(H (Ta)) <
e
2m
By Induction Hypothesis I and II we know that for all j e |S] and i {0, 8Y
WEN: H
N H
Combining the above two expressions, we have
exp(H[[1)) (12) -
(N + m - e)exp( + exp(H1 (12))
<
V 2m
which gives
(1(()) (v²m)
2m 1))
(N (IN+ m-11 1)
(
(vm-1))
2m
log
(N + m-1)
3 log
N-1
By (68a), we have
H
(T1)
(e Ti)
4mS
Combining the above two expressions, we have
12mS
(N (N+ 3-++-1p m- 1)
(v (v---))
22m -1 ))
T2 Ti &
log
176
C.1.10
Proof of Lemma 10
For notation simplicity, we let n-1 (8) From (49a), it gives that
A (t.A)
1
Σ [=[S]
Σ
(aft))
(((h)))) 2 + (1 - a s(f))s
(t,b) )2 + (aa (1(2)))22
n =
2S
25
1
Σ
max (t.k)
aij
Σ
aij
(1,k)
(1
(t.a)
a (t.ki
*(i),s
)2
m
jE[S]
esige
1
Σ
E
(1
((,b)) 2
n-
5
JE[S]
and similarly, (49b) gives that
A (LAA)
E
Σ
a
(t.k)
2 (1- 2 (1-aa)) )
m- m=
25
JE[S]
(efs). 190
1 Σ
(1 (t,k) (2
<
E
زز )
n- 9
5 jE[s]
Thus we have when >> T2 1:
Lurain ( (0())) =
Σ
A (1)b) + 1 (c.
2 Σ
Σ
E
(1
((.))) 2
11 =j E
S k-
1=
[
42

---

where the last inequality follows from our choice of Ta (see (65)).
Since H is strictly increasing after iteration T2, we have
exp(H exp( ( +
3 > T2 11
((.K) 32
ijj 2
1
2m
from which we deduce (76).
C.2 Proof of Theorem
For any non-leaf node i, we let (1) and c2(1) denote the left child and right child of node i, respectively. Let (1) ei (i) for any that's not a leaf node, and we define č1 (g) (n2n)
We also define p(t) p(i) for any i that's 's not the root, and p(r) = p(m1) = 0.
Throughout the proof, we suppose Assumption and Assumption 5 hold, N is a sufficiently large constant, v Nm and e v 1/poly(NN is sufficiently small.
Loss simplification and gradient computation. Letting
(b(b] == (²)) (1: di), (()) (di + 1: :1:d),, g((k = 5 (1: d1), g(k) = (d) + :d1d,,
2(6) (2d 1: :21 d2),
g(k) (2d) 1: :1:21 d2),
and
A (k) 0
1= - ET Pian (b) - 1 A/ai = ET 2
1
T Pain ( -gt)) (k) = ET Poan
2(x)
((k) 3
2
2 2
1
(134)
we can rewrite the loss function as
2m
Ciaiin
( (8) A( + A
(135)
Recalling À = (ao, a1..........................aa ai....,ss) as) E (8+1) and S = (s}, sb), we further define the following
matrices:
U = À B
= G ((++)) lE (1,2), 22x2 == C35 e 2x
(136)
U3 = St 22x2
whose entries are denoted by U1a and Vi.ig, respectively. Define the following gradients for all ee [2m] and e (1, as
2 (k)
AA. (k)
(k)
80t.i.5
(t.k)
(k)
(137)
2
2
VVi.
since A and A (k)
depend only U, and A (k) depends only on Vi.
We define the following attention weights:
N exp (a Bia; az + ay B2aj + s B35.) (p.qa)))G..)) En (((k-)) Ve
(k)
(138a)
exp aT) B1aj + af Braj + 8 B380
Nybhal exp (a Cia, + az C2a, + 8 C38r)
((K) (p.g.))....)) En oe N((-)) exp a Cia + C2a; st ST
(138b)
where u, 10, p.q.j E (0, ........... 1.... S). and N (k) P-q.u N (k)
p.g.a (T) denotes the number of times
(ap; Gg. sw) appearing in the columns of ((k-1) ( E(k- (T) (set (((k) = 0 if the combination
(ap, ag. sn) does not exist in E(k-1)). Then we have
Vk [2m]
(k--1)
10,1,2)
43

---

and
(K) oo 1,
Σ
((k)
(p.q.o)((..))
=1.
(139)
We simplify the loss function in the following lemma, whose proof is given in Appendix C.2.4. Lemma 11 (Loss simplification). Suppose Assumption 3 and Assumption 5 hold. For A x we have
(k)
Wk E [m]
A (k)
1 (2) 1) (1- -11-appaaa
(0) (((
(K) 0f (aG))I
2
E
(a
+ + +((aaio))
(R)
gE[S|/b1
(140a)
k = m 1:
A
(m++)
E
1 1
(1- (0.00)..1) -00000,0) (0..1)...))
-1-
a (K)
(2)
1-1
Σ
(..00)
(k)
+ a (((a)))))))))
$E(S]MB
(140b)
m 2 k < 2m:
A
E
1|
(k)
53 (1-
a1- (k)
(K)
2
--1
Σ
(0
(K)
+ (R)
(p11))0000... +
(E(S)M
(140c)
For (k) we have
1 8
vk E [m] : (k)
Σ
E
1/
01-
(4)
(4)
2 j-1
(0
+ (K)
(10.mmm1))(mm.111))) )²
2
(0)
Σ
(k)
(10000))0) + (k)
2e[s]
(141a)
k m + 1: A (m+1)
Σ
5
1im in{ (1 (22e)
(k)
2
(m22))))))
(.2....6.1))
2
(k)
+ a (0 1)6.1))
Σ
(1.66))))) + a (k)
(k)
((.666)))))....) + (oP0))))))))
(141b)
m + 2 & lo < 2m : A
Σ
E
Lin =(c(())) = (1-a a
(1.0000)).)
(A)
(1.000))))0..) - -...0)(13)))
(k)
(1)
)
2
iJEts]
(200.m.00)....))) + (k)
Σ
(A) (R)
+ (
(k)
+ (0000.00
)"}
)²)
(la.0160.000010.0
<c[s]
9/4
(141c)
44

---

For ! we have
k
A
Σ
8
38 34(1 (1-
Σ
3(k)
( +
2 *EE
Σ
3(k)
(.....))...) )
2
qe [s]
(e[s]
(142a)
k 11 A (m+1)
1
s 2
E
1(m1
i) (0((1-
Σ
((h)
2
(k)
Cc a.)...))))
}1
2
j=1
(E[S]
qE[s]
(142b)
1 8
ib 1((11))
2 )
}}
m + 2 k& 2m:
A
2 E
Σ
lim1
(e)). qE[s]
Σ
((k)
Σ
[(k)
(a, (2.
ge[s]
(142c)
Here, n-1 denotes the node index of the embedding x=1 (k-1)
We also provide a simplified expression ( 68i.2.3 given Appendix C.2.5.
(k)
(2(h) in the following lemma, whose proof
Lemma 12 (Gradient expression). We have for all i (O...., S). j e [S]. k e [2m], & far, u):
d1 (k)
21 (k)
(K)
Σ
1,1
OU1.1
(p. da (e)
a
(n...)).(.))
(K)
(a.u,0) da (L...))((..)
(k)
((.g.)).(..))
(b)
1
Σ
(how.))
(k)
(143a)
se[s]
wef0.13
a
(M)
(k)
Σ
a (6)
Σ
(k) ((.i,w)(..))
252,
a
2U2,5.5
do (K)
(p.e.u.e)
(p...u).(..))
A (k)
(k)
(1
11
1
Σ
(k)
(143b)
(p.u,v) do (k)
(p...))((mm))
(p.i.))....))
rejo.....3 me(0.1)
and any u, E. {0, 1}, have
2 (k)
(b)
(k)
Σ
(k) (p
Σ
853,0,0
a ((n
UUs,u
(p.4.)) da (p.q,a) (.6.))
(k)
refo Ae[s]
Σ
5
(A)
1
Σ
(k) (ross)) (143c)
(k)
(p...u)((..))
(p.9.) (p.g.))(..o)
reto Ae[9]
where ue 1 - u. Similarly, we have
(k)
(k
(k)
(k)
(-
((k)
Σ
g(k)
821,i.j
V
(g(kk
(p.g.a))(L)))
(pla.n,o) piz
(p.g.a).(..))
2
we(D,1)
aA (k)
+
Σ
(()
Σ
(144a)
ag(k)
(g.u,))
(2...))((..)))
sc[a]
weto.13
(k)
(k)
(k)
822,1.1
(k)
Σ
(A))
Σ
((k) ((rd.w))(i.))
VV2..
(K(
(K)
(pro.n.0)
(n...))
.e.0.....)
45

---

(k)
(1-
Σ
g(k)
(p..a))((.))) 1
Σ
(
(rossnndtt))) (144b)
03(h)
(p,a,))
(p.i.))(i..)))
rc(0....53
wc(0,1]
and for any u, {0, 1}, we have
(k)
Σ
(
Σ
8 ((aa))))
023,0,0
(Pai))
(.
reta
Se(8]
a (K)
Σ
g(k)
1-
Σ
(r.))))(6.)) (144c)
(...a),(..)))
(p-9-5)
(..0))
ret0
sc(3]
Define
E EyE[s] ((6)
ke [m + 1]
(R)
(a)))...)). 1.1)
(145)
LSss 3
ke (m + 2,..., 2m)
(.0)))
E gess (a (0)....)(..
(k)
ke [m 1]
W1
(146)
Eyels ( (1)
A E {m + 2, ...... 2m)
which represent the attention put on the tokens with stage embedding 8f and sb respectively the k-th reasoning step. We have
Ak [2m]
w
(R)
+
w
(R)
(147)
We also define
((k)
ke [m + 1] R {m + 2,...,2m!
3 (k)
Ke [m + 1] Ee {m + ,,...,mm, (148)
(k)
0) g(()
(K)
Po
90
()
10,m
for all (0...., S). we define
((k)
E [m + 1] be fm + 2..... 2m}
ke [m + 1]
(k)
P25 (e000066m11)
(K)
Pi.s
g(k)
he (m + . 2m} (149)
and for any e [S]. define
(k)
(()
(t...))) e [m + 1]
9i (.).).
(k)
kE [m + 1]
Pi (p(0)..00..mm...)))))
E {m + ,,...,,m..
ke {m 2,...,,m. 2m} (150)
We further simplify 8z( (1 (1,2,3)) in the following lemma, whose proof is given in Ap- pendix C.2.6.
(k)
Lemma 13. We have
R
(k)
2
w
(k)
(8)
(k)
n_ j
ke [m]
Vje [
8z1,0,
m
E
wi
(k)
w
(R)
Po
(A)
90 (A)
n_
ke {m + 1,..., 2m)
(151)
2
E
Wi
(k)
w
(4)
(k)
P1.
+
(8)
(1)
R- j
ke [m]
82
mip
(k)
wi
(6)
+
(6)
(4)
(A)
n- j
ke E fm + 1.... 2m3
(152)
46

---

E
w
(k)
2
(8)
9
(8)
(4)
n- = j
ke [m]
Vijj e [S]
(R)
10g
E
Wi
(k)
(R)
P (4)
n- 1
ee fm + I,... ..... 2m)
;
(153)
and
2E
W
(k)
Wi
(k)
A: E [m]]
3,0,0
(k)
-8=3.1
1,1,111
(k)
2E
00
2
w
m + 7
(154)
0,
A: E {m + 2,.... 2m}
E [m 1]
8 (k)
02000,0
-82g,1,0
2E
W11
(K)
2
w
(R)
{m + ....., 2m)
(155)
We also have
Vi.j E (0
S):
1,,,0
(k) 8z
022,0.J = O.
(156)
We next give lemma that simplifies the gradient expressions of (h) and 8y, (8) for 1 = 1,2,3. Lemma 14. For any j [S], we have
R
a (1)
(0000000000 - (00.)))..))))) (0.3.))
(1)
(1)
(1- (1- (1)
(0m.000))))))) (...))).)))
(1)
+ (1)
(((0m100)63.)) )² (1- (1-a (1)
(0..1.0)...) (1-a
(1)
10000000000000 - -..... a (10))00))
(100 + + +05....1))))
(1
Σ
(1) (P))).))).....
(
(0 + (1) 12m
(.......))
ve(STetimm)>
(157)
For k E (2, .2....,,1y m., we have
8x,,0,0
5
1
a (6)
(0.0..)) (1
10.11.))
(8)
10 13.0.1)
(K)
(K)
(K)
(k) (p(ng Jan 1).(.))
1 a (k)
(k)
(p( h.gm 1))
(0. 1).0.1)
(U..).0)....
(0,m 0).15.1 + +a (K)
(p(ngm (...).
0. 1).()))
(p(ngm Jigm (1.0))
a
(0.0gm (D)..)..)).(s.a))
(K) -
(PG)).0).....
(Pb))))))...)) (116)).))))....1
(k)
(k)
(10)).)))))))
(k)
(a (fo0.).0)...1) + (K)
Σ
(4(0)
(Glg)) + (2100))))))))))))))
(a
(a (1..))))..)))) + (u.gm ().g...)...))))
(k)
n1 = 5
(158)
for k m7 + 1, we have
801,0.1
(m-1)
_ -
5
1
a (0.00).11)
(m+1)
(2, ()).)...
(m++1)
(0.1.)).0.3)
1m+1)
(1 (1-0 (0.000..)).C..1)
(m++)
(m++1)
(1).gmm1).(G..))
(p(ngm) j.ggo 00,(1.1)
(m+11
(0..ggm. 11.(4))
1m+1)
a (p(ngm (agkgg))ngmmm).... ()((i)) (1-a (m+1)
(m+1)
11-0/00.0)...n ((..1.0)...1)
(0.ngmm
(m+11)
(m+1)
(m+11)
(m+11)
(m+
+ (m+1))
1
(0....)....
(m+1)
1m+
(m+1)
(m+1)
(m+1)
(p(ngm))) 0).(.1)
a (0,0g. 11.3.1)
(ringm Janm 0).6.1)
(0.1)))).6.)
(0.ngmm 1).6.1)
Σ
(c(m+))
((p())a0)) +
(m+1)
(m+1)
(
(m+1)
+ + (m+1)
n n- =
gangm. )
(159)
47

---

for k E {m + ,,...,2mi. we have
601,0.3 s
(a
(0.m..0)..0.0 + a ( 1..000
a
(1-a
(0,0gmm 13.0003
(k)
(00mm,m1)):0000
+ a (plng)) (0.mm)m)mm0)....))) D).G.0)
00,2m
(10110)) (10,.1000..00
- a (0, 1265,0)
(1-a
(K)
0 (8)
(b)
(k)
(k)
(K)
(p0)).00660011
+ + m0.0
a (k)
(2)
(0,m1 01.1200
mm..mi...aa 1).10))))
03,400
+ a (0mmm.1300)))
(8)
(k)
(A)
(R)
(K) a00,ngm 13,6500
1.0.mmm1MMi0
a (p(ngm higm 0111100)
Σ
(a
(a
(atom) + a (0.0gm (000gmm11..0.00 (1.5..))
(K)
121
ec[s]
(160)
For all E [m]] we have
(K) 8yi,0.)
E
(a (0)
00..m. + ((..ngm.))..... (
(4)
(2)
-10.n.m 11.01.1)
(1)
5
+
a
(P(....)..... - (o0)..)..).....
(k)
(K)
(P0))..)).G.)) + +0000000000000 (o 00.01.) + a (0.
Σ
( (4)
(oma)))))) (a (000m.)))...))))) + 00
(K)
1- =
(9[8]
(161)
for k = 17 + 1, we have
5y(m+))
1
E
(0.1.0))6..)
(m++1)
(000a0 a).6.a) 2 (0.10))(a))) (011111100
(m+1)
(m+1)
(a.mmm1))...))
(m+11
5
(m+1)
(m+1)
(m+1)
(12)).0)))
(1g. (xg).0))6..)
(m2.11))6611)
(ogg) + (m+1)
(m+1)
(m+1)
(m+1)
+ (m+1)
(m2.02(1g).00).63) +
(o2.5.1)6111)
(0.03)))))
Σ
(a (m+1)
(1911))))))) + * (0 (..6))))))..))
(m+1)
(m+1)
(
(0..%)) + (m+1)
(m+1)
1-1 1-1=
(162)
and all k {m 2.... ,2m), we have
8y1,0.
S
1 Σ
1 1m = ( (2/0..
+ +o (k)
ag (10..gm.).cc.0))) )
(1 (1- -0 .0.6..0
(0.0gm (00.2mm.).....0
(8)
sE[s]
+
a
(La100,00000000 (Loz())0))00 - (1..1).1..0)
(k)
(000
(1 + (k)
(k)
(..2..,)....) (...).0000)
(8)
(0.m..0).... + a (k)
(k)
Σ
(a00
(1 (g02000000000000
(k)
(a
+ (K)
KE[S]
(163)
For all E [m]] we have
(k)
1 R.
(1- (K)
(K)
(
(o + (K)
5
48

---

Σ
(a)
(2.61) + + (100)))))))))
(1 + (8)
(000))) n =
(164)
for k
L, we have
62255
(1
a (0..0)....) (all)))-(llln (616).......)
(m+1))
(m+1)
(a
(<o...0..0.))) + ((11)))))
(m+1)
5
Σ
(1(())..00)....) + + (......))))
(m-=1)
(m+11
( (aos.0).01))
((G6)).. 1- =
(m++1)
2c[3]
(165)
for all & {m + 2.... 2,,), we have
8x 2x30
(6)
1 E
(1 -
(k)
- (
- (2)
(4)
5
Σ
(k) (p(2)..00) + (K)
(210))).) (2 + (4)
(o
+a00ss....... m==
vc(s)
(166)
For 1, we have
8y2.3.
(1)
SE
(1 a1
- (om.emmm " (1)
"(orroaaa (1- (1)
(19..00).00000.. (0..)...))
(
S
(1)
a ( a (ooorr (1 (1-a 1--(6rrs.0..a (...)...))
a (0.11.0)...)
(1
a
(1)
a (1)
(011010000.)) (iG))o))))) (160)0006.)) + (1)
(1)
(1)
(1)
+ + f O0
(3)
(10)))))...) + (0)
(
Σ
(a
(1a10))))) + ①)
(1000,) + (
(
n- =
vPb())
(167)
for all ls E {2, m]. we have
5y2.1.1 = E - -11-a a
(l0......mmm - (oo..n) (
(00 (nthsoraa. (1 (1- - (0)
+ (MaathsanM + a (8)
(K)
(110))44)00 + + (k)
.elooss...... (1 a (11)).0)00111
(1 - (1- 1-0000000:100000000000000000 (01
(8)
(1p0),.))) (1 (aranam)
(6) (600).0
(8)
(A) (A) +o (6)
(8)
(k)
(100110)).)))
+ (aMAA.G
(8)
(100) + (k)
(GGG) + (4)
(allaaaaoo.. (p6)..0)... + (k)
(k)
(am.m..)))))
Σ
(4)
a (4)
G2000.0600 + a arta...)))
(aammoo.. (
(K)
+ (k)
a n =j
(168)
for k
1, we have
(m+1) 1
1 - (1
(m+1)
(m+1)
(m++1)
(m+1)
5y333 5
(m2.91(99))006...
(.a. (...).))....
49

---

(1 (1-0 (0.22.0.(.3))
(m+1)
(m++1)
(( (m+1)
(0.10);(111) + +a (m+1)
(m+1)
(m+1)
(m+1)
(0.mgmm3)).6..))
a (0.1D).G.1)
0 (0.12.0)(61))
(m+1)
(m+1)
(m+1)
(a
(m+1)
(m-1)
(12.01(1)))0))66311
(02.62(22)0))6111 (m2.)))66.))
(122:11
(m+1)
(m+1))
((0...0).6.))
+
0 (o2d11)6.))
-
(m+1)
+ a (m+1)
(m+1)
(m+1)
(m+1)
(0, (26,))
(0.51.0)...
Σ (m+1)
(a..1ta).00..... + (m+1)
(m+1)
(0...00).31)) + (m+1)
(m+1)
n-1 =1==j
(5[9) 9m
(169)
for E fm + 2,.., ,2m,, have
8y2.5.
S
1
Σ
E
- 11g
1 (1gge-m
a1-a -1-0af (B16000))
(N)
(6)
(1..1))610))
(K)
(6..1)
(K)
1E[8]
(K)
(1.1).00))
(1)0) (K)
h +000 (K)
(K)
(8)
(1-a (4)
(K)
(P0)).000000
(1.1).60
+ (1 (1c1(0)030000)))
- - (.....))))))
(1e20t)00000111
(k)
(0.00.0)
+ + (1c2(()0000
(N)
(K)
(I...)..3
(611000011100
(k)
+
(k)
(k)
(K)
(K)
(N)
(Leas) + a (0100.)))
Σ
(a (k)
+ G...0))
(...)))))))
R
(a ((..)))))))))) + + (4)
+0000000000
m-1 =
(170)
For 1: E [ma]. we have
5x = -6x 0,
(171)
and
5
(1
a0 (4)
(P0)...00.66))
(K)
(0)
Σ
a ((a)).0))...))
[E[S]
Σ
(k)
(6)
(b)
(2)
(2)
(c)
1- Σ dE[s]
(0) ((..)..0)..11
ge (E(8)
-1-1 (P0)..00...
(1 (1-a!0 (6)
(1)
(1)
Σ
( (P.)).
rE[s]
Σ
(N)
(il + (b)
Σ [E[8]
n-
(
(172)
For m we have
(m+1)
-55
3,1,0
(m+1)
0,
(173)
50

---

and
80.3,0.1
- -,,1,111
(m)
3,11
1
(m+1)
(m+1))
(m+
(m+
a (0..0).(..))
a (0.1.0).111))
5
0E[8]
Σ
(n++1)
(m+1)
(m+1)
(m-1)
1- Σ
(m+
oE[s]
(1 - (0.10)..1)
(m+1))
Σ
(m+11)
ge|S]
(o++1)
(m+1)
(m++1)
(m++1)
n- n-1 =
qu[S]Ma)
ge[S]
(174)
For E {m + .. 2m), we have
3.0,001
(M)
0x3,,11 0,
(k)
(175)
and
8x. = -8206.1.0
5
E
(1
(1-af
a
(P0)).00..00000
(k)
(8)
(2)
1-
Σ
(P(.))..0)......
(0)
DE(S]
Σ
(4)
(k)
+ (.11.0)......0..0)))
+
(b)
(2)
(K)
1- Σ
((0))..)).G..0
(A)
(E(M)
RE[S]
(1
(1)
Σ
vE[s]
(4)
(K)
Σ
n- =
yc[8]
(176)
For ke E [m]], we have
6y3.0.0
(k)
-8y3.100
(k)
(177)
and
(K)
- -5y(h)
(1-
5 F.
(1
(160)..0).00.))
(K)
(... (0)
(k)
(o0hnoosoa) + +
Σ
(k) (...)..0))....
dE[S]
(k)
+ (K)
(6)
1- Σ
(k)
qE[s]
(0
(4)
(4)
(a
(omttt)))) + +
(2)
(u(S](Pp())
GE[S]
( o).(ia)) f n. jo 1)GA))
(k)
(
a 400.00)))
(k)
51

---

Σ
+
(0 + + (4)
(M)
(..66. 1366
(4)
Σ
n- =
(E[STMP())
(178)
For m + we have
(m+1)
(m++1 0,
8y3,0,0 -dy3,1,0
(179)
and
0y3,0,1 111
(m+1)
(m+1)
1 E
(m++)
(m+1)
(m+1)
(.2...).6..)
(o2.).0)))))).G)
(v2.))())))))
a (m+1)
-
(m2.02 (mgg))2(m2).)))))..)
(m+1)
1-
Σ
()..
(m+1)
[E[S]
+ (foo..))...))
(m+1)
+0 10,0ggm
(m+1)
(m-1)
Σ
(m+1)
9c[8]
(m+1)
(m+1)
(
(m+1)
(m+1)
((..116).0))6...)) + (a..2()).))....)
+
gE[3]Mm2)
1
Σ
(m+1)
(Ft))..00.G.1)
(e[S]
+ + a (m2cs)).6.)
(m+1)
(m2))((())))))..)) - 0 (m-1)
(m++)
(m+1)
Σ
(m+1)
[E|S]
(m+1))
(a.cstt).0)....) ((....).........) (
(m+1)
(m+1)
(a.pto).))....
(m+1)
Σ
(m+1)
T-
(E(3][mm)
(E(
(180)
For e {m + 2,...,2') we have
833.0.1
(k)
-833.111
(181)
and
8y3,0.0
8y3.1.0
1
Σ
E
1
[fgk-m
il{
(1-aat
(K)
(122())00)00000 - -a........))
(K)
(K)
S
xc[8]
-000200000.c..00 (1-
(a000
(G.00000))) - (A)
Σ
(K)
(E[S]
(4)
- - ().0.)))). ((0m..)).0..
(K)
(6)
1- Σ
(la)...))....
(k)
qE[S]
(a
(oaat.)))). + (2)
+ (o00))))....0))
(0[S]]()
(0
1- Σ
(E[S]
+ +1 (1-a(0) (4)
(- (...0000)0 (6.201000000000
(k)
f...)...00))
(k)
(K)
Σ
(3)
(
52

---

(K)
(K)
Σ
000.gg 1105001)
mm 126000
96/8]
(aa110000)000) + (1 (6[S]M4)
Σ
(o
# (1pt))........ a (ara).
(k)
Σ
(4)
n =
[E[S]
(182)
Lemma 14 can be verified by straightforward calculation, whose proof is thus omitted for conciseness.
C.2.1 Proof outline
Let the iteratio-vvarying versions of (136) be
u1 À B A
Ug!0) A B Ae R U.!) 3 3 2x (183a)
((() A c AE (8+1)xx2
V2 A c Ae Sx2 Vy 5 c/ 5 R222 (183b)
and the index starts from 0 for U/ and vi( when 1,2 (i.e., ULo ag B aj. Ví.o. af (
az). Then we have for all / [3]
2m
Vi 10, 8),j [S]
((++1) 11(1)
Σ
(1.b)
+ (y((.h1 (1.b)
(184a)
2m
(++1) v.(1)
82 1,,,5
(t,K)
(184b)
k=1
Therefore, is possible to analyze the dynamics of 07(0) and v,(0) separately.
Similar as the backward case, by symmetry the training distribution, we have for any e N. i [S]. i * (iا) (1 2) are equal to each other, U( (1 1, are equal to each other,
and U,00, equal each other. Note that dg only appears in (k) but never appears in (1),, thus U13.0 2.j.0 U:4 0 for all N. The same holds for (t) 1,5 1, 1). Therefore, could
define the following pattern:
p (t)
!1)
a
0 11
1
12
u!!t)
0
11,1
0
V1.2
(t)
11
(1)
(E)
(0)
14 (0)
1.1
VL.
V1,2
V1
(185)
and
0
P2
ub (t)
0
1/
0
b(t)
(t)
2,2
1/2
(t)
v(())
0
2,1
2 (8)
((()
=
V2
(()
1/2,2
(6)
(186)
0
2,1
(t)
2
0
2,2
(1)
1/2
(1)
((t)
12
(9)
for some (t)) B(0)
(())
11 .12 (9) 1/2
(1) (0)
1/ 1,2:
(1)
(1)
2,1
(0) (1)
2,2
E R. We set a
b(0) 1.
Regarding U and (!) by Lemma 13 and Lemma 14, follows for & (I, 2]
Vt N: [t -5c(1 and 8e.0 -5(( 5c11
53

---

Thus by our initialization (Ug, V/3 = 0), we have
(1t))
3,0,1 = - U31111
and
Ut(t)
= -
3.1,1
3,0,0
3,1,0*
(tt)
(t)
= (0)
3,0,1
= 3,
and
3,0,0
3.100
Therefore, we could define
US p1 61
(1)
(b
82
(())
U2
(t)
(4)
(8)
(187)
b2
for some 6 (0)
Ci C2
(1) (1)
R. where p1 P2
(1) (t)
defined in (185) and (186).
C.2.2 Training dynamics of C1, C22 C3 Training dynamics of Phase L.-C. Define
T = max N cxp((()) & v2N
(188)
We show that the following key relations hold during Phase I.I-C. Lemma 15, Assume the learning rate satisfies 1) E Nm When [T]] we have un monotonically increases, and
1.
(0) 12 z
in
(189a)
SN3/22
b9
(0)
+
log
))
(t) 112
(189b)
N
llog N
log N
log
1 =
122
102 (0 = 0
122
122
O
11
(8)
(189c)
N
N
N
2
(1)
> 1/2.
(189d)
The proof of Lemma 15 is given in Appendix C.2.7. By (189a) and (188), we know that
SN3/2 log N
T x
(190)
"
Training dynamics of Phase I.2-C. We let
T = max IE N: exp(n!)) & NJ.
(191)
We will show the following relations hold during Phase I.2-C.
Lemma 16. Assume the learning rate satisfies n S Nm When t (TP, ... T have μlz (0) monotonically increases, and
(1)
(2
Ttt TT)
11
42
(192a)
SN
log² N ug
600) 6001100 11 = 1 O
(t)
(192b)
N
|1 = o
(MOR/
log N
N 166
| = (1(11g
(log N'
1u22 =
log N
N
N
(192c)
10 1/4.
(192d)
At the end of Phase I.2-C, we have 02 (T
[0.25, 0.26], and c (T
> 0.48.
54

---

The proof of Lemma 16 is given in Appendix C.2.8. By (192a), we know that
SN log N
T T
(193)
Training dynamics of Phase II-C. We set
N exp (0.4742 + N
TØ max teN:
(194)
V 6m
We show the following relations hold during Phase II-C.
Lemma 17. Assume 60 > 0 is small enough obeying co x 1/poly(N) and nj E Nm t tE (TZ. ,T3), J. any e E (0, col. we have μ2 monotonically increases, and
1
When
12 H2 122)
(1)
(1 12
3/2
(195a)
5
(1 112
(1)
1 θ
1
p1
(1)
(195b)
1221 δ
12
(0)
11
12
(4)
10
Ō
12
(0)
(195c)
N
2 [0.24. 0.26],
(
((+1)
C2
(1-1)
1
(195d)
N
The proof of Lemma 17 is given in Appendix C.2.9. By (194) and (195a), we have
TØ T x S 3/2
N
log
(196)
It's easy to from the proof of Lemma 17 that P1 keeps increasing after Ty and (195b), (195c) and (195d) hold for all t >/ Tς. We give the following lemma to mark the convergence of ΣAT1 A! whose proof is given in Appendix C.2.10.
Lemma 18. After t > T + 1, we have
2m
(197)
3
C.2.3 Training dynamics of Bi, B2, B3 Training dynamics of Phase I.1-B. We define
rP = max t E N: exp (n5*) & N...
(198)
and call the period & < T as Phase I.I-B. The following lemma gives the training dynamics of U1, U2, U3 in Phase I.I...
(1) Lemma 19. Whent t E (0, 1.... I. we have (1) ((0 pi decreases and then increases, and
monotonically increase. 6 Pi first
(1)
n N
M1 2
(199a)
6 < and exp (1 -1-33 26 )
(t)
z
N
when 6()) O.
(199b)
m
(t)
(1)
(t)
v
(1) = δ
#
11
(1)
= δ
(
= ô
(199c)
N
N
N
exp (a (11-a (11aa())11)) $ N.
(199d)
55

---

The proof of Lemma is given in Appendix C.2.12. By (199b) we know that b(7T*) (199d) we know that (T > 0. Further, by (199a) we know that
x= 0. and by
NS NSlog. N T x
(200)
7
Training dynamics of Phase I.2-B. Let
T = max {t tt T exp (11 26l2) (1 )((") 3 CoN..
(201)
where Co is a sufficiently large constant. Without loss of generality, we assume T.2 > T We call the period I E (TP, T + 2.... T..1 as Phase I.2-B. The following lemma formally describes the training dynamics of U,, U2 U3 in Phase I.2-B, whose proof is given in Appendix C.2.13.
(1) Lemma 20. When (2P.. TT + 2, 2.... TT3,, }. we have Hi b1 Pi monotonically increase, and
(0)
n T
11 - 11
(1)
(T()
(202a)
NS
exp
(1 (10)))41)) (1)
s N 1
(202b)
0 & (1) $ 2
(202c)
(1)
(0)
(1)
(1) = δ
Mi
(0)
M1
(e)
P1
P1.1
111
(202d)
N
N
N
At the end of Phase 1.2-B, we have
(1 (1-a (1-a(T3))) (3
exp
= N.
(203a)
by (T
01 01) (T9))
(203b)
By (202a) and (203b) we know that
NS Nogog N
T -TB s
(204)
Training dynamics of Phase II-B. We define
N 2m 1
T = max
tEN:
(205)
exp (0.46]( (0)
+ N + 2m
6m
and call the period t E 1 TT,8 + 1,... TTg) as Phase II-B. We'll show that at the start of the next phase (i.e., shortly after T. exp (1 - a (t) n1 increases to, and stays at O(N), indicating that when training long enough, (4) will approach 1, ((()) will approach and stay at 0.25, and 6(t) is larger than 669) The training dynamics of U1, U2 U3 in Phase II-B is formally given in the following lemma, whose proof is given in Appendix C.2.14.
Lemma 21. When t E (I , T. + 1,.... ,T,"), J, we have 11 (1) (0) (t) #1
monotonically increase, and
(1 - T5) 1
1 (1)
M1
(TF)
(206a)
S
m
exp (1- (t (1) N,
(206b)
(1) [0.24, 0.26] (1)
(206c)
(t)
(t)
(4) δ
N
1,1
(e)
P1
N
V12
(
N
(206d)
56

---

At the end of Phase II-B, we have
a (T5) [0.99, 1),
(207a)
11 z b2
(T5)
(I3)
+ Ö
1 N
(207b)
By (206a) we know that
S
3/2 log
N
T T x
(208)
After Phase II-B. It's easy see from our proof that after Phase II-B, ji slowly), (206d) and (207) of Lemma 21 still hold. After Phase II-B, the loss reaches as indicated by the following lemma, whose proof given in Appendix C.2.15.
(9) keeps increasing (but
Lemma 22. After T
have
2m
A A!AA)
2c
(209)
C.2.4 Proof of Lemma 11
Note that when & [m]] we have (K) ngm-k44, and hence by (9), when &t [m]] we have
A
Σ
(k)
(k)
2 On
(1.2.0)). 1.1) + (1, 02i41
a (b)
a (K)
Σ
(k) (agnma--1 1.1) On
O),(n on
(0, 11. 2.1) Onam
1=m-k+t
8
(n-1 = o1 E
(1-0 11-- (P0...0). - (k)
(100).0))... aj
2
Σ
(a
((l.........)) + ((lg))..)) 0g
(210)
By Assumption we have ao, a as are orthonormal, and hence
1 s
1n In (1 (1-0(0) (1) (N) (PO)..))...) ((oosoo0..
A
(A)
2
}1
a (ppl))-.0))..... + a ((0.D.
(211)
sE[ST|bi
This gives (140a). The relations (140b) and (140c) can be verified in a similar way.
For A (k)
when & [m], x1
(k-1)
angm & = ap(n.)); and hence by (9), we have
2 -1
A (k) 1 y
Σ
(12210))) + (k)
(k)
(k)
Opin 1)
1..22)+1..000--111)) a;
2
- (aftt) (k)
+ (000000..)) () +
(K)
2000001 Σ
(k) (1 Ony!
1=m-kt2
1i- 11n-1= = /
1
E
(1-
((0)..0).GG.)
(8)
-000)..6...0...... a. (a
(k)
(a/om...0)....))
0 (0 1).(.11) ao
(k)
57

---

Σ
(g) (%aano.. + + (1.600)))) +
(M)
((Pp)). Oq
(k)
g
(212)
Since 4o, 01, are orthonormal, we have
A
(R)
1
8
1mm = ih (1 -appsoo..n--uoooo.....
(6)
(a (k)
(a + a (00.gmm.))((..))
(e)
2
2
}}
21
E
(k) (1..2(0).00)...) (11.610)))))...1)
+
(R)
+ (op0o)))...)
(213)
(E(S]6))
This gives (141a). The relations (141b) and (141c) can be computed analogously.
Finally, for A (R)
have when kE e [m].
A
1 E
sb
(
(1.20,0) + +..22++1.000 81
2
((A)
((A)
((M)
(00.m.0))(mm.11)))
(0,02m,1)((n_1111))) ().
(n
(1-111)) 85
t=m-k++ k+2
2
1
8
1
E
1mm 51
2
(11-
8 [E[S]
Σ
(l....))
2
+
Σ [E(S)
( (P..)..0)0.1111
(214)
This gives (142a). The relations (142b) and (142c) can be computed analogously.
C.2.5 Proof of Lemma 12
It's easy verify by straightforward computation that for any b {0, 1.... S), c [S]. [S).
u, U.E.E. E (0. 13. k [2m], we have
da (K)
(k)
(1- leisr a (omes)t))) sefo.11
(k)
if p, j, if # p; j. otherwise,
day (((aa.. an
(p.g.)).(k))
(p.g.))((.)))
Σ
sc[s] ((ha,s))
(k)
(215)
se(0.1)
(k)
(1 (1- Σ re[s] a (k)
((a))))))))) (k)
if g, c = j.
da (K) (p.g.u).G.v)
(...))(())))) zc10,1 (k)
dag Bbd ac
(p.q.)).Lo))
2e(0.1)
r(ts (ra))))
if + c; j. otherwise,
(216)
and for any b, {0, 1}, we have
da (k) (p...u).L..))
(naa).( (K)
(k)
(1-Σ re(o.... a (k)
sc(s] (oa)))))
if u, E, if # u, otherwise.
(p.g.))((nn)
Evero... (k)
(217)
ds, B Sc
Then by the chain rule and the above three expressions, have (143a), (143b) and (143c). Also, (144a), (144b) and (144c) can be computed analogously.
C.2.6 Proof of Lemma 13
Using (142a) and (144a),
compute that for any i [m]]
(2(h) =S*
E B (1 w
(k)
(K) do (1 40 (A)
(k) + w 90
(N) (k)
(1
(8)
%0
Po
58

---

+ ((-
w(()) (wot))
% ) (0b
+ Po ) w/th)
(wo" (1) - pb
(
(k)
+ po p
n- = 5
(218)
By using the relation
wh
(k)
+ wi
(k)
1, we can further simplify the above expression as
(k)
(k)
1k E [m] : 821,0,5 =
2
(R)
wi
(k)
90
+
Pu
n-
1
(219)
S
w
(k)
wi
Similarly, we can compute that at the turning point where k = m + 1,
(m+1)
+ 90 (m++1
,,1,0,3 = s
2 E
(wfm+tt)
w/
(m+1)
Pu
1-- =
(220)
(m+1)
wi (m+1)
w)
and for any & E {m + , 2m}.
12.,j
S
2
E
[1i
(R)
w0
(k)
(k)
90
(221)
Wg
w1
Combining (219), (220) and (221), we have (151). The other expressions in Lemma 13 can be computed analogously.
C.2.7 Proof of Lemma 15
We prove the lemma by induction. First note that (189) is true when = 0 by our initialization. For induction, we assume (189) holds for all time steps & t (0 & t & T 1). Below we prove they hold at time + 1.
Step I: showing (189a) holds at step t + 1. By (184b) and (186), we know that
21m
(t+1)
112
1
Σ
(-52
(t.bk)
(222)
Thus to track the growth of jlz and show (189a) holds at step + 1, we need estimate its the terms in the gradient expression (151).
We start from bounding the iteration-varying versions of
%
Ph
for k E [m]] and
P
+ 9
47
(t.k)
for & (1,...,2m) 2m.. Note that for
o ()
when A = 1, we have
(t.k)
40
1,
(223)
w (t.Ac
where d0 ((.k)) and w ((.k) are the iteration-varying versions of (148) and (146), respectively;
when &o E (2..... m + 1). we have
(t.N)
exp
(1 + 2 112 + 22,2
(()
%
w (1,8)
exp (1 02 )uz (t) + v/²3) + (k- 2) exp
(t)
(t)
(2) (1)
exp (v()) + (b(0) c ol (3)
2 + 22 + 62 112
exp (1 + 02 142
(1)
log
N M exp (1 + (10)) + (k: - 2) exp
1+0
(1)
(1) c22 112
(1 ol)))
+ exp
(1 NN
"log² N ))
exp (12
3 1
(Lag**
log N N))
")
1
(224)
2 2exp
(12
(1)
+ k 2
N
k'
59

---

where in the second line we use the fact that when t & Tp, we have
log N
exp (x) =
= -(t)=+++)1+000
N
where r = (1) 12,11 12,2 (1 b()) )p2 (1) (note that 22 x log N by our choice of TT);
when & E {m + 2..... 2m), similarly, we have
exp (1-4) (()))(²)))
9o
(t,k)
11:
log? NN) N
wr(t-k (1.8)
exp (1 - ci yo (1)
(m 1) exp
-ce 41)) k/2
exp (1 - e
(4)
= = 1
log² N
exp
(225)
N
exx
(4) [12
m77 -
where the last step follows from our induction hypothesis. Besides, turning
Po (6,8)
(5,6)
when & [m]]
exp (11 - -(-)))
(6,8) Po
11 0
log N
ung (1.A)
N
exp
(1 -2 )(0)))) (N - 2) exp
(-cf0 11/ + exp (DL) - )/
(m!²)
= =(1+0
log? N
N ))
exp
(226)
2exp
(u2 N
VN
where the last two steps follow from induction hypothesis and the choice of TY:
when & = m + 1, we have
Pu (t.m+1)
1 N ))
log² NN)
exp
(2
(1) (1)
1122
o [1,m1]]
-N - 1) exp
--2 112
(1) (1) + exp
(2 - 2 )
(
exp (2 (t)
- -(1+0
log² N
(227)
N
N + exp 2412
when & E {m + ,,,,. we have
2m),
Po (1.b)
log² N
))
exp
(1 + 02 (1)) (t)
1 1
wig (t,k)
N
(N 4+ h: - (--m)esm mxp
(t)
+ 3exp (1+cl")
exp (2s5)
= -(1+ O
log" N
(228)
N
(1)
3 exp
u12
+ N 4+ k
Together, (224) and (226) give that for any k [m]:
(t.k)
(t,k)
9n
(229)
w [t.k)
(t.k)
VN
Besides, from (223) and (224) we can see that when k 1:
(t)
(t.1)
(t.1)
log" N
exp
H12
4 Po
1-
1 O
1,
(230)
(t.1)
[t,1)
(t)
w1 w
2 2exp
(12
+ m
60

---

and when l E m2]
(1)
(1)
40
(f,++))
(Log* N) ))
exp
(42
exp
112
1
2
wi wi (1++1)
(c,R)
N
2exp
M/
+ k
2 2exp
112
(2)
+ m 1
VN
(231)
where the last inequality follows from the choice of TP, which guarentees 12 x NN, and by (226) and (227), have
(t.
(c.++))
Pb
Po
(232)
(t.k)
(t,m+1)*
wo
w
The above two inequalities imply that for any k [m]]
(- (L.k)
(t.k)
Po (1.m+1)
(l.m+1)
(log* N
%
Po
+ 9o
O
(233)
w ((..)
(t.kk
w (1.m-1)
wi 11.m+1)
N
w
By comparing (224), (225), (226), (228), we have
(1.k)
(t,k)
(1,k')
(e,b")
tke [m]. R' (m 2,..., 2m)
40
Pb
Pb
0.
(234)
(t,b)
(1.A)
(2,1")
(1,
w
w
wi
wi
(k))
(k)
(k)
We next bound the iteration-varying versions of the terms
wwo
wi
and
w/
in the
expression of 82),0,; (c.f. (151)) for l [m]] and & {m+11, 1,.... 2m), respectively. By the induction hypothesis, have
When 1,
w6"
(t.1)
(11 O
log" N N
(N - 2) exp --2
(0) 11
+ 2 2eepp
((1 (1-c8 (1)
exp
(2 + 0 114 + (N 2) exp - |b
+ 2 eex
(1- -=("))²
(235)
exp ((2 + ( )((())
log" N N
(6.1))
+
exp
(2 + cla dl'))(()) + (N 2) exp
(-ch 112
(2)
2 2exp
(11- (t)
)u²) (2)
(236)
When E (2....,m
log² N)
+ 0
NV
(N - 2) exp (- ((6))) + 2exp (1-c(()))
2exp
(1 + cl )u5²)
+ (I - 2) exple (000444 ( + (N 2) exp (-c²") (16))
2exp
(1- 1--c()) (t)) c2 (1)
(237)
log' N
witth) kt.k)
1 o
N
2 2exp (11 - + (h- 2) exp(c
exx (1 + (()))
(Ic 2) exp(c"") (6/ (N - 2) exp (-c5" 2exp
(1 -
(11-22))) (6'))1 (9)
(238)
When = m + 1, we have
log² N
(1im==) 1 0
N
61

---

(N- 1) exp (-d()1(") + exp (2- )1(²)
2 exp (1+ (1+c))) )a{))
+ (m 1) exp(c (///)))))))))(9)) + (N exp (-dh" w!!))
(2- cf )2()
exp
(239)
log² N
1 + O
.V
2exp (1 + 02 )²') ( - (m - 1) exp(c22 ( 2²2
2 2exp (1 (1+c2)) ) + (m 1) exp((c) 50 //²) + (N - exp (-c3" + exp
(2 ce )4
(240)
When ke E {m + 2,..., ,2m], we have
loog N)
wfftht
1 o
N
2exp ((1 (112))) (m e)p (-cf² pp²)
(N 1 k - m) exp (c")114
122
(t)
3exp
(1 cf )u{²) + 2exp ((1 - cl") )x3")
(m exp (-c(² m1²))
log N N
2
< max
(241)
(2c)²) (0) 1)x2²) )
N exx
(1 1 O
llog² N))
N
(N + k m) exp (c (1)
1/2 exp o )()))()))
(N 4 + m) exp
(d (d")(²))
+ texp (1 + cf )xb") + exp
(11 (1t)) )د²²)
(m - exp
(-cf) (1²))
(242)
By (237), (238), (239), (240) and the fact that exp
(uz (1)
VN (c.f., (188)), compute that
for any s & t:
2 exp ( (0("")/²²) (1)
(unt.t)) 2 (wo(.at) =
Vk e [m] :
1 + N exp (a + 2c (2) )ub)) )
2
1 + & exp (1 + 20 )um
(243)
x/ exp
(1 + 2(()
2
(ug (uflmmtt) 2 (u!
(f.m+))))
1
1 + & exp (1 20 )ub )(("))
1 + 2 exp
+ (1+2cl")) )ut²)
(244)
which gives
(wg (wgg)))) 2
(w0 (1))
N
(245)
(o (
(w) (()mtt)))
(1 + 20 (2) )²²")
exp
With the above expressions, could derive the following lemma, whose proof can be found in Appendix C.2.11.
Lemma 23, Under our induction hypothesis, there exist absolute constants C, C" > 0 such that when 1) & Nm.
C we have
wo
(1)k)) 2 ((.k)
wi > (1 (1-
O m (m) (w (ftmmm1)) (t.m-1)
wp
(246)
62

---

By Lemma 23, (233) and (151) we know that
82 821,0.3)
$21.0.
(m+1)
E
(8)
(1)
(247)
SVN
Lemma 23 combined with (245) also implies that
exp (11 + 20 ) & mN.
(248)
By (243), (248) and our initialization we know that for any & [m]]:
2 ((.k))
(t.k) Wi
1
E
(249)
N
Moreover, similar as how we prove Lemma 23 in Appendix C.2.11, using relations (237), (238), (241), (242), we can show by contradiction that
(wff (cft
2 (e,6')
log N
ww (()))))
2 (t,k)
Vk E [m]] k' E {m + 2..... 2m}-
wy
wi
(250)
Recall (151) gives
(t.b) Vj [S]
82((.)
2 E
wo
(t.k)
2 (1,b)
wi
mi MTT
(1.b)
(t.k)
1- = j
k E [m]
E
wi
(t,k)
2
w
(2)b)
pb kc.
So
(t.k)
n-1 $
k (m 1,..., 2m).
w6
wy
Thus by (250) and (234), we have
Vj [S]
(t.k)
z
log N
(-62(1003
(2A))
Vk [m]] k' E {m + ...... 2m).
(251)
N
Combining (247), (249) and (251), we have
(t.k)) (1.2) wi
1
2
E
(252)
SVN
N
SN3/22
and thus
2m
(t 1y
f12
((+1)
222
P12
(t)
SN3/2
(253)
where we use the induction hypothesis in the last step. This shows (189a) is true at time t + 1.
Moreover, by (252) we know that #2 (t+1) μz
Step showing (189b) holds time + 1. By (151), (153) and the induction hypothesis, we have for any & t:
m++1
(a.k))
O
log N
m++
Σ
82(.003
(254)
N
and when A e {m + 2, .,...,20 2m.).
(2.6) = -8%
(255)
By (251), we have
Σ
2m
(s.k) = 1
log N N
Σ
2m
(-5z (-
(256)
Combining the above expresion with (184a). and summing over 0, 1, l, we have
log
)) (((+1))
60+1) (1+1)
+
(257)
N
63

---

This shows (189b) holds at time t + 1.
Step 3: showing (189c) holds at time t + 1. For any i,j e [S] and m = J. by definition (c.f. (150))
we have
Ak E [m + 1j:
Ap(t) (..10
(t.k) ((.,))
Thus we have
when k = 1:
(t.R)
p(e)
0.
(258)
w1 [t.kc
When & E (2...,, m + +1.
2/1 w(th)
(t,k)
-1/3/ E :0...................... k 23, R1m-1 =/m = 5}
Op(1)
w (1.k)
exp (1 + cz cd"))x?"))
= =1(p(1) = iin =} (1=0(
20000" [oog">
N
exp (1 + )=² (t)
+ (k- 2) exp (c2)) (b5 + exp (1 (L+ ux
+ +1131e e t0,..., k-- 3), A2m-1 = i| (K)
= =
log² N
))
N
(1) (1)
exp
C2 l12
exp (1 11² (t)
+ (k- 2) exp (c")((]) u/ + exp (11 + c2 )x(")
$ (p(i) = jim
5} 90 (k)
+ 1 e (0..... ... k 3), Mym i/m 5} 1
(259)
(6)
In addition, we have
When &e E (1,...,m)}
(1,8)
(t.k)
(t.A) (s.K)
Piς + P22
L E V(T), is not Ieaf (14) = j} Pl.s + P2,
(1,b)
w6 (t.A)
w
= (1) = (2,)) = (1+ o (o0g m) )) (N - exp
cxp (1 c
112
(1)
exp
(-c2 11
(t)
(-c2 (t) (2)
11
+ eexp
(1- (2) )112)
+ e V(T), is not leaf, (i) # (1 = =s (1+0
log² 2 N
N
(() (4)
2cxp
c2 12
P1
(N - exp (-c2)
(2) (1)
11
+ 2exp
(1 - ( (()))(()) (1)
(22)
$ ((e :/ (t.k)
= j}
Po
(1,k)
+ 1iie V(T), is not leaf, (i) # 5l (4,4)
j} 1
[t,k)
N'
W/
(260)
When k= m 1,
Piái + p2, 11)
(t.m+
1li e V(T), i is not leaf and root (1,+1)
-1 = j]
w (t,m+1)
log" N
))
2xxp
(-c(() H2
(t) (1)
1 + O
N
(N - 1) exp (-o8)
(1) (1)
+ exp (22 C2 )p1 (1)
U1z
$ e V(T), i is not leaf and root|n (t.m+l)
(t.mmm( 1)
=5} '''
(261)
64

---

Thus we have for all k E [m 1]:
(2,k)
(C.R)
(1.6)
((.k)
(152)
$
2
E
(u(
((.kk)
wi
(1.k)
p(s)
Pi + P22
5
wi (t.k)
(2.h)
w
(t.k)
(t.k)
2 N
5
E
(w0 (1)
wi
(2.1)
90
(t.R)
(r,R)
j
+
N
2
E
(w)
(t,8))
w/ (1.1)
wi
log N N
(-52(,053
where the second line follows from (259), (260) and (261), and the last line follows from (151), (224) and (227). When & {m 2..... 2m,, similarly, we can compute that
(t.A)
log N
N
82 k2)
1,03
(251) log N log N
N
(-ba((0))
Ak' E [m]]
(262)
Combining the above two expressions, we have
2m
log N 2mn
Vi.j E
[S]
Σ
(t.b) Lig
Σ
(-52(,033
(263)
N
This combined with (186) and (184b) gives for any i,j E [S]. i + j:
22,
(++1)
<
1221
17
2m
121,
821
(t.k)
x
log N
P22 +
(6)
n nog
-82(,0.3
= log N (t+1) (264)
N H2
N
N
(++1)
<
P2
(1)
+
1J3
(t.k)
log
N
N
22
(8)
n log N N
2m
(-02(,00)
=
log N (++1)
N P2
(265)
where the second inequality follows from (263) and the induction hypothesis. Analogously, we can show that
(++1)
log N (++1)
2.2
N
P2
(266)
Combining (264), (265) and (266), we know that (189c) holds at time + 1.
(++1), 1/2 by contradiction. If
Step 4: showing (189d) holds at time + 1. We show 62
(1+1) 1/2,
(267)
1 since 1) S Nm. by (245) we have
Mm=1 (ug ((,))) 2
wi (w(tk)))
(t,k)
mN
1
2 llogg N,
2 (uo(e.m (((ttmt))
(wutt (uffmmt)))
Nm
exp
(1
20 (++1)) yu2
(++1)
which combined with (151), (154), (251) and Lemma 23 indicates that for any E [S]:
2m
2m
82g11 2
(1,A)
Σ S )2((103)
(268)
This indicates that
de (t-1)
(1)
Σ
2m (t.b)
82g,1,1
(t)
Σ k=
(t.k)
2
1,0,J
where the last inequality follows from (268) and our hypothesis that c2 (t+1)
> 1/2. The above expression
contradicts (267), thus we have C2 > 1/2.
65

---

C.2.8 Proof of Lemma 16
We prove by induction that (192a), (192b), (192c), (192d) and
Ekm 2k=m+2
(f.&) wi (oftth))
W/
log² N
x
(269)
wff (uflimett)) 2
N
wi
exp (1+ (1+20/²)
(
2 N
(270)
hold for all E
Base case. When / = 7F, (192a), (192b), (192c) hold. By (189d) we know that (192d) and (270)
hold at time TF. By (189d), (239) and (240), we have
(TT mm
2
(T 1)+1)
[1]
(271)
(2c²* (TT)
1) M2 2111))
exp
and by (242) and (241), we have
(TF (P))))
x))
(T..))
w
w
2 22c (2(2 1) p4n ((155)
(T))
(272)
N exp
Thus we have
2m
(1)
(TT x)) 2 (TT x)
exp (2ch (2c -1 fa/
(TT)
(299)
2k-m+2 wi
wi
(TF m)-1))²
2 (TF,++)) m+1)
exp
2 (2 (T9) -1 M12
(TT))
w/
(202 (TF) + 1) H2
(2:5)
mexp
(248) log N
(273)
N N2exp
[4c] (cc(T()
(T)
(Tf) (Tf)
'12
N
where the last inequality also uses the fact that (1(() (187)
Gi M2
(1) (t)
> 0 for any 4, because by (155),
we know that V, strictly increases with By (273) we know that (269) holds at time ( = TT:
Induction. Assume that (192a), (192b), (192c), (192d) and (269) hold for all time steps 8 E [TT t]
(TF & TØ 1). Below prove that (192a), (192b), (192c) and (269) hold for t 1.
Step 1: Showing (192a) for + 1. By (227) and (224) and the induction hypothesis, we have
Po (1)
cxp (2x_) (2)
1
1
2
1 δ
1 δ
w [t,m+1)
N
N
1 + exp
24127)
NV
3'
(t)
90
(1.
1 δ
1
exp
112
1 δ
1
))
1
(274)
Wi
[t.m+1)
N
2exp
2
H2
(t)
+ m
N
2'
and thus by (151) and the above expression, we have
11.m+1
2 10
2
(t.m+1)
= S
w
n- =
(275)
By (224), (226) and (151), we have
- _
11)
2 O, and by (269) we have
log² 2 N
Σ
8 (t.k)
8=1,0.j
(276)
N
Amm+2
Thus we have
(I+1)
z f2
(t)
18
(t.m+1)
(277)
66

---

By (239) and (240), we have
2 (l.m+1) (uflmmtt)
exp 2 (1)6 )
1
when (0)
C2
> 0.5
w
wo
exp 2 1- 2o )(²))²
when (1)
C2
< 0.5
(278)
and when C2 2 0.5, by (237) and (238), we have Wk [m]]
w (t.4) 2
(t,8)
N2
w
N2
(t,b) wi =
and
exp (2 (2050 + 1 mn)))
W (uffmmt1)) 1))
2 (1,m+1)
exp ((20g) + +3) n1
(279)
Following the same argument as we show (248) in Phase by leveraging (245), here by leveraging (279), we can show that when S mN and (f) > 0.5,
exp ( 22c (t) + m $ mN2
(280)
This combined with (278) suggests that when
cl 0.5, we have
(2)
((
wr (l.m+1)))
exp
4p2
w
(281)
m.
m
Thus by (277) and the above relation, we have
P12
1++1)
222
(TT)
(T))
1
z
n(1 1((11-79)
(282)
Sm
SNV
where the last step follows from our induction hypothesis. On the other hand, (280) suggests that
E 0.5 after w 2 m./.N.
(283)
After ( & 0.5, by (278) and (192d), we have
(t.m+)] wo
1
(284)
(t)
N
exp
11
Thus by (277) and the above relation, we have
1++1)
H12
(2
17
N
1 1t +
SN
1-119) TF)
(285)
where the last step follows from our induction hypothesis. (285) and (282) suggest that (192a) holds for + 1.
Step Showing (192b) and (192c) hold for + 1. Note that (269) indicates that
2m Σ
log² N
(-oz(,03 --52(tm))))
121,0.j
(286)
N
mmm++2
Following a similar argument as we show (189b) and (189c) in Phase in Appendix C.2.7 (Step and Step 3), we can show that (192b) and (192c) hold for I + 1. Here we omit the details to avoid repetition.
Step 3: Showing (270) hold for t + 1. We prove this by contradiction. Assume that
exp ((1 20t
+ 1)) 11 ((+1))
N
(287)
then when is small enough, we have
exp (1+
(1) 112 < 2N.
(288)
67

---

Then by (237) and (238). we have
exp ((20g) + 1 u12
1k E [m] :
(wb (
wi (((ck))
(289)
N
By the above relation and (278), we have
vm --1
((.k))
wi
(t.kk
mxxp
((3 - 2c(2")
122
(1)
mexp (4u2))
Z m,
(290)
2
wp (f.m++))
N
2N2
where the first inequality follows from (288). Thus by (154), we have
82g
3,1,1
2 S
(291)
indicating that
exp ((1 2c
(+11)
ui (
2 exp (11+
(1 2 N,
11z
(292)
where the last step follows from our induction hypothesis. This contradicts with the fact that (287).
Step 4: Showing (192d) hold for / 1. By (270), we have after cg
< 0.5,
Ek=1 w1 (u))))))
(t.b)
wi ((.m+1)
mN2
(293)
[w] (uf. ((f.m+1)²+
(8c (2) ug)
wy
exp
By a similar argument as in Step 3, we can show that when 17 is small enough, C leveraging the above relation.
(++1)
>> 0.25, by
Step 5: Showing (269) hold for t + 1. We can compute that for any h: E [m + 2, 2m].
(w((t()
w0
(t,k)
exp( (2 +1-40(") (u4))
when c2 > 0.5
(t)
(294)
exp (4-40h") 40(1)) )as")) N²
w1
w
(1.m+1)
when 02 & 0.5
(1)
When 02 > 0.5, by (280) and the above relation, we have
22m E2mm+z
(wf(tt))))) 1,)
w
m² N2
m
(295)
11) wa
exp
2p2 (p((tt)
N2
N
(0) 0.5, we can show (269) also holds at time / + 1 by contradiction. First by contradiction When C2
similar to Step 3, using (293) we can show that
Dimm+ wi ((+1,))
((+1,k)
Wi
S m.
(296)
2 (u(+t)
m1+1
w
This together with (151) and induction hypothesis implies that
2m
Σ
< (-02(108t)) (2((m))))))))
$ m (-ōz
(297)
k-m+2
By the induction hypothesis, there exists an absolute constant C such that
2m Ex=m+z (u(o.)))
(a.
wo
Cm²
(298)
10 (of(mm+1))
2 (s.m11)
N
wo
68

---

for E ,tY. We now claim
- (wf (w)t) (e+1,k)
wo
Cm²
( ( (ufttlmmtt) (e+1 1m)
N
w
which can be shown by contradiction. If
Ex=m+2 (wnettl)t) (e+1.k) [1 (uf(t)mmtt)
2
Cm²
(299)
(++1.+11)
N
then when vy is small enough, similar in the proof of Lemma 23 Appendix C.2.11, we can show that when learning rate x Nm we have
E (u((.k)) uu (u(tmmt))
(e.k)
w
Cm?
(300)
(1.m+1)
2N
wo
By (151) and (155) we have
023,0,0
(1.m+1)
m S
(-8z1,03
(297)
M.
2m
021,0.3)
(t.Ai)
(301)
2N
2N
indicating 4 - 1-4ch") - 4ct decreases from step to step f + i.e.
4c(t1)
(t+1)
4c
(et+1)
4c 4c
(302)
Thus by (294), we have
Dkm 2k= (w
w/
(1+1,)
2k-m+2
2m
(wftte)) (t,2)
wo
Cm²
(303)
w (u(tttmot))
2 (++1m+1)
10 (o(lm+1)) w
2 (c,m+1)
N
wi
This contradicts with (299). The induction is complete.
Step 6: Showing C2 (T5) <
ci (
0.26,
0.48. Recall in Step we demonstrated that (c.f.(283))
S 0.5 after ( 2 /NN.
(304)
After cl 0.5, using similar argument we show (291), by (293) know that before 02 0.26,
(()
2m
2m
023,1.1
(1.6k
(-62(.00))
(305)
indicating that (0) (t) 112 decreases approximately S times faster than μl2 increases, and will decrease below 0.26 before T
In addition, after 02 0.5, by (269) and (294) we know that for any i e [m + 2, 2m], (Tf)
exp
(4- 4c
4c
(T))
u2
(T2)
log N
(306)
N2
N
(1
(11)
This together with the fact that c2
0.26 and
exp
(12
N implies that
0.48.
(307)
69

---

C.2.9 Proof of Lemma 17
We first compute some important quantities that will be used in the proof.
Basic computations. When > TH by induction hypothesis and a similar computation as in Step 1 of the proof of Lemma 15, we have when is small enough:
(t,1) w
= exp
+0
M2
w
(s.1) 1.
(308)
exp (1 202 (8)
112
and for any & (2...., m,]:
w
(t.b)
= exp
+
112
N
(2c )x²"))
(t,b) =
(309)
exp
When m + 1, we have
(1)
w
(t,m+1] = 1, w1
(.,++)) exp
112
1
(310)
N
exp (11-2c9) 2222)) )bb
When & E {m + ...... 2m}, we have
w (1,k) exp
+0
M2
w/ (t,b) == 1.
(311)
22 H1 v22))
exp
The above expressions give that
w/ (
w/
(t,
exp
+0
N
exp 2(1 2(1+ )2² )a²")
k=
(312)
exp
+
exp( (1c5 )m!')
h E m........
and
wi ( (l. 2 (t,m+1) exp
H2
wi
+0
(313)
N
2(1 20("))(9 (1)
exp
and
Ak {m + ,,...,,m,,
(uftit)))
wo
(t:h)
exp
+
P2 N
(314)
[4e] 112
exp
Below show Lemma 17 By Lemma 16 we know that Lemma 17 holds for TG. We assume that Lemma 17 holds for all & & & (t e [T]] T - -]]). Below we show that Lemma 17 holds for t+1 1.
Step showing (195a) holds for t 1. Following the same computation as Step the proof of Lemma 16, we could compute that (277) holds here, i.e., we have
((+1) (C) (t.m+))
P2 z uiz 10021,0.
(315)
and
(tm+1)
2 (r,m1)) wo
821,0.5
nh
(316)
By (313) and induction hypothesis we know that when is small enough, we have
(w( (((tt)))))²
2 (t,m+1)
3/2
w/
exp (2(1 0.48) + 0.01) (22)) exp (1.05p51)
6m
(317)
70

---

This together with (315) implies that
((+1)
+ (
3/2
H2 2 H2
(318)
5
indicating that (195a) holds for t ++ 1.
Step 2: showing (195b), (195c) hold for / + 1. This step is almost the same as Step and in the proof of Lemma 15 and is omitted here.
Step 3: showing (195d) holds for t + 1. Similar to Step 3 in the proof of Lemma 16, could show by contradiction that when S mN
E 1
(f.A)
w
(t.b)
1
+ θ
(319)
W (f.m+1)
(t.m++1
N
wi
since by (154), if
Σ
wo
(n,A))
2 (e,a)
wi
(wff. (1). 2 (1.m+1)
wp
(320)
( will increase, i.e., lt+1) (1)
62
C2
and if
Σ
m
w (())))
2
wi
(c.k)
(wfl fffl.mtt)) 2 w/ (1,,m+))
(321)
2
n will decrease, i.e. C2
(I+1)
< <0² Combining (319) with (312) and (313), we have
(1)
(++1)
exp 22(1 - 20 (1+1))
112 ) (t-1)
exp
+0
= 1,
(322)
N
exp (4ch ((((t))) ) ((+11)
indicating that c2
(++1) [0.24, 0.26].
Next, we show by contradiction that
(++1
(1+1) -6 (N
By induction hypothesis, there exists some absolute constant C' such that for all [T] .
Cpoly(log N)
N
where poly N) stands for some polynomial log N. (155), (154) and (319) give that
2m
523,0.0
(1.k)
2m
Σ
w
(t,k)
2
wo
(t,k)
2m
8z 82
E
W
(t.R)
2
w0
(c,
(323)
N
Thus if
Cpoly(log N)
(1+1) (++1)
C2
N
(324)
then when E mx we have
Cpoly(log N)
C
C2
2.V
71

---

By (312). (314) and (323), when C' is large enough, we have
(1.+1)
(l.m+1)
-,00,0
>> 023,1,1
leading to
Cooly(log N)
1
2 2 c (1+1) o2
((+1))
02
N
where the last inequality uses induction hypothesis. The above expression contradicts (324), thus we have
9 7 (1+1) -
(++1)
N
C.2.10 Proof of Lemma 18
Note that by our definitions of wo
(k) (k)
wi and Lemma 11 we could rewrite 2 A as
Σ
IE
(2.R) wo
Σ
E
(w
(325)
A=m+1
By a similar calculation in Step of the proof of Lemma 15, Lemma 17 (and the fact that the relations in Lemma 17 still hold after time ) ), we have for all t > T + +1:
4k E [m]
w
(e.b)
N
V 6m
(326)
exp
(..47
(1)
+ +N
4k: E [m + 1, 2m]
wi
(2.
m +
(327)
exp
0.4774
(1)
+
+
6m
Plugging the above into (325), we have
2m
C.2.11 Proof of Lemma 23
We prove this lemma by contradiction. Assume that there exists the first time to & such that
k)) wi
=11-N) N
(1- C'm
) ( (ufoom=1))+
wo
m1).
(328)
for some absolute constant C". First note that when = we have
w
(0.8)
A: [m]
w
(0,4)
N+
ke [m]
k E {m + ,,...,,2m1 2m]
m+1
ke {m + 1,... 2m}
N+k*
(329)
which implies that (246) holds when O, thus we have to >
By (155) we have for any N:
2m
(t.))
(tak)
(1,6)
(t,+1)
823,0,1
E
wi
2E
wo
(330)
This implies that
2m
Σ
(k)
< m
22,00,1
2
72

---

and thus by (184b) we have at time to - 1,
2m
(-fo-1) 3,0,1 =
v-(ta)
3,0,1
dz k)
3,0.1
&
vitto
3,0,1
+
C
2N
(331)
Similarly, by (151), (153) and (184a), we have
(1t0-1) = Uiih
U(to)
+ +0
C'
(332)
NS
Thus by the expression of wh (t. and (c.h) (c.f.(237), (238), (239), (240), (241) and (242)), we know that
w1
vk E [2m]
WD
(.0-1.A)
= 1 [
w
(3o.k)
w1
(1-11k)
1 +
O
C
wi
(look)
N
N
(333)
and therefore, we have
Σズ k-1 w (1(fo-1.k)
wi
1,k]
C
D
(wo
w
11+0
(334)
W (fo ((fo=llm+11)
2
(1o 1,m+) 20
(((o.mmt))
wo (fo.m+l)
By (328) and the above expression, we know that when C' is small enough, have
Σ wh 1.1) (
(f0-1,88
w1
C'm <
(335)
(to (o(to-mm+t)) w
2 (to 1,
2N
By the above expression and (330) we know that
2m
Σ
82
(4o 1,k)
C'm E
( (wflo 2 w/
(1o 1,m+
(336)
N
Thus by (184b) we know that V3,0,1 (recall V3,0,1 (22112 will increase from to 1 to to. i.e.,
(ta) u2
(to)
(1o-1) (to-1)
C2
112
(337)
Furthermore, by (219) we know that
82 (1k-1.k)
1,0,1
2m
5
EE
2 (w (nflo-Lmtt)
(lo- (.m+1) w/
2C'N
S
82
3,0,11
(la-1,k)
(338)
suggesting could choose the absolute constant C' large enough that )
(1 + 20 (to-1)) (to-1) > (1 + 2c (to) (to)
)
(339)
(243), (244) and the above expression imply that we could choose the absolute constant C" large enough so that
k)) (o(lo.k))
(66 (u (w] (1(60-116))
(340)
Wn ()(tom--1)+ 2
w (nftommtti)
( (wfto--mmt)) { (10fto-mm
(341)
Therefore, by (328) we have
Σ
m
w No 1,))
wi
(to (.0)
<(1 ) N
(1-Cmm
C'm (ufto-Lmmtt)) (to (.0
(342)
This contradicts with choice that is the first time that (328) is satisfied.
73

---

C.2.12 Proof of Lemma 19
We first define
(1))
Σ exp x3
(k-1)
Biy-1 (k-1)
+ Vi B2y-1 + 2i -
()- (k-1)
(k-1)
(343)
where we let ii (-11) y/ is the denominator of the attention score at the k-th reasoning step.
((-1) ( be the i-th column of ((k-1), y(k-1), zik-1) respectively. ((k)
For any i.j E {0,................ (0,......) 53. we also define
200
(t.k) S0u
+ 5y(...)
(344)
L
and we define , ( likewise. Then by (184a) we know that v((t1) U1(h - n8!"
(345)
From Lemma 3 we can see that each 8 consists of some positive terms and some negative terms. We denote the sum of all the positive terms as 8(t,+) and the absolute value of the sum of all the negative terms 8("-)) Then we have
8/th 6 g(ert) ((t.-)
(346)
With the above definitions, we prove Lemma 19 by induction. In addition, we'll also show that after 6
b ui (4) (i.e., U.(..,1 monotonically increases.
By our initialization, Lemma 19 the above claim hold when / = 0. We assume Lemma 19 and the above claim holds time (t {0, 1,... . 1}) and prove it at time t 1. Step Showing (199a) holds time We break discussion into two cases.
Case 1: by < When o & O, by Lemma 14 and induction hypothesis we can compute that
((--) 2.j.j
(1,m+1)
(1 a (1.++1)
(l.+1)
( (0
(t.m+) (Gm+1) (o..)))))))) (......))..) |m1 =
E
(0.f.0).G.11
N exp (1 + 054
(1)
11
(347)
S (f(t,m++1) 1))
and
2.j.
(t.+)
5
E
(a (k)
(aat)ss. + (k)
(k) (K)
(a)
km2 2
(1 (1-a
(p0)..0)..1) (000M.. In- =j
(K) (b)
mN exp ((2 ((2- ((4) (10
(348)
S (r(t.k))
for any h ... m), where the second and third line is the second term of the expression of (m+1) given in (168). 5y255
By induction hypothesis we also have
14
)) (IN - exp (()) + exp (1b5 -00))))))))
((m++1) = [1+ δ
N
74

---

((aa() -(aa")-ba)) (m 1) exp (-b2) 11)) + exp (11 (11-b2 6l") )(
+ exp
Nxxp (/2 (t)
11
(349)
where the second relation follows from (199d) and (199b), and for all E f2.... .m} have
cxp 1 δ 1 (N) (IN- 2) exp (b4)) m1()) exp (00)) + 1)((()) exp (1-bb) a )(0)pp))
p((,A)
exp (10)0 a'') (I: 2) exp (- m(() exp ((1 (1- (6()){(
(0) (1) 14
N exp
(350)
(note that this relations holds during Phase I.I-C, even when 6/ 0.) Thus we have pl.m+1) = ((,k), Ak (2.-.., mm.
(351)
By (199h) we know that (t.)) 2,3. dominates (1,+) and thus 222.
exp ((1 262 ji
2.3
(0)
2JJ
SN
SN
(352)
where the second relation follows from (347), (351), and the last relation follows from (199b). (352) indicates that (199a) holds time + 1 if b2 < O,
Case 2: ( > When bl²) > 0, by Lemma 14 and induction hypothesis we have
(0)
(!,-))
2.
2.j.
m
(K) (P.)..)...)) a (01000000))) (pt))..))...)
(k)
(K)
s
h-
(1
(1-0 (PU)..0))).) - (K)
(2)
-aGGnDD. |--1 =
mN2 exp (1 31
(1)
(1)
Mi
eexp
1
(353)
S (p((.))))
SN
SN
where the first and second line corresponds to the first term of رy2 given in (168), and k
(R)
{2. ....... .m). and the last relation uses (350). Therefore, (352) holds at time + if b > 0.
Step 2: Showing (199b) holds at time + 1 if bl > 0, then 3,0,1
(t(+1)
> 3,0,1 We first show by
contradiction that if b (t-1) 0, then we have
exp (1 (1-261g ((++1) 1r (111)
N
Suppose for contradiction that ((+1) & 0 and
exp (1 - 25 ((12)) CN
(354)
By the induction hypothesis we have and there exists some absolute constant C such that
b2 0,
,t) Use (0. (0,))
exp (1 28 (8) )11 (3)
<
CN
As (0, .t).
(355)
75

---

By (354), when the learning rate 17 S mN. we have
1
CN
exp (11 26 (1)
(356)
2m
Using the induction hypothesis and Lemma 14, we can compute that
3,0,1 = B
(1
a (1..))1)00.1) - (1g..n(
(t.m+1)
(t.++1)
(t.m+1) (m...a(m)))0)0...))
(t.m1)
5
Σ
(t.+1)
|m = j
(00))..00.1111
[E[S]
a (1 (0..0)..1)
(1.m+l)
(1.+1)
(1.++)
(t.m+l)
E
a
(01.0)
()....)))
11 =
00
N2 exp (1
(4) (0)
pi
(357)
S (p(t.m+22)
where the second and third line corresponds to the fourth term of óy3,0,1 given in (180), and the (m+11
fourth line corresponds to the third term of 803,0,1 given in (174). By the above expression and (356) we have
N2 eeep ((1 + 6)2 14
CS (p(t.m+t)) exp (22-b))) b())
MN
233
3,0,11
(r(t.m+1)))
(358)
By (348), (351) and the above expression we have when C is large enough, 2,53 + 28 is <(t,+)
dominated by 62,55 leading to
(t.)) s(t,t) - 2 (5(e-) s(t,+) 2,j,J
3,0,1 $ 0.
(359)
This implies that
(1 (1-20f+t1) 21((+11 u4 (1tt])) v - (1122b m1t))
(360)
and thus by (355) we have which contradicts with (354).
exp (1- 2/((+ (1t2)) $ CN
We next show
2/2
(t+1)
2
by contradiction. If 6 (t+1) >
1. 2. then when 1) x
M. we have
O
(361)
N
We can compute by Lemma 14 and induction hypothesis that
mN exp ((2 1(4) 11
0 Σ 0y3.0.1
(353)
1
(362)
S (r(t.b))3
N
where the last step also uses (361). This implies that from step t to step + 1,
11 (I+1)
pi R ( (t+1)
(t)
(1)
76

---

and thus
(++1)
b2
(1)
$
(363)
This leads to contradiction.
Lastly, if (g 0. by (362) and (357) have
exp (2)
(1
3,0,1
(t)
(t.++)
P1
6,00,
indicating that -6 3,0,1 3,0,1 increases with after 6/ 0. Therefore, 3,0,1 3,0.1
(t,+)
(1+1)) U(0)
Step Showing (199c) holds time I. 1. This can be shown by following a similar calculation as in Step 3 in the proof of Lemma 15, and omit the details.
Step 4: Showing (199d) holds at time t + 1. If ((++1) z O, then by our choice of TP we know that (199d) holds at time L + 1. Below we consider the case when ((+1) < 0. By induction hypothesis and Lemma 14, can compute that
(((-) = 5 R
(1,m+1) (l.+1) (1-0 (0..0).11.)) ((GG)s..G.)) (0..0)...)) (1 (1- (1-0ffm (l,m+1) - (t,m+1)
(l.++))
(0.j.0)(aa)) (mgm J.1...))...)) (m = =j
N2 exp (336g) + +1-1110 -aa(t) u²")
(364)
S (r(t.m+1))
where the first two lines correspond to the first term of 8x1,0. We could also compute that
(m++1)
given in (159).
[ (a (t.m+)) (00..00.11111
+ a (0,n2m ).....)
(1,m-1)
(1-
(1-0(tm11) (tm+1)
(1-00000))011 (0..0).(.1)) (0.12m .)....) n1 =
(c.m1)
100 2 s
N exp ((365) + +2-22 20(()) 14
(365)
S (r(t.m+2)))
By comparing the above expressions and using a contradiction argument similar as in Step in the proof of Lemma 20, we can show that
exp (1-aa")/( $ N.
(366)
Step 5: Showing U1,05
U1,0.*
By Lemma 14 and induction hypothesis, we can compute that for all & E {m + 2, 2.... + 1 + 2m
50.30000
5
(P61..0).1001 ( a (4) (G..s..
1 IE
(1-a
(k)
(000).)))))) (8)
Σ
(lo)..0))
n =
(E[S]
(2exp(211)) Nexp (N (exp(-bbipp (m 1) exp(b1/s1)
(367)
(r((k))
and
(4) 6y3.0,0 = E (1 (1-o a (1.10).)))))))) a (.220.))) - (k)
1
(k)
(...)....))) (...)...
Σ
|
yE[S
(Nexp ((1 - by) m1) 2exp ((2 bn) (1) + exp ((1 b))m)) (N 2exp(m))
S (r(1))3
(368)
77

---

Recall that Lemma 14 gives that for all E E [m + +1]:
0z 3,0,0
(k+m+))
6333,1,0 = 0.
(369)
Thus by the above expressions
that
1+
(())
Σ
(t.k)
Σ
(t.k)
3,0,0
3,0,0
+
$33,0,0
k-m+2
is strictly positive, indicating that U
(1++))
(:(t)
3,0,0
3,0,0*
C.2.13 Proof of Lemma 20
We follow the same notation in the proof of Lemma 19. We first show (202) holds by induction. From Lemma we know that when f = TH (202) holds, We assume that (202) holds for all e (TH. TP 2,,..... .4] (T < < T We next show that (202) also holds for t + 1.
Step 1: Showing pi
(t+1) p1 and (202a) holds for t + 1. Similar to the calculation in Step of Lemma 19, we can compute that (note exp(x))) 2 N by induction hypothesis,
exp (1-265) ) )1 x N by our choice of TT))
(l.++)
N
Vk [m]]
(370)
(t,k)
(1)
exp
since
(t.++l) = N exppbb/)1 (
(((,)) exp (11 (640) (44)
Ak e [m]]
(371)
And compute that
((2 (2-12) (1) 11
mN exp
2,3.
5 (p((.)))) 3
(372)
2 exp
6120 Hi
mN² exp ((1 3000 R1
(373)
233
(r(t.m+1))
S (r(t.))) 3
g(t,t) By (370) and the above two expressions we can see that ) dominates 8235 and thus
22exp ((1 65") u4²) exp ((1 ((1-261") u1
(1)
((.))
(374)
2.j. 233
S (r(tmm+1))
SN
SN
where the last incquality follows from the induction hypothesis that 60 implies that (202a) holds for + 1, and 11] (t+1) > (
The above expression
2.
Step 2: Showing (202h) holds for + 1. By Lemma 14 and induction hypothesis, we can compute that
11,00 = 8yy.nn
(t.m+1)
S
E
1[ngt-m
(1 (1 (10...0..6.)))) + (00mm)m.).....)))))))) 00(11
(t.k)
(1- > (1-a (0.1.0)....) - (t.k)
(t.b)
((0..nm,1)).633
[S[
(t.k)
(1,6)
a (...).(.0)) (a
(1.A)
(111(0)0))) + (t.k)
(2.2(())))) + (1.)
+
(6.1(()00)
a (LCc(()))))G..))
(G.00...))
(EA)
+ + (t.k)
(0...))).
(0,1mm, (0.mmmm1))..0.0)
M1 = j
78

---

N exp (3365) +2- =2-2a(1)) pl())
exp ((0)) +2 a f11)) (N + exp (1- )
+
S (p((mmt))) 3
S (r(lmm)1))
(6)
(375)
And same as (364), here we also have
cxp exp ((36l) + +1-11-10 -a00 11
N
(376)
8 (r(t.m-11))
We show by contradiction that By the induction hypothesis, there exists some absolute constant C > 0 such that
exp (1-a4))af (1) $ N.
(377)
cxp (11 (11-a()) ( )1()) & &CN, As (T TB 2, .. ,t].
(378)
If at time t + L, we can compute that
exp (11) CN,
(379)
then when 17 x mN we have
CN
exp (1-a())af
(380)
2
When C' is large enough, by (375) and (376) we have the term (a) is larger than 1,0, and thus
8(t) 1,0.5 = 01,00)
(381)
Moreover, have
p((,m+1) = exp (1 (1+1/) 6/0 al1)) 111)) .(a)) = exp (1+ (1+bb9) m1)), wk E [m], (382)
and in Step we have computed that
N N2exp
((1 [1+ + b3t (1())
mN (1 (1+33b%))
3662") (49))
(0) 25.
+
(383)
S (r(t.m+1))
S (T(1.20))
(0)
(d)
By (375) we have that
N exp
(bg)) + a(.)
p1()
CN2 exp (1+ 669) -
aa()) a(t))
1e(")
(b) 2
z
2 C x (c), (384)
(r((.m+1))
S (f((.m+1))
where the last step uses (380). By (375), (383) and (382) we have that
(d) m2Vexp (1 [1 (6 m11)
(z())
S (f((m++1)) 3
(r(cA))3
(32) mN2 exp ((1 (1+ 36 (t( 11)1)
mN² exp ((1 36(²) 3a(() an(t))
(385)
()()m++1))
S (r(tmm+1))
where & [m] is arbitrary. On the other hand, by (375) and (380) we have that
C exp ((3b)) +1 +1-a(() al))) m4
(a) z
R C (d).
(386)
S (((lmm+1))
79

---

By (384) and (386) we have
1,0,3 2
-Co
2,9d"
(387)
This indicates that we could choose C' large enough such that
(1-a(tt)) 14 (e+1)
(1-a m4 (0)
(388)
Combining this with (378) we have that
exp (1- (1--((t)))uh l(tt1) & CN.
(389)
This contradicts with (379).
Step Showing (202c) holds for { Same (362), (357) given in the proof of Lemma 19, here we also have
mN exp ((2 (2+ b6(²) 111
(R)
3,0,1
093,0,1
S (r(1,k3))
(390)
and
2 exp ((1 (1+ + b/2
ue(
N N2exp (1+155) mn
g(.+) 3,0.1
(391)
S (r(tmm+1))
S (r(t,1))
where the last line uses (370). Thus following the same contradiction argument as in Step in the proof of Lemma 19, we can show that U,
((++1) & 1/2.
On the other hand, by our choice of T we have that
((.m++1) z exp ((1-b5)) (1))) 1410)
(392)
This together with (371), (391) yields
2 exp ((1 1/ m1(1)
(t,+)) 3,0,1 &
exp (6/2 14
S (p(t:)))
N2 cxp
+ 7b (2)
11
Nexp ((2 7b 19
(393)
()(t.)))
S (p((k)))
By comparing the above expression with (390) we have that if b thus it can also be shown by contradiction that (f+1) > 0.
< 0, 0 dominates (t.+)
0,0,1 and
Step 4: Showing (202d) holds for t 1, and 3,1,0 3,,,,00 (202d) can be shown by following
(t+1) ((()
similar calculation in Step 3 in the proof of Lemma 15, and U 3.1.0 can be shown by
following the same calculation as in Step in the proof of Lemma 19. We omit the details here. The induction is complete.
We now move on show (203).
Step 5: Showing (203a) holds. After exp (1- 21 (T?")
)M1 (T:7) =N, exp ((1-
(T2))
= N
(394)
we can compute that
y (T .T m++1 = exp (1- b (T
) ) (1((T))) (TP M) exp
(T: (
(T 31
Wk E [m],
(395)
80

---

and
(3- 6/20) 140)
mN exp (3- (3-76g) 762 (0) (0)
m
mN exp
M1
3,0.1
(t.
0/300,1
(t.k)
(396)
(r(t.)))3
(p(tmm+1)) 3)
N exp ((3 (3- ((3-2a 2a(")) + 622 (t)
pi)
3.0.1
+t,++) (1.m+l)
033,0.1
S (r(1.m-1)) 3
(397)
Note that when (14 = 1/4, the above two terms are equal in magnitude. If 03,0,1 (resp. 8 ((++))
3,0.1 (resp. o then by (ft=1) will increase (resp. decrease) to near 1/4. Thus by contradiction similar as in Step 3 we can show that when Cu in (201) is large enough, there exists & T22 such that
dominates ((t+t)
60 = (1 δ
VE (8,8+1, ,2
(398)
Step 6: Showing (203a) holds. When (394) is satisfied, by Lemma 14 we can compute that
01.00- = 5 - (1- a (t,m+1)
(t,m+1) (0.1.0)(11) (a()).
(t.m+1) (0.110)).1))
(1 (1- (0.1.0)6..1) (0,0g3m (0..gg.1)).G.1)) Im = 1..G..))
(t.m+1) (t.m+1)
N Nexp ((/5) + 2 a(0)
(T2))
Mi
(1+t)
010.5
(399)
S (p((mm+1))
In addition,
(1.m+1)
0100. =
=
Σ SE[S]
E 11n2-mm=
(1,k)
+ +0 (0.0ggm (....).....0)))) (1-
(1,)
a (0.10),(G.0) - a (0.nggm,)).... 1).(2.))
(f.k) (t.k)
+
(1).0))))
(e,k)
(6.c((0)0))(3.)
(t.b)
((..))...0))
((,2)
((t011000))000)) + (2,b)
(c,
(6.0((000)) + (C.)
(...))
(0...00)(11.)
(L.A)
+ + (0, (00.ngm,1).6.00)) )..1.0))
(1,k)
m1 = j
N exp (306)) +2- 2a(()) +2-2a(()) 11 (1)
exp ((b)) + (1))
u4 (1))
(N (N+ + exp (1)
S (n(t.m)))) 3
S (n(tmmt))
(6)
(400)
By comparing the above two expressions we can see that
(1- (tt) 1 CNN
for some small enough constant C > 0, then 01,05 dominates (t.++) causing (++ (t+1)
(1.-)
p1) (0)
And since ( (t+1) > (((),, (t+))) (++1) increases faster than (t) and could reacch O(N) before
t reaches 1:
C.2.14 Proof of Lemma 21
Proof of (206). We use the same notation as in the proof of Lemma 19, and prove (206) by induction. By Lemma 20 we know that that (206) holds for t = T.P. We assume that (206) holds for s E (T: T: + 1, ...... t}, and will show that it continues to hold for + 1. In fact, the relations (206b), (206c) and (206d) hold for t + 1, well as that 6 (1+1) ((+1) 6/4) p1 can be shown by 41
(t)
81

---

following a similar calculation as in the proof of Lemma 20. We omit the details here and only compute that (206a) holds for t 1.
Similar as in the proof of Lemma 20, we can compute that
exp
+
N
exp
(1 + 0/2 y14"
k [m]]
p((,k)
cxp
#1 N
exp
(1 b4))
k m k {m ,
(401)
(1 + )4")
exp
+
exp
Then following similar calculation in the proof of Lemma 20, we can compute that
(()
N exp
((1 (1+ b62) ui
23.
exp
+
P1
N
S (p((.m+2)) 3
(t)
N x2cxp/
( (4b/] - -2) (t)
Mi
exp
+
P1
N
S
exp
(4 x 0.24
2 0.01) pi
(33")
(205)
3/2
(402)
S 6m
when is small enough, where the last step uses the induction hypothesis, and thus (206a) holds for ++1 1.
Proof of (207). By (206b) we know that when is small enough, (207a) holds. For (207b), by computation could verify that (367), (368) is still valid here, and by which we deduce that when
(t)) is large enough:
(2exp (2nl")
N Nexp (41 (44))
(N (Nex) exp (-699111) (m 1) 1)exp((())4) (1)))
2nz >
5 3,0,0
exp
10
.V.
S (r(tdk)))
k-m++1
N exp
((1
b(t)
+ 2exp ((2 - ((0) (100 + exp (1
(N 2exp(u(1))
(r(ta)))
(3 110) (100) + exp 2 b(41 m4
(t)
exp
(401) exp
+0
(403)
N
(3 + 3b 11
On the other hand, we can verify by straightforward calculation that (396) and (397) are still valid, by which we have
exp ((3 b((0) (1()
A (1)
3,0,1 I exp
+0
(404)
S exp ((3 (3+ + ++bb")) Hi
By comparing the above two expressions and using contradiction similar as, for example, Step 2 in the proof of Lemma 20, we can show that
6
02
(t)
δ
N
C.2.15 Proof of Lemma 22
By Lemma 1l we know that for all l e [m], we have
1 s
= )1 (1 (1- (1,b
A (1.A)
E 1m-1==10
(P01..0))...) (.00).))
((.8)
82

---

Σ
(a
(t,K)
a (2.b)
gE(STEE)
E
m j}
(1-a (1.k)
(1,K) (00)..0)661
1-1
(N + 2m exp (6 (1(1)
(4)
Mi
exp
N
(1+ 6/ pl + (N + 2m 1) exp (2 M1
(2)
exp
(405)
where the second inequality follows from that fact that
n
(u) + unn)²
Vu, 20. 0,
and the last relation uses Lemma 21 (and the fact that the relations other than (206a) in Lemma 21 still hold after time 7.S).
Following the same computation we can obtain that when k = m + 1,
(N + 2m - 1) exp (b5²) 11
2
A (1.++)) exp
+0
11
N
exp
(1
((4)
11 (1)
+ (N 2m 1) exp
02 Mi
(0) (2)
(406)
when e {m 2.... .2m).
(N + 2m exp (640)147 (1) ((1+b)")
(1.,)) exp
+0
#1
N
(1(0) + (N + 2m 1) exp 1/2 Mi
(2)
exp
(407)
and for A (1.b) we have for all & [m]]
11 (2)
))
(N + 2m - 1)exp (1- 61 u11")
2
A((.0 exp
+0
N
exp (11 (60) (1 (N + 2m 1) exp
b
(e)
M1
(408)
(1)
)
(N + 2m 1) exp (1 (69²) u1
Alt (1.m+1) exp
+0
11
N
exp (1 (1)+
(110) + (N + 2m - 1) exp
b1
(409)
and for k E (m++2 + 2,. .., 2m},
(N + 2m 1) exp max (1 ",9, 11"")111"))
2
(1)
A(.hh) exp
+0
N
((++ ((4) u1 (N 2m 1) exp (max[1 1 (1))
(2)
cxp
(410)
and by Lemma 21, all (1. (E E fx, y). 1e E [2m]] can be upper bounded by the following:
(N + 2m 1) exp (1 b/2 p1
(2)
A (1,
exp (1 1/+b/// ((00
0.01 (1)) + (N + 2m 1) exp (1 (1- by pi
(1)
83

---

N + 2m
exp (0.24 x x 0.99 0.01)
+ N + 2m
N 2m
(411)
exp
0.46pp (t)
+ N 2m
6m
as long as c is small enough, where the last inequality follows from (205). This gives (209).
D Generalization Analysis
D.1 Proof of Theorem
Notation. Throughout the proof, we suppose Assumptions 3 and hold, and the model is trained for t 7 steps under the training distribution, where T is the same in Theorem Let
E(--1) = (E, ((-1)))
be the input at the k-th reasoning step. Define
34 (m d) y4 F(m) (d, 1: :d:1;))) Vie [Ñ m 1].
and
2(0) 0th) (1 d1),
g = 5th) + 2d1),
(k) = (²²) (1 di); y (k) = (h) (d) + 2d1).
Furthermore, let r(T) and g ( denote the root and goal node of T, respectively. Let p(i) be the parent of i in T, with p(r) = 0. Let c(i) = c(t; 7) be the child set of i, with c(0) = [r). We also let
0 = H
(0 = H1122
(412)
Then for all i.j E [S]. i # j. we have [ = e and H (t),
We first give the following lemma, whose proof is provided in Appendix D.I.I. Lemma 24, For any q E [Ñ + + m] there exists 15 such that
8
s 3 = 1, 8(a) > VOEiES.
Kgg 3q
3(0) (ap(s) 2
and
(413)
a,
Main proof. We begin with reformulating the test-time loss. Define
exp(μ)
(1)
= NNN)
(414)
exp(1) + (N -1 exp(v)+ +1
which is the attention weight of (0) y(1) = Og on ag- Further, define
Vk [m] a
(N+1+ (g)
(415)
Then (k) is the proportion of (k) in @(k) By our definition, we have for all & [m].
xk)) = (1)
Σ
((k)
(pp(i) ai
((0) = apt(s)
(416)
(pp-
Kigh (g)
and
(k)
Σ
g(K) 1.
(417)
sev(T)
1
84

---

Thus we have
m
Diest (ie (T; =
((k)
o(4)
2
k=
12
1 m
Σ
g/(h) (ap(s)
ai )
(1-a((t)
aph(o)
2
ape ((g)
LeV1T) Felple
1 m
m
ap10
apts) - (1-a!tt) ap
8,00 01 (1- (1-a(h)) (x0-1(9)
+
2
veVIT) (x)
(eV(T) itpk-1i0 (u)
(418)
The first term can be bounded via
2
Σ 8 ap1) ap(1) - (1-a (pt(o))
Σ
(o(tt) 2 + a1))
lEVITy
g + (1 (1- 2
|
Σ
E 2(1 - 11)
max sev(T) cap--1(0)
kept (()
Similarly, have
2
8 (1- a (g)
$ 2(1 a
icV(T)
2
Substituting the above two inequalities into (418), we have
Liest (T; 0(0) & 2m (k)
(419)
Therefore, bound the test-time loss by bounding
(kk] = 1 - (
(420)
We give the following lemma, where T2 is defined in (65). The proof is deferred to Appendix D.1.2.
Lemma 25 (bounding (())) For all k [m]] we have
Ñ + m
(()) & max
1,
20,
(421)
N + -
where
(xp(u(())
(65)
8 = 1
(422)
(N + m 2) exp(v(t)) + exp(n((t) +
2m
for > T2.
By Lemma 25 and (419), we have
N me
Crest (T; 0(0) & max
1,
. 1 c.
m
(423)
N + m 2
m
85

---

D.1.1 Proof of Lemma 24
We prove this lemma by induction on k. First, (413) holds trivially for 1 & q Ñ +1 ((rg. Byy is in E. Assume that (413) holds for q & s (N 2 N + 1+ m)). We prove that (413) holds for s.
For notational convenience, we let l - 1 N. and let
(k)
(7i
(k)
(k) 7=-1)
softmax (yik-I) BB(k-1))
(424)
Then by (5), we have
((0)
8
(k) 800)
(<pts)
=
((()
y
yi
==
where the second equality uses the induction hypothesis. Let 88.0) = Ei=1 (0 3/0
then we have
g(5)
[pp)) aj
(425)
Us
1-0
and by the induction hypothesis, we have ((0)
> 0, Vjj
E
|S|
and
S
Yi
(k)
(
(42
g
1.
(426)
f=0
This completes the induction.
D.1.2 Proof of Lemma 25
We prove by induction. When k 1, by (422) and (414) we immediately know that if N + m 2 >
N 1. we have 8(1) & 8. If + m - 2 < - 1. we have
(N-1) 1) exp( exp(z(t) +1
5(2)
(x)(x(()) (N m 2) exp(x(()) +1
Ñ
(N + m - 2) exp()(())
(x)(x(()) + N - 1 exp(vlt) +1
N m 2
(427)
Thus in either case, we have (421) holds for A: = 1.
Assume that (421) holds for all s [K] (1 < k & m 1). We aim to show that (421) holds for A 1.
Let B = B(t), where t > T2. Then (p)(s) = (()) = 1) And note that g(4k) is the proportion of (0,a,) a,) T in o(k), thus we have
( (+) a (8) (1 - (k) ((0)))())) yr.
(428)
For all ie ((T))=*()) we have
a. (2) < (1 (k) - -. 3 + a (k) v.
(429)
We define p (g) for all s [k.) and recall c(0) (r), we have
a (0) ( (0) +
Σ
g(t)
B
(k)
+
Σ
ap(()
(2p--169) (c)
12pk-T (g)
(a(0) Byr-2 + g(s)
x (g) (4)
Σ
(a)
JEc(f)
( (2)
86

---

+ (aa (a- - 086²) - (a()) (y))))p02 + g/6)
Σ
5
Σ
)
)m
jee(s)
"xp (0)
(430)
Note that
A S a (1 - (b) - B(+t) + a (1
(1 -a( - -86%)) + (1 a.()) - -g8"))(1 - (x) - -B.ha)
= (1 - )1 - ) (5) a (k),
(431)
and by Theorem we have μ > v and μ > 0. Thus we have
glt) BT(k) < (1 - (((h)) (1 - (1-) al)) ps (lt)a(
1 8(4h (5)) (2 1 a a
(432)
Combining (428), (429), and (432), and denote (+)) = (in fa()),, we have
exp (p)(s) (T((k)
(R+1)
exp ( (3>(0)) BT(
+ Σ
(evi)) exp (a. BT(k)) + E/-1 exp (D(N) Bi(()) (DO) + 1
exp (a(h) 1 + (1 (1- a(h) - -//4)
cxp ((1)) a + (1 - -(())= - 3,( + (N 22 m) m)xxp (1- 111))))(((()))) -(())))) v + (ace))²)) + 1
exp (ā(1)) + (a(t))2 1 8. (n- v)
exp (214k) + (a(A))² 2 - 1 3/ah) (μ (1) 11)
+N- - -2 + m + exx (- (-((()))3 v)
exp (( (ax(A) + (ai(4)) 2
-1 -1)(1- (11 ju-v)) vi)
(433)
exp (a(4k) + (5(4)) -1 1) (J+ -)) + N - 2 m + exp (- ((((4)22)) v)
where the second line uses the following fact:
(1 (k) -
-0(()1)())))))))))) ) -
(A) v x
1 1 (1-1((4) (()))())) 1r (k)
$ g(h) (215)) 2) ")
+ (al4) v.
(434)
Define
8(h) = 1 (()) mmax(6(=)) )
se [k]
then we have
(x()) + (a(k) 2 -11== (2 + alt)) (1 - - < 337(1)
(a(())²-
(435)
Combining the above inequality with (433), we have
exp ((1 - (11-3310 35(()) (n- -w)
11+1)
(436)
exp ((1 30(1)) (w-v)) + N - 2 + M + exxp (-(āāk))) v)
which gives
N - 2 + m exp (- (-(a(1)) v)
(11-1) = 1 1- (a+1++1) x
exp (1-388t) ) (nr w) + N -22 2 + mn + exp (- (a(k)))2) )) (437)
87

---

On the other hand, (422) gives
log (N + m exp(-v)) (1)1) )
(438)
Thus if exp((-)) & N + m, we have
2(N m 2) 2) -1)1
& log
(439)
If TTex((-V) 3 N 2 + m, we have
log (2exp(-)) ( -1) log (3(}-1))
(440)
By (74) and Lemma 10, we know that there exists E (m) such that for all E (0, (1))) after
t > T2, we have
& s & 5 5(u-v), v);
(441)
where the second inequality uses the fact that above two inequalities, we have
0, since exp(-v) > N 2 m. Combining the
(2 ((!-1))
#- UE 2og
(442)
Combining (439) and (442), have when (0, [m
2(N + m - 2) (1 (5-1)
E 2og
(443)
Note that
5 log 2(N + m 2) ( (1-1)) - as - 0
(444)
Thus we could choose cm N, m) & (m) small enough such that for all e (0, co m)]. we have
Ñ + m 5(1 v) & log2,
(445)
N + m 2
which gives
Ñ + -
exp (x (k) (px v)
& exp 6max 1,
N + m - 5(u 8(m- 22
(446)
where the first inequality uses induction hypothesis.
Define
N +
a = 2 max
1.
(447)
N + 2
Then
if N M 2 N m 2, we have and
N m 2 + exp(-v)
N4m > 1
if v 0
N m exp (- ((((()))2))
4m 2+
if
(448)
N +m 2+ 2+exp(
which gives
1 06
1
1 1
( - 1) (N + m - 2 + exp(-v))
1
(449)
25
8
N + m + exp (a(t))) 2 v))
88

---

if N m - 2 N m 2, have 2.
N+m- 1
1, and
N+m-2
2 1
+ m 2
a
$
N m 2 (1
+ m 1 8 -1)
(450)
N + m
0
Moreover, we have
N m - 2+ exp(-1)
24 1+mm-2+opp(- -v]
N+m-2
if < 0
N+m 24xxp( (-v)
N+
(451)
N + m - 2 exp
exp - (a(k)) v)
N+m-2
if R 0.
N -+m-1
Combined with (450), we have when N m
2 3 N + - 2,
1 N 2
7 -1) ( - 1) (N + m - 2 exp(-1))
1
(452)
a5
2 N + + ---11 -
2 + m + exp (- (a((k) 2 v)
By (449) and (452), we have the above relation holds for both cases. Furthermore, from this relation we derive that
exp(# v) (438) (N + 22--2 ex)((())
8 -11
(452)
(- (-(24t) v) )) (1) -1)
2
2
N m - 2 + exp
aδ
(446)
exp 33 (k)
(1 v)) N m - 2 exp (- (a(x)) 2 v) )) ad -1) (453)
Thus we have
exp (1 (1-35 37 ) (1)) 3 - 2 exp
(k)
(-(1))))
1/ 2) (1 -11)
2
exp (1-35" 36 (k) (M - -1)
2 ad
-1
N m- 2 + exp (- (ā(k)) ((x)))/)) 2
Ñ m + exp (-(p(())² v)
$ Edd.
(454)
(- (a(k))))) v) exp (1- (1) )) (W- w)
Ñ m 2 + exp
(m
Combining the above with (437), we have
g(++1) ad.
(455)
The induction is complete.
D.2 Proof of Theorem 6
Assume the tree T at test time is T with Ñ distinct nodes chosen from [S] and a path length m, and Assumption 3, 4 hold, and the model is trained for t l max[T TT) steps with the trianing distribution, where T:l and T are defined in (205) and (194). We let
E(K-1) (E, ((K-))
be the input at the k-th reasoning step, and define
2; = (m)) di.1),
yi = ( (d1 + 1: 2d1, i), Vie [Ñ m 1 17.
(456)
We let (bk) = ( E dd denote the output of the k-th reasoning step of the model at test time for all E [2m]. We define
((k) = (1 d,), g(k) = (d)) + 1: 2d
3(1) = otk) + 1 2d, + d2), (457)
89

---

and let x(h) = (1) d1), y/h) == 0 (d1 + 1: 2d1), 2(h) == (2d) 1: :1:2d1 + d2) (458)
denote the label of the k-th reasoning step.
By the same argument as Lemma 24, know that each reasoning step there exist { N
S,=1
$. (k) N+k
such that
i=1
N+b
N+b
(ROOT, (E00T. g(0)T)) = L
g(k)
(x) T
19:
g(8)
7
22
1-1
N+ k
N+
(K)
(b)
1, 1(hh) (k) 2 O.
(459)
Especially, we let k(x) denote the proportions of (x)) (( in (a(N), g(m)T), and let (K)
denote the proportions of (A) in (6)) We let
a = min p5 100
8(h) max {1 {1- ((()) ( 7
(460)
We define
4 (k) = x; B + y Bz + 2, Bzelk-1),
(461a)
( (1) Biy y (8) (xx(k=1) + (K) BBz
(461b)
o/()) x, Ciy 11-1) + y/ (2x(1-1) + 2 (12(1-1),
(461c)
w(() (((T (2 y (x) (2 2k) Cz2(k-11.
(461d)
We also define for some absolute constant > 0,
:: C(p +22),
(462)
where J4; M2 defined in (185) and (186). By Lemma 17 and Lemma 20 see that there
exists C > 0 such that for any i, k Ee [S]. and any p, E {0, 10,
a B1aj + a Brar + sp Bas < (k)
(463)
a Cia C2a 8p Caspl & (b)
(464)
We set C to make the above relations hold. We define
4k [2m],
g(4)
(A)
0
7 (b)
(465)
(14h max
0
(k)
02
= max
((0)
(466)
3(0) = min
((5)
(M) = min
(467)
sE[
and
N + 2m - 1
N
81 :=
82
(468)
exp 0.46pp
+ +N + 2m
6m
cxp (0.47p²) N
6m
where the inequalities follow from (205), (194) and that induction that
max B
T.
1.
We first prove by
N + 2m
N
Vk E [2m].
7(t) < max
81
8
max
82.
(469)
N + 2m
N
We define
N 2m
a == max
(470)
N + 2m
90

---

When & = 1, we have
exp (14
011)
(468)
max
N
1
81.
exp
Σ
N+ exp
H4
(1)
N + exp (0.46
N
2m
(471)
where we denote the index of pi and the second inequality follows from Lemma 21. Similarly, we have
N
82.
(472)
Thus (469) holds for & 1. Now we assume (469) holds for k (1 = k 2m 1), and prove it for k + 1.
We have
(3(h)) (1) (1 (1- (3400
2
exp
g(k+1)
exp
(3(4))
paa
(1 (1!0) ) 1/
2)
+ Σ
N+* im1
exp (3!
2
14;
(N)
+ (1
3(00
(473)
indicating
N + k
81 (k+1)
(474)
N k 1 + exp 0.46 (3(h)) 11 -
1- (3(ht)
Note that
101-11111 (1- 3(b)
( [ (1+))())
1- (3(+)
3((A)
< 20(k))
&
20d1.
(475)
where the last inequality follows from induction hypothesis and the definition of a (c.f. (470)). Plugging this into (474) we have
Ñ + k -
0 (k+1)
(476)
Ñ - 1 exp 0.46H1 2adi (2v + 0.46p1)
(a)
On the other hand, by the definition of 81 and 82 we have
81 -1). P2
1
1 -
111 log
log
(477)
0.46
0.47
82
and
81 &
82
(478)
6m
6m
Also observe that log 04 04. Thus could set é (m,, N, m) small enough
such that (a) (476) is smaller than log for any E (0, 61]. Thus by (476) we have
N + k
Ñ -
N + k -
8 (k+11)
2
& 2 max
51.
N + 2m 1'
Ñ 1+ exp
0.46mm
N N exp
0.46p1
(479)
Similarly, there exists e2 e2(m, N. m) small enough such that
Ñ
(k+1)
22
max
82.
(480)
N
91

---

We let co = min(c1, <2), then (469) holds for h + 1.
Following similar calculation in the proof of Lemma 22 and Lemma 18, show that
(R) E ((4))
A (k) & (1))
(R)
8
(481)
2
and thus
{
| 2mm
2
N
Lesst (T;0) < 4 4max
N 2m
N
1
E.
(482)
92


Reviews:
[{"review_id": "Ih1I2VMPzR", "reviewer": "Reviewer_Avae", "summary": "This 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.", "questions": "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?", "limitations": "Please refer to the Weakness part in the previous section.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "*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).", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer Avae\n\nThank you very much for your insightful comments and your appreciation to our work!\n\n>**W1: The backward pass process constructs a transformer adapting copying mechanism, which is well studied in many theoretical analysis papers [1][2].**\n\nThank 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.\n\n>**W2: In the construction for forward reasoning, each position has two nodes and a flag, which seems too artificial.**\n\nThank 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.\n\n\n>**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?**\n\nThank 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.\n\n---\n\nReference:\n\n[Brinkmann et al. 2024] J Brinkmann et al. A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task."}, {"review_id": "fiBGKs9IGK", "reviewer": "Reviewer_WaDb", "summary": "This 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.", "questions": "Please find my questions intertwined with the weaknesses above.", "limitations": "No.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "#### 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.", "was_revised": true, "final_justification": "The authors have shown that their can provide more clear arguments that are convincing, but the paper itself still has many more examples of unclear or vague arguments that make it almost impossible to validate all their theoretical statements. I hence increase my score slightly and have increase the scores for significance and quality, but can not accept.", "rebuttal": "## Response to Reviewer WaDb\n\nThank 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.\n\n>**1.No proof is given for Theorem 1**\n\nWe 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.\n\n>When Assumption 1 holds, the existence of B satisfying (6) follows from Lemma 1. Letting $\\alpha\\rightarrow +\\infty$, it follows that \n$$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)}$).\n\n>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]$.\n\n---\n\n>**2. regarding the existence of parametrisations and limit of the softmax in Theorem 2**\n\nThanks for the question. The proof of Theorem 2 is correct. We clarify the reviewer's concerns as follows. \n\nRegarding 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.\n\nRegarding 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:\n\n- **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$.\n\n- **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.\n\n>**3. mathematical derivations of Theorems 3 and 4 do not go anywhere...strong indication that modern generative AI was used to construct those \"proofs\"**\n\nWe 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.  \n\nFirst, 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.\n\nThe 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.\n\n- **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.\n\n- To make the logic transparent, we'll include the following proof outlines in Appendix B to summarize the key steps for Theorem 3:\n\n\n    - 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.\n    \n    - 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\\}$.\n\n    - 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\noff-diagonal elements $H_{i,j}$ are strictly decreasing and satisfy \n    $$-\\frac{H_{j,j}^{(t)}}{S}\\lesssim H_{i,j}^{(t)}\\lesssim \\frac{\\log N}{N}.$$\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 \n    $$H_{j,j}^{(t)}\\gtrsim\\log\\left(\\frac{m}{\\epsilon}\\right).$$ Furthermore, the training loss $L_{\\text{train}}(\\theta^{(t)})\\leq \\epsilon$. \n\n- Similarly, we'll also add the following proof outline for Theorem 5 in our revised paper. \n\n    - 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.\n\n    - 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$.\n    \n    - 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}.$\n    \n    - Theorem 5 is obtained by combining Lemma 17 and Lemma 22.\n\nWe 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."}, {"review_id": "ogfFF3q5z5", "reviewer": "Reviewer_LK9v", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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?", "was_revised": true, "final_justification": "The author has addressed my concerns in the rebuttal. I updated the final score to 4.", "rebuttal": "## Response to Reviewer LK9v\n\nThank 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.\n\n>**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.**\n\nThank 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.\n\nBased 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.\n\n>**W2: the experimental settings and the tasks are weak...**\n\n- **On the \"weakness\" of the task and setting:** \n\n    - 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.\n\n    - 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).\n    \n- **On the lack of evidence for head specialization and coordination:**\n\n    - 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.\n\n    - 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.\n    \n    - 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.\n\n>**W3: no discussion on multi-layer transformers**\n\nThank you for your comment.\n\n- 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.\n\n- 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.\n\n>**Q1: regarding experiments that compare the results of single-head and multi-head transformers**\n\nThank you for this question. We speculate that you mean Theorem 5, not theorem 3.\n\n- 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.\n\n- 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. \n\n>**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?**\n\nThank 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.\n\n\n---\n\nReference:\n\n[Brinkmann et al. 2024] J Brinkmann et al. A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task."}, {"review_id": "VVQJiP5qze", "reviewer": "Reviewer_sUWo", "summary": "This 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.", "questions": "See weakness", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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?", "was_revised": true, "final_justification": "Thanks for the explanation. I think I don't have severe concern now. Thus, I'll keep my positive score.\n\nPersonally speaking, I like the idea and the research question of this work. Even if there might be some limitations, a large gap between tiny 1-layer model and large pretrained models, and the idea might not be able to be extended to general cases, the exploration towards the model abilitiy analysis and interpretation is still meaningful and necessary. Afterall, I understand that one paper may not answer such complex research question completely. Good luck!", "rebuttal": "Response to Reviewer sUWo(5)\n\nThank 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. \n\nWe believe our conclusions offer a strong foundation that can be extended to more practical scenarios, albeit with increasing analytical complexity:\n\n- **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.\n\n- **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.\n\n- **Initialization**: \n    - 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.\n    - 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.\n\nWhile 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.\n\n---\n\n[Yang et al., 2024] T Yang, Y Huang, Y Liang, Y Chi. In-context learning with representations: Contextual generalization of trained transformers.\n\n[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.\n\n[Du et al., 2019] S Du, J Lee, H Li, L Wang. Gradient Descent Finds Global Minima of Deep Neural Networks."}, {"review_id": "54oeQPtAJn", "reviewer": "Reviewer_WYm4", "summary": "The 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.", "questions": "* 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.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "While the learning-theoretic implications of the results are somewhat unclear, the studied task is very relevant to understanding reasoning of CoT models (for which we have little existing theoretical results on) and the analysis techniques appear novel enough to merit acceptance.", "rebuttal": "## Response to Reviewer WYm4\n\nThank 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. \n\n>**W1/ Q1 (a): regarding what is learned by the model**\n\nThank you for your insightful question. \n\n**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.\n\n\n- 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:\n    - finding the parent node: implement a \"lookup\" operation to find the parent of a given node in the edge list,\n    - chain operations: autoregressively apply this lookup to traverse the tree,\n    - identify termination/turning Points: recognize the root node to stop the backward chain and begin the forward one.\n    \n    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. \n\n- 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.\n\n>**Q1 (b): how can the tree structure and/or distribution be made more general?**\n\nThe 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.\n\nIt'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.\n\n>**W2/Q2 (a): regarding summary of our convergence proof**\n\nThank you for your valuable feedback! We will include the following proof outline for Theorem 3:\n\n- 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.\n    \n- 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\\}$.\n\n- At Phase II $(t = \\{T_1+1, ..., T_2\\})$, we show that $H_{j,j}^{(t)}$ keeps on strictly increasing, and the\noff-diagonal elements $H_{i,j}$ are strictly decreasing and satisfy \n    $$-\\frac{H_{j,j}^{(t)}}{S}\\lesssim H_{i,j}^{(t)}\\lesssim \\frac{\\log N}{N}.$$\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 \n    $$H_{j,j}^{(t)}\\gtrsim\\log\\left(\\frac{m}{\\epsilon}\\right).$$\n   Furthermore, the training loss $L_{\\text{train}}(\\theta^{(t)})\\leq \\epsilon$. \n\nSimilarly, we'll also add the following proof outline for Theorem 5 in our revised paper. \n\n- 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.\n\n- 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$.\n    \n- 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}.$\n    \n- Theorem 5 is obtained by combining Lemma 17 and Lemma 22.\n\n\n>**Q2 (b): regarding nontrivial parts of the analysis and extension to other CoT setups.**\n\n- **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.\n\n    - *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.\n    \n    - *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.\n    \n    - *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.\n    \n- **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.\n\n---\n\nReferences:\n\n[Huang et al., 2023] Y Huang, Y Cheng, Y Liang. In-context convergence of transformers.   \n\n[Nichain et al., 2024] E Nichani, A Damian, JD Lee. How transformers learn causal structure with gradient descent.\n\n[Ahn et al., 2023] Transformers learn to implement preconditioned gradient descent for in-context learning.\n\n[Yang et al., 2024] T Yang, Y Huang, Y Liang, Y Chi. In-context learning with representations: Contextual generalization of trained transformers."}]

Author rebuttal:

Meta-review: The paper provides a theoretical analysis of how one-layer, multi-head transformers can provably learn symbolic multi-step reasoning tasks through gradient descent. The paper shows that for path-finding on trees different attention heads learn to specialize and coordinate on the distinct subtasks within a single autoregressive pass. The motivation here is that this task requires a two-stage process of backward reasoning followed by forward reasoning. Overall, this work aims to provide a mechanistic explanation for the emergence reasoning in transformers.

On the positive side, the reviewers found that the paper addresses an interesting and important problem, providing valuable theoretical insights into how transformers perform complex reasoning. Reviewers noted that the mechanistic explanation of how attention heads specialize to solve different subtasks is a significant contribution. The theoretical analysis, which includes training dynamics and generalization guarantees, is solid. The reviewers found that it is a clear step forward in the theoretical understanding of multi-step reasoning in transformers, which is a long standing and important research area.

The reviewers identified areas of improvement, primarily concerning the clarity of the theoretical proofs and the simplicity of the experimental setup. Several reviewers found the proofs to be dense and complex, suggesting that adding high-level outlines and more intuitive explanations would significantly improve readability and accessibility. Some reviewers also noted that the problem setting, while useful for a controlled theoretical analysis, felt artificial and simplified. This included the use of a one-layer transformer, the specific structure of the model's output, and the idealized training conditions (e.g., using vanilla gradient descent without components like LayerNorm). Reviewers also initially requested more direct evidence for the claim of head specialization, which the authors later provided in the rebuttal. By incorporating this feedback to enhance the exposition of the theoretical arguments and contextualize the experimental design choices, the authors will produce a more polished and high-quality paper.

Based on these reviews, I recommend accepting this paper. The reviewers were excited about the paper's novel approach to providing a provable, mechanistic explanation for multi-step reasoning in transformers. They highlighted that the work offers an interesting and insightful viewpoint on the role of the multi-head mechanism and that the analysis is relevant to understanding the theoretical foundations of chain-of-thought, an area that currently lacks strong theory. While there were initial concerns, the authors addressed most of them effectively during the discussion period. I expect the authors to incorporate the reviewers' feedback into the final version of the paper, particularly by improving the readability of the theoretical proofs and perhaps adding more experiments.
Decision: Accept (poster)