Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization
Authors: Yu Huang
Zixin Wen
Aarti Singh
Yuejie Chi
Yuxin Chen
Venue: NeurIPS 2025 poster ()
Keywords: transformers
chain of thought
length generalization
deep learning theory
non-convex optimization
Primary area: theory

Abstract: The ability to reason lies at the core of artificial intelligence (AI), and challenging problems usually call for deeper and longer reasoning to tackle. A crucial question about AI reasoning is whether models can extrapolate learned reasoning patterns to solve harder tasks with a longer chain-of-thought (CoT). In this work, we present a theoretical analysis of transformers learning on synthetic state-tracking tasks with gradient descent. Specifically: 1). We prove how the *algebraic structure* of state-tracking problems governs the length generalization of learned reasoning in transformers. In doing so, we formulate the **attention concentration** mechanism, linking the retrieval robustness of the attention layer to the task structure of long-context state tracking problems. 2). Moreover, we prove that a transformer can provably *self-improve* via a *recursive self-training* scheme that progressively extends the range of solvable problem lengths. We show that the model can achieve abilities outside the coverage of the base model in recursive training, different from prior theoretical works on self-improvement.
To our knowledge, we provide the first *optimization guarantee* that constant-depth transformers provably learn $\text{NC}^1$-complete problems with CoT, significantly going beyond prior art confined in $\text{TC}^0$, unless the widely held conjecture $\text{TC}^0 \neq \text{NC}^1$ fails. Finally, we present a broad set of experiments supporting our theoretical results, confirming the length generalization behaviors and the mechanism of attention concentration.

Transformers Provably Learn Chain-of-Thouge Reasoning with Length Generalization
Yu Huang" Upenn
Zixin Wen" CMU
Aarti Singh CMU
Yuejie Chi Yale
Yuxin Chen Upenn
Abstract
The ability to reason lies at the core of artificial intelligence (AI,, and challenging problems usually call for deeper and longer reasoning to tackle. A crucial question about reasoning is whether models can extrapolate learned reasoning patterns to solve harder tasks with a longer chain-of-thouge
(CoT). In this work, we present a theoretical analysis transformers learning on synthetic state-tracking tasks with gradient descent. Specifically: 1). We prove how the algebraic structure of state-tracking problems governs the length generalization transformers. In doing so, we formulate the attention concentration mechanism, linking the retrieval robustness of the attention layer to the task structure of long- context state tracking problems. 2). Moreover, we prove that a transformer can provably self-improve via recursive self-training scheme that progressively extends the range of solvable problem lengths. We show that the model can achieve abilities outside the coverage of the base model in recursive training, different from prior theoretical works on self-imprrvemmen. To our knowledge, we provide the first optimizution guarantee that constant-depth transformers provably learn NC¹ ccomllete problems with CoT, significantly going beyond prior art confined in TC unless the widely held conjecture TC &i NC² fails. Finally, we present a broad set of experiments supporting our theoretical results, confirming the length generalization behaviors and the mechanism of attention concentration.
learned reasoning
Introduction
Reasoning is a central theme of artificial intelligence [1, 2, 2J. Transformer-based [4] large language models (LLMs) achieve state-of-the-art results complex reasoning tasks via chain-of-thought (CoT) reasoning [5, 6, 7, 8, 9, 10, 11, 12], where the model generates intermediate steps before delivering final answer. Recent frontier models such as OpenAI-ol [1] and DeepSeck- R1 [2] typically produce long CoT traces at inference time, often elicited via reinforcement learning and/or supervised fine-tuning (SFT) that distills from longer chains [13, 14]. These advances have enabled improved performance on more challenging problems [15,, 16, 17, 18, 19], yet the mechanisms and limitations underlying CoT reasoning remain poorly understood, posing fundamental theoretical challenges.
Theoretical studies transformers with CoT have recently advanced along fronts: expressive- ness [20, 21, 22, 23, 24] and statistical learnability [25, 26, 27, 28, 30, 31]. On the expressiveness front, seminal work [20, 21, 22] showed that constant-depth transformers without CoT behave as shallow circuits and are restricted express the circuit complexity class TC of constant computation depth, whereas with OLL) CoT steps on inputs of length L, they can express log-depth circuits NC!, which a problem class conjectured to require inherently serial computation.' These
"The authors contributed equally, and the author order was determined by coin flip. Given the density of results. recommend the full version for readability: arXiv:2511.07372 For background on circuit complexity and a detailed review of expressiveness results, see Section and [32, 33].
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

results reveal that CoT reasoning equips transformers with the expressive power to solve inherently sequential problems. In contrast, there remains limited understanding of how transformers acquire such reasoning abilities during training. Prior optimization analyses [34, 35, 36, 37] were limited to simple, fully parallelizable tasks in TC that do not require sequential reasoning. This leaves a substantial gap between what transformers can express in principle and what they can learn through training.
Another key question. motivated by the success of reasoning models, is whether models can ex- trapolate their reasoning beyond the sequence lengths of the training data: a property known as length generalization.
practice, reasoning effectively over long chains of thought relies strongly on the model's long-context ability [38, 39, 40, 41], which can be severely affected by context rot, phenomenon in which model performance degrades as the number of tokens the context in- creases [42, 43, 44J. Even for synthetic reasoning tasks, empirical evidence on length generalization of transformers is mixed [45, 46, 47, 48, 49, 50, 51], and several prior studies reported limited extrapolation despite strong in-distribution performance. Architectural choices, including positional encoding and attention variants, influence generalization considerably [52, 53, 54, 55, 56, 57, 58). On the theoretical front, previous work established existence or statistical guarantees, showing that transformers can, in principle, represent length- length-generale generalizing algorithms or achieve favorable sample complexity independent of the length of the CoT [59, 60, 61, 62, 63]. Nonetheless, it remains unclear whether transformers can provably learn to length-generalize reliably when trained with gradient-based optimization algorithms.
In light of the aforementioned gaps in prior literature, the current paper seeks to make progress towards addressing the following two fundamental questions:
Research Questions
1. Can transformers, trained via gradient descent (GD), learn CoT reasoning to solve problems requiring inherently sequential reasoning beyond TC
2. Can the learned reasoning ability generalize to problems that require longer CoTs beyond the lengths of training data?
To address these questions in theoretically tractable manner, we analyze a minimally viable transformer: a one-layer transformer block with softmax attention and feed-forward network (FFN), trained by GD with positional encoding (NoPE). We study this model on synthetic state-tracking tasks, namely LEGO [45], which distill core LLM skills such as entity tracking, game-state updates, and code evaluation [64, 65]. This setup is tractable for analysis, while capturing the mechanisms needed for step-by-step computation via COT. We analyze the training dynamics with CoT on two LEGO task families with distinct algebraic action structures: simply transitive group actions, and symmetry group actions. By tracking attention patterns throughout training, we demonstrate how reasoning capabilities emerge and how length generalization is governed by the structural properties of these tasks. More concretely, our main contributions are summarized as follows.
Provable guarantee of learning CoT with length generalization in state-tracking. We prove that for the LEGO state-tracking task, a one-layer NoPE transformer trained with GD provably solves constant-length problems using CoT reasoning. Moreover, the learned transformer directly generalizes to problems of substantially longer length, when the group actions in LEGO is simply transitive. Conversely, for the canonical action of the symmetry group S, on Zn. the learned transformer generalizes only up to constant-factor length. We identify an attention concentration mechanism that dictates step-wise retrieval depending the task structure, leading distinet length generalization behaviors.
2. Recursive self- training provably extends the solvable reasoning length. When length gener- alization is limited (for example. under symmetry group actions), we introduce self-training curriculum that recursively trains the model on its own CoT traces, motivated by the empirical work [54]. We prove that this scheme can bootstrap the solvable problem length up to maximal allowable length after sufficient rounds of self-training, thus offering a theoretical guarantee of recursive self- celf-improvement.
3. Constant-depth transformers provably learn to solve problems beyond TC via COT. The first two results further establish that the model can learn the solution of state-tracking problems for son-solvable groups, which is NC' -comopeee and lies outside TC² unless the widely held conjecture TC a NC² in circuit complexity theory fails. Therefore, provide the first opti-
2

---

0.8
0.8
Cycli(c6)
Accuracy
0.6
(Graund-truth)
(Sei-libel
Symmetry(S5)
Seel-lbbblll
40 (Sh (Sbr
0.2
0.2
40
60
80
100
120
140
160
510
20
40
80 Length
100
120
140
160
Length
(a) Length generalization results of eyclic (Ca) vs.
(b) Self-improvement results symmetry tasks (Ss).
symmetry (Ss) tasks.
Figure 1: Empirical results of length generalization on LEGO tasks with different group actions. (a). Transformers generraieee to solve significantly longer CoT tasks for simply transitive (cyclic) group (Theorem 4.1), while generalizing poorly for symmetry group tasks. (b). When direct length generalization falls short for symmetry actions, a recursive self-training scheme that train on the model's own longer CoT traces bootstraps the solvable problem length (Theorem 5.1). The dashed lines indicate the training length.
Full attention
Full attertion
10
10
Overy postton
7-
6
4
3
0.2
0.0
0.0
9
10
Key position: prove 1:1], anc[O:t-3]
positiont predu1: 1, anr[D:1-11
(a) Cyclic (Ca) task
(b) Symmetry (Sa) task
Figure 2: Attention concentration at convergence for LEGO task with length L = 5. The heatmap places the query clause index on the y-axis (keys on the x-axis). For a task of length I., the LEGO sequence prior to the final answer clause has length 2L; focus query positions L + 1 to 2L. corresponding to answer clauses Zans,o to Zans Two diagonal bands in the upper region indicate attention concentrating the answer clause Zans, and the predicate clause Zpred,E41 when the query Zans,E-
mization guarantee showing that a one-layer transformer learns to solve reasoning tasks beyond TC" with CoT, matching the expressivity result of [20, 22, 21] for iiner-Cool transformers.
Empirical evidence on synthetic LEGO tasks that supports our theory. We present wide range of experiments based on our theoretical setup. Our results corroborate our predictions on length generalization for different group actions, showing clear separation between the two algebraic structures the theoretical setting. Moreover, we empirically demonstrate how recursive self-training effectively improves reasoning length. The attention concentration mechanism identified our theory is also supported by our experimental findings.
Background
Computational complexity [33] has been employed to characterize the power of neural networks [66, 67]. Historically, circuit complexity has been used extensively study the power of neural networks [68, 69, 70, 71, 72, 24, 73], due to the structural resemblance of Boolean circuits and neural networks with threshold gates. In nutshell, circuit complexity evaluates computation models by the size, depth, and gate types of Boolean circuits that implement them. Below we provide a brief introduction of the basics of circuit complexity.
3

---

2.1 Circuits and Expressiveness
A Boolean circuit is a finite acyclic network of logic gates that computes Boolean function on {0, 13" - {0, }} for some fixed n. The gates with fan-in are the inputs, which are assigned one of the n Boolean variables. A circuit family {Cn, (=,)a computes a language by using Cn on length-n inputs. We measure the size of circuit by the number of gates, and depth by the length of the longest input-output path, respectively. For instance, "constant depth" means depth = O(1), while "polynomial size" means (C,) = (n²) for some integer c. Standard eircuit complexity classes are defined by restricting circuit depth, size, gate set and fan-in:
TC consists of constant-depth, polynomial-size circuits with unbounded fan-in {AND, OR, NOT} gates augmented with threshold (e.g., MAJORITY) gates. The cir- cuit class TC captures exactly the complexity of integer multiplication and division, and sorting [74].
NC¹ consists of log-depth, polynomial-size circuits over {AND, OR, NOT] with bounded fan-in.
This class captures exactly the complexity of recognizing all regular languages.
. P/poly is the class of all languages computable by polynomial-size circuit families.
These classes satisfy the following relations:
TC C NC P/poly.
Expressiveness gap of transformers via circuit complexity. One can analyze the expressivity of transformer neural nets through the lens of circuit complexity. Prior work [75, 46, 24, 20, 21] showed that the expressive upper bound of vanilla constant-depth transformer is limited to TC the class of problems solvable by extremely shallow, highly parallel circuits. In contrast, when equipped with COT, recent work [20, 22, 21] proved that (n)) intermediate steps enable transformers to simulate -Cmmcopee language. Further, [21] showed that 2-layer transformers with polynomially many CoT steps suffice to express arbitrary polynomial-size circuits, i.e., P/poly. Hence, under the widely believed conjecture that TC * NC [32], CoT strictly extends the expressive power of transformers beyond TC allowing them to solve problems that inherently require super-constant (here Olog n)) computation depth. A standard example connecting algebraic structure with circuit complexity is the following word problem, define by Dehn [76, 77].
Definition 2.1 (Word problem for a group G). Let G' be a group and let e denote its identity. For word w 91 ... Il G* (each Si E G), the word problem asks to decide whether J1 o I e.
Barrington [78] proved that the above word problem is NC -co-peeee when the group is non-solvable.
Theorem 2.1 (Barrington [78]). The word problem of every finite non-solvable group is NC!- complete. The canonical example is S for n 2 5, the symmetry group on n elements that encodes the permutations.
The word problem captures fundamental reasoning task: state tracking in finite world [65, 79, 80]. Given an initial state and a sequence of transformations, the goal of the task is to compute the resulting state. State tracking underlies practical LLM abilities such narrative entity tracking, chess move analysis, and code execution [64, 65], while exhibiting a rich connection to the circuit complexity theory. Therefore, it has become standard synthetic testbed for probing the reasoning abilities of language models, both theoretically [46, 65, 79] and empirically [45, 46, 80J. Motivated by this connection, we study the state-tracking task with varying algebraic structures.
LEGO for State Tracking 2.2
We focus on specific formulation of the state-tracking problem, LEGO (Learning Equality and Group Operations) [45], which was originally proposed as a synthetic task to study the reasoning behavior of transformers empirically, A typical LEGO instance in [45] takes the following form:
b b c c b,
== s,
-1,
Here, a, b, C, are variables, each taking a value (or state) in {-1, +1} in this example. Short expressions such as b are clauses, where- and denote actions: the action is applied to the right-hand-side variable's 's value to obtain the left-hand-side variable's value. For
2
We intentionally omit uniformity conventions here as this section is only meant to give readers sense of the classes we refer to and their standard relationships.

---

instance, from = and = -1, it follows that b = -1. Formally, the LEGO language is defined as follows:
Definition 2.2 (LEGO language [45]). Let X, ( be finite sets of variables, actions, and values, respectively, where each g 9 is map 9 : + y. The formal language LEGO(X, g, )) has alphabet RUGUUUUE (.) and consists of two types expressions (called clauses):
(1) Predicate clause x = g(x') specifies an action E Eg linking variables x, x' E X. X.
Answer clause assigns value y y to variable
A canonical LEGO sentence of length L with answer up to L' concatenates predicate clauses In 9n (1n-1) for n E [L] and answer clauses In Un for n E [L'] with - L:
91 (zv)
TL (LL-11
Zo yo
(I)
predicates
answers
which describes the chain of transitions:
21
96
gh, IL
starting with zo = yo-
answers V1. DE up to
For semantic validity, any sentence containing path In 9r (2n Zn 9n-k+1 (In for & E [ni] must satisfy: y/ = In o 9h-1 ..... 9n-k+1 (Vn-k).
In the LEGO language, predicate clauses encode transformations, while answer clauses encode observed states. Thus, solving LEGO sentence is exactly state tracking: compose the listed actions along the path and propagate the observed states predict the next answer consistent with the composition.
Problem Formulation
As introduced Section 2, the current paper employs the LEGO framework to investigate the reasoning capabilities of transformers. To set the stage, this section presents precise mathematical formulations of the problems to be studied in this paper. Before proceeding, we introduce the following notation:
Vocabulary. Define the vocabulary as V := ( where the blank token (blank) is a null symbol indicating the absence of other tokens. Vocabulary size. Let d | denote the (finite) vocabulary size.
To facilitate theoretical analysis, concentrate the asymptotic regime where d - Note that [. (9). | may depend on d, and we shall specify any required scaling assumptions as needed. Assumption 3.1 (Asymptotic regime). For language LEGO(X, G, Y) defined in Definition 2.2, we consider the asymptotic regime where both the vocabulary size d and the number variables | tend infinity. Assume [9| log Co for some constant Co [1, 100), and hence y and are much smaller than N' in size.
3.1
Data Distribution
We begin by specifying how LEGO clauses are tokenized, followed by a definition of the LEGO distribution.
Definition 3.1 (LEGO encoding). Each LEGO clause from Definition 2.2 encoded as fixed-length, 5-token tuple Z E VS More specifically,
For each predicate clause = g(x'), set Zpred := (x, x', (blank), (blank)) E vs: For each answer clause y. set Zans == ((blank), (blank), (blank), x, y) V5.
With Definition 3.1 in place, we can encode a LEGO sentence (1) to a sequence Z ZL,L
LL,L' (Zred,, Zpred,L, Zans,u. ... Zans
(2)
where pered.k and Zans represent the k-th predicate and answer clauses, respectively. For conve- nience, we denote LL,L' {pred, EeL U {(ans, f) the index set associated with the clauses LL.L' L L', we write the sequence simply as Z and the corresponding index set as I
5

---

To feed LEGO tokens into neural network, we first map each symbol to an integer index (tokeniza- tion), and then map indices continuous vectors via learned table (embedding). The following definitions formalize these two steps.
Definition 3.2 (Tokenization and token embedding). Each v E is assigned a unique index .d.. Denote by & E Rd the embedding vector associated with index i, and write
T(O) E e, = (v() for convenience. The blank token is assigned the zero vector, (r(((bank)) = 02 E Rd For technical simplicity, we assume that f eep v EEV. V) (blank) } forms an orthonormal set in R2 (note that this assumption can be relaxed results).
well-conditioned embedding matrix without affecting our
Equipped with Definition 3.2, we can transform the LEGO sequence encoding into vector embeddings that can be used inputs neural networks.
Definition 3.3 (Embedding of LEGO sentences). Let de := 5d be the clause embedding dimension. We define operator Embed: V IRd that maps clause embedding by
Embed Z) = (Rey Cpg: Ces) E Rde, for clause Z (11, , Us) E V.S.
Specifically, LEGO sentence L defined in (2) is embedded
ZL.L' (Z.prdd L. Zans, -..... 1:) (
where each column p E Rd (resp. Zans,e) is the embedding of clause Zredd (resp. aan,,t) When I'' = I, we simply write - = LL,L' for simplicity.
Next, we describe the distribution that governs the generation of LEGO sentence.
Assumption 3.2 (LEGO distribution D DL.L'). Consider LEGO G) defined in Defini- tion 2.2, and let L denote the sequence length. We assume that the distribution D of length- LEGO sentences satisfies the following properties.
1. All LEGO sentences Z ~ DL of the form (1) are encoded by Definition 3.1 into representation (2).
2. The variables æo, F1,... ,....,,,,,, Jz. E are sampled uniformly at random from X' without replace- ment.
3. The first value yo E y is chosen uniformly at random from V
4. The actions 91 92. . 9L E G are sampled uniformly at random from & with replacement.
5. The intermediate values 1/1, 1/...., V7 are computed recursively by = -e(()-1).
For any L' L, we define the truncated distribution DL.L' of sequences z containing all the predicates and the first + 1 many answer clauses, where zz.,!! is obtained by first sampling z ~ D and then removing the answer clauses Zans,t for all f > L'.
As can be easily seen, sequences sampled from DL or DL.L' correspond to valid LEGO sentences as defined in Definition 2.2. With a slight abuse of notation, we write Z4,L' ~ DL.L' to indicate that L.,L' is the embedding of a sentence I. drawn from DL.L'.
3.2 Transformer Architecture
In this subsection, we introduce the transformer architecture investigated in this paper. Towards this end, we first introduce a smoothed activation function that will be used in our network.
Definition 3.4 (Smooth ReLU). Define a continuously differentiable variant the ReLU activation function [81, 82] as follows:
2
x S -
29
sReLU(z)
E (-
0
lq
T
(1-
x <> 4
where q and g are some design parameters, Here and throughout, we choose q = O(1) to be a large, integer, and take e ((I/polylog(d).

---

Transformer architecture. In this work, we focus on an autoregressive transformer [4] whose block consists of softmax attention layer followed by a position-wil -aise feed-forward network, as described below.
Attention layer: Given LEGO sentence embeddings Z and indices j, k E IL,L' the attention from clause Zj to clause Eµ is defined through the softmax operator as
exp 2 Q kk)
Attny-*(Q.. Z
EPEILE exp(Z) Z,)
Since the model is autoregressive, standard causal mask is applied to ensure that the latest (answer) token attends only to preceding tokens (including itself). The attention output is
Attention( Q 2. " = ExETE, A'
Zk-
Note that in the standard formulation of transformers, the score takes the form Z, WQT WTTKK ZR instead for "query" parameter matrix w and "key" parameter matrix W Here, we fold W W into a single parameter matrix Q, resulting in a score zj QZ,, which can be viewed an equivalent reparameterization that simplifies analysis without changing expressivity [83, 84, 85].
Feed-forward network (FFN), The FFN with parameter W E R5xdxmxde is defined by
(W, = Σ Ere(mi) sReLU (W.,.rr X) + beg.r). Ai [5], j E [d].
re(m)
where X denotes the input, W..j.r E Rd are neuron weights, indicates the number of neurons, and bi.j.r E I is some fixed bias.
With the above layers defined, we are ready to introduce the transformer model used this work, which is summarized as follows.
Definition 3.5 (Transformer language model). Assume that our learner neural network F a one-layer decoder transformer block composed of an attention layer with no positional encoding (NoPE) as well as an FFN layer: / FFN o Attention, with parameters W Bx and Q E Rde Formally, for the i-th token position and the j-th vocabulary index, we have
Fr (i(ZLL1) FFN., Attention Q 2.) E RR, Vi [d]
(3)
We interpret F(Z') five logit vectors {Z) C Rd, each parameterizing the distribution of the i-th token of the next clause. Recall that V is the vocabulary with [ and 7 V - denotes the index map. Given encoded LEGO sequence 7º (71,..., Ze) with embedding Z
the model's predictive distribution for the i-th token of the (e+ 1)-th clause is the following softmax:
elF,
Pr. =
Zt) /
Σje{ elF(
V.
Now we can sample the next clause Ze+1 by sampling from the product distribution
7 (7241,1, Ze41,5)
pF. pF,
in an autoregressive manner.
3.3 LEGO Task via CoT Reasoning and Training Objective
With the LEGO distribution and the transformer model in place, we now formalize the task within the LEGO framework. We view solving length-L LEGO problem generating sequence of CoT steps, where each step E produces an intermediate result used to predict x; = ye, ultimately leading to the final solution Fı, yı. To be precise, we define the following reasoning task.
Definition 3.6 (Reasoning tasks T4). We define a family {T" that captures the ability to solve sequential reasoning problems. For each L 2 1, task T measures the model's 's accuracy along

---

the chain from step 1 to step L:
(F) Σ
EzenDe
I ( (mPe(1Z2,1)) 1 [1 Zama (5)
Ost'<L
where pr is induced by the model from (4).
Clearly, Accz(F) ? [0. 1]. We say that task Th is solved if Accz (F) = 1. As L. grows, (T+) LEN+ poses increasingly difficult state tracking challenges. At step I/, the model conditions on the partial transcript ZI.L' and predicts the next answer L/ns. To enforce step-by-ste generation of the CoT trace, we define the following training objective.
Definition 3.7 (Next clause loss). The training objective for Th with L 2 1 the next clause loss
Loss L (F) Σ
Loss (F),
(6a)
ISL'SL
where Loss T.I!
(F) (F)) = E2t.MNDUM [- loggorr pr anns, 2L.L'-1)]
(6b)
Here, pp is induced by the model F from (4). We also define the per-token loss
Loss L. (F) = EEzL,
log Pr, (Zane, 24.
This is a teacher forcing style CoT objective: the model is given the ground truth answers for all previous steps and is trained to match the next step's 's answer [86, 34].
Further, we adopt the following initialization for training.
Assumption 3.3 (Initialization). Let F' be the transformer network in Definition 3.5 with parameters W and Q The attention parameter is initialized to zero: Q = Od, xd. The FFN weights are initialized randomly and independently as (( ~ (O,, of with o = d-1/22 The biases are i.j.r
trained and instead fixed at b,,,. og log d for all i, j, r, chosen to keep most sReLU units active at initialization. All random draws are independent across indices.
Learning CoT on Simply Transitive Actions
In this section, we present our first main results for the case where the group action G on y is simply transitive, which is isomorphic to the action of the eyclic group C, on Znn.
Assumption 4.1 (Simply transitive group action). Let ) = {0. 1, , By 1} with Ty [2 log d), log ]. Assume the group & acts simply transitively on y: the action is transitive and for any y1. V22 E y there exists unique g E G such that g = y2- Equivalently, for any fixed y E y the map g + g y bijection from & to V |9] = | nly.
Our first main result demonstrates that, for such a simple task the model obtained via Algorithm I for short-chain tasks 7 and T successfully generalizes to significantly longer tasks.
Theorem 4.1. Under Assumptions 3.1 to 3.3, 4.1, C.1 and C.2, for some constant 0 < c² < 1, ny < m < log² d, the transformer model F(T++Ta) obtained by Algorithm with learning rate 7 (aly and stage and iteration Ti Ō n(
Tz = Ō poly(d) satisfies
900
I. Direct short-to-long length generalization:
Accr. (1th+Ta)
2 1
poly(d) for every L S O(d
(7)
i.c., F((T+Ta)) which is trained for task T and T generalizes to solve the tasks T $ L.
2. Attention concentration:" given 2 with E 10, have
(Ti+Ts) (-pred + Attn (Tt (Ti+Ta)
ans, 2 1 0 de
Attn
(8)
3 For readability, we abbreviate Attn (Q(9, by Att omitting explicit dependence when immaterial.

---

Atenam,1
Att p
Retrieve (Attention):
91 (To)
2/2
(T1)
yo
y1
query
Apply group action (FFN):
1/2 = 92
Figure 3: Illustration of how the model solves the LEGO task: given Z the goal is predict y2.
Mechanism of CoT for state-tracking. Given the current sequence z with intermediate steps up predicting the next state y ge+1 (yr) requires two steps:
(i) Retrieve the correct action ge+1 from the context clause Zpred and the current state ye from the answer clause Zans,i
(ii) Apply the group action, that is, compute the next state ye+1 = (
It is well established that attention can implement content-based retrieval [87], and that FFN can represent the group operation [88] (see Figure for example). Algorithm decouples learning the attention (retrieval) and FFN (action) components to simplify the analysis. For task T the transeript Z1.0 contains only the two relevant clauses, Zpred, and Zams,o, without useless contents. Fixed uniform attention (Q initialized to be zero in Assumption 3.3) therefore suffices to expose both clauses to the FFN, and we optimize the FFN to learn the group operation. The high accuracy for T in Theorem 4.J indicates that the FFN has indeed learned to apply the operation correctly. For task T2, with the FFN already trained, the attention layer only needs to learn to route the correct context the FFN input. The attention concentration result in (8) confirms that the learned routing pattern is correct. Figures 2 and 4 also demonstrate the attention concentration patterns empirically.
How does attention concentration induce strong length generalization? As we increase the chain length in T for l > 2, the FFN layer remains largely insensitive input length since the learned group action is lucation-invariant. By contrast, the attention layer is affected: more irrelevant clauses appear, retrieval must
over longer contexts, which risks diluting attention the relevant clause. Theorem 4. guarantees that training on short chains already yields attention concentration with error Odd-c). ). This "purity" allows the model tolerate dilution and maintain high attention on the relevant clauses for chain lengths up to O(d"). Technically, this concentration arises because the query matrix Q learns to locate the same variable 2; that appears simultaneously: the third token of the context clause Zgred,6+3 and the fourth token of the answer clause Zans, This co-occurrence furnishes a strong, consistent signal that enables robust retrieval across longer chains, which will be elaborated in the proof overview in Appendix C.2.
5 Learning CoT on Symmetry Group Actions on Zn
We now turn to the case where the action group G is isomorphic to the symmetry group, under Assumption 5.1. In this case, the problem is NC -complete for ny 5.
Assumption 5.1 (Symmetry group actions). Let y = {0, 10. ny - 1}. We set 9 = ())),)), the symmetry group of all permutations of 31 so that | = ny.. We let G act on y in the natural way and write g y (or g(y)) for the image of under g G. We assume ny O hence [9] = nty.! | polylogd.
log log log log a and
log log d
Challenges of learning CoT for the symmetry task. When the model tries to retrieve the group element 96+1 from the correct context clause p there can be other context clauses whose group elements also send Ye to ye+1; we call these distractor clauses. In the symmetry case on ), each pair () admits (Ny, 1)! elements mapping to j. the fraction of distractors is substantial. By contrast, in the simply transitive setting each pair has unique element, so distractors are unlikely and can be ignored. Attending to distractors still produces the correct next answer, so training may converge with insufficient attention concentration. This weaker concentration makes the attention layer less robust dilution longer contexts. Hence, for this harder setting, directly proving (1(1) length CoT generalization from constant-length training is difficult.
Self-improvement for reasoning length extension. Recent empirical studies [89, 90, 54] show that length generalization can be bootstrapped via model self- self-improvement: models training on their own

---

output can bootstrap their capability to solve longer problems. In particular, the work [54] motivates a recursive self-training scheme for the symmetry task. To perform recursive self- self-training, we adopt the greedy language model as data annotator: the greedy language model Py induced by the network F is defined by
1, if Z/ns, = 2 (r(Z)
Pr (2a +/+1 12
otherwise.
Now can define the self-annotated LEGO data distribution:
Definition 5.1 (Bootstrapped LEGO distribution). We define D as the LEGO distribution in L S L' is given recursively by sampling the prediction ans, DF 1-1 1 S i L' from the greedy language model pp.
Assumption 3.2 except that the answers Zans,
Definition 5.2 (Self-training loss). Given a (fixed) model F and length L, The self-training next-
dlase-predicction loss is defined by replacing D'.L' with D' (Definition 5.1) in (6):
Loss
L.L' (F) 2 E. Nu. ND.
D. [- -loopp((anss. | (ZLL'11)
(10a)
L.
L.L' A E EZLL F.i
log Pp, (Zans 2 for E [3].
(10b)
We present main results, establishing that recursive self self-training scheme can provably bootstrap the reasoning length for the symmetry LEGO task.
Theorem 5.1. Assume the distribution D induced from LEGO(X, G,Y) satisfies Assumption 3.1, 3.2 and 5.1, and assume the transformer network satisfies Assumption 3.3, C.I, C.2, and ny < m < log² d. Then for any S k < loga ( the transformer F(Tx) trained via Algorithm 2 up to length
LA 2 and T3 O poly(d)) satisfies:
7
1. Constant-factor length generalization: F(TO) is able solve T Latt with Lk+1 2k+1
ACCL
(F(TT))
(11)
poly(d)
2. Attention concentration: given Z The with E E (o.....,.. Lx - -1}, we have
(Tt) where ē is some sufficiently small constant (smaller than 0.01).
Attn ans, ans.d--pred.--1 Attn (1x)
ant. ans,E-ans,.
(12)
At convergence for the current task T (at time Tg, the loss has fallen below 1/dF: Algorithm 2). (12) confirms that attention concentration is still insufficient. Nevertheless, Theorem 5.1 shows that while this level of concentration cannot withstand the dilution from much longer contexts, it is sufficient for doubling the length. Consequently, model trained progressively on T for = 1.2...., 2 generalizes the more challenging task of length + yielding the following corollary.
Corollary 5.1 relloimmpoveeenttt for |e'l-length reasoning). Under the same assumptions as Theo- rem 5.1, letting K é(log d), for any length S [X), the model F(Tx) trained via Algorithm achieves
1
Accy (F(Tx) 2
poly(d)
Significance of the result. Note that { is sampled from X' without replacement (Assump- tion 3.2), the longest feasible chain scales with the variable size: L+1 < |x| O(d). Thus our guarantee attains the best possible length in this setting. Corollary 5.1 also demonstrates that the transformer can be trained to solve a task beyond TC" with linear-step CoT, matching the expres- sivity result of [21].a While prior empirical work reports self-improvemmnt in practice, theoretical guarantees, especially for transformers and length generalization, have been scarce [91, 92, 93). Theorem 5.1 provides, to our knowledge, the first rigorous evidence that transformers can bootstrap their reasoning via self- training without additional supervision.
There is a caveat that we do not study when the embedding dimension is logarithmic in the problem length, although it is possible if allow over-complete basis for the tokens in X' only.
10

---

Acknowledgmentt
The work of Z. Wen is supported in part by NSF supported part by the Alfred P. Sloan Research Fellowship, the NSF grants IIS-2218713, IIS- 2218773, and CIF-2221009, the ONR grants 400014-22-1 22-4-22---333- 1-344 and N00014-25-1 1-2344,, the Wharton AI & Analyties Initiative's ' AI Research Fund, and the Amazon Research Award. Z. Wen thanks Yuanzhi Li for initial contributions and helpful discussions on this project.
and DMS-2134133, Y. Chen is
References
OpenAL. OpenAl ol system card. ArXiv, abs/2412.16724 2024.
[2] DeepSeck-AL. DeepSeek-RI: Incentivizing reasoning capability in LLMs via reinforcement learning. ArXiv, abs/2501.12948, 2025.
[3] Francois Chollet, Mike Knoop. Gregory Kamradt, Bryan Landers, and Henry Pinkard. Are- agi-2: A new challenge for frontier ai reasoning systems. arXiv preprint vrXiv:2505.11831 2025.
[4] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems, 2017.
[5] Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. Chain-of-thouuht prompting elicits reasoning in large language models. Advances in neural information processing systems, 24433777 2022.
[6] Lucie Charlotte Magister, Jonathan Mallinson, Jakub Adamek, Eric Malmi, and Aliaksei Severyn. Teaching small language models to reason. arXiv preprint arXiv:2212.084110 2022.
[7] Xuezhi Wang, Jason Wei. Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203. 11171, 2022.
[8] Karl Cobbe. Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, and Reiichiro Nakano. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14.,14,,,, 2021.
[9] Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the few-shot paradigm. Extended abstracts of the 2021 CHI conference on human factors in computing systems, pages 1-7, 2021.
[10] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35:22199-22213,, 2022.
[11] Aman Madaan and Amir Yazdanbakhsh. Text and patterns: For effective chain of thought, it takes two to tango. arXiv preprint arXiv:2209 200.02209907666 2022.
[12] Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang. Dale Schuurmans, Claire Cui, Olivier Bousquet, and Quoc Le. Least-to-most prompting enables complex reasoning in large language models. arXiv preprini urXiv:2205.10625 2022.
[13] Yiwei Qin, Xuefeng Li, Haoyang Zou, Yixiu Liu, Shijie Xia, Zhen Huang, Yixin Ye, Weizhe Yuan, Hector Liu, and Yuanzhi Li. 01 replication journey: A strategic progress report-part 1. arXiv preprint xrviv:24100 18982, 2024.
[14] Yingqian Min, Zhipeng Chen, Jinhao Jiang, Jie Chen, Jia Deng, Yiwen Hu, Yiru Tang, Jiapeng Wang, Xiaoxue Cheng, and Huatong Song. Imitate, explore, and self-improve: A reproduction report on slow-thinking reasoning systems. arXiv preprint axXiv:2412.09413. 2024.
[15] Aviral Kumar, Vincent Zhuang, Rishabh Agarwal, Yi Su, John D Co-Reyes, Avi Singh, Kate Baumli, Shariq lqbal, Colton Bishop, and Rebecca Roelofs. Training language models to self-correct via reinforeement learning. arXiv preprint arXiv:2409.12917 2024.
11

---

[16] Yuxi Xie. Anirudh Goyal, Wenyue Zheng, Min- Yen Kan, Timothy P Lillicrap, Kenji Kawaguchi, and Michael Shieh. Monte carlo tree search boosts reasoning via iterative preference learning. arXiv preprint urXi:2405.00455/1 2024.
[17] Tianyang Zhong, Zhengliang Liu, Yi Pan, Yutong Zhang, Yifan Zhou, Shizhe Liang, Zihao Wu, Yanjun Lyu, Peng Shu, and Xiaowei Yu. Evaluation of OpenAl ol: Opportunities and challenges of AGI. arXiv preprint arXiv:2409. 18486 2024.
[18] Tianchen Gao, Jiashun Jin, Zheng Tracy Ke, and Gabriel Moryoussef. A comparison of DeepSeek and other LLMs. arXiv preprint arXiv:2502.0306 2025.
[19] Cameron R Jones and Benjamin K Bergen. Large language models pass the turing test. arXiv preprint arXiv:2503.23677 2025.
[20] Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36:70757-70798, 2023.
[21] Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint arXiv:2402. /2875, 1, 2024.
[22] William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint TrXiv:2310.07922. 2023.
[23] Lijie Chen, Binghui Peng, and Hongxun Wu. Theoretical limitations of multi multi-layer transformer. arXiv preprint arXiv:2412.02977 2024.
[24] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 2023.
[25] 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-2788 2024.
[26] Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in sequence to sequence tasks. arXiv preprint wrXiv:2204.02892 2022.
[27] Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Hal- crow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. Advances in Neural Information Processing Systems, 37:78320-78370, 2024.
[28] Juno Kim, Denny Wu, Jason Lee, and Taiji Suzuki. Metastable dynamics of thain-of-thought reasoning: Provable benefits of search, rl and distillation. arXiv preprint arXiv:2502.0/694 2025.
[29] Xinyang Hu, Fengzhuo Zhang, Siyu Chen, and Zhuoran Yang. Unveiling the statistical foundations of chain-of-thought prompting methods. arXiv preprint arXiv:2408. 14511, 2024.
[30] 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.
[31] Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thoughtt. Compositionality through in-context filtering and learning. Advances in Neural Information Processing Systems, 16:22021-220446 2023.
[32] Heribert Vollmer. Introduction to circuit complexity: uniform approach. Springer Science & Business Media, 1999.
[33] S. Arora and B. Barak. Computational Complexity: A Modern Approach. Cambridge University Press, 2006.
[34] Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. arXiv preprint ar arXiv:2410.08633,02 2024.
12

---

[35] Kaiyue Wen, Huaqing Zhang, Hongzhou Lin, and Jingzhao Zhang. From sparse dependence to sparse attention: unveiling how chain-of-thoon of-thuugh enhances transformer sample efficiency. arXiv preprint arXiv:2410.05459, 2024.
[36] Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gradient descent with chain of thought. arXiv preprint arXiv:2502.21222. 2025.
[37] Ruiquan Huang, Yingbin Liang, and Jing Yang. How transformers learn regular language recognition: A theoretical study on training dynamics and implicit bias. arXiv preprint arXiv:2505.0099 2025.
[38] Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of iims with long context reasoning-in-a haystack. Advances in Neural Information Processing Systems, 37:1065 10-1001191005544 2024.
[39] Yuchen Yan, Yongliang Shen, Yang Liu, Jin Jiang, Mengdi Zhang, Jian Shao, and Yueting Zhuang. Inftythink: Breaking the length limits of long-context reasoning in large language models. arXiv preprint arXiv:2503.0660 2025.
[40] Zhan Ling, Kang Liu, Kai Yan, Yifan Yang, Weijian Lin, Ting-Han Fan, Lingfeng Shen, Zhengyin Du, and Jiecao Chen. Longreason: A synthetic long-context reasoning benchmark via context expansion. arXiv preprint 2rXi::2501. 2025.
[41] Wang Yang, Zirui Liu, Hongye Jin, Qingyu Yin, Vipin Chaudhary, and Xiaotian Han. Longer context, deeper thinking: Uncovering the role of long-context ability in reasoning. arXiv preprint awv.ariiiiv......... 17315, 2025.
[42] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024.
[43] Kelly Hong, Anton Troynikov, and Jeff Huber.
Context rot: How increasing input tokens impacts LLM performance. Technical report, Chroma, July 2025.
[44] Anthropic's s Applied AI team. Effective context engineering for ai agents. Technical report, Anthropic, 2025.
[45] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04:3 2022.
[46] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transform- ers learn shortcuts automata. arXiv preprini urXiv:2210.10744 2022.
[47] Samy Jelassi, Stéphane d' A'scooli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
[48] Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint 20. 16028. 2023.
[49] Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Transformers can achieve length generalization but not robustly. arXiv preprint 0rXiv:2422.09377 2024.
[50] Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, and Eran Malach. Universal length generalization with turing programs. arXiv preprint arXiv.2407.03310, 2024.
[51] Changnan Xiao and Bing Liu. Generalizing reasoning problems to longer lengths, In The Thirteenth International Conference on Learning Representations, 2025.
[52] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 044-2482224225 2023.
13

---

[53] Mahdi Sabbaghi, George Pappas, Hamed Hassani, and Surbhi Goel. Explicitly encoding structural symmetry is key to length generalization in arithmetic tasks. arXiv preprint arXiv:2406.01895, 2024.
[54] Nayoung Lee, Ziyang Cai, Avi Schwarzschild, Kangwook Lee, and Dimitris Papailiopoolos. Self-improving transformers overcome easy-to-hard and length generalization challenges. arXiv preprint arXiv:2502.01611 2025.
[55] Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowyez, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length gen- eralization in large language models. Advances in Neural Information Processing Systems, 14:35566385555 2022.
[56] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv:1803.02133 2018.
[57] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-677 2020.
[58] Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, and Ronan Le Bras. Faith and fate: Limits transformers on compositionality. Advances Neural Information Processing Systems, 36.70293-70332, 2023.
[59] Annie Marsden, Evan Dogariu, Naman Agarwal, Xinyi Chen, Daniel Suo, and Elad Hazan. Provable length generalization in sequence prediction via spectral filtering. arXiv preprint arXiv:2411.01033. 2024.
[60] Kartik Ahuja and Amin Mansouri. On provable length and compositional generalization. arXiv preprint 1xxix:2402.048755 2024.
[61] Xinting Huang, Andy Yang, Satwik Bhattamishra, Yash Sarrof, Andreas Krebs, Hattie Zhou, Preetum Nakkiran, and Michael Hahn. A formal framework for understanding length general- ization in transformers. arXiv preprint arXiv:2410.021400 2024.
[62] Noah Golowich, Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. The role of sparsity for length generalization in transformers. arXiv preprint 25X2:25502.1692. 2025.
[63] Nirmit Joshi, Gal Vardi, Adam Block, Surbhi Goel, Zhiyuan Li, Theodor Misiakiewicz, and Nathan Srebro. A theory of learning with autoregressive chain of thought. arXiv preprins #xXiv:2503.07932 2025.
[64] Najoung Kim and Sebastian Schuster. Entity tracking in language models. arXiv preprint arXiv:2305.02363, 2023.
[65] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.0888 08819, 2024.
[66] Gail H Godbeer, John Lipscomb, and Michael George Luby. On the computational complexity of finding stable state vectors in connectionist models (Hopfield nets). Department of Computer Science, University of Toronto, 1987.
[67] Hava T Siegelmann and Eduardo D Sontag. On the computational power of neural nets. In Proceedings of the fifth annual workshop on Computational learning theory, pages 440-449, 1992.
[68] John S Shawe-Taylor, Martin HG Anthony, and Walter Kern. Classes of feedforward neural networks and their circuit complexity. Neural networks, 5(6)97719777, 1992.
[69] Ian Parberry. Circuit complexity and neural networks. MIT press, 1994.
[70] Wolfgang Maass. Lower bounds for the computational power of networks of spiking neurons. Neural computation, 8(1):1-40, 1996.
14

---

[71] Jifí Síma and Pekka Orponen. General- purpose computation with neural networks: A survey of complexity theoretic results. Neural Computation, 15(12):777-2278. 2003.
[72] William Merrill and Ashish Sabharwal. Exact expressive power of transformers with padding. arXiv preprint arXiv:2505. 2025.
[73] Zixuan Wang, Eshaan Nichani, Alberto Bietti, Alex Damian, Daniel Hsu, Jason D Lee, and Denny Wu. Learning compositional functions with transformers from easy-to-hard data. arXiv preprint arXiv:2505.2366 2025.
[74] Ashok K Chandra, Larry Stockmeyer, and Uzi Vishkin. Constant depth reducibility. SIAM Journal on Computing, 13(2):423-439, 1984.
[75] William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant- depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-855 856, 2022.
[76] Max Dehn. Über unendliche diskontinuierliche gruppen. Mathematische Annalen, 71(1):116 144, 1911.
[77] Max Dehn. On infinite discontinuous groups. In Papers on Group Theory and Topology, pages 133-178. Springer, 1987.
[78] David A Barrington. Bounded wiith polynomial-size branching programs recognize exactly those languages in nc. In Proceedings of the eighteenth annual ACM symposium on Theory of computing, pages 1 1-5, 1986.
[79] William Merrill and Ashish Sabharwal. A little depth goes long way: The expressive power of log-depth transformers. arXiv preprint arXiv:2503.0396/ 2025.
[80] Belinda Z. Li, Zifan Carl Guo, and Jacob Andreas. (how) do language models track state? ArXiv, abs/2503.0288 2025.
[81] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:20/2.008816, 2020.
[82] Yu Huang, Junyang Lin, Chang Zhou, Hongxia Yang, and Longbo Huang. Modality competi- tion: What makes joint training of multi-modal network fail in deep dearingg(proovbbll) In International conference on machine learning, pages 9226-9259, PMLR, 2022.
[83] Yu Huang, Yuan Cheng, and Yinghin Liang. In-context convergence of transformers. Interna- tional Conference on Machine Learning, pages 19660-19722, 2024.
[84] Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. A theoretical analysis of self-supervised learning for vision transformers. In The Thirteenth International Conference on Learning Representations, 2025.
[85] Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. Multi-head transformers provably learn symbolic multi-step reasoning via gradient descent. arXiv preprint arXiv:2508. 08222, 2025.
[86] Namgyu Ho, Laura Schmid, and Se- Young Yun. Large language models are reasoning teachers. arXiv preprint arXiv:2212.10077. 2022.
[87] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma. Dawn Drain, Deep Ganguli, Zac Katfield-Dodds Danny Hernandez. Andy Jones, Jackson Kernion. Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Cireuits Thread, 2021. https//transfoomrrccrruisss http://///////aaeeooooo////eee
[88] Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang. The expressive power of neural networks: A view from the width. Advances in neural information processing systems, 30, 2017.
15

---

[89] Avi Singh. John D Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Xavier Garcia, Peter Liu, James Harrison, Jachoon Lee, and Kelvin Xu. Beyond human data: Scaling self-training for problem-solvine with language models. arXiv preprint arXtv:2312.06585. 2023.
[90] Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahem, Miaosen Wang, and Chenjie Gu. Reinforced self-training (rest) for language modeling. arXiv preprint arXiv: 08999, 2023.
[91] Audrey Huang, Adam Block, Dylan J Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jordan T Ash, and Akshay Krishnamurthy. Self-improvement in language models: The sharpening mechanism. arXiv preprint arXiv:2412.01951 2024.
[92] Yifan Sun, Yushan Liang, Zhen Zhang, and Jiaye Teng. Theoretical modeling of Ilm self- improvement training dynamics through solver-verifier gap. arXiv preprint urXiv:2507.00075. 2025.
[93] Yuda Song. Hanlin Zhang, Carson Eisenach, Sham Kakade, Dean Foster, and Udaya Ghai. Mind the gap: Examining the self self-improvement capabilities of large language models. arXiv preprint arXiv:2412.0666 2024.
[94] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6980, 2014.
[95] Zeyuan Allen-Zhu. Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers. SSRN Electronic Journal, May 2025. https://ssra. com/ abstract= 5240330.
[96] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 167063, 2024.
[97] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint 2XXi:2339900071,, 2023.
[98] Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens. arXiv preprint arXiv:2402.1375:. 2024.
[99] Lingrui Mei, Jiayu Yao, Yuyao Ge, Yiwei Wang, Baolong Bi, Yujun Cai, Jiazhi Liu, Mingyu Li, Zhong-Zhi Li, and Duzhen Zhang. A survey of context engineering for large language models. arXiv preprini arXiv:2507. 13334, 2025.
[100] Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. Advances in Neural Information Processing Systems, 37:85867-888 85898, 2024.
[101] Yuan Cheng, Yu Huang, Zhe Xiong, Yingbin Liang, and Vincent YF Tan. Transformers provably learn directed acyclic graphs via kernel-guided mutual information. arXiv preprint 1rXiv:2510.25542, 2025.
[102] Zeyuan Allen-Zhu and Yuanzhi Li. Forward super-resolution: How can GANs learn hier- archical generative models for real-world distributions. arXiv preprint arXiv:2106.0266 2021.
[103] Zeyuan Allen and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pages 977-988. IEEE, 2022.
[104] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self- supervised contrastive learning. In International Conference on Machine Learning, pages 1112-11122. PMLR, 2021.
[105] Zixin Wen and Yuanzhi Li. The mechanism of prediction head in non-contrastive self- supervised learning. Advances in Neural Information Processing Systems, 35:24794-24809, 2022.
16

---

[106] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37833 37736, 2022.
[107] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint (rXiv:2402.144 14735, 2024.
[108] Alessandro Rinaldo. 36-755: Advanced statistical theory, lecture 27. Lecture notes, December 5 2016. Scribed by Xiao Hui Tai.
[109] Martin Raab and Angelika Steger. "balls into bins" simple and tight analysis. In Inter- national Workshop on Randomization and Approximation Techniques in Computer Science, pages 159-170. Springer, 1998.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes],
Justification: Main 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 be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses the limitations by acknowledging the specific assumptions and settings in Section 3 and appendix under which their methods and results hold.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to 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 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
17

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
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 impor- tant 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 a complete (and correct) proof?
Answer: [Yes]
Justification: For each theoretical result, the paper provides the full set of assumptions and includes a complete and correct proof the supplementary appendix.
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 in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged 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 supplemental material.
Theorems and Lemmas that the proof relies upon should 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 whether the code and data are provided or not)?
Answer: [Yes]
Justification: While the main contribution of this work is theoretical, provide experimen- tal details in Appendix A to empirically verify and support our theoretical findings.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to 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 or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary 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 to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
18

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should 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 the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide details in Appendix to enable reproduction of our synthetic experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (kttps://nips.ca
for more details.
While we encourage the release of code and data, we 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 (e.g., for new 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:
//ipp.cc/pppl wttp:cccccppccccccceccceccccccce herbllic/guidesgee gui /ooeeSubmm ssionPolicy) for more details.
The authors should provide instructions on 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 baselines. 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 versions (if applicable).
Providing as much information as possible in supplemental material (appended 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 (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify all the training and test details in Appendix necessary to understand the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to 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?
19

---

Answer: [NA]
Justification: Error bars are not applicable in our setting; instead, we report results averaged over a sufficiently large number of independent runs to reduce variance and provide stable estimate of performance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence 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).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is 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 range (e.g. negative етхог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or 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 small- scale synthetic and can be reproduced on single GPU; we include the necessary implementation and hyperparameter details, and not require specialized compute resources.
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 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 into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /tte::////eeiieee https/ccccbbcccbecicecec ctniddeeelee
Answer: [Yes]
Justification:
This paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS 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 (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
20

---

10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
This paper contributes to advancing the field Theoretical Machine Learning. We do not foresee any immediate societal implications arising from this work that warrant specific discussion in this context.
Justification:
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 societal 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), fairness 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 a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is 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 unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or 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
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: This paper does not use existing assets.
21

---

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.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular 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 be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
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 is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions 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 sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals 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 the requirements of your country institution) were obtained?
Answer: [NA]
22

---

Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If 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 adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not 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 it is an important, original, or non-standard component of the core methods this research? Note that if the LLM 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: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: anttps//neurrps.. cc/ he/////eeeennee22222252 /LLM) for what should or should not be described.
23

---

Appendix: Experiments and Complete Proofs
Experiments
25
Discussions
27
Proof Overview
28
C.1 Learning One-Step Reasoning
29
C.2 Learning the Attention Layer
32
C,3 Significance of the Proof
35
D Learning In-Context Retrieval of Variables
35
D.I Preliminaries
35
D.2 Induction Hypothesis
36
D.3 Gradient Lemma
37
D.4 Growth of Gamma
38
D.5 Group and Value Correlations Are Not Large
D.6 Off-diagonal Correlations Are Small
40
D.7
Non-target Correlations Are Negligible
D.8
Convergence
Learning Simply Transitive Actions
42
E.1
Preliminaries
42
E,2
Phase I: Feature Emergence and Competition
46
E,3
Phase II: Feature Growth and Cancellations
49
E,4 Phase III: Convergence
54
E.5
Auxiliary Technical Tools
56
F
Learning Symmetry Group Actions
57
F.1
Preliminaries
F2
Phase I: Emergence of the Feature
61
F3
Phase II: Feature Growth and Cancellations
63
F.4
Phase III: Convergence
67
G Learning the Attention Layer: Simply Transitive Case
68
G.1 Gradient Computations
68
G.2 Some Useful Bounds for Gradients
70
G.3 Stage 2.1: Initial Growth of Gap
G.4 Stage 2.2: Continual Growth of Diagonal Entries
88
G.5 Stage 2.3: Decrease of Gap and Convergence
94
G.6 Proof of Main Theorem
105
H Learning the Attention Layer: Symmetry Case for Short-Length
105
H.1 Gradient Computations
106
H.2 Some Useful Bounds for Gradients
107
H.3 Stage 1.2.1: Initial Growth of Q
111
H.4
Stage 1.2.2: Convergence with Small Wrong Attention
127
I Recursive Learning the Attention Layer: Symmetry Case
138
L1
Preliminaries
139
1.2 Reducing the Wrong Attention
143
I.3 Proof of Main Theorem
150
24

---

Peruuadion of predictte 1o
Fall attention
attontonn
0.5
04
Jrowe
02
00
eny pootion aoli-alaadssaa
postion predicd andbacs
Perrutation preditate 39500
Fall itterttio
umeentinn
08
02
Predicat porsioe
Kay posttoe predlc amiii.1
postion prediss., a1ibbs111
(a) Ground-truth permutation
(b) Co
(c) Ss
Figure 4: Attention patterns of the same trained model Figure 2, evaluated with randomly permuted predicate-claus positions. Column (a) gives the ground-truth permutation. Column (b) and Column (c) show the attention heatmaps for the simply transitive and symmetry tasks, respectively.
A Experiments
In this section, we conduct synthetic experiments to verify theoretical claim.
General Setup. We adopt experimental setting closely aligned with our theoretical setup. Data strictly follow the LEGO distribution the wve-token-per-caaee
format defined in Definition 3.1
and Assumption 3.2. The simply transitive task uses the cyclic group of order 6 (denoted Cc). and the symmetry task uses the symmetry group on five elements (denoted Sz). The network is one-layer, decoder-only transformer with two attention heads and a FFN block. Training optimizes the next-clause loss in (6a) using Adam [94] with learning rate of le-4. We train for 300 epochs to ensure the training loss approaches zero and the model converges. For evaluation, to measure the transformer's 's ability solve CoT reasoning tasks, eschew the teacher- forced accuracy in (5), which conditions on the ground-truth answer prefix, and instead evaluate final-answer accuracy after autoregressively generating all intermediate answer steps without teacher forcing. This better mimics CoT reasoning and is more challenging because errors can accumulate during self-rollout. For computational efficiency, we report the probability that the model predicts the value of all answer steps correctly, following [21].
Length generalization for different group actions. Figure 1a shows that when trained the short length 1, = 5, the model exhibits strong length generalization, achieving near-perfect accuracy at much longer lengths, for example, up to L 160 on the simply transitive task, which corroborates Theorem 1. By contrast, on the harder symmetry task the generalization weaker, yielding only constant-factor extensions expected.
Self-training improves reasoning length. To extend the solvable length on the harder symmetry task, we adopt recursive self-training curriculum inspired by [54]. We first train at L 5 with ground-truth answer supervision until convergence. At the next stage, we double the length to L = 10, the L.55 model to greedily generate answer traces (self-labels) for the L=10 data, and retrain these pseudo-labels. We repeat this doubling process for three stages (L = - 10 - 20 -> 40), so that the final model is trained self-labeled data at L 40. Figure 1b reports the length- generalization curve for each stage, with the dashed line (matching the curve's color) marking that stage's training length. Across stages, we observe least constant-factor generalization beyond the training length; after multiple rounds of self-training, the model achieves nearly perfect accuracy at lengths far exceeding the initial L 5, ultimately matching the simply transitive task's performance up to L 160. These results validate Theorem 5.1 and demonstrate the effectiveness of recursive self- frff-training for extending reasoning length.
25

---

1+
Dagoral nomalized position
L-30
Diogonal attenton
Ancower pea Mee
Aticcion boore
0.50
0.48
0.46
0.64
0.42
0.0
0.2
0.4
0.6
1.0
PWekaae posior
iedtcatepssin
Normalized predicate sosttion LI
(a) C
Diagonal vs. normalized position
0.62
16
Oiagonal antontion
11:00
Aroner
Prockoll postton
Phodoata postiio
15
Normalizad prodicate ocsition
(b) Ss
Figure 5: Predicate-claus attention for models trained with LEGO task length = 5 and tested with L. 10, 20. We extract the heatmap of the attention from answer queries to predicate clauses (the upper-left part of the full attention heatmap as in Figure 2), and line-plot the diagonal attention against the normalized predicate position. At longer contexts, the symmetry group task displays visible drop in attention to the target predicate Zpred (+1 at later predicate locations (larger E; note the absolute query index L+L also increases). By contrast, the cyclic group task maintains more consistent attention across lengths and positions.
Attention concentration. We visualize attention heatmaps at convergence for models trained at length L on both simply transitive and symmetry group tasks. Each attention matrix is averaged across the two heads and over 100 independently sampled LEGO sequences. Note that for a task of length L, the LEGO sequence prior to the final answer clause has length 2L. We focus on the query positions L+1, 2L, which correspond to the answer clauses Zans.u.... Zans,L-1 for predicting all outputs y1. ,,, VL.
In Figure 2, the two diagonal structures in the upper region indicate attention concentrating on the answer clause Zans. and on the predicate clause Zpred.E+1 when the query Zans validating the attention concentration principle highlighted by Theorem 4.1 and Theorem 5.1.
To ablate positional bias, we test on the trained model from Figure and only change the input format by randomly permuting the locations of predicate clauses. Results in Figure show that retrieval keyed to the shared variable rather than absolute position.
Figures 5 and probe out-of-length (OOL) behavior through attention patterns L E {10, 20). As the line plot indicates, the attention scores for both the target predicate and the answer clauses decrease as the task length increases from 10 to 20. This pattern matches the attention dilution predicted by our theory of length generalization. In Figure 5, for the symmetry group task, attention the correct predicate Zpred, 1+1 decays with position, becoming weaker at later predicates (larger f, hence a larger absolute index L+ By contrast, the cyclic group task shows comparable decay, with attention remaining fairly uniform across lengths and positions. In Figure 6, the pattern reverses: attention over answer clauses is more stable for the symmetry group, whereas the cyclic case is more sensitive to length and location. robust predicate attention attention on predicate clauses in the symmetry task indicates a non-robust learned attention mechanism, which may underlie its poorer length generalization.
lengths
conjecture that length generalization in the LEGO task hinges on aetticate-attention pattern. Accordingly, the observed length- dependet decay of
26

---

L+ 30
L- 20
Diagonal normalized postion
646
L-10
Diaporal atzentio
0.40
638
5.5
6:88
0.20
0.0
0.2
9.4
0.6
1.0
ABSWET
AntweY
positoo
Normalized answer position II/ 1L/11
(a) Cs
Diagonal vs. normalized position
0.45
16
Diogonyl antention
0:33
0.30
638
648
666
Atame po0don
Aramst poalorr
23
Normal
(b) Ss
Figure 6: Answer-clause attention for models trained with LEGO task of length L and tested with L 10, 20. We extract the heatmap of the attention from answer queries answer clauses (the upper-right part of the full attention heatmap, as in Figure 2). and line-plot the diagonal attention against the normalized answer position. In contrast the predicate attention in Figure 5, longer contexts, the symmetry group task preserves a stable attention for the answer clause, while the answer attention pattern of the eyclic group task is sensitive to the lengths and positions. Nevertheless, the sharp decay of attention scores for the cyclic group task did not result in visible performance drop for length generalization, as seen in Figure!
B Discussions
NoPE, recency bias, and length generalization. Empirically, standard positional embeddings often hinder length extrapolation, while NoPE has been advantageous considering its angth-generalization performance [52, 95]. Prior works [34, 35, 37] typically adopt fixed positional encodings, which tie the learned computation to the training horizon. Intuitively, positional embeddings inject location- specific biases that favor local neighborhoods, making longer inputs harder. In our setting, length generalization is possible because the model retrieves relevant information from long, unordered contexts by content (variables) rather than by position, as discussed in Appendix C.2 and verified empirically Figure This provides concrete architectural guidance for practice and identifies positional embeddings as plausible cause of observed failures to generalize to unseen lengths.
Local aggregation of tokens is beneficial. While NoPE is appealing for length generalization, lacks position- aware horizontal mixing within layer, which often hurts empirical performance. In practice, many strong models adopt RoPE (Rotary Positional Encoding) [96, 97, 98], and recent seminal work [95] introduces the Cannon layer, a short convolutional window that adds local residual links across neighboring tokens, to inject locality. Integrating Cannon layers substantially boosts the reasoning depth of NoPE transformers in synthetic experiments. Our clause-basee based data can be interpreted as mimicking the output of Cannon layer: locality is built directly into the data structure. enabling one-layer transformers to learn induction heads, which otherwise require at least a two-layer transformer express.
Context rot. Context rot is a phenomenon that the performance of LLMs drops as their input context grows even when the task is unchanged, which has been widely noted in practice [43]. A prevailing empirical view is that longer contexts introduce many distractors and other irrelevant tokens, forcing the relevant signal to compete with them for attention and thereby weakening retrieval. Our attention concentration perspective characterizes this competition at fine-grained level: we show that as irrelevant clauses accumulate, attention mass is diluted away from the correct clause, reducing performance at extended reasoning lengths. This offers simple theoretical perspective on the origin of context and why long-context training hasn't been effective in eliminating
27

---

Algorithm 1: Curriculum training for simply transitive actions
Input: Model ((() with parameters (W(O), Q(");; Learning rate fj; Stage snapshots Ti, T2.
Stage 1: Learning one-step reasoning (T');
for
Tj do
// Update the FFN parameter
( W v/V (
(M(--)));
Q (4) Q Q)
Stage 2: Learning two-step reasoning for length extension (T²); T2 do
for T
// Update the attention parameter Q
Q # t Q ((-1) 7
Ei=1 Loss 2. (F(t-1));
( = W(+1);
Output: Model ((+))
Algorithm 2: Recursive self-training for symmetry actions
Input: Model F(0) with parameters (W(O) Q (0)); Learning rate 7); Error degree E
(constant) T1, 72 Ö polyd ); Total Stage K.
"
Stage 1.1: Train FFN for one-step reasoning (T²);
for t = to do
// Update the FFN parameter W
(t) + W(t-1) v/V Losss ( (1)(+1));
Q (1) ((-1);
Stage 1.2: Train attention for length extension (T3)
for /
to + do
// Update the attention parameter Q
Q (1) <- Q (c-1) - V 2 2 (f(--)))
( W(+-))
Ti t
Till Stage K: Recursive self-train for length extension
for & = 2 to do - 2 (()) &- (T)-
// Stage to solve T
while Loss L,
(F(t-1))
do
// Update the attention parameter Q
(k))
6 1. 1
Q()) Q(1-1) 11 VQ L.2
Fi4)
L. (F(1-1);;
F(N),
w w(t-1);
To 1-+- L;
Output: Models {
the problem. We expect the insight from our analysis extend richer tasks and inspire practical mitigation strategies, e.g., context engineering [99].
C Proof Overview
In this section, we outline the proof ideas for the main theorem. Our training schemes in Algorithms and 2 alternate between two phases. First, train the FFN parameters W to solve the one-step task T Then, holding W fixed, we train the attention parameters Q to solve T and, for symmetry group actions, recursively 2 This mirrors the division of labor in our setting (Figure 3): the FFN learns the local update rule, while the attention layer learns to route and compose these updates by locating relevant context over long sequences.
28

---

Guided by this picture, our proof overview proceeds in two parts: (1) learning the one- mechanism for LEGO (T'): in-context variable retrieval (Section C.1.1) and group operations (Section C.1.2). (2) learning the attention layer: direct short-to-long generalization on 7 under simply transitive actions (Section C.2.1), and recursive length generalization via self-training on 7 under symmetry group actions (Section C.2.2).
Notations Let first define few notations to facilitate the presentation of the proofs. For each i E J E [d], r E [m/] define
Asgi (ZL.
Σ KETL,--1
ant-aast---k Wijsr Zu) + Ds
(13)
The quantity Ar.. is the FFN pre-activation, i.e., the to for token position i, vocabulary index j. and hidden unit r. According to (3), given Z the model's output at token position i and vocabulary index is
Fi (Z
E re(m]
AReLU(Ai (A (ZL.t-1)..
(14)
We denote W, A (Wa,.r (Wa.s.bedd...ee (e(...rm( as the FFN parameters associated with token position i.
Each Wig.r E Rld is written as a vertical concatenation W.g.r = ii..,......;; : Wig.r,l where
Wij..n. E Rd for E [5]: these five blocks align with the five token-type inputs.
Additional Technical Assumptions. Additionally, we introduce a couple of technical assumptions to be used throughout the proof. To control rare large deviations in the logits during training-time analysis, we adopt bounded-output assumption as stated below.
Assumption C.1 (Logit clipping). There exists Cn log d, for some sufficiently large constant Cn 0, such that each coordinate of the raw model output F; is clipped from above:
[ily f min B}
for all i.j.
This coordinate-wiss clipping is a technical device to control large-deviation tails and simplify the analysis of the dynamics; B can be chosen large enough to avoid interfering with the regimes study.
Moreover, to simplify the analysis of attention dynamics, we impose fixed block-sparsity pattern on the attention parameter Q.
Assumption C.2 (Block-sparse attention matrix). Let Q [Qp.llaeess] Rādx5d be partitioned into 5 blocks with Qp.g E Rdxdd We assume that
Qp.9 p.g = Odxd for all (p.g) $ {(4,3), (4, (4,4)),
i.e, only the blocks (4,3) and (4, 4) are trainable.
This block-sparsity pattern, which zeros out most inter-token attention, standard in recent theoretical analyses of transformer training dynamics [83, 84, 85, 100, 101]. Importantly, although sparse at the 5 x 5 token level, the two retained blocks (4,3) and (4, 4) are fully dense d matrices trained without constraints, leaving 2d free parameters and thus a substantive, non-trivial learning problem.
C.1 Learning One-Step Reasoning
For solving task T1, in stage (stage 1.1 for symmetry group task), we train W via the full loss Loss which contains the prediction loss across five tokens in the answer clause and W; is responsible for predicting the i-th token in the answer clause. Thus, there are three different types of prediction tasks here: the (blank) tokens (tokens 1, 2, 3), the correct variable I (token 4) and the action update y1 - 91 (vo) (token 5). Notice that the (blank) tokens are deterministic, the learning task is straightforward. We therefore focus on the learning dynamics for the 4th and 5th tokens, which involve in-context retrieval (W.) and one-step group action (W,). Across these stage, the attention layer is fixed and keeps as uniform attention due to the zero initialization the attention matrix, i.e., the input for the FFN layer is 1/2 Zpred,1 Zan,,
C.1.1 Learning In-Context Retrieval of Variables
Given the input 1/2 Zpred,1 + Zans,o for the FFN layer, the goal is to predict the fourth token in the answer clause Zans 1, which is the variable 21. Intuitively, the network should retrieve the occurrence
29

---

of 21 in the first predicate token 1r and copy it to the target position. Specifically, the FFN pre-activation for predicting the fourth token to be is
Aa.i.r (Z
W 4.j.r Zpred,1) + (W4,) w.j.r. Zans.a) + ba.j.r-
Using the 5-vector decomposition of W W.j..r and the fact that the embedding of (blank) tokens are zero vectors (recall e+(+) denote the token embedding vectors), we further obtain:
A4.jsr (z (Z - { ( (Wainn + (W4,in2 (r(11) + (W4inas ex(a)))
(15)
+ * (Was..)) (W W.j.r,, CT(n)) + ba.j.r
Therefore, the main idea of our analysis is track the training dynamics of (W W.J.r.p, ex/) for s' E. [d]. E [5]. and r E [m]. Letting E =(2') be an embedding index for a variable, our analysis shows the diagonal correlations (W..s.r11 ex) receive strictly larger updates than all other (Wa.j.r.p..e) This occurs because X1 co-occurs simultaneously as the first input token at pred and as the supervised target at Za which amplifies the gradient 4W.,.r;11 Cx) when In contrast, non-target coordinates coff-diagonals, wrong variables, value tokens, and group-action tokens) incur negligible gradients and remain o(1) throughout training, Hence the correct variable's diagonal signal becomes order-wise larger, and the active diagonal mass (W e,) dominately grows until the end of training.
T(x1).
Given the dominance above, the learned weights align that W4.s.r, points toward e, for 8 E (((). Substituting this alignment into (1.5) shows that, when j = =(x1), the red term (W (W.r(2)).r.11 (r(x1)) contributes dominantly to A4.r(2)) ).r (ZI.9), thereby realizing the intended in-context retrieval: the model copies , from the frist position of pred, to the fourth position of Zans,
C.1.2 Learning the Group Actions
For task T', the FFN input 1/2 Zgrdd. + 2 Zas, already contains the current value yo and the action Ø1, with distracting information. Accordingly, when predicting the next value y/, the role of W is to correctly apply the action 91 to the current value yo- In this section, we sketch how the model learns implement simply transitive group actions, and briefly discuss the symmetry case. We first introduce notation to explain what the model should learn.
Definition C.1 (Combinations, simply transitive actions). Assuming the group 9 follows Assump- tion 1. for each class index j T()), define the combinations
: (Er()) Dj, where D ((g' 12) E E r(d'6/)))
|6| = n for the simply transitive case. We call & = (g. E D3 a combination for predicting j = (g())). Hence, the goal of the model is to correctly identify all of these 6 = =g. E $
Analogously to (15), the pre-activation at the fifth token for class and neuron is
As.j.r (Z -0) (W..j..,, Cr())) + (W W.j...5, (r(w)) + other terms,
where slot "2" reads the action token 91 and slot "5" reads the current value token you We then define the following Ceature-magnituee notation:
Definition C.2 (V V-Notations), Given á = (g.8) E &, for the r-th neuron in the j-th class with j E T r())) and E [m], i.e., Wa.j.r E R.d we define
(g) = (W W..... eg),
Virr()) = (WoJ.n.sey)
and the composite feature magnitude
VØr = (V). Vi.,,, (v))-
Then Vy.r (e) is exactly the contribution of the input pair (g. y) As.j.r (Z1.0) for predicting j. Notice that for each class j. there are associated neurons [ 5.j.r 1= and we index neurons by the pair (j.r) to emphasize that neuron r specific to class j; an index alone has no cross-class meaning in this context.
)
The main proof idea is to track Vg) throughout training and show that it amplifies the correct correlations across the combinations while suppressing the incorrect ones. To make this concrete
30

---

and to clarify the roles of different neurons, we introduce the following neuron neuron-feature index set: V = [(., () 1 j E r()),, r re [m], @ E }
Here, (j, r) again refers to neuron for class j. as in Definition C.2. With this notation we can write (o) as Vy for 3/ = (j. r, o) E W.
Our proof shows that, given j. for each ø E 27. there is exactly one neuron r in Wa.j.r that is activated, denoted y. to learn Φ; that is, Vi.rg.y (6) will grow to a large value. Therefore, in total n distinct neurons will be activated to learn all combinations, i.e., {(j. gg.. (g. y)), Vj E =())), (g. E }3). The magnitude of remaining Ve with non-activated neurons will stay close to the initialization.
Our analysis shows that learning follows an implicit curriculum induced by the magnitude of features Vy at initialization:
VO
Vi,, E V.
Items on the left under this ordering are learned first, and those on the right are learned later. We denote by E' the learning curriculum: the ordered set of neuron-feature indices {(i, gggg (0, 3))) identified above, equipped with this order.
Then, for each s E E", the learning process follows its associated ordering, and when it is the turn of 3/ = (j. T g- (g. y)), the learning process is mainly characterized by the following two phases:
Phase I: Emergence of the feature Ve among other features. During this phase, V grows faster than any s/ * & with v < v², while not affecting the already learned predecessors vf E E* with w x tb, with growth rate
((++))
(1( + ((r)) (bj,r + ViO + o(n)))
(16)
This form permits the application of the tensor power method (TPM) [81] and explains the ordering induced by the magnitude of Vi at initialization, since TPM implies that slightly larger initial value leads to dramatically faster growth. By the end of Phase I, for any other feature
v * e with &i < v the growth of that feature is capped at Ō(ao). its initial magnitude.
Phase II: Growth of Vú and cancellation of incorrect features. After Phase I, the target feature Vi already has a relatively large magnitude; however, features the set {1/ = (i.Tg X¹) E V 8' = (g', v) or ø' (g. v') v') namely, the features in the same neuron (j. rg ) that share exactly one component of Φ = (g. y). may also grow as the shared component increases. Note that j is not the correct label for such o', and we call them confounding features. The key characterization in this phase that, due dynamics are coupled, the wrong half of the confounding features grows in the negative direction while the correct half continues to grow, and ultimately the confounding feature cancels out.
stationarity property of the gradients, although the
We show that for each E Σ* that should be learned, it retains its structure after its own learning process (i.e., at the end of Phase II) and persists through the final convergence while the other features E are learned. Specifically, we have the following properties at the end of training:
Vy (9), Vy (v) = B, Vi (g'), Varun (v/) ~ -B,
(17a)
(17b)
VY,,, (g) + Vir (1)) S o(1) and (g') - Vgr,,(u) S o(1) 4g v/ F y. (17c)
How does this structure perform group actions? For an input pair (g. y), let the correct answer be r(g(x)). Then the pre-activation for predicting is around B at neuron rggv- with all other neurons for j near the initial value Ö(oo). Thus the model output at j is around B. For all other predictions j'. there exist Φ = (g', y) and de (g.1/) such that j' = T(g' (1)) = +(g(y')) By cancellation of the incorrect features, we have Vy (01), Vy (02) S o(1). Since
B = O(lo d), this implies that the logit on the correct prediction j is very close to 1.
Symmetry Group Actions. The proof strategy for symmetry group actions mirrors the simply transitive case through the emergence, refinement, and convergence phases. However, symmetry actions create richer interactions because multiple group elements can map the same y the same j. which requires more nuanced control of the training dynamics and leads to different learned feature structures. To illustrate this pattern, slightly modify the definition of combinations used in the simply transitive case by introducing the notion of the fiber of a value.
31

---

Definition C.3 (Combinations, Symmetry Actions), Assuming the group Ç follows Assumption 5.1, define Fiber,,y := {g T(g(V)) = j}. This allows us define the combinations as
4 PP j E T)?), y E V). where Pi Fibber {y}.
Moreover, lp).yl = (ny - 1). We continue to call á o= (g. v) E G x y a combination, whereas an element of 4 is now a subset of combinations Ssyy (there are n such 4yyy in total), which includes all pairs (g. y) such that g sends y to j. and which reduces to the single pair (g. y) in the simply transitive case.
Based on these notions, the main difference from the simply transitive case is that previously the basic learning unit is each pair (g. y) E D whereas now it is the subset Pjyy of combinations. All combinations é within fi.y are captured by the same and unique neuron, denoted T' in the sense that the feature magnitude Vy.r,. (ő) is large for any & E fy..y. Accordingly, the learning curriculum E' is now based on the initial value of the ensemble of feature magnitudes: |py.ll Eeeps Virr,. Finally, at convergence, we have the following imbalance of feature magnitudes:
1
2B
Varr.(()) se 2B,
Var (2)
gg E Fiber
(18a)
Tly
2B
Vors,((g)) nl
V/rr,v(') -2B;
(18b)
Verr, VVr,, (3') s o1) and VYrY, (g') + (br,.()) S o(1)
Ag' & Fiber, 3 + y- (18c)
C.1.3 What Changes for Longer Tasks?
FFN is length invariant. The feed-forward network (with weights W) only acts on the attended linear combination at the current output clause, rather than scanning the sequence to retrieve informa- tion. Retrieval from the context is delegated to the attention layer. Hence, as sequence length grows and positions shift, the FFN computation remains the same mapping on its local input, making length- and position- position-invarianee invariant. On the other hand, uniform attention becomes increasingly diluted as irrelevant clauses accumulate, so selective attention and learning attention layer are required.
Desired attention patterns for predicting each tokens. To predict the fourth token, the model only needs the variable c from the predicate clause 2pred.8+1- If the attention mass on Zpred 8+1 does not vanish with L,, the same pattern learned in T with Wa applies directly. Predicting the fifth token is harder because it requires two retrievals: (i) the group action from pprd,,+4 and (ii) the current value from Zans,é, followed by applying the update via W5. Thus robustness over length hinges on maintaining both attention links. If the attention layer is robust enough to support the fifth-token prediction, then the fourth-toke prediction follows as a special case, since it needs only one of the two attention links to persist. Accordingly, in what follows we focus on optimizing the loss for the fifth token; achieving high accuracy there effectively equivalent to solving the entire task.
C.2 Learning the Attention Layer
Successful training on T shows that the model has learned the one-step update g o y. Turning to task T2, the remaining difficulty is routing: directing attention to the appropriate locations. Thus we train the attention matrix Q to learn the routing pattern and keep W fixed. In this part, we will show how training induces the attention concentration pattern highlighted in Theorems 4.1 and 5.1: given an input Z f_,f-1 the attention mass concentrates on AttDans, ---edd, and AttDans, --t-aann------ms..---
We quantify routing quality via the attention concentration degree
I.E. Cattn
Zz 1
Att.aans/----ree (Z) Atttnamsf-1--ansstt (Z (19)
which measures the fraction of attention mass not placed on the two key clauses, and the attention gap
L.E (ZL Att-aassr-1--prd/ (oerdd,((ZLt (z Att (Z).1-1)
(20)
5.
We use the notion of fiber to denote the left cosets gG,, C G. g E G. where y E y and Gy = {g E g ply) = y} is the stabilizer. Since the fiher of the orbit map fa(9) g(y) is the preimage of fy, i.c.,
fy (v') with y' = 7 ( E T())), this is exactly the set Fiber,,y-
32

---

Predicate look
Self attention
5dxis
Zrred
Zunsa
Zane
d
Res
Qus
4
Qua Qu
Figure 7:
The illustration of how different components of the attention matrix used to route the attention to the appropriate locations. The query clause is Zans, and the goal is to retrieve the correct action 92 from Zpred. and value y1 from the current answer clause Zans, [Q+.p) a,a will grow and dominate the learning dynamics for p (3, and T()). Thus, in this example, large [Q4.3] ((11).((()) indicates the large attention the predicate clause pred, and large [Q1.4] ((1)).((1)) indicates the large self-attention to the answer clause Zans
which captures how balanced the two target attentions are. Thus, effective routing corresponds to L.C small Sattn (high concentration) and small A L.E (good balance).
Under Assumption C.2, for an input Z L the (unnormalized) attention score from Zans,E-1 to clause Zk decomposes as
Z ans, 8-1 QZK
Z Z 1,4
Z Pans,6-1,4 11 Zr.
where Zμ =
-ZZK......... Zk.5] with Zu,i E Rd. By design, in the clause embeddings the fourth token of predicate clause and the third token of an answer clause are (blank). Consequently, we further have
Z Q
D f- Eredd
Z Zans QZans, Zans,e 1,4 Q4 Zans ans, M,
(21)
which means Q4.3 governs attention to predicate clauses, while Q4.4 governs attention to answer clauses. An immediate observation from (21) that the desired allocation can be realized by growing the diagonal entries Q4.31s.. and [Q,, for E T(X') since æz-1 appears as the third token in E and as the fourth token in Zans,--1 simultaneously. Our analysis of tracking the attention dynamics will show that these diagonal coordinates indeed receive asymptotically larger gradient magnitudes than all other entries due strong co-occurrence signal, and the training dynamics dominated by their growth. See Figure 7 for an illustration.
For notational simplicity. we will refer to the relevant diagonal entries [Q4.3ls. and [Q4,4ss. simply as Q1.3 and Q4,4 below. The remainder of the proof quantifies how the growth of Q4,3 and Q1,4 simultaneously drives the concentration degree L.E toward zero and controls the gap A 1.f ensuring that the FFN consistently receives (ge, yt-1) and thus outputs the correct ye.
C.2.1 Simply Transitive Actions
For the simply transitive case, we analyze the gradient contribution at position = 5 on task T i.e., the loss Σ " We show that for 2 the attention concentration degree <atnn (for E E [2]) can be reduced below O(1/d")) for some constant 0 < c* < 1, indicating highly focused mass on the relevant clauses. When irrelevant entries of Q small, also have 2,1 2.2 since the number of irrelevant clauses doubles from l to f 2; hence focus controlling Cattn-
2.6
2.2
Stage 2.1: Growth of an initial gap. Early in training, attention is close to uniform, so given z we have the approximations
Vor((1)) + Vg.r(22) (30)
6
[A].. denotes the entry of the matrix A at the s-th row and s' -thh column
33

---

Az.j.((Z2.1)) = + (22) + V.,((b)) + (3.
By the cancellation at the convergence (17), for = all A lie in the small smoothed regime, whereas for é = 1 we obtain correct logit for 1/1 91 (ve) and a spurious logit for g2(vo) of magnitude about B/3. Consequently, LLLos is comparatively large and drives Q1,3 to grow faster than Q1.4 (increasing Q4,4 would also amplify the wrong prediction (92 (Vo)) and thus not reduce Loss ). An S2(1/ log d) gap emerges between the diagonals [Q1,3ls,s and QQ,4 yielding an early routing advantage toward right predicate clause.
2.2
is negligible, while - Losos
Stage 2.2:
Joint growth with a controlled gap. As Q4,3 increases, the weight Attn I 1-pred, becomes large, moving As, ( and (o so for = 2 into the linear regime. Gradients from = then dominate, and Q1,1 starts to grow to separate the correct 1/2 = 92(y1) ) from the incorrect 7 (g2())))) Throughout, the gap between Q1,3 and Q1,1 stays within [22 / log d), (()), so the attention gap satisfies A 52(1 log d).
Stage 2.3: Convergence and gap reduction. Continued joint growth of Q1,3 and Q1,4 con- 2.2 centrates attention near its ideal limit, making <attn small. We show 222 cannot remain above o(1) for long; otherwise an incorrect logit logit,, ( would acquire a stronger gradient and force Q4,4 to outpace Q4,3, which contradicts stability. Therefore, at convergence we have (i) 2.2
Cattn O(d==)) for some c' E (0,1) and (ii) 2.2 o(1); hence both Q4,3 and Q4 equal Coog d + (1), while all other entries of the attention matrix remain close to their initial values, where C 0 depends on c*
Strong length generalization. The key is that attention concentrates cleanly while remaining stably balanced. For Th, we obtain
0(1) CLe $ 0(1) L + 1/
$ 1 O(d" /L)
Hence the model tolerates Odd") irrelevant clauses: particular, if L = o(d"") then ct. = (1): and if L = O(d'') ) then e_. < TT5 which can be made small by choosing the proportionality constant T. = (dd") appropriately. Moreover, ALE 2.2 S o(1), so the two target attention masses remain balanced as I. grows, preventing errors from large imbalance (e.g., prediicii 7 (02())))
in Stage 2.3). Consequently, the correct logit satisfies 1 - logit, ( 1(xe) S poly(d) so T is solved with accuracy 1 - 1/poly(d).
C.2.2 Symmetry Group Actions
We now turn symmetry group tasks T and analyze GD updates with respect to the per-token loss 4.2
Loss (i.e., predicting the value token in Zans,2 from Z.,1).
The case J, = 2. The high-level picture mirrors the simply transitive case, but because multiple group elements can map a given y to the same j. the learned (g) and Vy.((y) structures are unbalanced. Moreover, as we discussed in the hardness part for the symmetry group in Section 5, there will be a non-negligible proportions of distractor clauses. These make harder keep a tight balance between Q1,3 and Q1,4 as in the simply transitive case. Nevertheless, we prove that both Q4.3 and Q4,4 grow, and the attention gap A 2.2 is controlled by feedback mechanism: exceeds a small fixed threshold (in either direction), some incorrect logit receives a stronger gradient, which pushes the system back toward balance. Consequently, after sufficient training: (i) catin < C1; (ii) A 2.2 < C2. for sufficiently small constants C1, C2; and since B = O(log d), we still obtain Loss 2,2 < 1/)oly(dd)
2.2
2,2
2 is already small constant, initialization for Recursive learning for 7 2 2. Because <attn
2² T satisfies aatin < 2 Fattn 212 (still small), and 2 2 < A g 1 Thus the attention pattern remains
close to that in 7-2-1 and = follows a bootstrapped LEGO distribution generated by the greedy model )) which coincides with the original LEGO source. In particular, y/1 = gl (yo) and
3/2 (2(1)) are correct. We can therefore reuse the convergence analysis from 72" to show that both Ettn and 2 decrease to a small constant, yielding stable, inductive concentration across recursive reasoning depths.
2*
34

---

C.3 Significance of the Proof
Our proof techniques are inspired by recent advances in understanding the dynamics of feature learning in neural networks [102, 103, 104, 105, 82, 106, 83, 84], which show how gradient- training induces useful internal patterns and representations. Building on these ideas, we analyze the dynamics of the FFN and attention layers to capture how transformers gradually acquire length- generalizable reasoning through CoT training. We conclude this section by summarizing the technical significance of our analysis for learning the FFN and attention layers.
Learning the group actions. Most existing optimization- basdd analyses of CoT [36, 35, 34, 37] hard-code the group action, e.g., parity, as a fixed FFN or an almost-linear map. In contrast, we train nonlinear FFN end-to-end to perform the group action at each step. with parity n = instance of our simply transitive framework. This formulation strictly more general and substantially more challenging than prior setups. Our proof shows that the model learns not only the basic action for simply transitive groups but also more complex actions for richer actions that is transitive but not free. Especially, we precisely characterizing how task-relevaie features emerge and spurious features are suppressed during the process. The technique used in the proof for learning the FFN on discrete combinations of data is of independent interest beyond analyzing CoT.
special
Learning the attention patterns. Prior work on transformer training dynamics has established the attention concentration principle as a key mechanism for solving various tasks, e.g., in-context learning [83], self-supervised learning [84], and graph learning [107]. In those settings, only a single token needs to be retrieved, so the attention matrix is learned to be diagonal (after a suitable change of basis), and single pattern suffices. In contrast, our task requires retrieving two different types of clauses, so the model must learn two distinct components (blocks) in the attention matrix, Q1,3 and Q1,4- As we show in Appendix C.2, beyond growing these blocks, we must maintain a delicate balance between them. This balance is crucial for length generalization across different group actions, introduces substantial technical challenges, and our proof provides fine-grained control of it.
D Learning In-Context Retrieval of Variables
In this section, focus on the learning process of W4 for the task T For this task, W. should predict the 4-th token in the answer clause Zans, which is the variable x1. We will show that, the network learns to retrieve the target variable x2 from the first token of the first predicate clause preed 1 to make an accurate prediction. Throughout the rest of the proof, we will omit the subscript for the expectation E, when the context is clear.
D.1 Preliminaries
First we define some notations for the presentation of gradients.
Notations for gradient expressions For each i E [5], l E [4].] E [d],r E [m]]. we denote
Ei (ZL,t11) 2 lzm....ee, logit, i (F, (F,ZZ,. 1
(22)
Aij.r (2
Att (Wijjn Zu) + bigjor
(23)
kEIZ,t-1
where logit,, (F, 2 are defined as
Pis (z
logit;,, (P, 1-)1))
Eyeta Eo (2 (2ire-))
Fact D.I. For any i E (5), j E [d]. E [m]] we have the following gradient expression:
-W Loss" E E (z (Z (.-es)RRLLLL' (A
Σ
Atttansse-1-xxl
KEIL,
For simplicity of notation, we will henceforth denote Ai,.r (Z by Aig. and Ex, (Z by E when the context is clear.
35

---

Given Z" we use Pt. to denote the appeared variables in the context clauses, i.c. 8r. (20,T1,.. 31,..,,, Jt). We write X X' for simplicity. Throughout this section, write [Fil] as Fi, for simplicity.
D.2 Induction Hypothesis
Proof Plan. The main idea is to track the dynamics of different types of weights W......p Specif- ically, we prove that for each j E T(2), there exists certain neurons E [m] such that the corre- sponding weights W 1.j.r:1 grow significantly along the direction ej, while all others remain small. Specifically, we proceed in four steps:
1. For j (()), define
- max (W rE/m] J. Ej) + on log d,
4.
to track the maximal activation associated with retrieving the correct variable token.
2. Establish rapid growth (early phase). Let A (11mm). and define the hitting time
Ti, - min tt 0: r 2 )
We show that for iterations Ti O(dog-2/n). the diagonal weights (Wl)) grow rapidly, causing [h enter linear growth regime. Simultaneously, the model confidently identifies the correct variable, indicated by logit, 4,r( o(()
4. ez)
(2)
3. Convergence via dominant neurons (late phase). For each j E 72()) define active neuron sets and their total activation
fr E [m]]
(Wi)
4.j.1 ej) 2 d]
rEAR
Ti, we analyze the refined dynamics, proving that the total diagonal acti- 4. for the weakest activated variable jr, eventually grows to O(log d), ensuring successful learning.
For iterations >
vation & (2)
4. Bounding non-target correlations. We finally show by induction that all other correlations (wit) 4.j.r.p- ex), including group actions, value tokens, off-diagonal tokens, and non-target variables, remain negligible throughout the training process.
Our proof begins by positing an induction hypothesis expected to hold throughout training. We then analyze the dynamics under this hypothesis and show that it remains valid along the entire training trajectory, establishing the claim at convergence.
Induction D.I. For T polyd all of the following holds:
(a). for E T(2), ((vv) < (w() 4..m. + p where (W(t)
1. is non-decreasing:
(b). E EG
(W4.m2: e()r) & Ō( 0 max (w())
(iii min (W
(w
MEAl) 4.²
(c). for ( y
max { (w()) Es), min (wl) (w) S ((00) + 0
rEAl
(d), else, [(W a e < Ō(oo) for any other j.j E [d].
Claim D.I. If Induction D. holds at iteration t, then for an input 1.0 we have
1. if T(æ1),
A 1 ( (..r2. e,) ((()) A. ((()))
Wil)
4.J.n ((())) z" +O(o))
36

---

2. else if \ {a1)),
1 1,,r 2 (wl) 4.j. ((a))) +
w
d..r. 4....5s. (ro)) + zz + Ō(oo):
3. otherwise, 0 s A S 3 (
Claim D.2. If Induction D.1 holds at iteration t, then for an imput Z Z
1. j T(11), logit (0)
2. otherwise, logit (1) 4.j = 0(2) (1 logit (6) ((21)
Proof. j T(11), by Induction D.1 and Claim D.I, we have
5 (Z v
Σ
Ag.r 1+ & (91) + +0 4.j
max +
(0)
4.5 )) + Ö(v)) + Om0g@mmgggg logd) d)
9
rE[m]
(t)
= (11) 1.
+
+ +(00) + 0
polylogd
For j E T(X' * T(x1), F (z (0) < ((0))
+ 0 max (+19
) else F1) s Ō(ao).
Combining them together, we complete the proof.
D.3 Gradient Lemma
Starting with the gradient computation from Fact D.1:
-W Loss's = B Ea (eL,U(Aaar))
Σ Zu.
kET
we first consider the gradient for j E T(x) Lemma D.1. For j T(X), we have
(a) for Wa.j...1. $ E T(X)
(1) if 8= j. (-V
w(o) Loss" ,e,) 1 logit (
) (
(2) 8 $ j. (-V a.g.r Loss ea) = 1 logit 1. RReLU (
(b) for W 4.j,2, = ((g) for g Eg
1 _ (1 logit |ReLU (A 1 (
V wio
Loss' ,ass",e))=a
4.2in2
2
logit sReLU (A (20)))000))))
(c) for Wa.in with p (3, 8 T(X)
(1) s = j. V-V ,e,) 1 logit (A [
(2) jj
(- W- Loss ee) .e,) - (1 logit ( sReLU (Ale) ( ((00)))))))
logit !)) sReLU (40)))))
(d) for W 4.j. s = -(v) for g E
Loss', ,ca) = [ logit sReLU (A1)..) (T
W(t) 4.r.b
37

---

logit " (A (
Move on to j & =((), we can obtain
Lemma D.2. For j E T2)), we have (a) for W W.....11 8 E T(2).
Loss .e,) ,6,)= = _1e
logit ② RReLL'(A S ( [=
w 4,g.r
(b) for W Wajr. 8 for G.
(-V w/o 4.2ir22
Loss' cca)
E 2-
logit (t)
CACLU [
(c) for W Wijp with (3,4], T(t),
Loss eg) = E
logit 1.
RReLU'AA (A1,))) (
w 4.mr
(d) for Wa Vaj.s 8 T(y) for EY
(-V w(t)
Loss' .es)
E 2
logit (t) 1.
sReLU' SReLU'AA m
D.4 Growth of Gamma
Lemma D.3 (Growth). Given (X),, suppose Induction D.1 holds iteration when 16 44 S
(1)
0.01 log d or (()) <
0.01 then it satisfies
m
p(++ = 4. + + )
Proof. By Lemma D.I, have
(-V With 4iir
Loss .e)) = B (1 2
logit SReLU (A 1
By Claim D.2, when $ s 0.01 log d, logit (t)
a lo 0f a)+d
< when ( and
combing with the fact that the event (7(x,) j} happens with probability X we complete the proof.
Lemma D.3, combined with the growth of the tensor power method, immediately gives the following corollary.
Lemma D.4. Suppose Induction D. holds for all iterations. Define threshold A ()))) Ler
(t) 1.j
2 A=, and Ti def O non (-)))) Then we have T 2 Ti, for every
Tis be the first iteration so that
j E T(X), for 1. 2 Tı, it satisfies ((t) 4. 2 A-
Lemma D.5 (Upper bound). Suppose Induction D.1 holds for all iterations < I, we have (1) < Õ(1) for j T(X).
Proof. We only need to consider the time 2 Ti. Notice that the gradient descent update in Lemma D.1 gives
Loss' ..)) 1 [1 logit (9git )RRLU (A)
Wit 4...n
Therefore, for sufficiently small 7,, have
$ ((+-))
1
E [ logit (2) ss ( + O(glogd) [1(t"1) AAl
(
2
(()
38

---

(0 + Σ
?
(1
logit (A [ + polylogd
1
2
rEA
When there exists T, s.t., maxjer(X) 11 Q(log d). by Induction D.I and Claim D.1, given an
input sequence 000 with (x1)) = =j = arg (E(()) T 4. we have
F (
Z LZ
4.. 1 0 (g)
$ TT) Ō((o) > 12(log d).
4.j
Following the similar analysis as Claim D.2,
(T) 1,3' (Z 10)
for other j' E T(X), and 1.j e -fllo' 5 d) Therefore, we derive
- 0
4
[o|
F (Z) S (1) for j' $ TxY), which implies logit
(T)
that for [ + 1, polyd
7
1 ( Oppolyd e e-Q(0og"' d)) Opplog m,
which implies ( 1.3 0. d) since @ < m logd"
D.5 Group and Value Correlations Are Not Large
Lemma D.6. Suppose Induction D.1 holds for all iterations then for any j T(X) and 8 = T E 9, we have
[W (0) (t())) 2((0)) { ( max (w())
Ag ey),
min (W rEA
9
Proof. By Lemma D.I, we have
(-V with) Loss', ea) ea)
(t) 1 e (1 )
SReLU' (A4..)) (A)) logit (0) RReLU' (A
(Anjir) 1
Clearly, the positive gradient can upper-bounded by 0 ( w the negative gradient, by Claim D.1, we have a naive bound
Loss ..))). Moreover, for a.g.a
sReLU' (A 1 {O (A ((1))=)/1="'
When / < Tí, by Claim D.2, we have 1 - logit (t) which implics
1)=r())) 2(1) and logit (1)
(iTT())) 0(2),
B (f ( 50(
1 )(
V
Wit 4ahind
Loss .e)).
Therefore, for t < Ti, we have
(w (t)
(TG))
s <Ō(p)) + +0(1g) (w
(g) 1.).)
For / 2 Ti, notice that by Lemma D.4, A (1)
4,3' + 0 for j' E T(2), thus for r' E (
sReLU' (A (A) ( < sReLU' (A ( m1=9"
Furthermore, ( (1)
(yAT(O)) (1)) - logit (1) b)-=(()))) which implies
(1) E logit
RReLU ( 0 |
1 ) )(-V W-V Loss' ,y-).
1.
39

---

Due to the arbitrary of r', we can conclude that
(w 2: - Ō((u)) +0()))
min (Wl) NEA
Lemma D.7. Suppose Induction D.1 holds for all iterations t, then for any j x(2) and 8 = T E V, we have
(w (1) 4g..s ((o)) S ((oo) 0(
1 1
max (wi))
Agg.n eih
min (wl) reAt.
Proof. The proof is similar to Lemma D.6, and we omit the details here.
D.6 Off-diagonal Correlations Are Small
Lemma D.8 (off-diagonal bound). Given j E TX),, suppose Induction D.I holds at all iterations tt, s T(2) j
[ (....nee)))) s ((o)).
Proof. By Lemma D.I
Loss ea) (ass,,e))=
_ E logit RReLU' 3 [r(o)=))
1
with 4.5i
2
Notice that by Claim D.I,
sReLU (A )
s
combined with Claim D.2, logit, 1.j (0) (.!) when T thus
E logit (t) RReLL (8) 1. (x))s B [ ) (1 logit
f. ReRLLL ( =s
0 d ( W....... Loss,es).
From Induction D.I, we have
(w (0) 4. ea)l 0( (wl) 4.)))))ss) Ō((o) s =012) -((o) Ō
Lemma D.9. Given j E T(X'), suppose Induction D.I holds at all iterations < t, we have
(w 4..mp Loss Loss", ea) s Ö(oo), for e (3, 4] and all s E T(X)
Proof. When = j, have
V Wa ,oss',e,)
,e,)
2 logit (A L((m)))
1 E
logit RReLU (A Σ [(
2
85)
Therefore, we can bound the above gradient in the similar way as the off-diagonal case, and obtain
d.j.rpp e))l s o
max (((())
(.).. e.)l Ō(op) Ō((o).
sEr(2)
When j.
W .c)
40

---

= (1- RRLUU D (A1,,)) ) (()))=)))()))= logit (2) sReLU
1
(Aj,)) " (A
((zo) (=o))=a ( jer (2)(())
Noticing that (7(20) = s, T (21) = j} happens with probability - thus the positive gradient can be upper bounded by 0(2) V-Vwa,.... Loss' ez) Furthermore, the negative part can be
upper bounded in the similar way as previous off-diagonal negative gradient. Putting it together, we complete the proof.
D.7 Non-target Correlations Are Negligible
Lemma D.10. Suppose Induction D.I holds at all iterations < t, for & T(X), for p E [5] and [d]
E
kw a e.) S Ō(oo).
Proof. By Lemma D.2, (W())
4.3',r.p ea) for p E {1,3, 4) and s r(t) can be bounded in the similar
way previous off dia-onag negative gradient.
We can observe that for j' & T(X), all the non-zero gradient on the different directions are negative gradient, which implies (w also non-increasing.
Ca) S (Wi)) m: = O(oo). Moreover, (1) & Ō(7o) is
4,8' inp.
For T(g).g E G, whenever (Wilt)
2, ea) reaches - 3u. we have 1 91 S -3p + 2
(aa) S 0, and thus (-V w(t)
(Ww Loss e.)) 1 logit 4 SReLU' (A 4.j.)) 1 which 4.i.n
implies (wl)) m,2 Cx) 2 -3p. Hence, (w()) m2 e.)I s Ö(oo). Following the similar argument,
we can prove the result for (w()) n,,. ex) for E T
D.8 Convergence
Lemma D.11. For 191 2 1D 21 log log
log log d), polylogd 2 m 2 13% <
log and sufficiently
small r) < -
polyd Induction D.1 holds for all iterations t < T polyd
Proof. Putting the results in Lemmas D.5 to D.8 and D.10, we can directly establish the results in Induction D.I.
Lemma D.12 (Convergence). For sufficientll large Ti S polyd we have
(a) Objective convergence: Loss' S polyy
(b) Successful learning of diagonal feature: T 2 Milog d) for any j E T(()).
Proof. Assuming for some sufficiently large constant n > 0, E[(1 - logit (.).) T(x1)= = =ja] 2
2 for / E (T1,TH + an+1 dn+"log" then by Lemma D.1, we have
(T+ 2l 4.j(0)
2 22 ) d log? d
(() 4.j(.) 2 2(log* d),
which contradicts with p(t)
S (1) 4,60 I Ollog" 1.5 d) = Ō(1) in the polynomial time. This
4-5(-)
implies after sufficiently large iteration I, we must have [[(1 logit!!) ((11) o s ( for j E (x)). Hence
Loss' = =El- log logit (t)
E log logit 4.j [=())==)
4.7(z))
j
Σ E(0(1)(1- ) 1+())==]
jEr(Z)
(logit (t) is very close to 1)
< <0 polyd
41

---

By Claim D.2, at the time of convergence, we must have $4 2 (log d).
E
Learning Simply Transitive Actions
E.1 Preliminaries
Definition E.1 (Combinations o, Φt)). Let G. Y' be defined as in Definition 2.2. For any pair (g.9) 9 x V. we call ó (g. y) a combination. We write Φ1 and Φ2 = y denote the corresponding components. We further define the set of correct combinations $* by
$ {Ø (g.v) j ((a(9,9))), Vj E r())
U jEr(D)
$
(24)
The combinations & E 3 are the ones that can composed to get ÿ (J) E y. We also define the set of incorrect combinations of for each j +(V) be
/ fo + 93)
(25)
For simply transitive action : x y, any combination @ E Φ) is incorrect pair of g and y.
That is, holds that any y/ * y and g' # g E 9 satisfy a 5/) aly. + (j). For each combination, we define the following neural-features:
Definition E.2 (combination features vb, Vr Nr* w/)., Given combination & (g.v) E Φ token index j E T()) and neuron index r E [m]. we define
Vy,r(v) (Wadnb,ey), (o) 1 2 (V)r. + Vyr(
(g) (W (.....ee.))
and call (Ø) the features for combination 6 in neuron (j,r) +()) x [m]. We write
zi = (j.r, r, o) and Ve = (Ø) to make the notation concise. Similar to above, we further define
{v (j. 7()) [m] x }
and T that contains the desirable features for each class j E T(V)):
T = (v) = (j, o E $3). V U
(26)
(1)))
and ! that contains the incorrect feature combinations:
T jr = (1> (r., o + Φ7}, Vt
U
(27)
(r)))r())) [m)
Events of combination appearance. Let us first define some useful notations. For a combination (g. y). write H. to denote the event when appears in the sequence '
Hos = [(11 vo) = $).
Hoi = Hg = = = g).
Hoz = Hy = {vo = =
Further, we write
H! 0.1 (6] + 9.90 v).
H! 0.2 [g] 9. yo # v).
H! H!, UH! 0.2
to denote the event where @ did not appear but its group element or value is the same. Finally we define a notion called learning curriculum, that sits at the center of our proof.
Definition E.3 (Learning curriculum). We define order on the set Σ, defined by the following process: Let Mg = V. At each [n]]. we choose
vb = arg max (0) DE
Let's write vi = (j, r, .) where (g. v). then we define E by excluding the following features:
Exclude the confusing combinations in neuron (j.r);
xt = = (t' (j.r,0) E E., = (g'.v') (g' =9) XOR (3/ =
42

---

2. Exclude the unselected combinations in neuron (j,r);
E =
= = (v' (1. E. = (g''0)) + g. + y)
3. Exclude the feature indices of the same combination in other neurons (i,r'), fr.
Σ = x = fo = (i,r' E Vr' + [m]}.
Which returns En+1 Ej (vv = ,,r','))) E Σ, either r' = r or の' = Φ). Iterate over the whole
set V, we will arive 2ng 2. This gives rise an ordered sequence "n and Et:
Ex = (v),, U2 .2....,, vig);
at : U
Dt
VEE
By our construction V (0) (0)
We write si -y v²' to denote that & is ahead of v in 2
4xi1"
Intuitively, 2n encodes the order at which the features V grow in magnitude and E contains all the unlearned features during the process.
Throughout the analysis of the FFN layer, we further make the following assumptions.
Assumption E.1. For rj. [d] [m] x [5], we fix ((()) 6.....p initialization for simplicity of proof.
for p E (1.3.4) at w! Giny
Assumption E.2. We use modified smoothed ReLU as our activation function. This technical assumption is to avoid many pathologies appearing the learning dynamics.
BB
x I B I ( - B, 50
x²
I É (-w, 0]:
sReLU(x) =
x /00
1000 9)
E (0,
2 0 1
1
E (e. B]
B - 0 1- 9)
x >> B
where q (1) is large even integer and o @(L/polylogdd)) (d 2))) B d-1 d-1+e"
Cllog d, C = (11) E (5, ( and X
E.1.1 Theorem Statement
We will try to prove the following theorem:
Theorem E.1 (Learning Simply Transitive Actions). Suppose F()) is returned by Algorithm I at t Tı, and let 51 d pr 62 कi.h, it holds that with probability o(1):
A. For all & E", we have (TT)) B O(11)
B. For any v E Etd, | (T)) < Ō62));
C. For any v E 2!? (T) < B + 0())):
D. For any v E Et,, v((h) < <(());
E. For dh (j,r, (g. E E*, 1n (g) v/t (w) < ((2))
In fact, we can decompose the statement of theorem E. into the following statement for every feature E.
Definition E,4 (Feature Shape). Let ő = (61, 52) be tuple of error parameters. Let zb E Σ* be feature in the learning curriculum. We say the feature reached feature shape (o) with error if:
1. F (8)) ((0 2 B -0());
2. Fo.2(62): For any & E 2x14 it holds that (v(0) < Ō(2);;
3. Fv.)(81): For any v 2t, it holds that (1) - B 0()1;;
43

---

4. Fv. (.))())) For any feature v/ E 2 it holds that V < ((())
5. Fv,)( Writing
(i (g. v)), then it holds that (1) v (v) S (
Obviously if we proved the condition (o) of all E E* are reached for ò (d |1 w Tr, then theorem E. is proven. We do this by following induction process sequentially for each & E Σ', and provide a guarantee that all features in Σ* reach the claimed convergence condition together at TT.
E.1.2 Induction Hypotheses and Phase Decomposition
To prove theorem E.I, equivalently, prove the convergence condition in Definition E.4 is reached for every feature v/ E h", we shall charaterize the dynamics of each feature and prove that they eventually arrive at the desired shape.
Phase Decomposition. Let define the following timestamps for different phases of learning each (i,,,)) E V U V. Let 1000 be small constant:
(a) Phase I: / E [0, T where T = minf! 2 0
V 2 d
(b) Phase I1.1: f E (T). 1 where 2 min 0 V (1)
2 2/2 log do
(c) Phase I1.2: E (T 2,1 22221 where T22 == min 2 Ellogit H. 2 1 +/
#
(d) Phase II.3: (T TT). where T m 2.3 Te + Ω (( and T is
2.2
defined as:
T2.3 = min t 2.2 F 62) holds, where 8 = d oo 82
(e) Phase III.1: E (T Ti.... where Ti, = T and thn is the last feature E This is the convergence phase where the rest of the features are learned and the feature converge to perfect shape.
4.3
(f) Phase I11.2: t (i], the end phase where all features are in perfect shape and have stablized.
Now we state the following induction hypotheses which naturally results in theorem E.I. Induction E.1 (Simply Transitive Actions, All Phases). For all S T1, the following holds:
(a) (t)) + bigjo 2 (()).
(b) T3 < D if v/ <x 1b for 1b, 10 Et, ie., the intervals ([T), T3) west
overlapping;
In order prove Induction E.1 and theorem E.I, we introduce the following feature based induction hypothesis to prove that Fs defined in Definition E.4 holds for some error parameters at = T
Induction E.2 (Induction for Individual Feature). Let db (i,r,)) E" o (g. v). att Tú the
following holds:
(a) At t < Th any g + g' + has Vg)( (v') S (()).
(b) During S T1 there are at most Ö(vd/n) iterations where logit
Vd conditioned on
Ho.
(c) Let u/f 3 E", the feature shape F (d 22 ) holds throughout t (T). Ti]:
(d) Let & > E*. then (((() < Ö(0)) throughout t [0, T3).
E.1.3 Technical Calculations
Let us recall some facts about gradient computation for W.
44

---

Fact E.1 (gradient computation), The gradient with respect to W Wijpp for i = 5.j E [d], [m]. p E and is
(- V Wasi Loss, Co) = | Es, (As.j.r)
Σ
(28)
2
kEZLO
where E5,j 1 u= Zem.1..
logits.; is the loss derivatives. We compute precise expressions
when j. p and v varies:
(a) when all j E r())), E (m).
2, for 9, let Φ (g.v) E t3 be the combination with di 9. we have for
1 ( (AR. Ho)
Wasnn Loss, Eg) = = logit,
E s sReLU (As.jr) (A
(29)
(b) when p =5 for E ) let (g.v) E $3 be the combination with d2 y. we have for all j E(()),r E [m).
V Wwsssnoooss,ey
= (1 logit (RR.LU(AA,,n))
E logits,,ReLUU' H
(30)
(c) For any combination of p. e not included above, we have due to our assumptions of the distribution and update rule:
Wo..irp Loss, ex) = 0
Definition E.5 (T -ootaaaiions.. To simplify gradient expression, we define the following notation: let s (rrr, σ) E V we write
r+.00) = (+) =
E (1 logit, 5. r--00 r
JsRcLU ,
(0) =E |
(31)
we can further define I' for each E G and E y: Let E G UY, we define
r+.(1) = (()
6,0 = - logit JReLU' )seL
jr.
r (t) -- (t) = E logits, s. SerittsssReeUS sReLU' (As.,))) Jay
(32)
This allows to define the following gradient condition:
Definition E.6 (gradient condition). At t < Ti, letting vi - (j.r. Φ) E V, we write (d) for 5 5>0 0 denote that the gradient update for Vi smaller than 6, formally:
(0) holds
r p_(()) (4)
jr..v
(6) holds C (D)
[Tt(1))
< 8
C (d) holds
Jrme 8
We further write
(6) holds
(6) A Cj.r.y(( (8) holds
C (6) holds
corgg()) (6) ^ C++ Crr.g (0) holds
C (6) holds
Cj,ng(0 (8) A C holds
Now proceed to characterize the lower bound of logits Fact E.2 (logit lower bound). We list some basic facts about the architecture Definition 3.5 here.
45

---

1. For any Z' supp(D'), we have
d-1
= min (1 (F
d d 1 + e
2. Suppose for some i [5]. Fi,j = o(1) for all j $ =()),, then
min logit, ( (1+ o(1)d d + ngeB = 0 F.PPZ
nyd
E.2 Phase I: Feature Emergence and Competition
We shall prove the following properties at initialization.
Fact E.3 (initialization range). At i = 0, the following holds with probability 2 1 11-000 o(1):
(b) S.j.r Gx) S Poaovloge For all j, r,p and V;
(a) (W (0)
(0) 7, where 22 m loga). ), for any pair vb x th E* in Def. E.3;
oo For Fact E.3b, Proof. Fact E.3a can easily verified from our initialization W iihp
we give a straightforward proof that every pair of v-(0)
v 3 E E" has a gap of 1:m==logd First note
that ((() of different & E " are independent and identically distributed the randomness of W due to the orthogonality of embeddings ev; E Then, by the basic property of Gaussian variable (notice that v0 ( is also Gaussian with variance 20,) (all though different pairs could be dependent), have with probability 1 ntm= log that their gap is at least 2 inmm log for each conclude the proof. C
&
pair. Then by union bound over O ) -mayy all possible pairs
We give another characterization of the initial activations.
Fact E.4 (activation magnitude). At 0, with high probability it holds that
A 1 = (1 + O 1
log ))pr O log d) <<
and thus (ReLL((A.,r)) (1))) sReLU' sReLU'(A.,r)) 0(w and Fs.j ( 19).
Proof. Combining Definitions 3.4 and 3.5, Assumption 3.3, and Fact E.3 gives the fact.
We establish some properties in Phase I. Lemma E.1 (Key properties in Phase I). Let th (i ().v)) E E*, and assume Inductions E. and E.2 holds & S T then
(a) ( fir (v) < O(/ O(//Vlog ff u
(b) 5. (Z) 2 for Z E Hoi
(c) logits.j (1/d) whenver H. happens.
Proof. lemma E. la is from both the initialization of weights Fact E.3 and Induction E.2a. lemma E. from lemma E. and Induction E. la. lemma Ee is due both Induction E. lb and Induction E.
E.2.1 Competition between feature combinations.
In this phase we compare the growth of different features using proxy. Lemma E,2 (Approximating gradient with proxy). Assume Induction E. I. For a fixed de = (j, r, ()) E
Et, define the auxiliary sequence
(0)
(0)
((+1)
(()) + SRELU (A Ino]
where in A (1) 5.j.r the weight Vo is replaced by Vo (all other weights are unchanged), Then for any
feature vb V. E, {14} and all < T (v - < /21/2-30
46

---

Proof. We need to prove the result by induction. Firstly we notice that the newly defined proxy sequence V monotonically increasing. Moreover, (0) is initialized the same as V (()) can use lemmas E.11 and E.12 to argue its growth by taking
It (-(t) + bej.er = Σ P=3,4
(Wa..r....o) + (Wa.j.,, F+41 2 2t + n
B( + 11)9-1)
n
And lemmas E. and E.13 guarantee V 2 d after Ō (p)))) steps. Now assume the bound
of difference between v/(() and Vi()) holds S T Since TT. have v(t) definition of T. By Fact 1, the exact update of Vi is
d oo by
((++)) (t) " (B[(( logit s.
SRRLU' (A (4) 5.j.r 11, Eloggit
2 SReLU g (0) s.jr ) xc.)
Compare the above with the definition of (t), first we notice that ()()) is dominating v for all iterations f S * Then, we can bound the difference between the two sequences by
v
(1(+1) (((+1))
[0 v (1)
SResL sReLU' (A (t) (. ||
+
? E
logit RRcLU' (A 1n
E logit RA
(0) (....)
We bound the three error terms for < T.
A. Activation perturbation: by the smoothness of sReLU' on [0, and that both pre- activations are in this range during Phase we have
[((ReLL (AG..) sReLU (1) (A (.rr)) S (V v EE [
(...)) )0 nol
where we have used the fact that As.j. 2 As.j.r because of the dominance of V.J.
B. Outside a set of iterations Ag (t St& T logit (t) z Va conditioned on Ho} whose
cardinality is no more than O(d tos /n) by Induction E.2a, we have logit 5. I d 1/2
d-1/2,
Hence for / & A
logit RReLU'AA (A
I O(d-1/2) (A... 11.]
Now we can sum over the iterations to get
Σ
Σ
E logit (8)
(A [n,]
+St,
ASt
so 0 ditas
:n(d" 11)-1 + 1/2 v
11
s {0(1/d/2-0)
C. Confusing events: before ((1) 2 1/2 loog d, we have logits,, = (1/d)) in Phase by lemma E.I. So we get
(d 11) (d |RRLL (1 u2) s O muv-2 mm(
Σ B| 5. SSt
< O(//Va)
d
So have that the accumulated errors from the last two terms are smaller than ((z/d /2-cc1 This allow us to bound the difference by using the more naive approach: we can assume the difference
47

---

started with 0(p/d1/2-o1) and now we only need to bound the following sequence for t - T
8e+1 b; + O(jō; JI
sReLU (AG)))) [oo) do S 0(q///1/2-)))
Now let 8' > 7) be some small parameter, we can do some slicing of time to, b1,t.... . io where min 2 0 (1) 8') and (1 + 8')" ((++1) I (1 + 8')nt1 u. This produce the
following bound:
log 0,+1 - log 6n
= Σ
log 1+)
UReLU'AA (.g.sr) (1o
Isisn sE(t,te11
5.j.r
Σ
Σ
(A
(log(1 + x)) S x when x 0)
Isi5a s(
5.j.r
s
Σ
r E|sReLU (A6g,)
ZEH, A(6 (Z)
min
Isisa
sEfte.tonl
Σ 13i5n
1
(V((a1))
(1 8' (1+8))'1
log( (1)) ((()pp))
/m)
log(1 +0")
< 2 2oog (/(1)//0 /4))
(log(1 - 8') < 2f 6' is small enough)
So when (()) E [d μr, 2d pll, it holds that 8, S de d0 S 1/2 This proves the claim.
Lemma E.3 (competition). Assuming Induction E.1, and let E", we shall have for any feature
2/ - E* or o Eg (vb) that v (TT) * I <Ō(a))
Proof. We can approximate the sequence Vo, using (1()) and straightforwardly use lemma E.12 to
compare vi with any v/ > v/s Σ*, fact, let's s compared the time they reach V 2 d oo. By
we can use lemma E. to get that the first time when ((() d og. we have (7 " 2 1710 +7,
lemma E.2, we can just compare Ti(() before = T3 now since by Fact E.3 we know ((()
Ō(a)). This
combined with lemma F.2 concludes the result. Moreover, we can also obtain bounds on all sb E v by discussing different u/:
For all ø' E <5* such that there exists x/=
= (j', r'', o') x the o' is learned. That means any z)
contains Φ') has gradient too small because of E.2c and therefore EEEs, 1 H.,1 s ((A and therefore wouldn't be learned.
All the rest of ø' E Φ- has their corresponding vi' > v E Σ* and therefore lost in the competition and has small growth.
All other & V Up lost the competition as well for the same reason.
This concludes the proof.
Proof of Inductions E.1 and E.2 in Phase I. Let v. Σt. Assume Induction E.1 holds for some t < 13 iy 1. then we shall prove that continues to hold for E [TSS T T
T1
Induction E.la: This is simple as all Vo for vb > the has updates approximated by (1(1) which is monotonically increasing, and has minimal update by lemma E.2. Induction E. lb: We proved in lemma E.3 that whenever & < vi E Σ*, T $ yr- This does not violate Induction E. lb so the non-overlapping preserved.
48

---

Induction E.2a: This is proven in lemma E.3.
Induction E.2b-c: They will be proven when we analyze the end phase of feature learning.
They are the corollaries of the end phase guarantee of the previous features.
Induction E.2c:
Induction E.le: We did not violate Induction E. either.
Induction E.2a-d will be proved in Phase III where we argue the end phase of each feature learned.
E.3
Phase II: Feature Growth and Cancellations
Below we present the induction hypothesis for phase
Induction E.3 (Phase
II). For all xo = (j. r, Φ) E E* and I. E [T, Te). the following holds
(a) Fort E [T 1. E 2 1 1 conditioned on Hg:
(b) For & T for all j' # j. we have logit (1) b.j' 5 5O(vd) logit 6,3"
(TT)
(c) For any ø' = (g', v') + oe V ()), (v)) E
(d) For t S 2 the number of iterations where ((++1) - (t) < 0 is bounded by O(1/dmnn)
E.3.1 Technical Lemmas
Lemma E,4 (gradient estimation). Assuming Inductions E.1 and E.3, let th (irr,)) E Φ
(g. v) and t E [T T.l. then
(a) At any 2 1 if v E (3 log d, log d), then r()) j.r 2 2(().
(b) Any of fge + E y has the following gradient approximation
r(t)
jin' = - E logit (0)
RReLU' (1 wl 1 =
r(t)
j.yy = -E logit <ReLU'(A (1) (...)) 11 H + +(w/nng)
(c) let ø' = =g', + De oj if (g. v) 2 B - o(1) and V (g' .// -B 2p,, we have
r
Σ
Q Aw dmg
logit 5. ReLL'AA (.. 11o
j.ngg' =
v (o D ( Bw-ww
and similarly for (1) jrry
Proof. Firstly (a) is simple to prove, because when v((() (1 logd, d, 3 log d), have logits,, In3 5 0 ), which from the calculation Fs.j S log d + O(1) guaranteed by Inductions E.2 and E.3.
For (b), we can estimate the gradient update of 1/ (g') where of # g. In fact
Vo =((t)
(t)()) r+)()) (t)
"jr.g'
j.n.
j.n.g
Σ E logit (A (I.
v'Ey
-ogiit
(0
sRRU 1 ( max max (0 w
where the last one is because for v/ # y. the activation gradient (A s.j.r) is either positive and bounded by 0(u9-1) negative and bounded by while logit,, Hgg S Similarly can also get the same for Vy.r(y') where y' + y:
(0)
-Vv,.
(t) j.r.y'
E logit ( 5.
(A
5.j.r Mlgg ((=nn))
49

---

To prove (c), notice that when ('r'''vv') < C- - 2/,, the part of T++(0 jing' 0 due to As.j,r exceeding the boundary - B. We group different ij * y E y by their feature magnitude and obtain the lower bound. The factor of x is due to the logit lower bound from Fact E.2.
dmg
Lemma E.5 (Feature Magnitude). Assume Inductions E.1 and E.3 holds, then for any &b E Σ",
(a) at t = T we have v (g) log d o(1), V 5ir (v) = log d + o(1);
(b) when C+ ( hold some å - (()/d1)) at some t 2 T we have vl9 that is, Fo,1 (d μ)) holds.
& 2 B Odd
Proof. Let's proceed the proof one by one.
Proof of (a) :This can be computed by comparing the gradient of (t)
j,r (g) which is (f)
jong
D+t(0) [ T and the gradient of jir (v) which is r r -
.())
r r-((t)
2.. Now since
j.n.
r+.()) p+) j..g jr.y
(1 0 RV (g) + v (v) 0(w)
1
and that
T = 0
(V (0) T = 0 (V
Vi
fr
one could see that Eigra1 (T()) - Jnn s Ō (2) and thus (a) holds.
Proof of (b): from lemma E.4a know that when C (6) holds for some small 6 = (
it is not because (1(1) dropped below 1 log d, thus when Inductions E.1 and E.3 holds, it can only be because (0()) > B -(d otherwise
(+) V
= [ logit ( [ XPr(Ho)) >
(if (if V
Odd p))
de
which concludes the proof.
Lemma E.6 (Gradient Stationarity). Let tb E², assume Inductions E.1 and E.2 holds at t [T) T.). and choose â >> Tw then we have the following guarantees:
(a) Define B = {t E [T T31 c (5) doesn't 't hold }, then |B S O my log2d
1, ny log" d); nd
(b) Define B3 {t [T T3l CJ (8) doesn't 't hold }, then |5, S O
(c) Finally, t T33 the feature shape Fo 82) holds with 51 d and 82
(" )
Proof. Let start with (a). First by writing vi (j,r, k). we define the following quantity
r = (z)
(33)
which sits at the central of our proof of counter-argument.Nov we have for any T 2 121
r(()
( (0)
(0)
Σ
ViF(6') (d')
((E)
(
(0)
Σ
(T
(0')
Σ
nD(()
jr
Σ
Σ
inn
2 .1)
sEfr?" .0)
50

---

We can further rewrite the I's on the RHS to
Σ
Σ
(t) Jind
fin
I
d'Engilo)
(
Σ
- (
p
wn +
Σ
( find
(34)
ME )
& #0EE
d'SCn;
Now we can start to estimate the RHS. Define P1 = fō' E x+ Ø' learned and P3 = {Ø' +
4 d' learned]. The sum Ed r can be decomposed and bounded by
- S
Σ n
Σ
T Ö(1du8-3)
MEEA
and then by the same methods bounding the gradients, we have that
Σ
Σ
logit (A (A
(1)
Σ THo
Σ In
8"4g'EG
V"
Blogit ReLLU'A
Hoy +
Σ
Ellogit sReLU (A (Ag.r)) H O(w)
8 #9
V'
r j.r.
+ ((w)
where ① is because for any incorrect feature combination o' = (g', & Φ5 and does not share
component with o, the activation of As.j.r conditioned on Ha/ is within [- ((x0-1]] and thus the negative gradient upper bounded by w w/n? and lower bounded by -((1--1). Inserting this back to Eq. (34), we shall have
(TOO Jond
Σ 3 .T)
Σ
(T(1)) Jirier
Σ
T jn
+ + -T )
('E4))(0)
2 .T)
Now we have arrived at the desired estimate of the update, finally, we have the following counter- iterations, letting T 1 be the last iteration which
argument: if c (6) does not hold for 22( By log" t (6) does not hold, we have
n1
m N
Σ
7I
j.r.
Σ
n + R(( T )mo) )
(For 1 BJ. ( 2 0)
t6 1 3 Br
My log d π1ō mp
Σ
n7 (
3 TNBB
2 0 ny log 2
9 176
(by choosing à > w)
2 (nn logg d)
By o' is learned mean there is feature e/ x/ E" such that v (i,r
51

---

which is impossible because the gradient would have vanished when T 2 ny log' would mean 36 E 27 such that 1V Jir (d)l 2 log" and it is forbidden in our setting, which proves the claim via contradiction.
1.5 d as that
Proof of lemma E.6b. Actually, by applying (a) combined with lemma E.5b, we know that for x it holds (²()) 2 B - O(d=) ud). From this iteration forward we shall prove that V/( stays above B p)) and can help. In fact, at / T, V (V) (-O(B), (()) for v E GUUJ \ {g. y) by Induction E.3, Say of + g. then due to lemma E.4, its gradient is approximated by
some T 2 T2 + 2
(())
(1) E logit ? sReLL' (.g..)) [x 1 O(w/m))
(35)
Let 0 2 w. we will find an iteration S +
2 + +0( Myy log"
that satisfy (o). By applying (a), we iters after T3) and before T
no dy
first get that C (6) holds for more than 0 (Bh 170
log
Let's assume this is the case, that C (6) doesn't hold for more than S2 "v log" d) iterations. In this 100
case, it must be that the gradient D (=(()) < -ō negative, since when the gradient positive it has ind
an upper bound (ww/ny) = o(d) and therefore cannot violate the condition C (6). Therefore the updates Eose 1 17,1 must accumulate to lower than
Ry log" d
(0) S r()) 22(
x np -S2 log" d)
r78
which is impossible in our setting. This proved the result for any 5 >> nyw. Proof of lemma E. Let d = n w, which is covered by lemma E.6a-b, can find an iteration where j.r.or S 8 ob E 03 \ [Ø). This implies we have found / such that
" (AoJ, 1 S niw = (y (d') O(w)) S 0 dew
(We lower bounded the logit by Fact E.2)
Thus we arrive at (V)((())) S O((() n do
Remember in the proof of b, v(t) = 1(V jir (0) +
v-ft (v)) B (d u) at t T + 2
d y
So we have for any o' = (g', E D \ {Ø).
dlw
2 (w イイ - < -
(t) 5.r (g) + frr (v)) + 0
< - B +
dL/34
Now we have verified that given = n there is a iteration 2 which gives the near sear-perfect feature shape (ō,, d2) with ô1 d μ, 62 2dde ) will hold until t = T2. For L 2 T,3 we verify F 82) in Definition E.4 one by one:
We shall prove that this feature shape
1. (d n)): By lemma E.6a shall find step after t 2 T
00 at which d nx)
(d μ)) is satisfied. In fact, the majority of the steps after T 0 will satisfy this. We shall prove that combined with the above feature shape guarantee, F (d 1)) will hold true until L = TT.
2. Fw,2(( 2d --):: ): From the above argument that applies lemma E.6b, we shall find a time step where Fa. (d μu) and 22.2(d11//5)) hold. Now we discuss the possible values of the gradient v E It.l In fact, let ø' = (g', v') E Φ \ {Ø), we have
-B Regime A: If V Vi.r < B - 24,, then there is no negative gradient for Vy.r(g') and Vy)r())) unless Vg y) or (g. s/) is positive. In fact, only one of the feature can be positive to have negative gradient, and in that case the gradient of ((t) j.r. v E {g', v/}
is pointing to the convergence direction V (g', v) + 0.
52

---

Regime B: If V (or) - - 2p,, then there is possibility of negative gradient for Vgrr(g)) and (3)) even when Vgrr(g'. y) or Vg.r(g. 3/) is negative, In this case the feature Vg 3) or Vg,r(g.vV v¹) might not go the desired direction to converge to zero, but given a 6 > n, r in(g', .y)
r (ar((.v)) cannot stay above 8, otherwise the term will dominate and revert the gradient direction. On the other hand, when one side has
(g, < 2((( n. dxz a--)) ), we have
(o) - (2) 22( do
24
which is back in regime A and therefore will converge back.
Induction over the above two regimes for all [T T3] showed that the feature will indeed satisfy F (6)) for all 5 > n x at f.
3. Fe.s( do } ) We need use both Fø,1 and Fu.z to get the desired result. In fact, we have for any 3/ = (j.r, (dd.8/))= E Dt.
P(0) (') + v (1/) + 1 (0) + Vi (V) + O(( n
7-1
because ( + frr (v)) and V (g') v (v) are bounded due to Fv.2-
4. (d p): The number of iterations [T. I is more than Ō (ل ) so by using
lemma E.13 along with the induction from the result at & shall conclude the proof.
5. Fv,5((( do --rr):: By using lemma E.7 and Fa.2 after 1 T23. we have the result. This proved (c).
Lemma E.7 (Symmetry of Features). Let ap (J.r, (g.9)) E E* and t T23 have that if |V S 8 for all & 1i,, then
| (g) 1l jrr (v) S (6)
Proof. For a feature 30 (f) E Φ (g. n). their update can be represented as:
(0) = Unn
Σ Rgg
+ (0) (g)
v'Fu
(v) v1)
Ug
Σ
Rg' au
v:(0) 3.r
(v)
9'
where Ug = Eose mi tlee j.n and Ry. is
Rg.,' Σ #Elgogtt
RReLU' (s)
5.
Therefore:
Σ (1) = Σ Us - ΣΣ
Rg.s'
Σ Ug
Σ
Righ + -
fir (g)
8
v'
PES;
s $tE0!
Σ (w) (v) = Σ Us ΣΣ Rgy
Σ
U
Σ
Rit + (10()) (v)
SED;
9'$9
ot
Thus
Σ VV
Σ (3) + O(p)
9'
53

---

Moreover, we have assumed
1 (g)
- /// イリ (1) + (0)) (A bg (v) (v)
Thus
Σ V V/ (g) - (my 1) (v) + (2₂0)) v
-nn 1)1; (g) + O
Therefore
fr (2) (g) - (00)
inserting = O((
which is from lemma E.6 proves the claim.
Lemma E.8 (Arrival time estimates). Assuming Inductions E.1 and E.3 for v E Er, we have
(a) 7 T s v
-))
(b) 2 TT I 0( 11tt
(c) 22. 2 S Ō 1
(d) T T I Ō ((
Proof.
Proof of (a): Once (tt) 2 1
og we can apply lemma E.13 to compute the number
of iteration needed to reach ((()
2 2
1 log d. In fact, by Induction E.3a-b we know that
((++1)
n11 - 0
vall RRRL +0(1)) v(()) +((((v(())00)
1
Since (33)
2 d ou. we have by lemma E.13 that T,2.1 T1. S S0( (0d
Proof of (b): Simply apply lemma E.6 with at+ai*
Proof of (c): Again by applying lemma E.6 with 8 = n m
Proof of (d): Follow the definition of T3-
Proof of Induction E.3. The proof of time steps are all in lemma E.8 and lemma E.6. We do not
repeat here.
E.4 Phase III: Convergence
We present the induction hypothesis in this phase.
Induction E.4 (Phase III, final). Let = (j,,,)) f" [TZ,Ti), Ti). the following holds:
(a) At t Ti,b we have Fx 52) holds with d₁ d 82 n²
(b) Art [T Til, have (61 52) holds with 5, d pt, 82 (n न We need a lemma to describe the logit distribution when all features are learned.
Lemma E.9 (Logit shape at convergence). Assuming Inductions E. and E.4. Let 1 E E" be the last E', then at = T the followings hold:
1. For & E $3. logit 5.j 2 1 - O()) conditioned on Ho
(8)
2. For & & T;. logit < (A/d) conditioned on Hág
(t)
54

---

Proof. Since v is the last feature in the learning curriculum Σ", all feature at t = T has satisfied
Fo 62) for d1 = d p4, 82 = dos ) Thus we can bound the logit by
e (
e B-o(1)
logits
(1) 3,
Po. (2) eB-o(l) d =1 -
(conditioned on Z E Hp. & D7)
and
logit
(1) 5.j
p(t) (2)
+ eB-o(L) + d 1 c = =02) (conditioned on Z E Hg. D7)
ed(1)
which is the desired result.
Lemma E.10 (Gradient Bounds, Phase III). Assuming Inductions E. and E.4 holds, and let ai- (j.r,) E E", then at E [T1,1: Til, the followings hold
(a) p+((0) 2 ((()) when v/ S d
(b) 1770 < 0 dry for any guy;
(c) For ø' E $ \ {Ø), we have
- A(Xw/n2) E (- ((WD/dn3) {O},
if tV
E - 24, -w]
rt(t jr.
if ivvlt) inne
U
E - - 2/u, - + 2μ1
Proof. lemma E.10a is simple when lemma E.9a holds. lemma E.10b is basically redoing the
calculations in lemma E.4. lemma 1. is because when ((t) Aing - 24,, we have
RReLU (As.s, E[(1 logit AReLU'(A (1) [1 E 55
2 AsRcLU (v-(t
J.rder + p. + Pr(Ho))
And the rest is simply the property of s'eLU' at the negative boundary when some activations of To. crosses the boundary.
E.4.1 Proof of Induction in Phase III and theorem E.1
Proof of Inductions E, E.2 and E.4. We prove the induction by reusing the proof of lemma E.6.
Proof of Induction E.4a: The same as in the proof of lemma E.6c, with longer time. We do not use any property of T$ additional handling is needed.
Proof of Induction E.4b: When all feature is learned, that is, after /. = T3 for the last feature s E Σ', we have perfect feature shape and therefore have much better logit shape described by lemma E.9, We reuse the proof of lemma E.6a-b to get the following result: for any 8 w( A= ), we can guarantee that (5) doesn' 't hold for more than Ō ( dm,
iterations. And beyond that, we can just reuse the argument in the proof lemma E.6c to obtain the desired result.
Proof Induction E. la-b: They are correct in Phase III because there is no violation.
Proof of Induction E.2a-b: They proven in previous phases.
Proof of Induction E.2c: For predecessor feature v < v E Σ*, we proved in this phase that they maintain feature shape until Th, and improved their feature 7
Proof of Induction E.2d:
This is proven in previous phases.
Proof of theorem E. Since all feature in E have shape Fo 82) for on -d 62 1- the end of Induction E,4, we have proven theorem E.
55

---

E,5 Auxiliary Technical Tools
First we need Bernstein inequality for U-statistics
Lemma E.11 (concentration inequality for pseudo-U-staaistiess), Let I...... In be different symbols, and let m & be such that 0 (mod n). Suppose for some function h with [h| S M the random variables b(zi,;Figs.... Figs.....
Tim and h(zg. "..,....,, Iv, are independent and identically distributed
as long as {Ti, Eig. Tim n (ag. T.g....,, BV 0, then the pseudo-U Coaddo-U-statiin Utsttiistic
Un.
Σ
h( Tigs Ti Timm
054 ca
satisfies Pr([Un,m E|U,. 2 t) S
m.32
Proof. The proof is the same as in [108].
We present two lemmas related the tensor power method.
Lemma E.12 (TPM in expectation). Let p > 3 be an integer, and a 0 satisfying μ " a Aog d. Now suppose & ~ [0, σ) a Gaussian variable and there are two sequences Vt, yt defined by the following update rules:
T++1 x4 C,EEE + 24 for some C. O(1):
* Mt+1 $ yt v 1E ((5 + ] some S = (
Then if To. yo 2 μ. vo - yo 2 Q polylog(aj) for every 2 such that Ty E [d0.01 xos 0(1), we shall have yt S (ppoyylog(d).
Proof. Let 5 > 0 be such that 8 >> # and ā o(1). Let Tú be defined the first time when Tt (1 6) 20. and min [i] + 8)'20 2 A}. Now we can compute a growth lower bound for
ε, := at each step:
T4h 2 Sn
Σ
C x1)0 (E y()
2 ST; ET;
Σ
C,B(l( + 21))--1)
(E(T,TiH))
2 ET + ST, Σ
CrEel(E + xx)r (1 +1+ 8) T1
(E
2 sn + Ehg 01 + +))'2
xo
(1 8742200
8
(1+ (1 + 8)2 ee,
Therefore let i' be such that (1 + 8)
v have that
8
ETp (1+
(1 (1+ I' 2 (1 o(1)) > o log
(36)
At t Ta have Vr, < Fr, I (1 o(1) To. After t = Ty, construct surrogate
sequence Ti 2 ye for 1 2 Tv by initializing Iт, yT, thH and updates by
U++1 Jr C..
This update guarantees that 7. Vt because J. Rel(E y1) every step. Now by applying corollary E. to the sequence of 24, De using the fact that IT, Z Vr, + 1/ Ur, polylog(d)) ). can get the desired result for Jr and thus yt at To. i.e, when F, A.
1
Lemma E.13 (TPM, adapted from [81]). Consider increasing sequence 2t 2 0 defined by Tt-1 Ft + C)C 9- for some integer q 2 3 and C. O(1) 0, then we have for every A 2o.
56

---

for every 8 > 0, and every 1) E (0, 1):
5(1 + +0)-1
(1 8).0
)r
9-2))
(11A9-1)
log(A/20)) 1
Σ n 2
1
(1 1---899-111
A
log(1 + +6)
x/ x9-2
(1 11+8)% 8)9-2
))
(gg(o))))
Σ SA
n S
+
2
T0
log(1 + 6)
This lemma has corollary:
Corollary E.1 (TPM, from [81]). Let q 2 3 be constant and zo, yo = o(1) and A O(1). Le {z,, tc2 be two positive sequences updated as
F++1 xt C for some C. = (1));
Mt+1 = + SSCyt g-1 for some = 0
Suppose zu 2 S (1+ polylog(d) ). letting TN be first iteration s.t., 2 A, then yr, S Ō(9)).
F Learning Symmetry Group Actions
Roadmap. We begin by introducing notation and defining fibers, followed by a description of the learning curriculum and the key appearance events. We then break the training process into clearly defined phases and prove the main results for each phase: In Phase I, features emerge and compete; In Phase II, features grow and mutually cancel, with precise tracking of when different outcomes occur; In Phase III, features converge to a desired shape.
Let us restate the action structure of and y here.
Assumption F.1 (Assumption 5.1, restated). Let LEGO(X, G, be the LEGO language, where | ny and G is group acting on on y. We assume the (left) group action G ) 3 is transitive but not free. particular, we assume ny = | (w(1), O(log d)) and the group satisfy 1 for all y E y. è
ng | (polyyoog(d)
Remark F.1. When the action G x y + y is transitive, there is only one orbit y = {g E 9} g/Gy, all y Y. By orbit-stabilizer theorem, the stabilizerGy, at any point y has the same cardinality.
Remark F.2. The canonical action of symmetry group Sn on Zn satisfy Assumption F.I with the choice ny O
log log
a) that keeps ng! (pooyyoo(dd).
log log log
F. Preliminaries
We work with fiber-indexed features. We first define fibers and coppincombinations ip, then we define neuron features by its index vi (j,,,p) and finally set the curriculum and appearance events in this notation.
Definition F.1 (Fiber of values). Assuming 9 follows Assumption F.1. For each (J)) and y eV,, define the fiber of j at y as:
Fiber,,, == (g E 9 T(g y) = }}
which collects all group elements that send y/ to y = (j)-
Remark F.3. We denote Fiber,,, j.y be the that transport y y' - because it is the + Tg--39g33 y) over y E Y Algebraically the fibers are the left-cosets JgGyboeg, where G, := {g E g.g. y) is the stablizer subgroup. Moreover, have the following fact from the orbit-stablizer theorem.
pre- pre-image of the predictor map (.(+)
Fact F.1 (cardinality of the fibers). For every j E r())) and y V, [Fiber,,yl |/|/D' na/ny from the orbit-stablizer theorem.
Definition F.2 (Super-combinations). Let (j.y) (())) x we define the super-combinattins. which is a set of combinations:
Piv := Fiber,,, x (v) C g x y,
11.0 Fiberj 24. (1)
8 We ny ny interchangeably in the proof below.
57

---

Let Φ == Ujer(""; ve be the full set. For a fixed class j E T()), we denote the correct set to
be 3 := Uya and the incorrect set b := UytieeyPPa base pair @ = (g. y) belongs to Pj.y and only if g E Fiber,,y- To simplity the proof, we also define for each (Pyyy C D, the confounding combinations
Φ 3
U
Py'a
XOR #y
That is, the set of combinations that intersect with exactly one component with Piay Definition F.3 (Neuron feature indices). Let us index neuron features by 1b (i.,.) with & E Φ and (rr) E T()) x [m]. and define
#3},
T = (U.r,))) T))
And we further define
V : TQ) x [m] x +
='=
U
Vin vt
U
V
(1))))(0))
The contains all the neuron indices we care about, and u and / contain the correct and incorrect indices for class j.
V-Notations We call pair e = (g. 2) with E G and E y a base pair, or following the naming of appendix E, a base combination. For neuron r), we define
(g) (...,....)) (v) = (Wo.i...s Eg))
+V (v)
For a fixed value and class j. the fiber Fibery,, collects all permutations sending y to r 1 (j). For (i,,,) we define:
Vo = ng/ny gEFiberj,
Σ Vg,r(g.)).
V = max gEFiber;
Vg,((,8), V Fiber,,,.
min Vgr(g.
For compactness of notation, let us also write (x)===>) -x) (W s.j.n,, Cno) + (Ws.jr.1, which is a random variable depending on the randomness Zo. T1.
The definition of learning curriculum for the symmetry group actions is essentially the same with the simply transitive case, with slight difference of using the fiber-combination instead of the base combination Φ. We repeat here for reference.
Definition F.4 (Learning currieulum, symmetry group actions). We generate feature index set E over fiber-indexed features. Letting Σn = V, at each E [n]], we choose
(0)
max 4b; = arg (o eE. PE4;
Write v/ r.) we define the next iteration set Σ,+1 by excluding:
(1) The confusing combinations for oo in the same neuron:
Σ x ((.r, En E b/}}
(2) All other combinations in the same neuron (j,r):
> 2 D.² 2 (G.r, ,') 1 & (0) (+}));
(3) The same combination fi in other neurons of the same class j,r r' #
1 #
= F = {(irr' Pan)) E E, Grrr #r).
The next iteration is given by
E = En \ (I!+) Dt: A
U 1 1
(zth))
58

---

Eventually we obtain E- = (UI....,, vng) and Σ := Uves- Σge We write vf xx a (or simply vl) if v precedes v in E". More generally, also write < of for v, of E v if s E* and V Up Et is from the.
Events of combination appearance. These events specify whether a specific base combination (g. y) appears, only one of its component matches.
Definition F.5 (Events of appearance). For (g. v) Psy for some j E T(V) and y E V, we define
Ho = {91 = 9. 30 = y). Hy := [و] g), Hy == {yo y}. For the mismatched events H. the definition is a little different from the simply transitive actions: since P,y Fiber,,, j.y {v} and E Fiber,,, j.u. we write
(g) {9 = 9. yo + y).
HC(x)) == {9r & Fiber Yo y}.
When considering the event of all the base pair PPj.y also define the union events
Hop = U Har
H.. = HV) (U (Up- 26(0))
dEP
F1111 Theorem Statement
Now present the theorem for learning symmetric group actions.
Theorem F.1 (Learning Symmetric Group Actions), Assuming Assumption 5.1, and let F(t) be obtained by Algorithmm 2
at Tú, then the loss for the 5th-token is minimized, i.e., Loss, (Ti) 1
polyddj
and for each ab = (j,r, r, :)) E D", the followings hold:
A. For any & E sp, we have V/FA (b) [B - O(d n), B + O(H)];
B. For any vf = (j,r, E Ete and d' E sp', we have IV (d') s Ö(w Or more concretely:
]/ / (g') + + (3') s Ō( V(g'3/) E
2
C. For any == (g.v) E y the following relation holds:
( S xŌ
Ca C (w)
where Ca (ny 11/ny O(ny) due to choice of nanny ny in Assumption F.I.
D. For any v E 3 it holds that V &5' = Ö(H).
The assertions of Theorem F.1 described the feature shapes of the neurons, which we formally defined in the proof of learning simply transitive actions (Definition E. The main difference between the results here and Theorem É. is that Theorem F. IC described a different symmetry of features without the factor (nny).
Let us follow the proof in Appendix E and the feature shape for symmetry group actions. Definition F.6 (Feature shape). Let = (j,r, r, sp) E V, and let d (61, 62) be the error parameters. We say the feature v has shape (8) if the followings are satisfied:
1. (8))( V( B O(51) and 7( O(+)(
2. F Fv.2(02): For any & (ir,,p') st. it holds that |1 [V/00(0') (d') S O(62) for all o' sp::
3. (6,)(62) 11 (g) Cy jr (v) < O(8)) for any (g. v) E ip;
4. F For any feature & E 111 it holds that ((!) S Ō()).
Indeed, we aim prove that (6) holds for suitable parameter 8 at the end of training, for every s E * And this is indeed what the theorem states.
Remark F.4. A major difference between Definition F.6 and Definition E.4 is that now the conditions involve not just one combination of weights, but all combinations & E yp given a combination yp E $ This creates some significant difficulties for the proof.
59

---

F.1.2 Facts and Basic Calculations
now define the gradient notation and conditions.
Definition F.7 (gradient notation). Let (i r, p),, we write I and I' denote the different
terms the gradient computation. For any base combination @ = (g. y) E Piy where Pi.y E D we write
r+:(0 == E[(1 (
T_ Jnn = E sReeLt
When consider the gradient of (g,((9) or Vyr(()) for individual tokens g and y V use the following notations: let Φ = (g. y) E Pi.y, then we define
j.n logit s. RReLU (As,,)
p= E [lggi sRcLU' (As.j,)) [xt(ol
jng
1 r=-0t) := E , sReLU (ARLLL"(As,s)) R4)
Finally, we sum up over & E yp for the gradient of combination & = (j, rp):
r (+) E [1 logit )sReUU (AG.,))
r (0) E [logits B (As...)))
Let s (i,r, up) where 9 Fiber,y One can check that the defined r+ base pair I' by the following relation:
connected to the
p+.
Σ
m+ jr.g
p=(()_
-= +
P-)()
(37)
j..g
Fiberr,,
gEliber,,y
Definition F.8 (gradient condition). At S Ti and >> O, write for tp = (j.
Cpp (6) | (t) -
(*)
S 6
(6)
o, C (8) | (1) S
These conditions control the magnitude of the gradient of the feature at iteration t. Next compute the expressions of the gradient using
Fact F.2 (gradient expressions). Equipped with Definition F88, the gradient with respect W Wij.pp for = 5. E [d], r E [m). p E [5] and v E g U y can be computed by
(a) when E T()), E [m]. = 5, for E G, let Φ = (g. v) E Paa be the combination with Φ1 g. we have
V W. w.snn Loss, eg) =
(T+.(t) j.n T
(1))
(38)
(b) when E +())), E [m], p 5, for y E let Ps.y D7 be the combination, we have
W Loss, eg) (r+
(39)
j.r.y
(c) For any combination of not included above, we have due to our assumptions of the distribution and update rule:
VW Loss, ex) = 0
E.1.3 Induction Hypothesis and Training Phases
We define the following intermediate time-steps:
Definition E.9 (Phase decomposition). Define timestamps for v using the feature:
(a) Phase E [0, TH), where T := min(t 2 0: v(t 2 d mJ).
(b) Phase II.1: E (IT) 12,11 where 1 = min 2 0: v 2 1 log dggd.
(c) Phase I1:: E (I 2,211, where 1 minft > 0: 1 1/23.
60

---

(d) Phase IL3: E (I ,Tl,, where T = m ,2,1 + 0 (nddd )), and T is
defined by
T2.3 = min 2 0: Fi 82) holds, where 01 = uu, (((d=/))d)))))))))
(e) Phase III.1: t E (T Ti, where 71,1 = T and Wag is the last feature in Σ*
Phase I1I.2: t E (Ti,1, Ti], the end of the training where the feature shape has stabilized We give several induction hypotheses, each characterizing different aspects of the process. Induction F.1 (learning symmetric group actions). Lett Ti, The following holds:
(a) v bi.j.r 2 (p))) for all vo E E..
(b) T3 < < up x v E E'. Thus the intervals {[T.) T&) veE are non-overlapping.
(c) Let v x E E", then the feature shape Fur (d Jo. क--1) holds throughout t [T Ti]:
(d) For any v = Gr E V* There are at most Ö(vd/n) iterations where logit (t) 2 va
conditioned on Hip E [0, Tb)
Interpretations of Induction F.1. We explain what the inductions hypothesis mean here:
(a) This property simply asserts that the good features never shrink below certain level which is crucial for the optimization to work.
(b) This property describes the curriculum nature of Σ*: each feature Vi grows in different time intervals within 0 < t < T1.
(c) This property asserts that the feature e previous to db in the curriculum E* has stable and good feature shapes, which helps with learning the current feature.
(d) This property means there are very few iteration where the gradient of (g.r Φ E to is disturbed before T)mmpp is reached.
F.2 Phase I: Emergence of the Feature
We first prove some properties at the beginning of phase
Lemma F.1 (Initialization range for features). Under random initialization and Assumption F.I, with probability at least - o(1), the following holds
(a) For any E V, we have kw (0)
S.inn ea)l Ouu/ Vlgg for any j.r,p.
(b) Let 1 E E* a/ v and th v/ then v
vl0) Do a gap Do Mooo/polyloge
(c) We have (0) A (Z) 21)) for all j [d], E [m].
Proof. The proof is basically the same that of Fact E.3. The first one uses basic concentration properties of Gaussian distribution and the second uses a union bound with the anti-concentration property of the Gaussian distribution.
(c) is based on the fact that
(0)
(Z)
bs.j.r
Σ
Σ
(W
Sainp Zk.p) = μt + O Vlg
kET1. PE[S]
which concludes the proof.
Corollary F.1. For any up = (j, (.r.))== E E and E v², we have with probability 2 21-0 o(1):
Eg (sRLLU (A Hy.pl BassylEE,[stee E. (v) + (z)( + bs.orr)
gEpt
(
E, ReesRRRLLUUL (v (v) + Vyrr(z) bs...)]
v /nc/ny
Proof. This is true because of lemma F.la and Hoeffding's inequality.
61

---

Now we present an induction for phase I.
Induction F.2. Let e = (i.r,p) E for t S Ty the gradients satisfy r+t.(t) any (g. v) E
3.r. 2 ( nc
I
Irttttt i,ja
for
ny
Using Fact E.1 with a representative & = y) E Pia and the definition Vgr((0) (V(x.)))) (v),, single gradient step with step size 1) yields
((ttt) 1or (0) 1/
7
2 (
WWs..in sisina Loss, Cg) + (- V Vw, Wa Loss, cg))
" r (0
T r-.0) .(1)
(40)
2
jir.g jrw
Similar to the proof of the simply transitive actions, we define proxy term and argue that the trajectory of Ve is similar to the proxy.
Lemma F.2 (Approximating the gradient with proxy). Assuming Inductions F.I and E2 holds at t. Let vb = (j.r.9) :p) E for every o = (g.y)= E yp we define two proxy sequences V (v), u E {g.u) y}
by the update rule:
V (v) v00 (v) 3
(0) gr j.r.
where both (g) and V2r (y) start from V (1) V (v), frr (g) = v (g). And the I' are defined by replacing the features of Vyr((u), x E xp U (u} with Vyr (v). Now for all t s T7 we have
1 (b) 3.r (6) O(v/di - 30), where Vi0 (0) (0) (v).
Proof. Fix rb (j, (...)) (40) gives
1 (v) / (v)) _
(
1 (0) - 1/0 (+) = T
Since Induction F.2 holds at t, have that ( (v) T00 (v) > s ( nG (V I t+) jing jnz
ny ((
(g) T (g))
Moreover, since Egers implies that
r+((() by Definition F.7. This, combined with Induction F.I.
(v) - AU
(v) 2 Σ
pt.l) J.ng
2
RG 1 (g) (1)))) Vg sp
lly
_E
The above update bound imply, for all t S T the following property hold: -
j.r (g) [1) <
Sneemn ((Vi (v) (0 (v)) for all g E sp. Now, by the same induction procedure the proof of j.r Lemma E.2, we know that the sum
O O(mer 1
( (v) + Σ
V (1) + Σ
(10 fr (g)
S $O(1/d²-00 3ca)
gEpp
gEpt
And then by the above relations we know
12:10 (v) v20 (v) S O(1/d2-30)
Thus the conclusion holds.
Lemma E.3 (competition). Let up E E. Assuming Inductions F. and F.2 holds at &. S Th. Then at
t t T we have for any feature sb E D, 16' > wb v(t) W S Ö(p).
Proof. From lemma F22, we know that
)((+1)()) (v) = V y
Σ
E
(V (v) + (/(0) (g) b5.ir + ( '1H
9 SEpt
62

---

Suppose there is another s = (0', E" such that 2 > db, then we can also write a comparison based on lemma F.1:
(0) j.r (t') Do.
And therefore compute by absolutely bounding the average over (p using corollary F.1:
t
Σ
Eg
[( fr
(w) + J.r (g) + b6g. + ()rr(z)))
|
DEPt
1
E E (v) + bo.j. + (()(())))-1]
v Vno/ny
Similarly, for e/ (j', ,'' (²) and o' E Ex', we have
(+.
E ((
(1 frr (v') + (o + bs r + V))r(()))) (Hg.yr
lpti
DEPrA
S (1 Vng/ny
Ex (r) bs of + (x))
Now we can compare the proxies Vy. (y) + bs.j.r with (v') using lemma E. to get that when
(v) 2 d'i we have (v) - (μ)). Now due to lemma F.2 and Induction F.2, we have the desired result.
Proof of Inductions F.I and F.2 in Phase I. Fix +p = (i,r, E Dr. We check items (a)-(d) in Induction F.l:
Induction F.la: This is simple by combining lemmas F.1 and E2, where each Vj. (Ø)) is well approximated by a monotonically growing proxy (Ø))
Induction F.1b: This is proved later but here lemma F.3 showed that for any pair th < thl E E* that y T which precondition for Induction F.1b.
Induction F. Ic-d: These two are proved later in future phases and are not violated in phase
Induction F.2: We shall show that this is the case. Consider any feature g sp", iis gradient is clearly bounded by O(ny r initially before the feature Vg,((g) Mu)). this range &
the difference of its growth with the update of Vjrr (v) averaged over all g E sp is bounded by O(ny/n)) (1/(1) (y) + b5. because grows slower than O(ny /nG Vy.( and thus Hoeffding's inequality with variance parameter μ would work bound the difference between the averaged gradient and the gradient with (g)((g) in expectation, and thus we obtain a bound r+ > M(ng/ny (imnc/ny) " When Vg)(g) (24), the feature (y) has grown more than ((na/ny) }u and the growth ratio still holds. This proved the induction for < T
F.3 Phase II: Feature Growth and Cancellations
Below we present an additional induction hypothesis for phase II.
Induction F.3 (Phase II). For all features (j.r,() E Σ* and [T]. T3). then for any
g' $ vp y up², we have v for (g'). V e (-B-O())
E3.1 Technical Lemmas
Lemma F.4 (gradient estimation). Assume Inductions F. and F3. Let th = (i.r,95.v) E Er, let
(g.8) E Piap then E [T] 13), we have
(a) Fort E [T]. T2.11 we have r 2 ((PP((H))sRoIU((())
(b) (f jrr (d) (o) E (1 log d, log d). then (t)) j..
= r j.. - -r 2 ((Pr(H))):
63

---

(c) For any g' $ xp and any y' +
r
---E Ellgit (1)
RReLU (A (n mow ml = =O(ony/na))
jang
Σ E logit sReLU' (A G.j [ ((w).
Jay
gEpt
(d) Any individual component Vi (v), E & U' cannot fall below O(p), and T+.10) in.
2
Proof. Write vp (i.r, 9sa))
(a) By the definition of 7 we know there exist e E jo such that (0) S 1 log d. By the same arguments in lemma F.2, we can get that all e E satisfy 2 log d + o(1). Given that Induction F.I holds at conditioned on Hg, all other neuron (j', ) * (j. r) has activation smaller than Therefore F3, o(1), and logits,j' = o(1) for any j' # conditioned on Hg. So conditioned on Z E Ha. when (1 I (2))
2 3 log d, we have
logits,, s.j = e.,,((2) + (L S 0(d
(1) Now we can simply bound all r for v E op Up? to be smaller than Ö (2
jina Thus
the proof obtained by computing j.r.v and sum get (t)
(b) Similar to (a), here we also notice that PRRLUU(As,,j, (2)) because As.j. (Z) E le. BJ
is in the linear regime. Now we can see that the terms r+ j.r 2 ((Pr(H))) and p-
S
j..
Od (d)) and thus we have the result.
(c) For g' & Fiber, let y E y such that g' E Fiberr,,' then we have
p(t) j.n.g' B[(1 - logit,, s
Σ 5 v'
(As,)) Hogro)
Since H(g'.)) confusing, the negative term dominates; other appearances contribute most ((#ny//ng)a O(p9-1). Similar arguments hold for y'.
(d) This one is proven similarly to (a) but can adapt to any level of Vo)(()) for the corresponding à E o for v E +pt Up².. From Induction F.3c we know that all other feature (v, S (p)) for o' & 4p U2 So for any v E spt U y2 such that (u.( < Ō(u), we have that the combination & (v, v') +
\/90 has (o)( O(μ) and thus have gradient
(())
+.0 j...
- 2 r+.(0 O(ar1)
j.r,
=
Note that p+t00 either 2 A/d 0 at this phase. So can sum up the negative gradient j.n.
-0( for all iterations t
[T]] T3] (which is fewer than Ö ((77)))))))) )) to get that
the feature (v) >> μ.. This means As.j.r when Hy happens and thus r 2 0 for all E [T] T3]
j.r.e
Lemma F55 (feature magnitude). Assume Inductions F. and F.3 holds. For any feature yb (j.r,t) E Et, we have
(a) when C+ (5) holds for some 8 (nG=,0(1//")) at t T2.11 we have v-(t) (Φ)
2
frr
Old p), Vo E s.
(b) when CT (5) holds for some 8 (ozz,((\/(to)) at 2 T2.1 have (b)
[-B B - O(1), ((())) + for all øt
64

---

Proof.
holds with 8 v O(X/d=)), then there are two probability: Vi (o) for all Φ E sp is as small O which is not possible for å so small because of lemma F.4a; or that 1 ()) 2 B (d (i) such that Fs.5 2 B and thus the gradient has vanished for some events.
(a) If C (8)
(b) The logic from (a) can be applied here as well. Since the positive gradient is at most ((w) by lemma F.4b, thus whenever > w, the negative gradient conditioned on Hø, Φ E Φ bounded, which gives
Σ EEp
E s (As. $ 50 = (As.grr) I 6X/(dPr(H())
(because Fact E.2)
ViJ (0) $ B Ao न 4Φ E
d
which proves the desired result.
Lemma F.6 (gradient stationarity). Let v (j,r,)) e S", then any 6 w 2, the followings
hold:
(a) Define B == {t [T 2 c (8) doesn't hold) then |B3 - 0
log" "): 1/8
(b) Define B3 = {t E [T 2.3 C (6) doesn't hold) then [Bg - O
log"
(c) For some T2.3 > T + O ) ) have V (Ø) 2 B O p.) for all & i
Φ and [V (d') S (())))) ) for all E tt. moreover, this will hold until t T3
Proof. The proof is very similar in spirit the proof of Lemma E.6. We first define the following quantity:
r ((t) (v)
Σ
for
(g)
9E4
One can compute that the update of (() is simply
Σ
(t(tt J..
r p=((0) (t) jon
77
s' F'PEE DEp!
d'pLO!
The middle term is bounded by ((w) and the last term is bounded below by -ww), thus we have that the update satisfy
r (t+1) r(t)) 2 nr - O(7z)
Since I+,(0) t t T at
2 For any 5 > by lemma F.4c, we know that the above bound accumulate
r
Σ nō (1)=)
So for no more than 0 log iterations should have the that T()) > 22log d), which is no
impossible in our setting. This contradiction proves the result for (a). The proof of (b) is similar to (a) and the proof lemma E.6b. (c) also similar to lemma E.6c by using lemma F.5a-b with parameter X/de1 and parameter now, alonge with lemma F.7. We do not repeat the details here.
Lemma E.7 (Symmetry of features). Let v (j.r, Ex and 2 T we have that if F (6)
hold for some 8 0, then
(v (g) - C. C. (x) < 018)
65

---

For some Co Linglny -1) 11
(ny (oy--)+na/ny
Proof. For a feature & (j.r,pp) E E", their update can be represented as:
(1) V/5 (g) = -Ug
Σ
Rg.
Vtu
v1r (v) Viro (v)
Ug.
Rig.
gi
9 d'K.iberr,.
where Ug.= = EEgeet n ( and Rg is j.r
Ba Σ
nBllogt (4) g SReLU (1
Therefore:
Σ (1)
Un Σ
Σ
Rg
geFibery.y
géFiber,,y
geFiber,, V FV
Thus it is true for any Fiber,, and y E y that the following holds
( Jr (v)
Rg
Σ (V ()) 1 f.r
(g)) (g)
Σ
Σ
Rg.v
Fiiberr,y
gEFiber,,.
#EFiber,, s'
By summing both the left and right hand side with all the (Fiber,, 8) pairs, we have
E(v Jir (v)
1i (v)) + + E Rg',N- E(v (g) 1/ (g)) EnEy Eg'
Σ DEG ('
Σ
Rg.
vEY
BEV g¹ O'FFibers,y
GEG
Since
Rg' m Enes Ey Eytr(a Rgiy' we have
Σ (1) (v) - 1/ (v)) = > (v 3r (g) - v(
3.r (g)
DEy
gE9
Σ (( (2) Σ
((t) (g) + (())
yEy
gEg
Moreover, we have assumed
Vy * f.g E Fiber V (v) + (3)1 Ag' E Fiber
1 (g) + V/ -00),
Thus
9 (1 + yl (1D) 1)) (v) P'I( 1)
vlT + O(ng8)
gEFiber,,.
Therefore choosing whatever E Fibery,t. we have
9 (D)) - 1)) (v) = G
1/22
g y
9
(1 +
(g) + 0(1a0)) = (1)) - +
(g) + ((na))
y
Now by dividing the right-hand side by the factor on v(t) (v) on the LHS, we have the desired result.
F.3.2 Proof of Induction
Lemma E.8 (Arrival times). Assuming Inductions F.I and F3 for up = (j,r,y) E E', we have
(a) 22.1 - -T3 < Ō -d)=))))==);
10
(b) 2.2 2.1 < 0 1
66

---

(c) T 2 0
(d) T 2 0 p))=2)
Proof.
(a) By lemma F.4,, after phase within 0(1/(1((900)00-) have v (0)
1
2 log d for all $ E e.
(b) On [T' 72.11 lemma Fóa gives the rate O(d to /n) time for the positive grad to arrive at logits., S va conditioned on M4p.
(c) This again provided by lemma F.cc by using ngw, combined with lemma F.5b for the feature shape guarantees.
(d) This simply by the definition of T and that T n(d
Proof of Inductions F.I and F.3. Fix 2 (i, (Ps.)) and consider E [T Tél. then
Induction F.la: This is guaranteed by lemma F.4abb.
Induction F.lb; Because the time used is small between T3 and TV by lemma E88, we have
that Ty is still behind when this happens.
Induction F.le: This will be proven in phase III, do not use any properties that violate this hypothesis.
Induction F.1d: This is due to Induction F.3a-b and lemma F.6a with parameter å = 1/vd.
Induction F.3: For any & $ 4p, we have that they are either desirable features E j.y. for some different y' 2 y E D), or that they are the wrong features for the class i.e., d E Py for all y E y. The feature vi = (j. r, (y,s)) in the former case was out-competed by v (j, @) so their growth is bounded by (()) in the first two stage and negative in the end due the feature shape (d μr: (z )o-1 Combination @ in the latter case will be smaller as well due to lemma F.6c.
the end as shown lemma F.6c.
The time bounds follow from lemma F.8. This proves the induction in phase II for any feature Σ*
F4 Phase III: Convergence
We present the induction hypothesis in this phase.
Induction F.4 (Phase III, final). Let 1p = (j.r E E". [T3, Ti}, the following holds:
(a) At L E TT,,, we have Fw (61, holds with ō₁ = d 62 Ō θ({
Ti,1,, (b) Ar t E [Tı,, Ti), we have Fo (82, 52) holds with 51 = d 52 = Ō((=))-T)
We obtain the shape of logits at convergence.
Lemma E9 (Logit shape at convergence). Assuming Induction F.4. Let Unz =E E* be the last of E., then i [T Ti}, the followings hold:
(a) For & E $; logit
(t)
2 1 - O()) conditioned on Ho-
6J
(b) For o $ 4J. logits.) (4) 5.j
< ((A/d) conditioned on Hg-
Proof. The proof very similar to that of lemma E.9, we do not repeat here.
Lemma F.10 (Gradient bounds, phase III). Assuming Inductions F. and E4 holds, and let (j, E S*, then at t E [Ti,,, Ti], the followings hold
(a) (() 2 ((A) when (
B d μr;
(b) r- j.r.v
|= < 0 dry for any v GUy;
67

---

(c) For & E jyEV'jj.. we have
r+ Jin
- E
iVVl) fo
E
B 1000-20 -w]
(0),
E [ - 2u,-Bi
BB + 2pp
Proof. The proof of (a) is simple >> v which is the upper bound of all other gradient terms. (b) is because of lemma F.9b and that Induction F.4b holds which guaranteed suitable feature shape. (c)
simply combination of (b) and the lemma F.9a, and Induction F.4b.
F.4.1 Proof of Induction in Phase III and theorem F.1
Proof of Induction F:4 and theorem F.I. The proof is similar to the proof of Induction E,4, thus we leave out some details here.
Induction F.4a: We know from Induction F33 and lemmas F.6 and E.8 that Induction F.4a
holds at / = =T The feature shape (a ju, O ((((((((()))))) )-)) guaranteed a starting point for reusing the argument in lemma F.cc to guarantee that the feature stays at the same feature shape.
Induction F.4b: After t = Ti,1. which is t = 4g for the last eb E E We shall
prove that the feature shape Fo ne ((())=t)) holds for every v E Σr Since (a) holds for all vp E D", we have that at T 2 that the feature shape in (a) holds for all 43.3
1 (p) E Er. So by lemma F.10b and similar argument in lemma E.6c, we can prove that the feature (V)((()) S ((=T)) ) for any & E +t. which is F (d' μ.) is guaranteed by lemma F.10a and F
is guaranteed by F Fe. and lemma F.7. Fu.4 is simple any feature E lost the competition and are bounded by Ö(z). Moreover, once each Pr.u. T())), E y is learned the feature in Et. has too small gradient bounded by O 12) which will not have sufficient growth before t = Ti-
Since all feature vo E Σ* have shape (61, 52) for 51 = d u and 82 = Ö( at t = Tı, we have proven theorem F.I.
G Learning the Attention Layer: Simply Transitive Case
In this section, consider the case where the group operations form simply transitive group. According to Assumption 1, we assume that for any y1, y/ E y, there exists a unique g E g such that y1 y2- Without loss of generality, we let y = {0, 1, Ny - 1}, where my
[22[lg log d), log d].
We focus on updating only Q, while keeping W fixed. Combined with the attention structure specified in Assumption C.2, it suffices to consider the updates to the blocks Q4.a and Q4,4 only. We consider the contribution to the gradient from the position i = 5 on task T2; specifically, the relevant loss function is given by >²_ Loss? As W remains fixed in this section, we omit the superscript (t) in W and in all related notations that depend solely on W (e.g.. V,,r) for notational simplicity.
G.1 Gradient Computations
Notations for gradient expressions. We firsr introduce some notations for the gradients of the attention layer. For < f S L, given L and k IL,f-1 define
A
Σ E 5E|4
Σ rE/m]
sReLU' (Z Ca Zk), E [5]. (41)
For simplicity of notation, we will henceforth omit the dependence on Z Lhe-1 in the notation of Et when is clear from the context.
Fact G.1 (Gradients of Q). For any p, g E [5], we have
1.
-Q Loss
Vo Loss
where
t=1 G
68

---

V Qn.g Loss 4.8
E
Att 1-1
Σ
Zans,t 1.p T
KEIZ.
k'EZL,
Lemma G.1 (Gradients of Q1,3). Given s E T for the diagonal entry [Q1,3la of the block Q1,3 we have
Vo Loss
2.1
Σ
(Z
Σ
sReLU' (As.sr))
5,5
jEId
rE(m]
(Wo.j.r s.j.rr 2pred
As.jer
bo.s.r)
1 r(20)
Vai. Loss?
2.2
E AttDans 1-2ped,22
Σ (Z21(
Σ
sReLU' (As.).r)
5,5
jEid]
rE(m]
(Ws.j.r Zpred.2) As.jor
bs.s.r)
1 lozriar)
Moreover. for the off-diagonal entries QQisle,s with s + s', we have
VQ. Loss?
E Att Spredd
Σ
(Z200)
Σ
sReLU (Re.UU(Assr)
5,8'
jE(d]
rE(m]
(Wo...r Zpred.2) - As + bo,sir)
1 (==(()))
2.2 VQu. Loss,
E AttMaas, E Es
Σ
sReLU (Ae.s)
jc(d]
rc(m]
(Woj.er Zprod, As.err + bos..r) s==(())))
Proof. For = 1, due to Fact G.I, the diagonal entry [Q4,]sas with E T(t), the expected gradient contribution from QQ, Loss,
2,1 takes the form
E Att
(=) (rre,
Σ
Att E2,6,e) 11=7(
which is nonzero in expectation only when T(xg). Therefore, combined the definition of E in (41) we have:
VQL Loss
Atttaass00-ppred
(
(=)
Attnanst-1-xx22.. I
k
Attmons.0--peel
Σ Esg Σ
(ReLU'(An.jr)
jEld]
rE(m]
((W W.s.j.r: Zpred,1)
Σ
Attn f-1-k (W Zui) )
ke KETRR
E
AttDaans Atttaans.0-prr..11
Σ
Es,s Σ (ALLU'(As.jr))
jEld
rEmm]
69

---

(Wor.r D.red: 2pred.
As.g. + bossir) 1
Other quantities are computed similarly, and thus is omitted here.
Lemma G.2 (Gradients of Q4,a). Given & E (x)), for the diagonal entry (Qu,alas of the block Q4,4, we have
(Σ
Qo Loss 2.1
E Attaaans.0--as.00
(Z,
Σ
sReLU' (Asar)-
jEld
re(m)
((W...m s.j.rr Zans,u)
Asujo + bs.j.r)
10=7(00)
Qa. Vo,, Loss? 2.2
=B Att E Eg (Z2.1)
Σ
sReLU' (AsJ.r-
0,8
SEld]
rE[m]
((Wss.r Zans.1) Asder + (ossir)
I Ia=r(an)
Moreover, for the off-diagonal entries (Quale,s with s + s'. we have
-v,, Loss le,o" and
Vo Loss?
Att
Σ jEld]
Es (Z2.1)
Σ
RReLU (As...)
rE/m]
(Ws..,r Zans,o)
As.g.r + (s.s.r)
Proof.
The analysis is similar to that of Lemma G. and we omit the details here.
G.2 Some Useful Bounds for Gradients
In this subsection, we establish several useful bounds on the gradients the attention layer, leveraging the feature structure of the MIPP layer learned during stage 1. These bounds will be instrumental for the subsequent analysis. As established in Lemma E.7 and Theorem E.1, the end of stage 1, the network exhibits the following activation properties:
Sparse activations: For each j T)2), and any feature (g. v) 3y- there exists unique activated neuron r E [m]] such that, under event 9I = 9 and y/o = y, the following holds:
s.jrr B 0(8),
(Ti)
IT (g) V (Ti) (v) O())
A (Tt)
I 0(8-11) for all r' *
S.j.r
Cancellation of incorrect features: let [m]] be the activated neuron associated with (9.v) E Bj. Then for any g' + g E 9 and any y E ), we have:
y (g) + v(T) (1) < 0(0),
(T)) (g') + v(Th) (w) < O(6).
the analysis of FFN layer, we have pair (ō1, 62); in the expressions above, some terms should use di and others 82. For brevity, we write d = max (81, 82} and bound all such terms by d. Since both 6, and 82 are small, this simplification does not affect our analysis.
Notations for activated neurons. Since in the simply transitive case, the feature are disjoint across indices-i.e.. 8, 0 8,j- = for all j """ j'w denote the activated neuron corresponding to
70

---

(g. y) E 33 by rg.y. Moreover, let
21 Uj where 21, - fr (3g E B3,r Tg
In other words, the set of all activated neurons across all feature sets ß, for j E +()). Given z4,6-1, letting ((Z = Ugr 1ーー be the collection of all the chosen group elements in the predicate clauses. Similarly y Uye 1 Then define X, (Z frow (g.3) 8, A (g G(Z Vy E }. For simplicity, we omit the dependence on Z - in the notation of 2 when it is clear from the context. Equipped with these notations, we can summarize the above properties in the following lemmas.
Lemma G.3 (Properties of target feature magnitude). Given (g. 8, with j E +()), then, the following properties hold.
(Vi.r. (g) Vy B O(5), VVrro (g) VyL.a (v) S (()),
(42)
Va (g) + Visro (v')) S 0()), Visro (3/) <0 for all y + 4y,
(43)
Visrg (g') + Var (v) & 0(8), Virg,r (g') <0 for all g' + g.
(44)
& [V), 0()) for all r & 2z.
(45)
Lemma G.4 (Properties of irrelevant magnitude). 1f (p, $ (2) GU x >,, or j$ +()), then for any r E [m]. we have
WW...rpp e S Ö(ou).
(46)
The above lemmas give us some direct computations of the inner products between the weight matrices and input embedding vectors.
Lemma G.5. Let j E +))) and E [2]. Then any r E [m], the following holds:
(W s.j..r prred. = Vg.r(e) + Ö(ao), (W s.j.r, tan) (ye-() = ((o)).
(47)
(48)
Moreover, for j & r(V) and any k I and r E [m], we have
(W s,j.r Zx) Ö(o)).
(49)
Proof. By direct computations and the definition of Zpred, we have (Ws.j.,;eg)) + (Ws w.j.n, eze-1
(Ws.jrr
s.j.r Zpred, e)
(Ws.j.,,1,
Plug in Lemma G.4 and the definition of (Jt(, we obtain (47). The proof of (48) and (49) is similar, and we omit the details here.
Furthermore, we can establish some characterizations of the As..,. quantities, which are crucial for the following analysis.
Lemma G.6 (Characterizations of Lambda). Given Z with E E [2]. with an attention structure (AttD 1- 1aassst---kkkeet----- keT..--,,
(a) for j E for activated neuron r E A, we have
As.j. (Z2-811)
Att (ge) +
Att ams 11-aam,,e -V) (
F=1
F=1
(b) for E T(V).. any non-activated neuron r &f 21, we have
AG.j. (22.011) < 0(6).
(c) for $ T()). for any r e [m], we have
< <0((0)) (Z..
71

---

Proof. Recall the definition of As.j.r As.j.r(Z2..-) (Z in (23), we have
As.j. (Z2. = Σ
f-- (W5. Zu) + bs.g.r-
KEIZ.
Thus, the first part follows directly from (47) and (48); similarly, the second part holds by plugging (45) into (47) and (48); the last part is direct consequence of (49) and the fact bi.j,r do log d.
direct consequence of the above lemma is the following finer characterization of the activated neurons.
Lemma G.7. Given j E T(Y), E 21 Iy we have sReLU (As.ar) = O.
Proof. For j E 7()),, for E 2,, by Lemma G.6, we have
2
As.j.r (,...((22.1)) (2
Σ At (gr( +
Σ
ans -1--nns -1r + Ō(r)).
8=1
By (45), for for r E 21 Eg we have Vg.r( Vg, (Ye -10--)) - Hence
As..r - + ((0o) < -
which implies sReLU' (As...)) =
Now we are ready to further derive the gradients of the attention layer starting from Lemmas G.1 and G.2 and the properties established above.
Lemma G.8 (Refined expression for the gradient of Q4.3). Given s E (X)), [Q1,3 of the block Q4,3, letting j1 r(gi(y)) and j2 (g9(111)), we have
(, for the diagonal entry
VQiaLossh
E
3,8
(1 (ogiiit
Σ
sReLU (As,i.r) (V)r-(o)) As.r.. + Ö(6v)) + Ō((9)
re
logit sj
sReLU' (Asaar) (Vir(er) Aosr + Ō(6o)) + Ō(84)
H+her()))
rEů,
Σ jKr(3)
logit, J. Ō(08) 1l 1r(eo)=s
VQ Loss? 2,2
= AttMans.1-pred
(1 logitttt
Σ (RcLU'(As,zrr) (Virr(z)) As.jz.r + Ō(ou)) + Ō(69)
reii,
logits
Σ
SReLU' (Var(ee) As,,r + Ō(oo)) =0(8))
r
Σ logit,, Ō(08) Il Pr(ensos ((T(D)
Moreover, for the off-diagonal entries [Qi,alas with s # s', we have
- VQi. Loss? 2.1
E
72

---

logiigits)))) ( Σ sReLU (As.s.e))) (Vo,r(02) - As.s.. + Ö(0v)) = Ō(89)
logit
Σ
aReLU"(Asar)) (Vir((0) As..r + Ō(o0)) + Ō(88)
rEa,
Σ jKT({)
logits, , Ō(o)) , 1. D((0))=
2.2 Va Loss 5,8
Att
(1 (1-logits,j))
Σ
sRcLU (As,zar) (Vrer(o)) As.ja.r + Ö(ou)) = Ō(69)
reflyz
s.
Σ
AReLLU(Ass.)) (Ver(on) An.,r + 0(00)) + Ō(6))
ifinEr(")
rE
Σ }Kr(D)
logits, , Ō(08) 1l [
Proof. For 1,
for the diagonal entry [Q4,]e, with E T(2))
Vou, Loss
2.1
= E AttDaass,0 Opeed,
Σ
EEs
Es,g (Z20)
Σ
sReLU' (As..rr)
8,S
jEld
rE|m]
( (Wa..er Zprod,1) As.j.r + bs.j.r)
11 )omrtoo)
Att 0-pped
-l-loggts..j. s.jo
CReLU'(As,.io) (Vy..(9r) As. =O(o))
rE(m]
Σ
logit
(ReLL (Vr((9n) As. + Ö(on))
rE [m]
Σ
logit,
5J
Σ
<ReLUU(As.sr)
(A ((W..)r) As.j.r + Ō(00)) (((00==
(E)
rE[m]
E Attniass 0pppee,
(1 logits,h)
Σ
(AeLLU(As.ir))) (Vn.r(or) Ao..r + Ö(oo)) + =((0))
rE
logit
Σ
RRLLU((ss..) (Vy,((r)= - AS. + =0(oo)) = =0(89)
b.j
E+()))(M)
rea,
logit, Ō(o)))
(2)
where the last equality follows from Lemma G.6 and Lemma G.7.
73

---

for the non non-diagonal entry [Q1,les.r with * s' E T(X), the analysis is similar unless the condition that the gradient is non-zero only when = ((0)) and s' T(11). Thus, we have
Vu. Loss
2.1
E Atttonss.--pred
( Ez, jE(d)
Σ re(m]
(As,U"(Assrr))
(W s. 2)red
As.j.r + bs.sr) (
E A
(1 (1- logits.))
E
RReLU((Assv)))) (Vnrr(92) - AoJi.r + ((0o)) + =(0)))
r
logit,,
Σ
(AL,U((AS,.))) (V))( - As.ir + = = =0(84)
b,
(+()))())
rea
+ Σ logit, s. O(o8)) 1 (1(00)=
jET())
For 8 2=2,
for the diagonal entry [Q4,]s, with E T(2),
V Qu.a Loss?
2.2
E Atttlans -pprd,, Σ Es (Z
Σ
sReLU' SReLU(AAGrr)
8.S
jE(d]
rE[m]
(Wor.e Zprod.2) As.jer + bss.er) No=r(er)
E Attm p-ed
(11 (1-1ogits, llologits.. 3.12)
sReLU' (As.sarr) (Vi,i As.jarr + Ō(oo))
re(m)
logits.
SReLU (Var(92)= -A..rr = Ō(ou))
H+1EE)())
rE(m]
Σ logits
Σ
RReLU'(As,j))
(Ws.jr,, (A As...r + Ō(Go)) ((0)==
jEr())
rE[m]
Attmfass, 1=ppred,-
(1 logits,.)
(1-logits..)
Σ (AeLUU(As.nr))) (Viar(ee) As.ja.r + Ō(vo)) +((8))
rE
Σ logits,j
Σ
ReeLU((ssse (V2r( - Asar + =((o)) + =0(8))
(+1)
eei,
Σ jeT("))
logit, (igit,,((a8))
where the last equality follows from Lemma G.6 and Lemma G.7.
74

---

for the non non-diagonal entry [Qi,sls.r with * s' E T(X), the analysis is similar unless the condition that the gradient is non-zero only when = T(zo) and s' T(11). Thus, we have
(Σ E6,( SEld
Qu. Loss?
E Attmans.--prod.
Σ re(m]
(As.U"(As..))
(Wo.jor 2pred,1) As.j.r + bs.sr) (
E
1-tpeed
(1-1og') logits,t))
Σ
AResu(AAssav)) (Viar(91) Ao.aar + Ō(0o)) + =((8))
logit,,
Σ
(ReLU((As,s))) (v))( - AsH + =((o)) + =0(04)
b.
j+
rea
+ Σ logit, s. O(o)))
1 (
(+T())
Therefore, complete the proof.
Lemma G.9 (Refined expression for the gradient of Q4,4). Given s E Tx)), for the diagonal entry [Q1,1s of the block Q1, letting j1 (gg (y6)) and j2 (g1(111)), we have
VQ., Loss
E
AttDans,0-tass00
(1 logit
).
Σ (As.ii) (Vin=(10) As.j.. + Ö(oo)) + =((8))
reãj
logits,) ور
RReLU (Ver(vo) - Ao.,r + Ō(oo)) + +((8))
(+)ir(())
r
Σ jKr(3)
logit u. (ogit,,0(o0)) Il Dr(oo)=
V Vo Loss? 2.2
AttMans,1-tns...
(1- logit )
Σ (AC,U((As,sar)) (Vnr-(r)- - As.ja.r + =((oo) + Ō(89)
reai,²
Σ logits.
Σ
(As.5r(A,..r) (Ver(ur)= As..rr 1 Ō(Go) + Ö(8))
(*12ET())
rea,
logit,,, (((8)) 1 Dr(en)=o
5(T(D))
Moreover, for the off-diagonal entries (Q4ala.o with + s', have [-Vq,, Loss? and
Vo.,,oss; 5,5'
AttDans, >ans00
75

---

Σ
sReLU (ReLU'(As,srr)) (Virr(00)) As.ja.r = ((0v)) = Ō(89)
rea,2
logit Bن
Σ
AReLU"(As.aa (Vir(uo) As..r + Ō(00)) + Ō(88)
ifjsEr(D)
r
Σ (KT(D)
logits, , Ō(σ)) ,
1 D
The proof is similar to Lemma G.8 and we omit here. Notations for gradient decompositions. We shall define some useful notations further simplify the expressions of the gradient.
for = 1,
for [Q4,3las with E T(2), we have |- V Q1,, Loss 2:11 Is,- = NN,3,111 + Na,3,115 + +Na,3.1.. where
N,,,,,11
Att (1 - logits,)))
(50)
Σ sReLU'(As.I))) (Virr(01) Assir + +((oo) + Ō((0)) 7000)=
r
Ne,3,111 = Attmaans.- 1
logits.
(51)
i
Σ
(AeLU'(AGar)) (Var(9)) As =(0o)) Ō((8)) 1=20)=
r
N, = +E Att sprdd
Σ logits, 1 0(08)
(52)
)ØT(y)
- for [Q4,4ssss with E T(X), we have |- Qu, Loss 2.1 1e.0 = No.. + No + N. where
Na,4,1.1
Atttlass,0-aans (1 logits.5))
(53)
Σ
SReLU (Vo)) As.i.r +((0)) + +0(80) 7(00)=
r
No.4,15 E ant
logit
(54)
Σ
(Ae,UU(As..))) (Var(un) As. 0(oo)) Ō(87) 1-(2))=
r
N4 =E Attmaanss +anas
Σ
logit;, , 0(08)
(55)
for = 2.
76

---

for QQ1,3lass with E T(2), we have 1-VQ,,, Loss le,. = Ne,,2 + Ne, + NN where
Na,3,221
AttiMans, 2 (1 logits.t2)
(56)
Σ (ReLU'(As.a)) (V))r((2)) A +0(00)) + Ō((8)) ((01)==
Na,3,2,
E Attmans, 1-pped222
Σ logits,
(57)
(+
Σ
(As,,r(sssrr)) (Var((e) As,.r =bs:=0(00)) ((00)) + =0(8)) (-(2))=
r
N,,3,2,8 = +E AttDans,1 xpred
Σ logit,, ((
(58)
(KT(1)
for [Qi,le,s with E T(2), we have |- V Qu, Q1, Loss 2.21 le,e = Ne, + NeN + N. where
Na,4,2.1
E AttiMans, L-nns (1 logits,)))
(59)
Σ
(acLU((AG,aoo))) (Vizr(U1) AbJair + =((00)) + Ō((o)) 1 r((0)=
r
Ns.4,2,1 = EE Attmaans, anns
logits.
(60)
()))Er())
Σ r
(Assrr) (V).-(r1)) As..r =((o)) + =Ō(07) 1-
N, E Attm
Σ ((T(1)
(
(61)
Probabilistic Events. We conclude this subsection by introducing several probabilistic events that will used to simplify the characterization of activated neurons the subsequent analysis.
E,
91 + 92
(62)
Eg
a ge, (ye;) + gez (Be;), for any (l1, M) # (l2,6), where la E [2], (0.1)]
(63)
It is easy to see that E and & hold with high probability 1 - 0 (10gg)
G.3 Stage 2.1: Initial Growth of Gap
At the beginning of stage since Q has been trained for long, have the attention score is still close to the uniform structure. Therefore, for e. = 1, we have
1 As.j.r (22.-1) x Var(or) + + (..((o)) + 3/2 VVr((v)) + Ō(oo).
1
77

---

If 7 1 E E1,
for j J1, for r fn. only ror vo is activated since As. B and
As.ji 02 Aai) = - < -
for j = R - ((02( only ra vo is activated since Ao
B and
As.ii "1'5"' o1)
~ - -
for other j T()), we have As. S 3B for all E Ny thus no activation.
Moreover, for / 2, we have
(2.. ~ 4 If 2 E E2,
(0i + Vir((0)) + Vor((0)) + Vir((r) + O(oo)-
j E (T(ge(we) eae) only the corresponding Tee is activated in the smoothed regime since |As.j.or = O(0) and As.j.r = 2 < -2 for r E 21 \ frgeor };
for other j +("), have As..r 2/2 for all E My. thus activation.
Here, activation means that the corresponding (As.gr(As,.rr) is non-zero, which is crucial for the gradient computation. Based on the above observations, we can see that the gradient from & is relatively small since A is only activated in the smoothed regime. Thus, imitially, the learning process is dominated by VoLoss? Moreover, if we take a closer look at the gradient from l' = 1, we have
N,33, ~ -E AttD ans, 0 pred.
logits s.
CAReLU(AAs.APp (Vil (01) - A5 =Ō(on)) + +0(89) (((20)= 2 0,
Ne ~ -E Attmans logits.y(
(RR.LU((As,)..a0) (Vii vo (v0) - As + = = =(())) s 0,
since ViL.roz vo (Wb) 2 ((B) while V2 (0:) -B.... Thus, [Q4,3].. will have significant positive gradient while [Q.4]a., will have negative counterpart. This will lead to the growth of the gap between [Q4,a]s,s and Q4.4)=.e-
We formally characterize this growth behavior within this substage. the beginning of each sub- stage, we establish an induction hypothesis that we expect to hold throughout. Subsequently, we analyze the dynamics under this hypothesis within the substage, aiming prove its validity by the end of sub-stage. Due to the symmetry of [Qu,a]s,s and [Q,,a]e,e across E T(X), we may, without loss of generality, focus particular T22).
Induction G.I. Given T()), let T2,1,s denote the first time that (Q reaches 5 (1 For all iterations < T2,1, we have the following holds
log
(a) a, a.a monotonically increases;
(b) [ ell S $1Q(11 and (199 10 (((Q((aaa))
(c) for p E (3, 4), for E T(()) # [[Q S 00 1Q
G.3.1 Attention and Logit Preliminaries
We first introduce several properties of the attention scores and logits if Induction G. holds. Lemma G.10. If Induction G.I holds for all iterations < t, given imput Z2,6-1 then we have
1. for & = 1,
78

---

(a) Attn p-redd
0
3
Tog
(b) Attn ans pred 2 Attn (t)
ans,0 anns0 S Attn ans,
(c) Attn 0- Attn
(8)
S 0 (oga for + k' I(2.0)
a
2. for =2
(a) Attn ans, pred, | 0 (1otd))
(b) Attn
1- for k (pred, 2); (c) for k E {ans, 0), (pred, 1)}, k' E {(ass 1), (pred, 2)}.
(t) ans I Attn (t)
Attn (t)
Attn (1) 1-s
S :0 1 log d
ans,1->k
(d) |Atnn (t) ans, 1-ppedd Attn 1-tans. (
Proof. For l 1, given - according Assumption C.2, have dQ (ag
Attn 0-pee-
(0) e (Q.alte))) + 1 Hom + +e (1)
1 1 10 + 19 ( -1 (aggtt)))))
Thus, by Induction G.I.
0 1
(Q33)())))) - -10 (o03)))))) (Q. (00))) - (00)) S S0,
(0)
log d
which implies that 0 < Attn (
ans spred.1 S 1/4 0 ( (b) is straightforward since
1.3 (ao)) 1 S QQ Q (c
(c) is direct consequence of (a) and (b).
For = 2, given Z (a)- (c) are very similar to the above analysis, and then for (d), we have
Attn (4)
Attn (t)
ams.1-aann
O +e (llrttp) e 1 + dQlleeerr
2 o
(9) (19I - [Q
(1) ((...)(()))) (8) 1
where (i) is due to the fact that 1e-- ed S O(k - v|) when y small, and (ii) is due to
Induction G.1 (b).
Lemma G.11. If Induction G.I holds for all iterations < t, given imput 22,--1 then have 1. for 1, if E E1, then (a) for j j1, A1 S.h..r < -eforr E 2, Vrr
(b) for j = J _ T (92(1))),
1l9 S.j.i < -eforr Zn \ frewwot:
(c) for other j E T(V), r is not activated for all r E U, iee A!l & -2.
2. 8= 2, if E2, then
(a) for j E (r(ge(ee)) (eeee,rr)))) only the corresponding Tgr-u may be activated, with (1)
A.j.ro 1 O(1), while all other A & -2for E 2, frorgell b:
(b) for other E T T(Y), r is not activated for all E M, i.e. (1) E -
79

---

Proof. We only prove (a) for = since the other cases are straightforward. = (1 by
Lemma G.6 we have
Asdrocve Attn (0)
ans, predd Vererve se (01) + Attn ans, 1-ped, Vir ve! (92)
(t)
Attn
ans, tans, Virve ver (Vo) Attn (0)
(1)
ans, l-ns, Vir xet (V1) O(oo)-
Notice that since 72.t- E &2, have two V terms are positive and two negative with magnitude B = O()). Therefore, | S ( Toga
B = O(1).
Lemma G.12. If Induction G.I holds for all iterations < t, given imput 7 E &1, then have
(0) (a) A.jii 3/4 ----010 O(1);
(b) (6)) S A (1)
A Torivo
S O(1);
Proof. By Lemma G.6 we have
As.jv.
(4) Attn poep--rddddd Varoo (91) + Attn pred Vi (g2) + Attt Viroa (20)= Ō
By Lemma G. 10 and the cancellation in (44), we have
Attn (t)
ans 0 -pred,2 ViLTT vo (92) + Attn (t)
ans -ans, Visrn vo (36)
0 log 0 O log B.
Putting it back, and using the fact that Attn (!) I-pped, Visiron (91) 2 (B -(6)). we have
1 (B (()) (} 0 (og ) 0(6) -0 1
A5
loo d B Ō(oo)
Iv O
Moving on to (b), we have
A As.ji.rozv Attn (1)
ant (V Viroo no (g1) Vor 80 (91)) = Ō(oo)
+ Attn (t)
ans.0 2 Vo. (92) V ug v0 (92) + Attn (1)
a (V/ 11 n1. vo (3b) VA vo (v6))
(t) Attn appred (2B (6)) Attn anped-- (2B 0())) + Attn (0)
a O())
(2B (2B- 0(6)) o d 0(8) 0(1).
1
Similarly, we have
As.h Var
As.
Attn (8)
ans, ppred. ((B)) Attn (t)
ans ppeed, (2B 0())) Attn (0)
ant O(8)
Attn ppedd O()) AAttn (t)
a0 0(6) (8)).
Lemma G.13. If Induction G.1 holds for all iterations < I, given imput 2 then we have
1. for 1, Zf 2.6-1 &r, logit! ( (211 {h.Hi. logit! (() polyd'
6.11 logit
(t)
5.1
2. for = 2, - Eg. logit (1) 5.3 O1)) for all j.
80

---

Proof.
For & = 1, by Lemma G.11 11 and Lemma G.6, we have
(226-11-
Σ
SReLU(A (A (.1.))) S.i1
(1)
+ + 1)
(8)
+ +0( (polyyloge)
5.31
rE/m]
Fl)) (22.6-1)
Σ
sReLU(A (1d
A
+ +emm/q-1 - 1),
3. Vo
rE(m]
(221-11-
Σ
sReLU (0)
(A s 00 polylogd for urjii. 7()),
re/m]
F3 (2 S m ( for j $ T()).
Putting it together, we obtain
logit
6,11
1+c e 53
p(u plt)
+ (Eitsisteeo)) 503 + Ejtron e Fb pil
1 C "off-ro
(Ollog d) o Opaainga + O(d) em(m))
Fb.j
Thus, by Lemma G.12, and the fact that B Cn log for some sufficiently large constant Cn > we have logit (t)
018) +0 (0n//000000000))aaa 22(1). Similarly, we have
6.1
logit (4)
5.31
1te
b.3f b.3f.r00 v0
+A
6.6.1100010 + (log d) e 0l polyinga) + O(d) eO(mo3)
5.54
e010 +(1) (Ce//-3-o0 -2(1).
From the above analysis, it is easy to see that - logit (t)
- logit
S <0( (Cn/3-1)1ozd
6.5 5.8l
O polyd)-
For l = 2, by Lemma G.11 and Lemma G.6, have
(2 = Σ
SReLU(A (.s..) [om/a, 0(1) e(m//- 1)] for E (r(osuue))
rE/m]
r:0 (22.0-1)
Σ
AeeLU(A (Ag..r (
polylogd) for +()) (r(oww)))))seeeeee..
rE(m]
(22.-1) S m Ō for j $ r())).
Therefore, for any j. we have logit, (t) since rlg S 0(1) for all
In the following, we illustrate the activations on the non-high probability event. Lemma G.14. If Induction G.1 holds for all iterations < t, given input 2 then we have
1. for = 1, 2 & E,, then (a) for j j1. A frog wog. Sa B 1+
(t)
(b) for j * j1 E T()). assuming j T(gi (v)), then A Adray B 0(1 and
A < - r E a, \ (Tave)-
2. 8 = 2, 2- & E, then
(a) 9f = 92 ^ y0 + yr,
i. for = jz. 5.ja.
= 1 B 0(1) and A < foro E A \ (roznih:
81

---

i. for jorj=m =()2(V))), (0) 3.52--o2 so
1 + O(1) and SaJ < -e E 2g
(0))
to iil. for other j E T()), assuming ((2(1))), 1. s (S() and ! < -2
for E 21 frgzvt-
(b) if # 92 A V/ = y1,
i. for j = j2. A N.22 = 2 1 ((1) and 2.22.r < < E Aip \ (raan:):
(t)
i. for j r(1 (v1)). A (t)
A B O(1) and Aggr < - M
tr
10 ii. for other j E TJV), assuming j = (((())), S.j.pm
s 0(1) and A j.j.r < -0
for E 2] {raw)}.
(c) if 92 A yo y1,
i. for j2. A = fr J. A (8) 5.22.09- = O
i. for j # j2 E T))), |M 5 for all E ã,
(d) + 02 ^ yo # A (91 (vo) -92(v) V 92(30) g1 (wi)). |6 0(1) for all
EE Z
With the characterization activated neurons, we can derive the following logits for the non-high probability event.
Lemma G.15. Induction G.I holds for all iterations < I, given imput Z2.6-1 then we have
1. for = if $ E,, then logit (1) 5:J1 polyd
0
2. 2, 2f $ E2, then
(a) if 91 92 yo Vn logits. 2((1) for j {fs, ((2(0)))) 1 logits.ja
(t)
(1)
logit 5.7(9:(v0)) polyd-
(1)
2(1) for E L2z (()(1))) 1 logit (b) if 91 # 92 yo Vi. logit (t)
(6)
logit: (, (v1)) polyd'
(8)
(c) 91 = 92 A yo = v1 logit
(t)
0 (plyy)
6.32
(d) 91 + 92 vo + vn (91(v0) g2(v1) V (2(v0) 91 (vn) logit!") (()
O(2) for
all j.
G.3.2 Gradient Lemma
Lemma G.16. If Induction G.I holds for all iterations < t, given E T(X),, for (Q1t) we have
log d
Loss?
2.
(t)
d
8,8
Proof. By gradient decompositions, we have
2
Σ
2. Loss
= - Σ
Q43
5
3.3.E,x
[E[2
RE (i,
By Lemma G.13 and Lemma 1.1 is straightforwarr to see that (V(()) and thus we can focus on other terms. By Lemma G.11 and Lemma G.14, we have
IN10
0 ((poyaa)
N ( 3,1
E Attn (t)
ans predd (1- logit (1)
5.
82

---

vo (01) Abdi vo
=O(oo)) = =((8)) (=(20)=0))))
Attn (t)
spred, (1 logit
5.51
(1)
(Viirr (01) A
+ =0(0))) Ō((8) (r(zw)=a)NEf
1 2 2(11 ( ((B) + + 0 1 ( O (a) 0
polyd
log
0 log d
where (a) follows from Lemma G.12, Lemma G.13, Lemma G.14 and Lemma G.15, and the fact that (xx)) = 8 holds with probability 2 A s,3,2, can be upper bounded similarly.
Moving to A(()) 8,3,1,5 noticing that Vis ma (01) B + O(6) on ε we have
NO $3,1.
Attn ans, 0-sprd,11
logit 5.51
(10) (VV No (g1)
=((o)) + =((8)) (r(20)=))n
Attn (#)
logit (.((((0))
ams.0 ppedd
97 EJ
(VVio (0)) 6,r (w))
+ +0(00) + Ō(88)) (+(2))=)))8)
0 1 2) 2(1) ((())-((B)) ,( - O polyd
O O
log
log d
-
d
For A/t1)
8, only need control the negative gradient, since the positive part be easily upper bounded by O logd)
\( 8,3,2,1 2 - .
0(B) - Attn (t) ans, 1-pred,22
Σ
logit
(,r())()))
((v (Vrioa(i)))roo (92) As.j.r = =((o)) = =0(07) (1
[ 0 = )
e(B) o( .
0(1) O(B) O log d
2 (
where (a) is due to the fact that l (, 2(1) and logit (,7(0((0)
O(3) for other
y A yr, No. Putting everything together, we complete the proof.
83

---

Lemma G.17 (Negative gradient), If Induction G.1 holds for all iterations < t, given E 7 T(2), we have
2
2
Loss, 2.6
2 -0
V
Loss 2
logd
1
t=1
Proof. By gradient decompositions, we have
Σ
Loss
2.
Σ E
AOO
5.5
EE(2)
KE (7,61,50)
Similarly as Lemma G.16, (M(a)
(C)
s,4.2.8 = O( (olya) ) and thus we can focus on other terms.
s,4,1,88
By Lemma G.11, Lemma G.12, and Lemma G.14, we have
Ar(:) $,,11 + +A(())
8,111,5 = E Attn ans,0- (1 - logit (f)
6.11 )-
(VV (26) - (1)
- =((oo) + = (r(00)=
(64)
+ IE
Attn (6)
-an,0- (1 logit 5.51
(VVi (vo) (1
+ Ō(6o) + +((88)) (r(00)=
(65)
E
Attn (t)
Attmans,0-4asss logit (8)
5.5
(V (2o) (1) (1)
+ Ō((o)) + Ō(8%) 1 (1(20)=))
3.8 "22
(66)
E
Attn (8)
logit 6,r((())))
ans,0
##
((Vroowo))ie (vo) 6.7(000)))
(8)
+ +ō(0o)) + =((8))
(1(m0)=))(E]
(67)
First, considering the event (T(z0) 5) Es, we have
(64) + +(66
(a- (1-logit!)) (Vyrror (2o) - A10
Attn (1) E
5.51 =(()))+((0))
logit!! (t)
5.5 (Vsiroo (vo) (0)
=0(6))= (
(0)
Attn 0 (1- (t) (Vii (vo) - A.jiiroo + Ō(oo) + Ō(89)
logit
(1
logit 5.
polyd ) (Vir.ra (20) - (1)
+Ō(oo) +0(o))+018))) =5(8)) ((((m))=))8
3.31 uug
84

---

Attn 00t (a- (1-1ogit)))) (1)
EE
logit
6.j1
VVi (20)
Visiroz (v0) A
+ + ( Ō(03)) (1rce)==enes S
(0)
(b) E
Attn (0)
(.))) (a
logit
001) 001) + 0(08)) ( ((((0))=))e))
ans.0
Artt) 3.3.1.2
log d
where (a) follows from Lemma G.13; (b) follows from Lemma G.12 and Lemma G.3.
Notice that (65) 2 since Va (go) (1)
2 S(B), thus we just need to consider the
possible negative gradient from (67). By Lemma G.15, we have Σ Egggreg logit (t) 5,r()(u))) O polyd
on (T())) = }} ngf and hence (67) < A ,,3,1,5"
Moving to the gradient from e = 2, it is straightforward to see that M(l) 5,4,2,; is non-negative, and thus we can focus on the possible negative gradient from N (t) 8,4,2,18
A1 8,4,2, = Attn (9)
Σ
logit (t)
ans, ans -anss
Σ (A (V2r( A =0(o)) =0(8%) ) (
r
Attn E
ans, 1ans
Σ
logit! (1)
5.j"
()nEr())
Σ (ReLL'(A (V)r(()) - -.r =0(0o)) =0(8)) ((((e))=))883 rett
((t))
2 ) o O(B) ( ) O( 0
o log d
2,3,1,1
log
Putting everything together, and combining with the fact that Σ
Loss 2.
((N(() 8,3,1.)) from Lemma G.16, we complete the proof.
Lemma G.18 (Growth of gap). If Induction G. I holds for all iterations < I, given 8 E T(2)), we have
2
2
Loss?
2
+
Σ
Loss 2.1
2
S
(
Loss 2.6 8,8
1
t=1
Proof. By gradient decompositions in (50)-(61), have
Loss
2.
+
V
Loss?
3.
NO
8,3,Exx
-N(t)
S,4,f,K"
E=1
5.S
5.5
(E(2] KE (1.155))
Due to Lemma G.13 and Lemma G.15, (V((e)
e.p. 0 polyd) for p E (3,4 4) and f E [2], we can
focus the gradient difference between Q By Lemma 1. we have
(1) and Q contributed by other terms.
Attn 200.0 tans.
Attn
Attn ans.1 +ans,1
Attn (0)
0-pped
1-s.prd
85

---

Hence, for E [2]. we have
M( 43,0, a,,e, -0(0) (10))00)))1 O11) 2 (11)
-N
-(6//))
N (t)
3,3,1,
1,4,
-- Attn (t) - logit (1)
5.1i
((V (VVS (01) - 11
+ +0(oo)) + Ō((%)) (r(o))=o)na,
+ E Attn (1)
ans,0- logit 5,31
(Vo5 (20) - bil Po2.v0
-1
+ =((oo)) + =((89) ) (((8)))((((TO)))) [(((a))=
Attn ans,0 --pedd- E
logit 5,r(
((V (orr
(4)
=0(oo)) = Ō((0)) (7(m0)=))eE
E
Attn ans S-
l (.r()())))
VFuDEy
((Vrenn())rr (vo) - A (r) + =((0)) = =((09) (r(a))=s)nE
2 0 (
(a) 1
((B) 0(B)
o polyd
1 .((B) O
1
2 1 log d
log d
where (a) follows from Lemma G.3 that 20 (91) - 5.'i vo
I ~ii(B). and V oo (yo)
(1)
2 Q(B).
5.8g Fo2 va
AC(()
N N/11)
6,3,2,5
8,4,2,8
Attn (t) E
logit
ans, ans, 2
5,
)(hEr(())
Σ
(ReLU'(A (.rr) (V2.+(92) A + Ö(o)) + Ō(89) 1 (+(e))=
rea
+ Attn 1ns,1-tans.
Σ
logit!") (t)
6.j"
Σ
sReLU ( (V)r(1)) - A + +0(0))) + Ō(69)) 1 ((())==)e
rea
E
Attn ans 2pred,2
Σ
logit!") (t)
86

---

Σ RReLU' (.jr)) (Vir((02) A6 = Ō(00)) + =(88)) 1 (((2))==tNE
TEa
+ Attn lans,
Σ
logit
17
Σ sReLU'(A (.r. (Vir(v1)) Ab..r + =((oo)) + Ō(8%) (r(e))=))NE)
r
) ( 0(1) O(B) O log (2))
( 0(B) O(B)
where (a) follows from Lemma G.11 and Lemma G.13, which together imply that on the event &2. only a constant number of neurons activated and logit (?) S 02)) for all j * f2: and from Lemma G.15, which implies that on the complement event Eg. occurring with probability at most
Tog there exists at most one j # j2 such that logit (t) other j E +()).
5.3 2 (2(1) while logit s O polyd for
Putting all together, we finish the proof.
Lemma G.19. If Induction G.1 holds for all iterations < t, given E T(X), for [Q..., 2 U (( we have
2
2
Σ
Loss
2
Σ
Loss 2.
t=1
1=1
Proof. Notice that by Lemma G.17, when 1Q 2 D2 (o have 1.00 -0( log Hence
Attn ans, -ans,, 2 52 log d
Attn ans,1 ppred, + Attn ans, I
Attn ans,
which implies !0 3.f2.ro2 (Z 2 2 (odd) B 0(8) o already lies in the linear regime for
z E Eg. Then we have
Ar(e) 8,12, 2
Attn 1
(1 logitt (t)
(Vissoom - bida. + =Ō(on)) + =Ō(89)
1 ((((a))=)
2 (0(1)-((B) e ( 2 logd
Moreover, from Lemma G.17, the magnitude of negative gradient from other A terms can be upper
bounded by O ((Ar!)) 1.), from Lemma G.16, complete the proof.
Therefore, combining with the fact that Σ
V Q
Loss 2.
logd
Lemma G.20. If Induction G.I holds for all iterations < I, given + E T (3,4), we have
Σ
(e)
Loss
2.
0
log d 1
o(
Σ
res
Loss 2.
8.6'
8,8
Proof. The proof follows directly by combining the expressions from Lemma G.8, Lemma G.9, and Lemma G.16, along with the fact that the event (7(20) s, ,.) s'} occurs with probability
O (1).
87

---

G.3.3 At the End of Stage 2.1
Putting gradient lemmas together, we can directly prove that Induction G.1 holds for all iterations until the end of stage 2.1, where can conclude the following: Lemma G.21 (End of stage 2.1). Given & E T(X), Induction G.1 holds for all iterations T2.1. = then at the end of stage 2.1, we have
Ing a
(a)
( (og (llgdd).orr f.rrpporppprr E {3,4);
(b) [1.33e. (1) 2
1t logal
(t) p. S 50( ( for s' E T((X) # sfr (3,4]; otherwise, QQ O.
(c)
G.4 Stage 2.2: Continual Growth of Diagonal Entries
In Stage 2.2, the diagonal entries QQ,3le,s and [Q1,lee. continue to grow until they reach certain threshold. The analysis in this stage parallels that of Stage 2.1, but our focus now shifts to the gradients contributed by & = 2, as the logit at f = is already near-optimal and thus contributes negligibly the growth of Q4,3 and Q4.4-
Induction G.2. Given E T(X), let T22,s denote the first time that (Q4,3)s,s reaches 0.0001. For all iterations T2,1, S < 12,2,s we have the following holds
(a) Q4,36e.e Q s (1) monotonically increases;
(b) 1,3 1 E (( O(1): 0(1)
| & 0 ( for sf E T(X) + s; other
(c) for (p.4) E ((4,3), (4, Q =
Throughout the following analysis, instead of &g defined in (63), we consider renewed event &g for
{ (= 2:
Eg
91 +
92 A yo + y1
(68)
G.4.1 Attention and Logit Preliminaries
The proof in this part proceeds analogously to the arguments in Appendix G.3.1, with the induction hypothesis from Induction G.2 incorporated. Hence, we omit the details here.
Lemma G.22. If Induction G.2 holds for all derations E [72,1,5: (), given inpul 2 then we have
1. for 8 1,
(a) Attn ans,0 0 pred.1
3
+
2
log
d
3
C1
where > 0 is small constant;
(b) Attn (t)
ans O pred
C2 3
loga where C2 > O a small constant;
(c) Attn ans,0 0 2
l
S
Attn
(0)
ans,0 aans,
(d) Attn ans.0 0.s.0-+preed pred.11
Attn
Ω
2ns,0
log
2. for = 2,
(a) Attn (t) ans, pred,2
l Ing ( ca). where ea > O small constant;
(t) (b) Attn p-ppddd
Attn (t)
2ns.1
E
[ cs, 1
2 (ogg moreover,
Attn 1-tans0 Attn ans, I 0 0(2) where cs
(4)
0 is small con-
Stant;
(c) Attn ans, prod,1 Attn (8)
(4)
attinnns11-amg ans,1-+ans.0 < Attn (1)
att.nass
Toga
(d) Attn ans. Attn (t)
1 E
12 log c6
where > 0 is small con-
stant.
88

---

Notice that the constant depends on the threshold 0.0001 in Induction G.2. We choose the threshold 0.0001 small enough to ensure 2Cp(c1 + C2) < 1 cgCg and AcgCB >
Lemma G.23. If Induction G.2 holds for all iterations E [T2,1, 1), given impul 722.-1 then we have
1. for & = 1, if E E1, then
(a) for j - j1, (1)
Aaji.r -efor E Ai Himmooti
(b) for j - バ _ T (12(10)). A 3. < - E A1 \ gr
(c) for other j E T(Y), r is not activated for all E Ug ie, A0 < -
2. L == 2, 1f E E2, then
(a) for j j2. S.d2ir < - E A1 Afrrnn
(b) for j = 1/2 - T (92((0))) A < -eforr E 21 \ frowo::
(c) for other j E T(Y), r is not activated for all r E Ag ie, A < -0
Lemma G.24. If Induction G.2 holds for all iterations E (T2,1,s:)), 1), given input 2 then we have
1. E 1, for Z E E
(a) 10
2 3 +
l2 log ( B;
(b) (11) S (0)
(()
S 2(c + B2)
2. { 1= 2, for - E2.
(a) (0)
E [1(1), 4cs B
5.ja.Toz-11
(b) A (1) E [1(1), c6 B] and A (0)
(1)
2 2(1).
5.32.
Lemma G.25. If Induction G.2 holds for all iterations [12,,, 1), given imput 2 then we have
1. for l = 1, 2f E logit
(1)
2 2
5.11
5.11
2. for 2, 2 - E 1- logit (t)
= 22(1), logit (8)
= 0 (
5,22
5.3%
Proof.
For l = 1, we have
logit (t)
10)
10 6,3 maz O(d) -e
(1) b.ff
1 +
1+ 2 (d) 22- 2012)
pCnle
where the inequality (a) follows from Lemma G.24 and the last inequality is due to the fact that (c1 + 02) is some sufficiently small constant s.t., (1 (1-2c) -2ca) ) 1/polyd.
For l = 2, we have
logit
(()
(t) -0.
1 +
3i3hiaa2 vo + O(d) 32:
1
1 Odd) 1
di-1csCa
where the inequality (a) follows from Lemma G.24. Similarly, we have
1
logit
J/
16) 50..9 +O(d) - Di3h
-A(t)
1 + 333
BB =0
1 dl-coon
< 1 + 0(1 + +Od)
89

---

Lemma G.26. If Induction G.2 holds for all iterations E [T2,1,5 t), given input z then we have
1. for 1, if $ E1. then (a) for j = j1. A, = 1r91-wo.. A bJs. = = B 0
[ 0z) B,
(b) for j + n E T()),, assuming j ((gi())), then A (0)
o2)),
a B)BB B and As.jr < -eforr 2, (ravzl-
2. l t= 2, ( $ Eg, then
(a) if = 92 A yo + yl,
i. for j = j2. 2.j2.
(2)
L + 2(1), + 2cs) and (1)
S,222 < -0 for
r TE 2 \ (rgavn)i
i. for j 7(92(y0)), J22ioo2
(8)
2(11 and A < -eror r
á \\ ::
ii. for other E T)),, assuming = = (92(1))) A S.j.To E (1(1),, cg and 3.j.
-e E 2g (r
(b) if + 92 A 10 = y1,
i. for j j2. 2.2.
(t)
( + Bcs) and 3. & -0 for
r E Az \ (rgzvn):
i. for j = ((91(11)), (4)
-90
2 2(1) and A < - T
2 ::
ii for other j E T()), r is not activted for all r E Mj. Le, ! < -0
(c) gi 92 A yo y1,
i. for j2. 2, (1 (1)
6.J2. B+
i. for other T(?), assuming j = T (g2(1)), A S.r E [9((1), cB B) and A Ggg
r E 21 f.
Lemma G.27. If Induction G.2 holds for all iterations E [T2,1,s; t), given imput 2 then we have
1. for 1, 2f $ E1. then 1 logit (t)
0 1
5.11 polyd
2. l = 2, if 2f $ Eg then
(a) 9 92 A ye + v1.
logits.) 5.12
(1)
2(1). logit
(1)
logit (t)
5.j2
(,7())))))
polyd
(b) di + g2 A ve yı. logit
(e)
2(1). logit
logit (t)
l (.7()((1))
5.j2
5.32
polyd
(c) 9f = 92 ^ yo = yı. logit(
(4)
O
polyd
G.4.2 Gradient Lemma
Lemma G.28, If Induction G.2 holds for all iterations E [T2,1,s t), given s E for (1...
we have
Loss
2.
= 0
log d
d
6.5
90

---

Proof. The proof similar Lemma G.16, but we need shift our focus By Lemma G.23 and Lemma G.26, we have
[-
Loss
A ,,3,2, E Attn (0) a pprad,2 (1- logit (1) 5.12
(VVirro (92) (00
- =((0)) = =((89) (1(2))=))))
Attn (t)
ans, (1
logit
(4)
5.32
(( (02) (0)
+ +((0)) + =((81) (+(e))=))ne)
1 (a) θ 2(1) ((1))-(B)--)) + -01) 0
log d
θ log d
where the inequality (a) follows from Lemma G.25 and Lemma G.26.
Furthermore, for Av(e) 2,3,2,60 have
(1) N
2,3,2,
< Attn ans.
logit 5.12
( (VVi (92) - S.j.r + - =0(oo)) = Ō((0)) (r(o1)=af(e)
-
+
Attn (t) E
ans,1 1-pred2
Σ
logit! (t)
5,j"
jTtEE)())
Σ
(A (0 (Vir((9)) ( +((oo) +0(8%) (((0))=)n8;
rea
< <0 o - (B) + +0( 0(1) O(B) of
1 (oog S o
[W and ( 2,3,2,1 can be upper bounded by 0 ( (laad) Lemma G.16. Thus, we complete the proof.
Lemma G.29, If Induction G.2 holds for all iterations E [T2,,,,, given s E T(2), we have
Loss? 2. = - log d Σ[- Q
Σ
d
E=1
Proof. The proof follows similar analysis to Lemma G.19, and we thus omit the details here.
Lemma G.30. If Induction G.2 holds for all iterations E (T2,,,;;), t), given r(2), have
t
2
Σ
Σ
V 134
Σ
Loss
Tal..
1
Σ
Σ
Quyy Loss 2.8 1,0 )
Tog
V=T2,1.
1=1
91

---

Proof. Following the analogous analysis as Lemma G.18, we have
2.6
Loss 2.6
Σ
A00 6,3,Ex
-N
s,4,6,
t=1
1=1
tE[2] Ke (1,1.8))
Meanwhile (W((t)
s.p.f,ii ( (olyaa) polyd for p E (3.4) and E [2]. we can focus on the gradient difference
between
1.3 and Q contributed by other terms.
For 1, since Attn (t) A/(0)
1-peed 2 Attn a- and thus it is straightforward to see that
(1)
Ar(e) ,,3,1,1 4,4,1,i 2 Furthermore, we have
Ar(e)
Ar(1)
s,4,1,8
-
Attn ans --pad, logit
(2)
(Viisr (91) A 6.31 vo
+ +0(0))) + +((8)) (r(20)=
+ E Attn
0-n.0-aans.0 logit 5.31
(VVisron (vo) - 5.5. 02:000
- O((o)) + =((8)) 1 (r((2))=))a
Attn
logit 5.7()((0))
0-
((Vren (01) A (o
+ =((0))) + =0(0)) [+((0)=))ne
E
Attn (1) ans 0-an,0
logit (4)
(,7()
VTWEY
(( ((Vra (20) - A (,)()n =0(0o)) + +0118)) ((((a)
(0) ((t) (PCnler+es) ( ((B) - 0 polyd .(B) 0
log d (1+2Cn(er+oa)
log d
where the inequality (a) is due to Lemma G.25 and Lemma G.27.
For 2, since Attn(t) 2 Attn ann. and thus it is straightforward to see that (v(()
(5)
Ad(e) 23,,, 2,1,2, 2 O. Moreover, we have
(t) (t) $32,
$,1,2,8
-E Attn (t) ans,
logit (t)
5.55"
(Vigrro (92) -1 5 -o2
= + =Ō(8)) (1(2))=0)
(69)
+ E
Attn (8)
llggi!!
1-tans.
5,16
92

---

(Visroo (v1) (v1) [/
- =((o)) =((89) (T(2))==)
(70)
Attn (t)
logit (t)
prod,
Σ sReLU' (...)) (V))(12)
A + =((0)) + =((8)) (((()))=)EEE
(71)
r
+ Attn (C) ans, 1
logit (t)
Σ sReLL'(A (A
(Vir((n)) 1 -A0 + =((o)) = = (
(72)
real
By Lemma G.25, we obtain that
(69) (70) 2 -0()) ( (--ce)) O(B) 2 -0( log d ((==cocn))
By Lemma G.26, for Eg, we only need to consider the case that 91 92 A V/ # V1. and we have
(71) +
2 -E
Attn 1-ped_
Σ
(t) logit (,7())1))
logit
(( ((Vr (92) A (4)
+ =Ō(0))) =((8)) (F(e)) (o
(d2(v)))
+ E
Attn ans,aass1-
logit (,((x(0)))
((VrLo2())) (31) - (1)
+ =O(0o)) =0(89) 1 (or (12(V))))
2 ( 0(1) O(B) o (1ngg) -0
Putting it all together, combining with the fact that and (cl + 02) are sufficiently small, we finish the proof.
Lemma G.31 (Lower bound of gap). If Induction G.2 holds for all iterations [T2,,, t), given 8 E T(X), we have (0) [., 2 Ω (oga).
Proof. Letting ī denote the first time that (Q [Qli S Q43 J..0 1Q2 (22...)) 10.0)) which implies that
(2a.))
2
Σ
Q Loss
Σ
Loss 2.0
5,5
T=
Hence, by Lemma G.30, we have
1 s.s and [Q/ 1s. reaches 2(1). Thus, we can have a refined lower bound for (71) + (72) Lemma G.30, and obtain:
+ 2 -0 polyd
93

---

E Attm (1)
1-peed logii (,7(02()))
(?)
(VVrooz(ooD) (92) 5,7 (22 + O(00)) + Ō(87) (
E Attn (
1-ans.. logit (1)
(10))
- =0(0o)) (Vrroo(loo))roo (v1) A (oa
= Ō(69) ( (r(())= (2
--0 1 Ive
( (p(oy -
0 d ((B)
0
log d 2 0 2
where (a) follows from the fact that on the event (gr 92 yo # vi), have A (1) b.da.
Aoja--o2-v 2 ((B) once (0())sss and [Qf.llss reach constant magnitude, and consequently, the
(1)
logit satisfies logit!") (t) Therefore.
6,7(02(0)) 50 (F1(1))
2 Σ -Vv-awaa p=T+1
Σ
Loss 2.0 3.8 I..)) 2 -0
log d d2a)) O(1),
which means that [Q - 12
52
log
(
log d
logd)
Lemma G.32. If Induction G.2 holds for all iterations [T2,,,tt), given + 0 E T(X), for p E (3,4) (3,4), we have
Σ t=1
Loss
2.
5 log d
= 0
Σ
Loss 2.
(6)
Q
s,Nt
8,8
G.4.3 At the End of Stage 2.2
directly prove that Induction G.2 holds for all iterations' t until the end of stage 2.2, where we can conclude the following:
Putting gradient lemmas together,
Lemma G.33 (End of Stage 2.2). Given 8 E T(X), Induction G.2 holds for all iterations T2,1 S
t < T2,2, = 0 d
log then at the end of stage 2.2, we have
(a) [Q ((11) for (3;
(b) O 19ll E
[(( (nga) 0(1):
(c) 1. s O ( for s E xxxX) + s, and other [Qu.pea,s'
G.5 Stage 2.3: Decrease of Gap and Convergence
After rapid growth of diagonal entries in stage 2.2, we now focus the convergence of the attention
(t) and logit matrices, and the decrease of the gap between [Q4,3sss and [Q(!)] [Q,,lss.. Recall that
Lif cattn (z (21-
Att.tttttnsstt... (- 1--ppee, (z 12)- Attn Attnnnss--1-oms--1111 (Z
L. (Z) Atttngs.--1-peddd L.1111
(Z (ZL.-11).
Throught stage 2.3, we will focus on the attention gap A (ZL, instead of the gap of the attention matrices. We abbreviate. ( (ZZ. and ( as Eattn and A LE for simplicity.
L.t
Induction G.3. Given e 2 (oloo).for T(X),, let T2,3,a denote the first time that E 2.2 T((1))
s] c. For all iterations T22, S tt T2,3, we have the following holds:
94

---

(a) e4,3le, and [Q.,dlss monotonically increases O(1);
(b) 2 2 0 for any 7, with E (1,
(c) Attn ans, I- S 0.5 for some small constant č1 > 0;
(1)
| S o [ for of E T(X) + other
(d) for (p,9) E ((4,3), (4,4)),
[b4 = O.
G.S.1 Attention and Logit Preliminaries
Lemma G.34. If Induction G.3 holds for all iterations E [T2,2,s) 1), given Z2.-1 then we have
1. for =,
(a) Attn O-aans0 2 2(1), and Attn
Attn (0)
--pred2
2. for 8 =2
(a) Attn (1) ans, 1 (11), Attn (0)
ans, 1-pedd Attn (9 ans 1-ans, Attn (0) ans, 1-
(b) Attn (1)
1-a,n-+ Attn (t)
ans, ppred, I OH))
Moreover, given Z and corresponding Z Attn (t) ans,0--pred (Z2.0) 2 Attn (t)
p-eed (22.1).
Lemma G.35. If Induction G.3 holds for all iterations E (T2,,,stt t), given input Z then we have
1. for = 1, 2. E E, then
(a) for j jj, A 5.1.. < for E A, \ Mo
(b) for j = a T (92(100) ((g1 10
B < - E An \ vo
(c) for other j E T(),, assuming j = (y)) for some y + yo. then A < - ffror EE +E (m
2. l t= 2, 2f 22 E E2, then
(a) for j j2, A1 Addir -fror E Eg (rognn):
(b) for j = 1/2 - T (99((00)). 6.31e < -eforr E 2 \ (ro
(c) for other j E T()), i = r(z(())) for some E V, then S,g. & -e r E 2 \ (T
Lemma G.36. If Induction G.3 holds for all iterations E [T2,2,ar 1), given imput Z then we have
1. 8 = for Z E1,
(a) (0) b
-11-26 - 2011 1 0(8) 2 + c1 )))B:
(b) A (1)
0.11 v0
-a( (t)
A 5.31.r92-90
2(Attn (t) ansed--pred.1
Attn (t)
ans D-pped--peed.2) + 0(8) 2
2(1 + 02)
(c) for y + yo, A (t)
(2Attn ans.0 - -1) + 0(8), which is only
(2Attn (0)
5, (vi): (t)
activated if Attn U >
2. L = for 2 E Eg,
(a) (11)
2,2
B + 0()) 2 4cs B;
3.j2.-92-91
(b) A11) (1) A B + 0()), and
(Attn (8) = ans, 1-sans,1 Attn (1)
1 B + 0(8);
(c) for y + yo, U1, A a.r
(2)
2Attn (0)
1-ppdd -1 0(6), -(0) which is only
(,7()x(1))
activated if Attn
(0)
1.
1-speed
95

---

Lemma G.37. If Induction G.3 holds for all iterations E [T2,2,s t), given inpur 2 then we have for 1, 1 E E1.
1 logit (t) 5. (1)) logit (8) 0
5.51 d22 (Attn (1)
0-ppred) Attn (8)
D-ppeds On
2. for = 2, e E
logit
(1)
0
1
5.15
Attn (1)
Attn (1)
Co -1--22.203
moreover,
2 min (2(1)
1 logitt
5.3 5.
do (1-2
Lemma G.38. If Induction G.3 holds for all tterations E [T2,2 1), given inpur 22,-1, then we have
1. for = 2f $ E,, then (a) for j ji. An fr yo A 5,j1. = B =B+
(b) for j * j1 E TD'), assuming j = (v)) for # yo. then (1) 5.ji.
- S.j.-a 2
2Attn ans, ~ans. B Z QBB) and (1)
(t)
S.j.r < forr M \ (rargh
2. f= 2, 2 $ Ez, then
(a) if = 92 A 10 + yı,
i. for j = j2, D.juoon = (1 (attm)
(t)
< -e r E B 066) and 6,ja.
(1)
Az \ (r bi
i. for j = (19(1))). Aa
A a.r
(t)
2( 2(Attn ans
Attn (1)
ans, 1-mns B = 0(6) and (t) b.i.
& forr r 21 \ tro
ii. for other j E +()), assuming j = ((9(())) for some y # yo: Vı, G122fa2
Girory = 2Attn (2)
(1)
1-a,n- 13 + (6) and A < -eror E 21 (r
(b) if + 92 A y0 = yr,
for j j2. 5.j2.
(1)
21- (atim)
+ 0 and 5,jo. < -e
712 (r (F ):
ii for T(9 (B1)) ((92())).
A (0)
(t)
2((Atnn (t)
Attn (t)
ans, pred, BB B
ans, p
(0) 6.Jo.
2(Attn ans, +ans. + Attn (1)
ans, 1- 1) +
S.j.roz-3
where 192-ū is only activated if Attn (t) ans. 1.2 Asr - &
2 (Ta gl
ii for other j E T(Y), assuming j = ((g(()) for # yı, G. then raav is only
activated if Attn anm 1-pee_ and
(t)
D.jiPoo 2( (1ttn ans, + Attn (1)
a B
6,12 22
(1) < -eforr 2y \ (roww)
(c) if = 92 A yo = yr.
for j = j2. 2 = (ra-vi }. 10 A6.j..m2 = 0
it for other j E r(Y), assuming j =(g2(v)).
A (1)
(0)
= Attnn (t)
Attn (8)
ats 1)B 0(8).
1-5,n
96

---

and A < for E EX, (rgrz)-
Lemma G.39. If Induction G.3 holds for all iterations E (T2,2,s) 1), given input 2 -then we have
for = = 2 & & then logit 5.51 0
0
polyd
2. 8 2, tf & Eg, then
(a) if = 92 A + v1. logit 5,7(0((0))
Attn rass.1 -Attn (1)
sama, (m..0OB
(b) if * 92 A y0 y1. logit
(,r())
Attn!!) (1)
are,1 pped, Attn (2)
poodd
CD
(c) gf g2 yı. 1 logit (0)
3.j2
Attm
-200
AActo (6)
1-ann
On
polyd
G.5.2 Gradient Lemma
Lemma G.40. If Induction G.3 holds for all iterations E (T2,2,s,)) t), given s E F(r (Qu,ala,s. we have
Σ
2
Loss
2.
2 12
elog d ([1-2:]CBB
(1)
4.4
Proof. By gradient decomposition, we have
Σ
Loss
2.
E Σ
Arti
Q
S.L.
1=1
[E[P] (F,
Firstly, for A1)
s,4,2, by Lemma G.35 and Lemma G.38, we have
A/0
Attn (4) ant (1 logit (4) 5. 5.
5,4,2,
(Vi (11) 100
3.j2. + +o(oo)) + Ō((8)) (1(m))=))8
+ E
Attn (t)
ans, 1ans,1 (1 logit (1)
5,12 ).
(( (11) 3.52%"02
+ =(0o)) + =((8)) (r(m))=))88
(a) 0 1 (2))
min (2(1).
dC (C2-))-2-22))1)
2: [1(((1)) =
2 clog d d(1 (2)1-2))Cn))
(73)
where (a) follows from Lemma G.37. Noticing that Ar(e)
s,4.1, 0 and W (t)
3,4,,al for l E [2] is
sufficiently small, to provide a lower bound for Σ
( Loss
2.6
we only need to focus
on the negative gradient from (() 5,4,1,5 and (() 5,4,2,1
2,4,2,18
Attn (8) ans.1ass11- 1anns, logit 5.12
E
97

---

(Vo2roo (31) -
- +((00) + =(08)) (+(2))=))
E Attmaass. (0)
Σ
logit (t)
ans, 1-aans1
(HieT(())))
SReLU' (...)) (v)) - A + =(00)) + =((8)) "(+(())=))nEs
r
2 E Attn lans. logit (6)
(,7()1(21))
(.
= =((89) ( (Vr(o 21 (v1) - (.
(74)
When E
CB (1 (1-2 ((x1)= = o) 1, Lemma G.37 and Lemma G.39, we have
(m)) ( ( (
(t)
-pee ) (1 log)t )1 (7(2))==)nz))
(t)
d2 (Aten!!) (t) Attn (t)
2.2 During this time, cattn 2 9(1), thus we can lower bound Ao(()
2, by Ar(1)
-0
8,4,2,6
Ar(t) 2,4,2, which implies that the negative gradient from Ar(t) 2,4,2, is dominated by the positive gradient
from Nr(t) 8,4,2,
When E Cn (1- 2.2 ( s) 2 1, since 2( ((Attn".)
ans, x-red, - Attn (e) ans, 1- ) 2
(t)
(1 - 2r)22) Cn, we obtain
logit (t)
((m))))) (r(mm)) S -1
-1- logit (1) (>(2))=)))8)
(75)
For the event (g) # 92. 30 =y1). A (0))) is still in the linear regime, we have
A (1)
(1 2Attn ans, 1 B 0()) e
(2. (v.))
which implies
Fattn 2 1 - 2At spred, 2 12
0 B
(76)
Hence, putting (76) back to (74), and putting (75) back to (74), we can lower bound (r))
e,1,2,4 as
follows (t)
$,1,2.4 2 -0(3) (y(())
2,1.2 If (t)
(5:7()(11))) falls into the smoothed regime, we
can upper bound sReLU' (t)
(.(( by O Can and then similarly, we can obatian
(/(1) 2,42 2 (2) -/(0 $.4.2.1*
Following the analogous analysis, the negative gradient from (+ 8,1.1 can also be dominated by the positive gradient from 2.1.251- Hence, we complete the proof.
Lemma G.41. Inductioo G.3 holds for all iterations [T222,s 1), given s ((X), for Q4,llas p E (3,4), 8 # & T()), ? E (2), we have
V R
Loss?
1
V 00 Loss 2.
98

---

G.5.3 Non-negative Gap
Lemma G.42. If Induction G.3 holds for all iterations [T2,,,0: t), then at time t, we have A 2, 2 for any 2 with & {1,2}.
Proof. Let T denote the first time that E[ ((x)) = =] < a, where a
polyd
Following the analogous analysis Lemma G.30, we have
Loss
2.6
+
2 Loss
Σ
ArTT 8.3.
-NT)
5,8
EE|2]
(E
We can ignore the negligible difference introduced by AYTT)
s.p, for p E (3.4) and E [2].
(T) For l = 1, since Attn ant, 2 Attn a-n-aan and thus it is straightforward to see that
N 8,3,1,
-N 1, 2 Furthermore, we have
NV(TT
-
1,3,1,2
1,4,1,
-
Attn ( p
logit 5..
sReLU' (A (Ad
(V (01)
A!T)
+ O((00) + Ō(89)) 1 (r(xo)==)]8)
(77)
+ E Attn a-a,o-
logit 5.
sRRLLU (A (7)
((Vii VV (v0)
(
+ Ō(0o)) + +((88) ((((2))=))(8)
(78)
bidf T22
E
Attn
Σ logit (,r((((1))
'RRLL' sRRLLU (A
((V (VVrOL(W))) (01) (())
A (1))))))))) + - =((8)) (1(m0)=))nef
+ E Attn 0ans,0
Σ DF
logit (,7()((1))
sReLU' (T) RReLL (A
((VroonwW)rn (16)) A
+ =o(0)) = 5(89) ( (r(=o))
(1))
By Lemma G.39, obtain that (T) ( Lemma G.37, we have
is dominated by (77) + (78). Moreover, by
8,3,1,6 1,4.1.5 is
logit logit (v)(
(2 2
2 Attn (?) -xpred, Attn (T) 0-predd CB
C
(79)
Attn (P) Lapoedd Attn (P)
1 ,)
For f= 2, we have
(T)
((T))
8,3,2,
2,1.25
99

---

E Attn 1-ppod2 (1 logit (
6,2
((Virro (92)- -
- O(0o)) + 0(8%)) ()(e))=03nE)
E Attn
ans.l -oggit
(T)
3.jo
(VVerror (11) (T))
- Ō((o)) = =0(88)) 1 ((()))=o)n
((aclogg d)
(1 logit ()) (=()))=
(80)
where the last inequality follows from the definition of T.
N
(T)
(2)
$,3,2,58
s,4,2,8
Attn (T) logit (T) 6.22
sReLU' (A
E
an
(o
(Vilero 20 (92)
A
(T)
= + =((8)) (T(2))=))
(81)
+ E
Attn ans. I-ans, logit 5.32
SReLU' (A
((Visre (11) A (T)
+ O((o)) + Ō((7) (+(())==))
(82)
Attn E
logit!
(T)
ans,
2
Σ
(ReLU'(A
B.g.)
(92)
- =o(o)) + =0(09) ) (((r))=)fã
(83)
r
logit! (7)
+ E Attn ans. 1-ans
)+b))(0))
Σ sReLU
(2) (g
(v)) 15 + Ō((o) = =0(0)) (((0))=))(s))) (84)
Notice that by Lemma G.36 and the definition of T, we have
(81) (82) S (oa E (1(1 logits.))
1-(2)==0
which is dominated by (80) due to the choice of a. Furthermore, by Lemma G.38, for Es, we only need to focus on the output logit case 91 = A yo * y1, and obtain
(,7(02())) from the
(83) + (84)
2 E Attn ans, 2prdd2 logit logit (,702(00))
sReLU' raer
AReLU'(A (
100

---

((Vroo(00))) (92) (7)
(0.)(22(10))))0020 =0(00)) = = (
+ E
Attn allinans.1-mmm logit (T)
sReLU' (2)
(A (15 6.
(,2(0)))))
(Vrlea(n))) (v1) A (T)
3. +0(oo)) + Ō(89)) (1oo))==s)((o===mm)mm
By Lemma G.39, we have
logit logit (92(0))) (2 =0
(T)
d (t((Att am0 Attn are 200,0
(3) 1200,
Cp
Since E/422 r(1))= = s) aa where a is sufficiently small, then combining with (79), we have
logit (10 (z = 0 logit (,7()2(1)) (Z..).
Therefore, since E happens with probability O a) we obtain that
(83) + (84) 2 - ( log d
(N119 5,3,1,5 v.4.1.6)
Ar(e)
Putting it all together, we can conclude that when E[ 7(x1) = reaches below o, the gap in non-decreasing direction is guaranteed, i.e.,
Σ
2/ Loss
2, Loss 6,5
0,
t=1
4.3
8,8
which completes the proof.
G.5.4 Upper Bound for Q
Lemma G.43, If Induction G.3 holds for all iterations E (T2,2,s) 1), given E T(X), then at time t,
we have [ S ((1) for E {3.4].
Proof. Denote the first time that 14,als.s reaches 2(log ++ d) for some small constant e T. Then by direct calculations, have
Attn (t) ppree,2 Alllog++ dj
eflllggl=d) {
(t) Attn 1
Attn ans,
(t)
50
Moreover, by Lemma G.37, we can simply bound the logits follows:
1 logit 50
HB122-1 for Z E E1;
cCa (1-2Atn 12ttn -pedd.
1 logit 50
CC 2Attn (1)
iped,1 -Attn -
+0
0 (Cn/2-1)
for 2 Eg.
2 Attn (1)
1100111 Attn
Cp
101

---

Thus, by focusing on 3,11. and N,,1,1,1, we have
Loss
2.8
S 0 (d(0//2-1)
.0 clllogi-e (callos--r))
B BB )
(
log d elllogh d
8,8
which implies
+0 +0
Olog d
+ +0
polyd log.logge
1.3
4,3
elillgt tc d) dB/2
Q 1,3
elllogt+ d) d0/2
- (log!++o))
G.5.5 Attention Upper Bound
Lemma G.4A. If Induction G.3 holds for all iterations (T,,,,st), given T(X), then time t, for any 2 1 we have Attnaas, (t) -xppp- S 0.5 + č₁, where č1 > is some small constant.
Proof. Let T denote the first time that ||tttn ans,1 ans,1-1ppeee -ppred2 ((x1) = s] reaches 0.5 + where
(()
(T)
ē is some small constant s.t., 28 Cp 1. At T. have Attn ans,0 0ns,0->predd ( (T)
(T) Attn ans, pred, 2 ((x))=s moreover, Attn 1-nns and Attn a is still at the constant level.
(T)
For = 1, by Lemma G.35 and Lemma G.38, for E V {y0), we have only Ter has been activted to the linear regime for the prediction T ((9 (v)). Furthermore, obtain
logit (1) 2 log 2200n lugd d + O(d) (1 logit (T) VE)'(Bo)
Σ
log d. 22002 logd
6.51
= o(1)) (1 logit (T)
3,51
Thus,
[
V Loss 2.1 Q
Attn (T)
(11 logit
(VViroo logit))) (T)).
v0 (91) (T)
A
S. = +((8))
ans, ans,0--pred,
Σ
logit (,7()1(0))
((rieetosra (0)) - (T)
(x))) +0(00)) + 0(08))
MEY'(an
[ 1- logit
( 00633) L
polyd
< Attn ( ans.0 1pree,1
( (11-1 ((2Attngm.0--peed2
Ons,0-ppred,2 B Ō(oo)) + =((8))
Σ logit (
(22ttn an.0
ans,0 vans.0 B + O(0o)) + +((8))
(
(,7()))))
y
polyd (1 - logit 5.51)
-(68)) 1 (((0))= <
For -= 2, for 2 E U (gi = 92, yo + yr we have
logit
3,r(92(1))
>
log d e e log d + O(d) (1 logit (T)
log d 200' log d
5.jo 5.5
- logit (T)
= (1 o(1) (1 logit (T)
logit 5.1% 5.3%
102

---

Otherwise, we have
Σ
logit logit 5,r(ga(1)) = (1 - -(1) (1 logits logit (2) 5,33
5.32
Therefore, similar to é = 1, we obtain
V
Loss
2.2
< 0.
Q
8,8
Combing the above two cases, we have E
-t=1
2, Loss 5
0. It is also clear from
8,8
Lemma G.40 that Σi=1 [- V once it reaches 0.5 2 ट.
Loss 2
2 0. Hence Attn
1-pee-+ cannot further grow
G.5.6 Decreasing Gap at the End of Convergence
Let T' denote the first time that B T(x1) s] 3e, if T(x1) S 00
(d ())
log
then we can let T T and stop the training. Otherwise, we have E| T
2 (d2.01 63
Following the similar argument as in Lemma G.40, we have the gradient contribution
log
from obtain
is dominated by the gradient contribution from l 2. Thus, focus on l 2, and
(2)) Loss?
8,8
Attn ans, I pped,
(1--lggi (ogit( ) ((Virree (92) - (T) N.ji
# =Ō(oo)) + Ō(89)
polyd (1 logit 5.52
(T)
Ō(0g) (r(m))=o)
E Attn (T) 3ns.1-tpeed.
Σ
logit 5.7(0((0))
sReLU' <ReLL"(A (A
((Vroolwi)) (92) A
A, (o))) +0(0))) + (((9)) 11 ((((e))=))8))
Attn (t) E
ans, 1,pred22
logit 5,3"
Σ RReLU(A ( (Vi,((e)) - A + =((00)) + =((8)) (1(1))
r
(1 l ) (eelogd)
Turning Q1,1. have
Loss 2.2
((1 (Igitg (VVerrormn (v1) (
A + =0(o))) + Ō(89)
E Attn ans, -anns,1
1 (1 - logit (T)
6.52) Ō(a8) [1((0)==0)
+ polyd
E
Attn 1
logit! 5.r(ga())
sReLU' (M
103

---

((rrer(w)) (01) n
A =o(00)) = 0((0)) [
(85)
E Attn (0)
logit (1)
ans,
5.5"
Σ
sReLU (As.r) (Ver(vr) 1 + Ō(0o)) + =((8)) (r())
(86)
r
For (85), have
(85) 2 E Attn (T)
(-logit
logit 5.12
ans, 1-ns,,1
min ([(log d). (eslog)).9((AAtttt)
(Attnn ans, -pped,
Attn (7)
1ns. log d) d))= log a) (r(20)=0)(8)
2 E (1-1ogith
logit (T)
8.12 -2(2)
min (1llog d). 22
1xpped Attn 1-amn, log logd) d) 9- log logd) } 3 110000000)88
Moreover, for (86), we have
(86) 2 -E Attn ans aans.1 logit (.1()(
(T)
sReLU' (A (A.ri
((Vroo (wn) T 5.
+ =((o)) + Ō((9)
(
-E E
logit
( o log
((clog
(7(20)
5.52
where the last inequality holds since A (T) (gr (vn)).
Attn Att 1-pee
Attn Ins, B
(c log d).
Since 2.2 [(x))= = 2 22 (dd. c) o2)-0
we have
logd
[[-
(85) > 0.01
V
Loss
3,8
and thus (85) > (86). Thus, for 1. if the attention gap does not decrease the level of (1 d
O
Q4.4ks.s will start to dominantly increase while [Q4,ss,s will not change too
much. On the other hand, if the gap of attention holds, then [Q4,3].,+ 2 [Q4,4s..s. have
0(1)
2,2 <attn
Attn ( ans, 1-pred22 Attn 1-s.n-+ans.1
1
O(1) + e. 9
+ 10
0(1)
3e e.
O(1) + 2e Q4
This implies that, we can find some time between and T2,3,s. s.t., the gap will decrease to o (21-01,) log d ) We denote this time T and stop the training.
104

---

G.5.7
At the End of the Training
Putting everything together, we have that at the end of the training, we have
Lemma G.45. Given T(X), T =6 d if (f < O ) we have
(a) Attention convergence:
Eatn I O(c) for l E [2], and A T(x1) = s) s
(d c)
0
logd
Qip (T2a..)
(T23.a) Oup 0,,
(b)
2 Rllogd) for (3,4] if = o' E T(N), otherwise, 4,
S
(c) Loss convergence: Ez Loss 2.
polyd'
6.6 Proof of Main Theorem
3.3, C.I, C.2 and 4. for some constant 0 < c² 1, ny < m < log" d, the transformer model F(T++T) obtained
Theorem G.1 (Restatement of Theorem 4.1). Under Assumptions
by Algorithm I with learning rate 7, poly(dy and stage 1 and 2 iteration Ti Ö 71o0|0
T2 = = poly(d) satisfies
No0
ACeL ((T
poly(d) for every L S O(d"),
(87)
ie., F(T++T2), which trained for task T and T2, generalizes solve the tasks T PPS I.
Proof. By Lemma G.45, at the end of Stage 2 training, we have Q 0
(T2....)
2 Milog d) for all
p E {3,4) and r(X). Therefore, for task Th with I. < O(d"), where 0 < c" < in Lemma G.45, we obtain
a constant depending on the value of
0(1)
L.E attn
= =(1)
(72,3,>)
Q
Q1,4
0(1) .Ll) L
2.2 + c
Moreover, we have
L < A1211
Attn (12,,..)
Attn (12,3,.)
ans,1 S o(1).
1-ppee2
These together guarantee that
0(1) d d e1) 0(1) d - (l()) + elllogd) poly
< (1(T") 1 - logit,, (, (p(T)) z((.0)
which implies
Acc (F(TT) 2 poly(d)
H
Learning the Attention Layer: Symmetry Case for Short-Length
In this part, we consider the scenario where the group operations form symmetry group. Specifically, following Assumption 5.1, we assume that y = (0, 1,... My 1). and let G be the symmetry group o' where ny simply transitive case, we restrict analysis to updating only Q, specifically the blocks Q4,3 and Q4.4-
for y, with order |g] ny. @(polylogd)
log log log logd). Similar to the
Throughout this section, we focus on the simple task T2 and analyze gradient descent updates with respect to the per-token loss Loss? Given & E T(X)), let Loss? 2,2 -E log an,
105

---

(2 s].
Due to the symmetry of [Q1,3le, and 1Q1, across E T(X), we may.
without loss of generality, focus on a particular E =(2') and analyze the corresponding loss Loss in what follows.
H.1 Gradient Computations
We start with the gradient computations for the attention layer.
Notations for gradient expressions. We first introduce some notations for the gradients of the attention layer. For S & L, given Z ZL.t-1 and I I define
E (Z).
Σ (Z jE(d
Σ rE(m)
sRcLU' (A...(Z i.j.
s.j. Zk),
E [5]. (88)
For simplicity notation, we will henceforth omit the dependence on 24,6- in the notation of Et when it is clear from the context.
Fact H.1 (Gradients of Q). For any p.q [5], we have
Qu Loss
-V
Loss be
where
8=1 1E
Van. Loss
L
Σ
EL
Σ
Attnam,
KEIL.
K'EILE
Lemma H.1 (Gradients of Q1,3). Given s E T(X), for the diagonal entry [Qt,3las of the block Q4,3. we have
Q Loss?
2.2
= E Attm
Σ
Es (Z
Σ
sReLU' (As.sr)
8,8
jE(d]
rE(m]
(Wa.j.r Zprod,2) As.ger
bs.s.r) M
Moreover, for the off-diagonal entries (Q4,3aaa with 8 + s', we have
VQ.. Loss? 2.2
E Attmans. 1-spred1
Σ (Z22)
Σ
sReLU' <ReLU"(Asar))
jE(d]
re(mm
(Wos..m Zpred, Asgi + bus.))) ( (==
Lemma H.2 (Gradients of Q1,1). Given s E T(X), for the diagonal entry (Q1,s,,s of the block Q4.4. have
Vou. Loss 2,2 = AttMaas, 1--ans,
( (Z2))
Σ
sReLU' (ReLU'(Asar)) (Asar)-
BELd
rE[m]
(Wa.... Zans.)) Asgg.r + bs.j.r)
1 La-r(zn)
Moreover. for the off-diagonal entries QQiale.o with # s' we have
VQ.. Loss?2
AttD 1-ran,00
Σ
Es,3 (Z21)
Σ
sReLU' (As..rr
6.0"
jEld]
rE[m]
106

---

for some C. (n,)): (Wo..r. Zans,o) As,ir + bbssir) 1o=(())))
H.2 Some Useful Bounds for Gradients
In this subsection, establish several useful bounds on the gradients of the attention layer, leveraging the feature structure of the MLP layer learned during stage 1.1. These bounds will be instrumental for the subsequent analysis.
Recall that for E T())) and E V, the fiber Fiberjy. and the set of feature combinations for predicting y = -- (j), are defined as
Fiber,,y j.y = (g E G ((((v) = -}, 8y=
= (Fiber v) E 29 }
As established in Theorem F.1, at the end of stage 1.1, we have
Sparse activations: For j E T()), let Φ = (Fi y) E Bj. then there exists exactly one activated neuron r E [m] such that when 9I FFiber j.y yo = y happens:
A
(Tra) 2 B - O()),
C. (C v (g) - O(6)
bgrr
((Ta.) < 0 Agg
Vr' Fr.
Cancellation of incorrect features: For j E T())), let é = (Fiber,,, j. v) E B, and let the r E [mi] be the activated neuron, then for any g $ Fiber,,g- and any 3/ + y E y. we have
+ ( (v') s 08) and
V ( (0')
v( (v) s 0()).
J.r
In the analysis of FFN layer, we have pair 5 (81, 62); in the expressions above, some terms should use ō1 and others 82.
Notations for activated neurons. We denote by Tiay the unique activated neuron corresponding to é (Fiberz.y; y) E }y. For any Fibers.y. we also write yg.y for the same neuron rj.y. Note that T y T92. for distinct Φ1 92 Fibbrr,,. Moreover, define
2 P UsEr(r) My, where 21, A (rew 1 E J)..
In other words, 21 is the set of all activated neurons across all feature sets 3, for j E +(V). Given z_.(-1 letting (ZZ..111) U(g Y be the collection of all the chosen group elements in the predicate clauses. Similarly D U(ye 16=1 Then define (Z Tgg g Fiberryy
( (gE } For simplicity, omit the dependence 2 in the notation of a when it is clear from the context. Equipped with these notations, we can summarize the above properties in the following lemmas.
Lemma H.3 (Properties of target feature magnitude). Given j E T()) and y E ) then for g Fiber the following properties hold.
Varr.. (g) + Virro (v) 2 2B 0(6),
|Cav (1) jr (v) - v,(0) (g) < 06)), jr
(89)
Vi.Tery (g) + VA.T.. (v) < 0()), Vi.ro (V') < for all y + y.
(90)
Vyrrar (g') + Vg.r,., (w) < 0()), Vyrr, (g') < 0 for all g' & Fiber
(91)
(V),( [V)+(v)] < 0(6) for all r $ 24. Lemma H.4 (Properties of irrelevant magnitude). (f v) $ (2) (U X,, jr $ T(V), then for any r E [m], we have
(92)
(Wa eu)) $ Ō(oo).
(93)
107

---

The above lemmas give us some direct computations of the inner products between the weight matrices and input embedding vectors.
Lemma H.5. Let j T()) and E [2]. Then for any E [m]. the following holds:
(W b.jors pred e = (ger(ge) O(ro), (W s.j.r a t-1, = (vr = O
(94)
(95)
Moreover, for & r())) and any k I and E [m]] we have
(Ws.j.r 2 ((a)).
(96)
Furthermore, we can establish some characterizations of the As.j.r(Z.. quantities, which are crucial for the following analysis.
Lemma H.6 (Characterizations of Lambda). Given Z with (k) then (a) for j E TJ')), for activated neuron T E A, we have
Asggr Σ Atttans.1--pred.f (ge) Σ - Ō(o)).
(b) for E for any non-activated neuron $ 21 huve
Ass.r 06).
(c) for $ T(V). for any r E [m], we have
As.j.r
$ ((o)).
A direct consequence of the above lemma is the following finer characterization of the activated neurons.
Lemma H.7. Given j E r()), for 2y Ay we have sReLU' (As.g.)) 0.
Now we are ready further derive the gradients of the attention layer starting from Lemmas H.1 and H.2 and the properties established above.
Lemma H.8 (Refined expression for the gradient of Q4,3). Given & E ((X),, for the diagonal entry [Q,,3a,s of the block Q4,3. letting ((92(1))), we have
VQ.a Loss
2,2
E
AttMaans, 1-prdd,22
8,8
(1- logits
Σ sReLU (As.z.r)) (Viar(02) - -A..j.r + =((0)) + Ō(89))
r
logits
sReLU' (ReLU (Vir((0)) - As..r + Ō(oo)) + Ō(6%)
(+LET(())
r
Σ logits, Ō(o3) Il Dr(enson
j4r(1)
Moreover, for the off-diagonal entries [Q1,3la.ar with + s', we have
VQ. Loss 3,0"
22
E Att
(1- logits.j))
Σ
sReLU' (Asjzrr) (Ver-or) As.j..r + Ö(oo) + =(88))
108

---

s.
(Σ
SReLU'(AsaN) (Ver(o)) As + O(00)) =0(89) r
ifizE00)
- Σ logits, Ō(o8) ll [r(2))
J4r())
Lemma H.9 (Refined expression for the gradient of Q1.a). Given s E T. for the diagonal entry [Q4,4ss of the block Q4 letting (19(1))) we have
Loss,
2.2
E
Attnaans1--aass.
8
(1-logite logits.ja
Σ
sReLU (As.2ar)) (Vear(11) - As.jer = =( = =((8))
r
logit
bj
sReLU (Asaar) (Vir((1)) AoJrr + Ō(00)) + Ō(89)
HAiET(())
r
logits, Ō(o)) Dr(onsas
(ET(D)
Moreover, for the off-diagonal entries [Q1alaws with + s', we have
Vo,,, Losss
2,2
E Att
(1- logits.z)
Σ
sReLU' (As.ear) (Viar(20)) - As..a.. =((o)) # =((87)
r
logit bj
sReLU' (As.grr) (Vir((0)) As...r + Ō(00) + Ō(88)
(#1EET(1)
rea
logits, Ō(o)) 1. I=(a:))
(KT(D)
Following the above caleulations, we further obtain the gradient summation of Q4,3 and Q4,4 as follows:
Lemma H.10 (Gradient sum of Q4,3 and Q4,1). Given E T(X), letting ju ((92(11)), we have
VQ Loss
2,2
+
Qa.a Loss
2,2
E (1 logit
sReLU <AeLL'(As.j.r))
rea
AttDans.1-and -ans,0 (30r Att 1-pred, Vi2r (g1)
+ (1 Att 1-ans,
Attn pred As.ja.r + Ö(oo)) + =(87)
logits.j
(AeLU'(As,.)))
(7()))()))
EEa,
Att 1-1ans0 (yo) + ALtmans,1-ppre. Vii(
- -(1 Atttigns.--am Attnyns.1-pred.l As.a.r Ō(ao) + +0(8))
109

---

l. ( + 1t (
Σ
2) 01(8)) 1r(en)=
Notations for gradient decompositions. We shall define some useful notations to further simplify the expressions of gradient.
Lemma H.11. For any E T(t), we define the following notations for the gradient decompositions:
1. for [Q1,3les we have
Vo,, 2.21
E [N,,3322 + Ne,3,2,5 + +Na,3,2,85s) where
N,,3,2,1 Attmans, I 1-pdd, (1 logits..t)
(97)
Σ (Ac. (V)).(92) AsJerr + =0(0o)) =((8)) ((=1)=s
r
N4,3,2,8 Attillass, 1pppdd
Σ
logits,
(98)
Σ
(Ac.UU(As.r) (Vrr((2)) Asy.r =0(oo) = =Ō(0%) (()))=s)
reat
N,,3,2,0 = + AttMaass 1
Σ (TT())
logit,,
((0g) (ogits,5((o8)111
(99)
2. for (Qu,a we have Qa.. Loss
2.21
[[ + Na,4,,8 + NN,,2,,a), where
No.1,2,1 Attmanss1--ans. (1 logits.j2)
(100)
Σ (As.j.r)
(Va,r(V1) AsJe.r Ō(oo)) +((88) (m))==
r
Na.12,8 Attnaass 1-nns.
Σ logits."
(101)
(ANEr())
Σ
(As.r) (Vir(11) Ass.r + H(0o)) + Ō(()) (m))==
eea
Na42, = Attn mn..
Σ logit,,, (agg) ((1)==)- (
(102)
jØT()*)
3. for the summation of [Quals,, and [Q4,dla.s. we have
- VQ Loss? 2,27
+ [
VossLoss?? = =E E[Na,2, + +N + Ns.2,tt.. where
Na = (1 - -logit,,.)))
Σ sReLU'
(As.jar)) sReLU'(As.j.r)
real,
AttMons.1-ans.0 (y0) Atttaas.1--peed V112((91)
+ (1 Attt 1-tans, Attnan..--ped. As.j2.r + Ō(oo)) + +((8)) 1 ((0))=
Na2.
Σ
logit b.
Σ
sReLU (Ao.jr))
ither())
Attnnass, +ans,0 (vo) + Attn -pprrd, Vgrr(ai)
(1 AAtt 1nnnann.1- Att 1-tppd) A + Ö(Go)) + +0(8)) 1,)))=
110

---

2,. =
logits. iAtt + 1 j4r()
Σ
Probabilistic events. We conclude this subsection by introducing several probabilistic events that will be used to simplify the characterization activated neurons in the subsequent analysis. We first define some events that may contribute non-trivially to the gradient of Q1,3 and Q1,1
81 yo yı 91 (ve-1) + g2(ve-1), for all & & [21]
(103)
&2 91..9 92 (iber 91 (v)) + se(vi)
(104)
&3 you 91 (11) # g2(v1))
(105)
& yo # yı g1 (ve-1) = 92(ve-1), for all ee [21]
(106)
We first observe that event & occurs with high probability 1 1-0 2.). and serves the primary regime of interest. In contrast, events Eg and & each oceur with probability @(1/n,) and correspond to instances of initial prediction ambiguity, where certain incorrect classes may exhibit disproportionately large logits. Similarly, & corresponds to instances of initial prediction ambiguity, while occurs with probability (1 Furthermore, we define the following events, which yield negligible gradient contributions to Q1,3 and Q1, as they do not lead to significant confusion among the incorrect predictions:
&g yo + 9n (11) = (e(v)), (un) + se(vn)),
(107)
Eg you = y1 91 = 92(1))
(108)
Here, Eg occurs with probability (11mn,),, and Es occurs with probability @(1/n²). Together, the events U,etoe&; form a partition of the entire sample space.
H.3
Stage 1.2.1: Initial Growth of Q
We define the following notations:
I.f.
(Z
A Att E- 1--pped (Z (Z.) Atttmasss / ( 1. 1.
Gattn
1. (. Attnaans 1--ped, 1-e-1) Attmann.6-1--ams-- (Z
L.f We abbreviate aatn (Z2,t-1) and A. (Z as cattn and A LE for simplicity. Since we only focus on the imput 2 in the following analysis, we will omit the notation related to the length, i.c., we use Eattn and A when the context clear.
Induction H.1. Gäven E T(x), ler T1.2.1. denote the first time that [f(aatnn T(x1) s] S 0.4 For all iterations < T1,,, we have the following holds
(a) [113.. + [Q S 5O()) monotonically increases;
(t) 2.d otherwise [Q.t.]] 8,8"
(b) for p E for E T(x) #
[Q s
0;
(c) for any sample Z we have Attn (t)
Attn 1-sped 2 -0
log log d); -1(1og log
1-pred2
(d) for any sample 2 have Attn (t) constant e₁ > 0.
ans, 1-ppe-2 Attn (1)
1-aans11 S Ci for some small
H.3.1 Attention and Lambda Preliminaries
Lemma H.12. If Induction H.1 holds for all iterations < 1,then we have
(t) 1. Attn ans, I + (ttn
ans. E {00.4 + ((t),
2. Attn (0)
Attn (0)
Ō()).
1-ppeed

---

Lemma H.13. If Inductiio H.l holds for all iterations < t. then given Z E E1.
1. for the prediction J2. have
Atttn ant Attm antmaass1--nns.)) 2B ans Attn (1)
(e)
(t)
ans, I-s.rdd )
2B
O()).
nly
2. for the prediction 8/2 ----((p (g2(1))), we have
(0)
= -Atttn
ans, 1-prad Attn (8)
(1)
att1aas..-)ams..1 2B + Ō B + O(0).
12-000
-ly
3. for the prediction T(9 (Vo)), we have
A (t)
(#) (Attn 1-pedd Attn ant, 1-ans
(A)
2B + (Attn ans,
Attn 1 --ped,)
2B
+ 0()).
furthermore, we have (t)
5 (,7()1
3 11 (1)
4. for the prediction 7(g (11)), we have
A (0)
-Attnn (#) ans, I-ns,,
Attn (e) Att 1-s.ed)) nl
2B
+ +0 + + +0(6).
5. for other j E r()), if there are and y. s.t. 91, 92 E Fiberj,y (notice that y # yo, y1 for εx), then for such a j, r M, \ (raww) cannot be activated, moreover, we have
A ( S.j.f
Attn (1)
B Attn Attn ans.1-) 2B + 0 ( dd O(0);
(
1
else, none of r A, can be activated.
Notice that for ε,, the neurons mentioned in 1-4 should be different neurons, while the predictions in I and 3, or 2 and 4 can be the same in some cases, e.g., 91 (31) g2 (yo). In all cases, except for the neurons mentioned above, i.e., (e (r
se'e ,) (which may not be activated), all other neurons T E U,, (Trggere-s) \ (rgrre-)) cannot be activated.
Lemma H.14. If Induction H.l holds for all iterations < t, then given Z E Ez, we have
1. for the prediction j2, 215 (1g22m) cannot be activated, moreover, we have
(1) 5.j22.7o2 Attn (0)
ans, 1-spred22 Attn (1)
ant. 1-s,1_+ans,0 2B
2B + 0(5)
+ (Attn ans, 1-ans,
Attn 1-pped
Rl
2. for the prediction 8'2 T r E Ay (rggwvod (notice that in this case Tgr-wo Tgs va) cannot be activated, moreover, we have
= -(Attm (t)
anst 1-peed- + Attn I-pred Attn 1-ans
2B + 0(0).
3. for the prediction ((9)(1))), r (()) have
(w)) (ramn cannot be activated, moreover,
A (A)
6.7( (11)))or
Attn ans, 1-tans Attn (1)
ans,1 1-pred,2)
2B
+ 0()).
ny
112

---

4. for other j E =()), if there are and y, s.t,, 91.92 92 E Fiber,,y (notice that y # Yo, y1 for
E2h, then for such j. 2 (Tazu) cannot be activated, moreover, have
(8) b
Attn ans, Attn ans.sss. 2B + +0 ( d.
B O()):
else, none of r 2i can be activated.
Lemma H.15. If Induction H.1 holds for all iterations < t, then given Z Es, have
1. for the prediction j2. E 21 \ (rorn cannot be activated, moreover, have
Attn (t) ans,1 2B + O(6).
(vi) (rrm cannot be activated, moreover, we 2. for the prediction T (( (1n))) r (vr)) have
A (o (11))) Attn (t) 1-ppedd 2B + 0(6).
3. for other j - g(y1). where 91,92 & Fiber we have
ans, l Attn (t) vI (Attn
(0)
(2)
1-p.ed-
2B
fl
+
(d-
d
B
+
moreover, if there exists s.l, 91 E Fiber jy (notice that $ y1 for Ex), we have
(8) Saji (Attn (1)
Attn 1-ans,) 2B + Ō
+ 0(8);
1-ppee2
d. (d.
besides, other r E Ay cannot be activated.
Lemma H.16. If Induction H.l holds for all iterations < t, then given Z E Es, we have
1. for the prediction j2, r E Aip (rowm cannot be activated, moreover, we have
JaJa-aa Attn (1) 1-ppeed2 2B 0
2. for the prediction r(92(1))), r Ar(as(o))) fr cannot be activated, moreover, we have
3.2 Attn (8) 1-peed2 + Attn (1) 1ns,e-- Attn
11
(1) ans, -anns. L-s,n,, 2B 0(8).
3. for other j E T()), if there exist j and s.l. 91,92 92 E Fiber then for such j.
a, (Fare} cannot be activated, moreover, have
2B + 0 ( ( 0(8);
(2) S.j.rov
(Attn (!) I
Attn (1) ans,
d. d.
else, none of E Ag can be activated.
Lemma H.17. If Induction H.l holds for all iterations < t, then given Z E Es, we have
1. for the prediction j2. E 2i (rawm cannot be activated, moreover, we have
Attn 1ns,1--pred.2 2B + 0(8).
6.ja.ramm
2. for the prediction T(92()0)), r ( (ro2 cannot be activated, moreover, we have
A (8)
Attn (1)
1-spee2 Attn (4)
1-an1 2B + δ B +0())
5.52. 20
d Buy
113

---

3. for the prediction r(gi(y0)), r (02(2())))) fr cannot be activated, moreover, we have
A 1, (Attn 1
Attn (e)
ans, 1-aanss)) -2B
(Attnn (c)
anttmass1I-ansp
Attn (1) ans,1 (ppedd)
2B
+
n
4. for other E T()), if there are and s.t., 91, 92 Fiber,y jy then for such a j. r A (oruut cannot be activated, moreover, we have
(t) S.j.roz-
(Attn (t)
1- Attn anttaanss--ans.]]
2B + Ō ( 0(0);
B
dd .d.ngg
else, none of By can be activated.
Lemma H.18. If Induction H.1 holds for all iterations < I, then given Z E & have
1. for the prediction j2. 21iz (Farvi) cannot be activated, moreover, have
(1)
(Attan (0)
1-ppee2 Attn (I)
2B + 0()).
2. for other j = g(v1), where & Fiber,, we have
G = =(Attn (e)
lans, Attn ans, pprdd
2B
+
0
B
+ 0(0);
d. nly
moreover, if there exists y + U1, s.l. 91, 92 E Fiber,,, we huve
(t) = (Attt 1-pprdd Ba
(t)
Attn l-an ) 2B 1
B
0(8):
d d.
besides, other Mj cannot be activated.
H.3.2
Gradient Lemma
Lemma H.19. If Induction H.1 holds for all iterations < t, given E T((),, we have
[- V [
Loss?
2.27
+
V
Loss?
2,2
Q
5,8
5,5
min
22
d
B
(1 logit
logit 5,
5.12) T = =.,5 >0.
}
d. nl
Proof. By Lemma H.10 and Lemma H.11, we have
V Qu.
Loss? 2.21 ]..- + +[- VQ,, Los,,
...
B[Ns,25 + +N + NN.2.1). Based ommmemma emm H.18, we can first directly bound the term An0
4,2.4 as follows:
E (t)
< Ō
1
8,2,5
polyd
In the following discussion, focus Ar(1)
2,2, and consider two regimes for the 1-spred22 Attn 1 ( = s] s S. and (ii) Attn ans, -ans.0 T(x1)
8,2, and Anl
gradient lower bound: (i) |[Att (1)
E Atttnns..-pper
(e)
n. The proof strategy is analyze A
and (t) 4,2,8 under different event conditions in two regimes.
1. For regime (i),
114

---

(a) For Z E Ei, by Induction H.1 and Lemma H.13, we have an immediate logit upper bound for j T()): logit (0) S ( Then by Lemma H010, we have
((t 4,2, 1c
E
(1 logits.
Σ
SReLU' SReLU'(As5.. (A5.)
r
Attn
Varr((10) (16)
Attn (1)
ans, -opred, (91)
ans, ans, -ans0
Attn ams,n-
Attn p-pred) ) + -(00)) + Ō(89) (((e))=s 1e
By Lemma H.13, obtain
B (1 (6)
(sRCLU' 5..2 3.22."02
Attn (1) 1-ann.0 Visroa (Vb)
Attn (1)
ans, I-pped,
Varroz (91)
(1 Attn (9 Attn (0)
(1(0)
+ +O(oo)) + +((89) (=1)==)
l-predd 5.J0i
B
0 E RReLU (A 6. (T(11) = s
(109)
On the other hand,
(1 - logit! sReLU'(A (A
Attn Attn aan,0 Vi (20) + Attn (8)
I-pped, Vi2i701 (01)
[-(2))= Ll, 1l (x0)=
2 -0 d. nly
E sReLU' (A
6.2i (n ([(()))=00) =,
which implies that E No.
1e,
2
B
SReLU' do 17((1))
&
Ar(1) O. Moving to 2,2,00 by Lemma H.10 and Lemma H.13, we have
Aett)
1
(1) logit
(Σ (AcLL"((ss..r))
E
5.j"
((())T())
r
(ttt
1-aas Vy.((u)) + Attn ans. 1-pred, Vgr(
(1 Attn (
Attn (!)
ans, sprad, A + Ō(oo)) + Ō((7) (e)))))
1-a.1-
2 - ny
d2
(b) For Z E2, by Induction H.l and Lemma H.14, we can also obtain crude bound of logit: logit!)) < O(2) for j + j/- Then (00))
4,2,3 can be bounded in the same way as
115

---

(109), and we have
8,2,8 Es 2 20 d. d E
(RCLU'AA J. )/((21) = 8, 8,82 20.
Moving to (v
8,2,50
for T(9 (v1)), we have
logit (t) S.j
sReLU' (As.sromm)) ).
(0) Attn ans I--ans. Viran (2o) Attn ans,1-ppar +pred,1 Viron in (91)
(1 Attn (0)
Attn (0)
a 1 + -(00)) + +((88) 1 (0))==
B ny .2)
where the last inequality due to the cancellation of the term Attn (4)
1
Visan (yo) Attn (t)
1-spred Visrn v (g)) and the fact that
(,,((1101))) (ttnn 1 Attn 1
(1)
2B
+ 0(0)) 0
ny
T4y
for j ((2(()) if there exists & yo, y1 s.t., 91 = 92(U)
(e) -when E8 Attnnt ann.1- Attm ans, 1- T(x1)
s] S 0.01, by Induc-
tion H.l and Lemma H.14, we have logit logit 5.7(9((1)) < 0( llogit (t)
5.3% and
logit (0)
5.44 -(1). Then
E
logit
(1) 5.j
(sReLU' (
Attn (0)
>ans,0 Vi.ran (Mo) + Attn (1)
ans, ppred.1 Vir
(1
Attn (4)
1-aans attann,,
Attn (t)
ans, 1-edd. Andrro + =o(0o) + =Ō(8)) 1-
B
0
By d²
On the other hand,
logit E
( (o
5.
Attnn ans,
(t)
Vigeron vo (vo) + Attn (t)
Vig vo (9))
1-pppd
(1 Attn (0)
Attn (1)
ans, 1ppod,, 5o - +((ov))
A (1)
1
+ +((89) 11 (r(e))=) 16
B 2 2 ny d
where the last inequality follows from the fact that
Attn (t)
ans.1-+ans.0 ViS vo (20) + Attn (t)
ans, 1 Viger (91)
(1 Attn (2)
1-nnns Attn (2)
ans, 2 A 5.16-
116

---

Attn
ans ppred,1 (2B
24
0 Q(B).
01220011000
when [[Attn ans aans,
Attn (t)
T(21) s) 2 0.01, by
1-pee-2
Lemma H.14, we have A A
(8)
cannot be activated. Furthermore,
logit (f) 6.25 (ReLLU (1)
Attn (e) ans, 1-aan,0 Vis v0 (yo) Attn (8) ans, -pred,1 Visro vo (91)
(1
Attn 1-s.1 Attn (2)
ans, ans ppddd vo +((00) +((8)) (
20.
Putting the above discussion together, we have
A 2,2,i E 2
-0
B n
da
(c) For Z E Eg, by Induction H.1 and Lemma H.15, we can first have some facts of logit:
logit 5.jo = 211), logit (4) (1))))()))) = 2(1)
(t) S logit!
polyd for other j.
1
Therefore, we have
E
Nett)
(t) 2
1
E
(1
logit
Attn (t)
ans, Vissaa (21) Attn ant, 1-pped V Vo (91)
(t)
(1
Attn (t)
Attn (t)
prdd 2 2 Po
- =Ō(o)) + =0(89) (((1) 1
ans, 1-rans,1
B
My
(r(e) Moving to 2,2,00 we have
vr (1)
E
&
22,
2 logit (t) (x)))
(e) Attn - Vrign 11 (v1) + Attn (1)
I Vi
(1 Attn (1)
Attn (1)
ans, -pprod ) (t)
+ +((0))
5. (.r(0r(0))))
+ (((8)) 1 -
- -0 ny B
d. dd-polyd
2 52
ny -0 d. polyd
B
My
(d) For 2 E4, by comparing Lemma H. with Lemma H.14, we can directly bound Arl 2,3 in the same way as &s, where the only difference is that we do not need to
117

---

consider T(9 (31)) for Eq. Thus No 2 Len
0. Moreover, it is clear that
Nitt) 8,2 De,
(1 (llogitt (t)
(sReLU bubiii "02
Attn (+) ans, 1-ans.0 Vi (vo)
Attn ans,1-ppedd Vi (91)
+ (1 Attn a-ass- Attn (1)
A Jori
+ =0(oo)) + Ō(09) 1l [((er)= 1e
1-pred
0.
where the last inequality is due the cancellation of Attn (9)
a Visiroaiv
Attn (t)
ans, ans,1-ppee ppred,1 Via.Nou. (91). and the fact that
(1) 2.
Attn (e)
pred, 2B + 0(0) 2 Q(B).
(e) For 2 Es U &g. by Induction H.1, Lemma H.17 and Lemma H.18, we can derive
the following logit condition: 1 logit 5.j2 logit (t)
(1)
5.j polyd for j2- Hence,
can simply bound Ar(t) and Ar()) as follows:
s.2.68
1.. IE
[
||
B diy
1
E
S
polyd)
for E (5,
Putting it all together, have for the regmie (1),
V 0
Loss
2.2
+ [[- V
Loss 2,2
Σ E
( 2.2
2
02
B nl
(E(,,,))))) (E[6]
2. In regime (ii), the analysis follows structure analogous that of regime (i). The primary (4)
distinction lies in the fact that, under the main event &,, the term 5.j2.r92291 guaranteed remain within the linear regime, thereby serving as the dominant component driving the overall gradient.
(a) For 2 E by Induction H.1 and Lemma H.13, we have
A (t)
2 181
E
(1 logit
Σ
sReLU' (A
r
Attn 1ns,
(12r(
Attn (t) ans, --pprdd V1or((1)
(1 Attn Att
Attn att 2) ) + Ö((o)) + =(87) [-(=))=
By Lemma H.13, we can obtain
(1 logit b.22 ((
Attn ans,1- Vinsor (vo) Attn (1)
ans, ans, Ispred, Viaarsu (g))
+ (1 Attn
Attn a spred, A 5.j2. + +0(on) + +((09) 11 ( 18
(110)
B ( E (1 logit (1) (0 = 3
118

---

On the other hand,
(1 - logit 6.12 SReLU'AA (t)
(A )
Attn lanso Vi no (30) + Attn (8)
1-pedd Viairo (01))
[=(x))= Le (vo) (20)) o2(v))
B
2 --0
)
E (1 logit
5,22 [T(1)) [=x)) -
d nl
which implies that IE
A Ne
2 Ω ()) B (1 B
logit ( = 2,E1 20 0.
Moving to Ar(1) 6,2,58 by Lemma H.13, we have
for j (() (vo)), rgr
(ttn (t)
ans, +ans,0 (30( + Attn (t)
1 Vgr((11)
- -(1 Attn (e)
-1--Attn ans,
Attn a ppred,)) A + -(00)) 2 0.
for j = (111) r = Trr due to the cancellation of Attn a-s,1-
(yb(( + Attn (1)
1-peedd (91( we have
(1) Attn ans, 1-aans0 Vg,r (36) Attn (1) ans, Ippred, Vgr(
(1 Attn (+) I-sans,
Attn (1)
ans pred A + Ō(oo)
(t) -(1 - Attn ans, I
Attn ans,1 ppred 5.g.
for j =(92(80)), r9g vo
(Attn ans, -ans,, Vg + Attn 1-pped (g,(
- (1- Attn ans.
Attn ans, pred, ) (1) + Ö(on)
2 (1 Attn (1)
ans, 1-ann.1 Attn (2)
(1)
ans, .prdd S.j.r
for j = 9i (v), where By + No MI. s.t., 22 (v) = 91 (4), we have
Attn ans, 1-ans,0 Vy.r + Attn ans, V9.((11)
A (1 Attn ans, 1-ans.1 Attn ans, pred, ) A + Ō(oo)
2 - Attn L-ans.. Attn ans p A600..
Putting them together, and upper bound Ejthero logit!") by 1 - logit (1)
5,52 we
have
(
8,2,1
1
E
(1) logit!)
Σ (AG..r)
((())
r
(Attn
Vy.((0) + Attn ans, 1 predd Vir(gi)
- (1 Attn 1-ans,
Attn I -
1 2 A + Ō(on)) + Ō((3)) (
119

---

E
(1-llgit
5.10) (1 Attn (0)
1-ans,
Attn (8)
1
(A (,7 (un Dir b,r (p2((()))) (
max v7
where the last inequality due to the fact that A (4) (20
(o +
Ō d-n.
for + yo. U1 s.t., g1 (v) = 92(y). Notice that
max { {1!0 11
( xn (.2()2()))) }
Attn 3ns,1-+ans,1 O 1 2B,
max
Attn ans, 1 spred,2
while
(110) E (1 logit (0)
.tt 1-mm 2B (x1)==0
(+)
5.02
E (1 logits. (1)
(1 Attn (0)
ans, 1ans1 Attn 1pppd. ((x))=s
(t)
5.j2
Since by Induction H.I, Attn ons,1-+pred,
Attn
ans, tans, Si << 1. thus we have
E
(0) 8,2
LE,
LE,
E
Ar0o
Ie,
2
d
B
(1
/ = =
(1)
(b) For 2 E E2,
N()
can be bounded in the same way as (110), and we have ,,2,
E
An 8,2,3
( &a
B d. nly
E
(1) ) [()) s, &g 0. ?.
Moving (1()
s,2,ii*
for (gi (81) )), we have
E logit (1) 5.j
sRcLU' RRRCLU((Assipgaa (As.j. )-
Attn 1-s.1--ans.p
(()
Virerr (vo) + Attn (1)
ans, 1ppedd Vi (91)
(111)
(1
Attn (0)
ans. 1-5an-
Attn (1)
ant 2) A + Ō((o)) + =0(0%) [-
(1)
B EE (1
6.j2 T logit
(8) logit
s, 8,82
where the last inequality due the cancellation of the term Attn (t)
ant ams_1->ans.0
Virrn (yo) Attn ans, ppedd Viraa. (gi) and the fact that
(t)
6,76)1 = Attn ans.1
Attn ans, 1-ped-) ny
2B + 0(0) < 0
nly
for 1/2 = ((9((8))) = 7(pn (vo)), clearly, we have
E
logits)) 3.35
(sReLU' (A()) 5.56.
(ttt
ans, 1- Visron (1o) + Attn I I
(1 - Attn (t)
Attn (t)
(e)
+ =((o)) + Ō(89)
a -anns
a 5.5h
120

---

1 7 s 2 0.
where the last inequality is due to the fact that
Attn 1-ped_1 ViS ir
(01) (1
Attn l-ans.
Attn (8)
ans, A (1)
Attn (t)
ans, 1predd
2 Attn (t)
ans, + Attn (t)
)
Attn (4)
1-anass ) 2
ans,1
2 Attn ans,1-ppred,1 (1- (1-2(c1+0.55) 2B 20. 0.
for j = (92(()) if there exists y * yo, 81 s.t., 9n g1(v) = 92(
when E Attn ans,1-+ans.. Attn (t)
((x1) 8) S 0.01, by Induc-
ans, 1-
tion H.l and Lemma H.14, we have logit! (t)
(,r(0x(v) < 0 logit 5.13
Then
logit 5.j
(SReLU 5.j.vn
Attn (f)
ans, +ans, VisTay (Mu) + Attn 1ns,1-+pred.1 Viron
(1 Attn I-s,
Attn 1- Aodror + =((0o)) + +((88) (((0))=oll l
o ) (111).
when EE [E[Attn (1)
I-ans, Attn I
(t)
T(x1) s) 2 0.01, by
Lemma H414, we have A
S.j.r cannot be activated.
Putting the above discussion together, we have
E
(t) 2.2
Le,
+ E
[ (0 8,2
1e
Ω
B
E (1 logit b. (T())) s, S,
(1)
d ml
(c) For Z E Eg, by Induction H.1 and Lemma H.15, we can first have the following logit bound:
logit!
5.j S
polyd (1 (1- logit!!) 5.92 for j + f2. (9)
Therefore, have
(t) 1
E
E
(1
logit
(8)
Attn (t)
a Visi (v1)
Attn ans, Visiron (91)
+
(1 Attn 1-ans_
Attn ans, 1- 3.f2.ro2m
11 (1)
+ + +0(8)) ( 18
Ω
B
R (1
logit (8)
o. ( = E3
My
Moving to A/(()
8,2,1 we have
(0
8,2.8 1sa
121

---

2 logit (t)
(an
Attn (4) ans, Sans, Vilan(m))) (11) + att (o(a(())))))))))
(2)
(1- Attn ans, -anns,
Attn (t) ans, 1pped,)) A A. (,0(0:(01)))rr + + +
((22)== E
0 My d polyd
ny
E (1
logit () = 8,83
Thus,
E
S 8,2,5 Es
+ E2
Ml 2,2,
1sg
52
B
IE (1 (1- logit 51)) s, s,Es
Thy
(d) For Z E Es, we can bound the gradient a manner similar to regime (i), and obtain
E
+ E
Ad $,2,
E
2
(c) For Z E & UE, by Induction H.I, Lemma H.17 and Lemma H.18, can derive the following logit condition: - logit 5,j2
logit 5.j
S
polyd E (1 logit (.. =
for j # j2. Hence, we can simply bound Artt) ,,2, and Artt) 9,2, as follows:
E
E
Ail
cm s,2,4
E
(t) s,2,ic
0
B dny
E
(1
logit (1)
5. ( s,E, for momme (5.6).
polyd
Putting everything together, have for the regmie (ii).
Vol
V
Loss
2.2
+
V
(e)
Loss
2.2
E
2,2, 18
8,8
6,8
ic/6]
2 N
E (1 logit (8)
6.2o ( = 8,
Lemma H.20. Induction H. holds for all iterations < t, given 8 E for (Qtple, (3,4). # E TX)), we have
[-
V Q.0
Loss
2.2
S 5o 12) Q Loss
H.3.3 Bounded Decrease of Attention to Related Context Clause
Lemma H.21, If Induction H.1 holds for all iterations < t,, then for any sample Z we have
Attn (8) 1-spr-1 Attn (4) ans,1-+pred,2 > -0( log d log
log d
Proof. Denote the first time that E Attn (0)
1np_e--
Attn ans, sprdd (((1)
S
(Q(t)
-Ω2 log d'log d) logd
as T. Notice that [
ma d
for E (3, 4). thus for any
0.0"
122

---

sample z satisfying T(21) = 8, at time T, we have
( Attn p-ped-- Attn ans, S -52 log d dog d log d
(
Based on the the gradient compositions from Lemma H.11, we have
VQu. Loss, 2.2
E [[N s,3,2 + Na,32288 + NN,,,,,03 and we will discuss N,,,,2, for A. samples Z Following the similar argument as Lemma H.19, we can first directly bound the term Nx 5,3,200 as follows:
{i, (i, idi} on different
(T)
S Ō = 1
E
8,3.2
polyd
for Z E εg, at time T, since the neurons for predicting j2 cannot be activated, and (T)
logit... 5.j
S ( for # f2, thus we can naively bound the gradient on the event & as
follows:
Σ
E[ 8.3 lel so Bl + 088).
(1.3)
for z E Eg, similarly, S.je.r is not activated, and thus we can focus on the term 3,,3,2 and specifically, the prediction of j/ (gg(U6))) By (98) and Lemma H.14, we have
A (T)
Attn (T)
E
1e
2 -F
5,3,2,0
ans, 2-rred2
Attn 1-ans,
Attn (T) ans, prod,
Attn ans, I-ppod, 1 =
(T) logit
5.3 ((x1) = 8, E2
2B ny
O((g).
(112)
ATT) for z E Eg, by Lemma H.15, we can mainly focus on the term 2,3,2,5 and the prediction of (()(1))) in (T)
8,3,2,5 since logits, (,7()1(11) - O ( Hence, we have
E
8,3,2,8 +
E Attn 1
(113)
$33,2,58
(1 Attn ans, 1-peed-
Attn ans, 1-pre_) |T(xI)= = (1 (1-0(10g)) 1 2B Tog Ny d + O
for 2 E E4, the negative gradient can be bounded in the same way as (112), however, the probability of & is order-wise smaller than & and &g, which can be neglected. Moreover, for for z E &g U Eg, the overall gradient is also negeligble since logit 5.j2
(t) is very close to
1.
Putting it all together, we have
Loss
222
2 E
Attn ans 1-aped222
+ Attn ans, I
Attn (7)
Attn (7)
ans, T(x1) =
ans.1
B
E
logit 5.32
[((1) = E
ny d
+ E Attn (T)
ans,1
123

---

(1+ Attn (T)
1pp Attn (T)
1nd. ((x1) = (1 (1-0( 0
2B ny
= 1o)))
log d
Notice that
1 + Attn 1-ppedd
Attn ams, 1-pred_
Attn 1-ppe- (1+ Attn ans, 1-nss
Attn ans.
2Attn 1-ppeed Attn
ans, 1ans,
If 2Attn (
ans, 1-pee_ Attn ans, I < for some small constant > 0, s.t., log d
cB
< 1, then
A 92
o(1), and (112) is dominated by the positive term (113). Else, clearly, (112) can be cancelled out by the positive term (113). Therefore,
(T)
S cB. Hence logit
532
lox
Loss
2.2
s Fly
B
ri 4.3
8,8
As consequence, together with the growth of Q Lemma H.20, we have Attn I-s larger than ( log logd loglog log d
+ Q and nearly no change of [Q 8, in
Attn
1- must start to decrease and cannot be
H.3.4 Attention gap is small
Lemma H.22. If Induction H.1 holds for all iterations < t, then for any sample 2 have
Attn ans, 1-peed Attn
ans, tans,1 S c1,
where 6j > 0 is small constant.
(e) Proof. Let T denote the first time IE Attn 1-pped Attn (t) ans,1-aans.1 [xx1) =
1.0005 log
2B
Q
for E {3,4), thus for any sample Z satisfying Jain
Notice that
S O
7(x)) 8, time 7, we have
Attn (T)
1.0005 log d
1-pee_ Attn
2B
Based on Lemma H.11, we will discuss the gradients of Q4,3 and Q4,4 on different samples Z
Since (T) P. A (T)
5.jh."w3 w2 is guaranteed to activated to the linear regime for 2 E E we
only need to focus the main event E,.
From Lemma H.13, at time T, we have
23 2B
logit 5.32
((d) (1-
logit
7 5.52
= (1-0(1/080000)) ) logit! (T) 5.52
1.000
Then by Lemma H.11, we can obtain the gradient on the event E as follows:
((T)
633,,t + (T)
6,3,244 + A
2,3,2,811 1e,
E Attn (7)
1-pred= (11 (1-1og
(
(Vinronu (9z) - ( - + =Ō(ov))
(T)
5.j2.
11-0
0 00.0005 )(1
lo (t)
5.j2
(Visror (92) (92)
A
A.Ja =((o)))
124

---

- =((0)) (((o))=
Attn (2) E
(1-0((po00n)) ))
5.32
( (A
(T)
ans, upred,22
3.j2.-oz v
+ +0 [0.0000 ) (1-logt))
logit (.)))
(o..eoroo 2(02) A (8)
(.pe. ))
Br(en==s
on the other hand, we have
E
(wit)) +NTT)
6,4,2, a,1,2, + a,4,,,00)) 1
+N7))
)
Attn (T)
(11- logit
(Vearoarm (11) -
+ =((oo))
ans, 1-pred
Goda v1
(1-0 (o.000s
(1 (1- (1-logits,) logits?)) ) (o (311
+ =((0))
+ Ō(08) 1 (((0)==
Attn ans, -ppeed.2 (11-0( (0.00005 ))
logits.).
5.40
Vil (31) + G32i
(T)
(T)
o( d0.0000) )(1 (1- logit 5.f2 5.22.
(x())==
Since at time T, have Attn (2)
Attn (T)
ppred
Attn ans, -ans,
1ns.
Attn ans, 1
(T)
1.005 logd 2B
we have 5.li Va2
< -Vghro (11). Therefore,
Loss
+
Loss
2
E - logit 5,j2 [(1)) =
which implies that [Q1,1l,s will grow faster than [Q1,3]a and thus the attention gap cannot be further increased.
H.3.5 At the End of Stage 1.2.1
(0.1
Lemma H.23. For all iterations & T1,2,1, = O nyd llog
+ 0
we have Induc-
log
tion H.l holds, and time T2,,,1,s have
(a) 1 (1
[ (T..))))
15. 2 I(1): 1,1
1.3 1.50
(b) other [[4, b for (3,4), s' T(x)
most Ö(3)-
Proof. The existence of 1. can be directly obtained by using Lemmas H.19 to H.22. Furthermore, [Q44 (Ti221))
S.5 2 21(1 can be guaranteed since Lemma H.22 implies that Attn ans, a-rans, 2 0.3 e +0(1/d) > 0.2, which means that 14,4 (T1....)) s.. should at least grow to a constant level compared
(Ti,2,1.))
2
11 (T.a..)) Ja,er1 ((1/d) with p E (3, 4).
We will handle (Q (T.2....) ls.s by means of proof by contradiction. Suppose that Attn (TI.2...) ans, 1->ans,
0.4 - C. where ī
logd is sufficiently small constant. Then denote T the first time that 8/7
125

---

[
(t)
8
0.4 22. Notice that [
(1)
0
for
8,0
P E {3,4]. thus for any sample z satisfying T(xi) = 8, at time T, we have
Attn ans, 1 2 0.4 - 27 + O(1/d)
If Attn ans, p Attn ans,0 2 2g, then
- for 2 E E1, A 202 has already been activated to the linear regime. Furthermore,
Attn (T)
ans ppred,2 Attn ans, a ans,0 - 0.2 + 2č - 0.2 I 22, which implies 1
logit 5.
=1 - 011) Thus, by Lemma H.11, we have
(NIT)
+AX (T 8,3,2 2,3, 8,3,2,365
B d
E
11
2 22
while
(N
+MII)
B
R
8,4.,1
+ 8,422
2,4,2,06
S <-50
d
for 2 Ez, 53 = (Attn
1-pedd + Attn ans
Attn Lan,, .)2B < (0.4+ 20 - 0.4
5 lng
Thus logit (1)
0 2
Hence by Lemma H.ll, we have
(v 6,3,2, + +N7))
+N(T)
2 200
B
2,3,2,8 +
18
833,2600
dd:/0Tg
while
(Not2 + -VT +
S -
B
2,2,66 +
s,42,88
16
nyd
Z E Eg, use the following naive bounds for p (3,4]:
(M
ArTT)
AxT)
s 0 B
IF
s.p.2. + s.p.
8.p.2,08)
dry
Putting them together, combining with the fact that the gradient contributed by &g U Eg U is negligible, we can conclude that
Loss
S2
B d
Loss 2.2
- 0 d
which implics Attn ans, +ans,1 cannot further increase above 0.4 2T.
If Attn ( ans.1--pred,2 Attn 1 < 2g, we shift our focus to the comparison between
event &g Es.
- for 2 E E1, we have
E
(Nr(T)
8,3,2, + +N7))
+NT)
1s,
-
9,3,2,8 2,3,2,10
while
E
(Mr())
9,1,2, + ()
8,4,2,8 +N(()
2,4.2,50
$ zō(og).
- for Z E Eg, similarly as previous case, we have
(N1(T)
8.3,2, + +
2,3,,, + +
( (
B
6,3,
18,
2
126

---

while
(N107)
8,1,2,t + +17))
+ +N(TT)
S 0
E
8,1,2,8
0,1,2,606
113
di d2/00
Z E Ez.
(7)
-A6
- 4ppB= = =o()), hence, we have
E
8,3,244 + +
2,3,2,6 + +
+,3,,,2,001 Lss
2 l
B dryy
while
EB
(N
9,4,2
+N(T)
8,2,,8
A
2,4,2,00
1sg
00
B log log d dry log d
Putting them together, again we can conclude that
Loss
2.2
2
2,2 Loss
I 0 B log log d dry log d
8,8 ny
which implies Attn
cannot further increase above 0.4 2T.
ans anns,
Consequently, this leads to a contradiction, and Attn (.1.). ans, -anns, < 0.4 - T, where č == log
Then it
8B
would follow that Attn (T1.
pred,2 2 č, and thus [Q]
(T))2
I 2 ((1).
1.3
H.4 Stage 1.2.2: Convergence with Small Wrong Attention
Recall that Loss EE log Pr aZ..,225 (2.m)
for T()).
Induction H.2. Given E T(X), let T2,22,s denote the first time that Loss decreases below
0(c1-3+301a)2n) For all iterations T22,1, St < T1,2, we have the following holds (a)
[
(1)
monotonically increases;
8,8
(b) for e (3, 4), for j E T(x) s. [[Q S S0(
(c) for any sample 2 have Attn ans, ppredd2 small constant C1
Attn ans, 1-nans,1 for some sufficiently
1.005 log > 0;
B
(d) for any z have Attn ant, Attn (0)
1-ppp- min Attn (t)
1
22 ) where e2 4B
logd
0 some sufficiently small constant.
H.4.1 Attention and Lambda Preliminaries
Lemma H.24. Induction H.2 holds for all iterations (Ti,,, then for any sample 2 we have
1. Attn (1) ans, -pred + Attn (2)
4 0.4, (t))
ans, 1-ans. E
(8) 2. Attn predd22 Attn
ans, - 2 2(1);
3. Attn p-pred Attn (t)
1 < (2).
Proof. In the following, we focus on the main events E,, Eg, and Eg, which correspond to cases where confused wrong predictions occur. We denote
Attn (t) ans, a 1pred Attn (0) ans.1-+ans,1 c+ Ö
If Attn (1)
ans, 2 2 Attn 1 then
127

---

If 1 E E1, we have
logpr (Zan,Z,,,225 (2.1))
< (1) max (Attm Attn (0) -m..)) )B (Attn (1) Attn (0)
1m )
Jog (Aten (1)
Attm (1) sars. .o.
S θ 201 (Aten (t)
Attn! von.0)
1prod.2
(=2(2-20m)))))
< θ
where the last inequality follows from the fact that Attn 1
p-op 2 211 - -2c)
If z E E,, we have
log Pr (Zas, 2. (2.1)) - O(1)
e (Attmii (0) 1-ppod +Attn (1)
-ppedd) Attm (!)
1mm. . -Atttn (5)
get ppo.. AAtt (0)
1-mm000 2B
max
log d- (Aten (1)
1ond Attn (8)
sars. )B
2eotc))BB B- (Aten (c)
-pass Attn (t) sans bne.o)28 )2B
4
2 ((=-30-a)
< θ
If Z E Eg, we have
log PF ( 2,5 Z (2)1))
(tttn 1ppddd -Attn!!) -ppd, dog d- Attn -ppes 2B
$ (1) max
2( aa -)B e 212 22)
< (1)) max
If Attn (t)
ans, p-rdd S Attn (9) l-ans, then
- Z E E1, we have
log (Zas, ((.1)) < O11)
log (Aeem(
(ct (1)
< +
e 2r B (Attn
Attn 2B
Aetn
2(
9 2 2 20-2--10
< + è
where the last inequality uses the fact that Attn (1) ans,1 tans,
Attn (t)
1- S C-c2
implying Attn (f)
ans, 1->pred 2 3 g.
If 2 E Ez, we have
log Pr PwpprrrZZaaa.... (Zass. (((.1) < O(1)
128

---

Attm ma. -pprd22 AAttn pmat, -Attn an, 1-sne, )BB (Atta (6)
m. 1ppd22 AAtt 1-tare,0) )B
max
e Jon d- (Attm
Attn (1)
2nd, -aa0
min1 20rol.2
(2xx(B)))2B- (on](ca)B--AAttm... 22ro1.2 Attn am, -
+
(2 20- (axx2-mxx((rs)) θ
If ZZ E E3, we have
Po (Zas. Z (2(21))
e (Attn mn.. 1prod.1 Attn (1)
mo.1 12 Augd-Attn?2. 2B
< (1 max
S (1) max
21 -0- -112-200 -0)B
2(
Putting all cases together, if S ge,, we must have
2.2 2
2 e-22222
20 toam ) )
Loss
which contradicts the definition of 7,,2,2,3 Therefore, it must hold that c fo for all / < 2,,,,
Lemma H.25. If Induction H.2 holds for all iterations 2T t). then given Z E
1. for the prediction j2. we have
(8) N. - -Atttm 1-predd Attn (
(0)
2B 0 My
B +0()).
Attmaas.1--ass.u)
2. for the prediction J/2= (=((22(U))) we have
= -(Attn (8) A (
ans, ppeed22 Attn (2) ans,1-+ans.1
2B Ō
B
+ 0(8).
nyy
3. for the prediction rgg (un)), Toi to cannot be activated.
4. for the prediction T(g: (y1)), we have
(1)
Attn (e)
Attn (t) 1ns,
2B ny
+ + +
B
1-aans1
5. for other j E r(Y), if there are and y, s.t., 91, 92 E FFiber,,v (notice that y + yo, y1 for E1). then for such j, 2, (Tory) cannot be activated, moreover, have
(2) S.j.ror
(ttn (1)
Attn (1) ant 2B + +0( idd.ng O(8);
ans.
else, none of E By can be activated.
Notice that for ε,, the neurons mentioned in 1-4 should be different neurons, while the predictions in and 3, or and 4 can be the same in some cases, e.g- 91 (y1) = g2(yo). all cases, except for for the neurons mentioned above, i.e., Ue (o geevr_ (which may not be activated), all other neurons
r E U1,e (Br(ge-pr>1) \ (rgeve-1) cannot be activated.
Lemma H.26. If Induction H.2 holds for all iterations E [Ti,2,1, T2,22,,), then given Z E Eg. we have
129

---

1. for the prediction j2, r 21 (razz11 cannot be activated, moreover, we have
(1)
= -(Attn (+)
1ns,1-ppreds Attn (2) ant..-ss.. 2B + Ō + +
2. for the prediction 1/2 = Zg frr (notice that in this case gz Tg gyy) cannot be activated, moreover, we have
(Attm ans. + Attn 1ps,e--pred,1 Attn
2B 0
5. 20
ans, -anns
3. for the prediction ((91(1))), r (a))) fromm } cannot be activated, moreover, we have
A (4) (ar (11)))r
(Attn
(8) ans, bans, Attn (0)
2B + 0()).
ans,1-+pred.2) nly
4. for other j E r)"), if there are and y, s.t., 91,92 E Fiberj,y (notice that y + you y1 for E2), then for such a j. M, (Tazy) cannot be activated, moreover, we have
(8) S.j.aaz-
Attn (t)
anttnans,1-ppae 1-pped-
Attn ans, 2B + 5 B d. d..
O(8):
else, none of 2, can be activated.
Lemma H.27. If Induction H.2 holds for all iterations [T1,2,1,a, T,,,,,,),, then given 2 E Es. we have
1. for the prediction j2, E Az (rggze1 } cannot be activated, moreover, we have
Attn 1-ppee- 2B 2B 06)).
(v1)) froim cannot be activated, moreover, we 2. for the prediction (1)(1)))), r ((r)) have
A (, Attn (4)
ans, 1-eedd 2B + 0(0).
3. for other j = g(y1), where 91.92 & Fiber we have
G.d.ov = (Attn (Attnansss-aa ans.1
Attn pred,
2B
+ Ō
B
+ 0(5);
ny
d ..y
moreover, if there exists s.l. 91.92 E Fiber,, j.y (notice that + for Ex), we have
Ao = (Attn
(t) 1ns (Attman.1--pe222
Attn 1-s,n,) )) 2B + 0
( d.
B
+ 0(8);
besides, other E Ay cannot be activated.
Lemma H.28 If Induction H.2 holds for all iterations E [71.2,1.; T2,2,2s), then given Z E Es. we have
1. for the prediction j2. E Uz (Tar09) cannot be activated, moreover, have
A (1)
Attn (t)
1nppe-5 2B +
2. for the prediction (92(20))) r Al(o(()))) (oo2 cannot be activated, moreover, have
A 3.
= Attn (t)
1 + (ttn (t)
ans, I-sped, - Attn 1amm)))) -2B+000
130

---

3. for other E T()), if there are j and s.t., 91.92 92 E Fiber,,y then for such a j. r E
Aj fnaant cannot be activated, moreover, we have
(8) Gditag
(Attn ans, anpee-- Attn 1-s...s))))) 2B + +0 ( +
B dd
+0())
else, none of Bg can be activated.
Lemma H.29. If Induction H.2 holds for all iterations (T,,,,,, T1,2,,s),, then given Z E Es. we have
1. for the prediction j2. E A1 \ (ra2-11 cannot be activated, moreover we have
(1 GJ2iraze Attn (1) 1-ppee+ 2B 0(6).
2. for the prediction (g2((30)), r A((xx(1))) Uor cannot be activated, moreover, we have
(1)
(Attn (t)
ans, spredd2 Attn (8) ans.1.s 2B + Ō B +0()).
=
d Byy
3. for the prediction (gg (yo)), none ofr A((gs(o))) can be activated.
4. for other T(V), if there are and y, s.t,, 91s 92 Fiber, then for such a j. r
A (rozvt cannot be activated, moreover, we have
(4) A A.j =
(t) a-tteass1-pedd2 1-ppeed Attn ans, 1 -sans,) 2B + Ō B - +0(8);
dd
else, none ofr M, can be activated.
Lemma H.30. If Induction H.2 holds for all iterations E (T,,,,,,: 1,2,2,s), then given Z E Eg.
we have
1. for the prediction f2. E M2z (riz-m) cannot be activated, moreover, have
(1
(1)
(Attn (1)
ans. 1-peed2 Attn antpass1-- ) 2B + 0()).
3.j22 v1
2. for other j == g(yn), where 92 & Fiberjan we have Tg cannot be activated, moreover, if there exists y # y1, s.l., 91, 92 Fiber we have
(t) A.J.-a
(Attn (t) ans, 1 1-red,22
Attn (0) Attn aas.. 2B + Ō B ) +0();
dl d.
besides, other Ag cannot be activated.
H.4.2 Gradient Lemma
Lemma H.31. If Induction H.2 holds for all iterations t [T1,2,1 T1222,), given T(X'), if
Attn (e) ans, 2 (T(1)) s] we have
2.2 V (2) Loss?
+ [ V (1) Loss? 2.2
8.8
5,5
Ω
E (1 logit (1)
= [xx1)
= s, E
+
B nyd
3
E (1
logit (1) 5,12
(22)) (((1))
) = s, Em
5,12
m=
Proof. The analysis follows the similar idea as Lemma H.19, while Induction H.2 ensures that E [ttnt (t)
1 - Attn ans.aannn T(x1)= s] 2 2(11 -. which falls into the regime (ii). Thus for the main event ε,, the term A (4),, j2, rg2 yi is guaranteed to remain within the linear
regime. We can first directly upper bound the term (t 2,2, is by Ō(o)).
131

---

(a) For 2 E &1, by Induction H.2 and Lemma H.25, we can obtain
E2 N
5,2, Ne
E (1 ---ogitss.)) ((
(t) (AAttn
ans, a-ns V32. (vo)
Attn ans Va vs (0))
+ (1 Attn ans, 1-ans,1
Attn
ans, I I-o 2) 3.22 + =0(o0)) = =Ō(08) 1- 12
(114)
2 d
(1 logit ( = S,
Moving to w/(t) s,2 by Lemma H.25, we have
for j = ((g r vn due to the cancellation of Attm 1-mmam (yn
Attn ans, ppred, (9)( we have
Attn (t)
ans, ans l-nns (yb( + Attn (t)
ans, predd Vgr(
(1 Attn ans 1-ans,1
Attn (1)
ans, pred, 2 A + =((oo))
2 -(1 - Attn (1)
Attn ans 1 Agg.r
(c)
1
for j = (022(10)) == 092
(Attn (t)
ans, Vir((U)) + Attn (8)
I-ped- (g1r
(1 (1-Attn ans, I Attn ans, 1- 2 b.jir + Ō(oo))
2 - Attn (6) ans, l-nns
Attn ans, pped, A.g.r
(t)
for j = 91()), where By + yo, U1. s.t., (2(2) = 9 (v), we have
(Attn (ir(()) + Attn ans, pred,1 Vgi(
(4)
(1 Attn ans, bans,1 Attn ans pred, 2) A Ō(Go))
2 2- Attn (0)
ans, banss
Attn (
ans, 1 pred, A6.,r"
Putting them together, and upper bound Eipierr> logit!!) (1) 5. by 1 logit!!) (t)
lojit
5.j2" we have
E
Lei
8,2,81
logit!
Σ
SReLL'(As.sr)
5,
j+
r
(1) Attn
Vir( + Attn (t)
ans, 1-pred, Vir((11)
ans,
(1 Attn (0)
Attn (0)
- (1(0) + Ō(0o) + Ō(8%) (0)0
1-s,
E
11 logit()) (8)
(1 Attn (
attaams..-sas.11 Attn (0) ans, p-ppdddd))))
(9
max
132

---

where the last inequality is due to the fact that (1) O a
(4) 6.7l0a(0)).
+
for y # you yr s.t., 91 (v) (2) Notice that
max { A (0) 3o.91
of (o
max
Attn (1)
1-p.e--
Attn (8)
1-ans 0
2B,
ny
while
(114) E
1 logit (1)
Attn (1)
ans, +ans,0 2BB (()))sl)
3.jo
E
[1
logit (4)
(1 Attn (8)
1ams1 Attn
(t)
ans, 1-ppod, (()))))
5,32
Since by Induction H.2, Attn
Attn 1-mm S <
1 2.
thus we have
ans, 1-ppred22
10
N 10)
L1
10 +
2
B
E logit (1)
5.32 T s, 8,
2,2
(b) For z E Ng 8,2, can be bounded analogously to (114), yielding
A0U E
(1) 2,2,1
1
2 0 B
E
logit ( [ s, S, 2 0.
d nyy
Moving to No(t)
For j = T91 (y1)), we obtain
E logits.)
sReLU' (AS.LU"A
Attn (t) ans, >ans, Vi (30) Attn 1ns,1-- Viran (91)
(115)
- -(1 - Attn (4)
ans, 1tans, Attn (1)
1ns, A = Ō((o)) = =Ō(07) (5(5)))ITAaa) 1
2
B
E((1
logit
5.jo
) ((x)) = 8,
my d
where the inequality follows from the cancellation of the term
Attn (e)
ans, >ans0 Virra (yo) Attn (t) ans 1-red. Virn m (g1),
together with the fact that
A (1)
AAtn
(0) ans, 1-tans,1
Attn (0)
1tt.nd)).-)prrdd))) )
2B
0(8) s 0 ng
B
S.r(s (vi)).romm
nly
For 1/2 = (92(2n)) = ((91(0)) it is clear that
logit 5.14
sReLU' (A
Attn (t)
ans, Vis (vo) + Attn (t)
Vilii (91)
(1 Attn 1-ans,
Attn
+ +o(oo)) + +0(09) 1 (0))= 1
ams, 1-pedd) G.
0,
133

---

where the last inequality is due to the fact that
Attn ans, 1 1peed Vigiru (91) (1 Attn (9)
Attn
ans, -spred, A 6.j. "61000
ans,1-+ans,1
Attn +predd
1 2 Attn (4)
Attn (t)
(t) Attn
ans, lonas, 2B
ans,1 -preed
ans,1
Attn monss11pp 1-peed
(1)
1 1--2(c1 (a
0.2) (cal+002))
2B20.
For j = ((9((1) with some y * yo: 11 such that 91 = 92(y). we distinguish two
cases:
If
E Attn (E) ans, 1 Attn ans.1 [((1) s]
(t)
4
then by Induction H.2 and Lemma H.26, have we have A, (t)
(
(- 6 + Attn
I )BB 2 (- + fc)22 = ci B, which implics
Σ logit 5,r()((1))
(1) (n) < 5 (1/))) logit (0)
5.34
Consequently,
logit SReLU'(A (1
(ttn ans ann
Visron (vo) + Attn (t)
1-pped- Vi
(1 - Attn 1-s.n-
(1)
Attn (1)
) A.J..ror Ō(0o)) + Ō(89)
ans, 1-pred
(x:)))= 1 1 (0(())=2(())
S 0 (( (115).
Attn ans,1 1 T(x))= = 2 C
Attn
ans, -ans,
4
then by Lemma H.26, the activation
(4)
cannot be triggered.
Putting the above discussion together, have
Mel (ey E 22 Ne, +
+
(M0 s,2. De
B d
E (1-
logit
(9) 5.52 5,52 ((=))) 8.
my
(c) For 2 E Es, the gradient admits an analysis analogous to that in Lemma H.19, In particular, by Induction H.2 and Lemma H.27, first obtain the following logit bound:
logit!
(t) 5.j
<
polyd (I (1-logit.)).)
i j2, (91()))).
Therefore,
(()
E
8,2, 1ea
E
(1 logits()
(- Attn 1-ans- Visira (21)
Attn (8) ans, I-pred, Visiro 1 (01)
+ (1
Attn (t)
Attn (t)
ans. 1-pped) 1 A + Ō(oo) + +((8)) (((0)==
1ns.1-
B
> Ω
E logit 7(x1) = =8, E3
Rg .d)
134

---

Moving to 2,2,60
we have
IE (((t)
Ls
2,2,5
2 logit (1)
aAt 1aan00 Vi(on (01) + Attn (t) 1-ppeed V V (())))))))
(1 Attn ans, +ans,1
Attn (t)
pppodd I 4 4 (g =Ō(00) =((8)) 1-(x))= Es
ans,
-0 d- polyd
B
E
(1
logit (0)
6,jo (((x)) E3
Combining the two bounds, we conclude
A10 5,2, |ss + +
L 822, 82,
Lea
2
B
E logit 5.j2 [(x1) = = Es
Gy d
(d) For Z E &g, in analogy with Lemma H.19, we obtain the following crude bound:
E
M 222
L
+ E
[ 8,2,8
N. 2 0.
(e) For Z E ES Es, by Induction H.2, Lemma H.29 and Lemma H.30, obtain the following logit condition:
logit ( 5,22
logit (1)
E (1-1iogtt 5.12 r(x1) E i 52-
polyd
Hence, Arlr) ,,2,2 and Artr)
,,2,8 can be bounded as
B[NO 5,2 1e..
B B[A 2,2, Nem_]
so B
F. logit
= [(x))={,,, E.] m me (5.6).
dny polyd
5.j0
Moreover, for Z Ec, if
(t) Attn ans, 1-
Attn 1 > 1
2
then A11 s.2,; = 0. Nevertheless, due to the above order- wise bound, this observation does not affect the overall analysis.
Putting everything together, we have
V
Q
Loss
8,8
+
V
Q
Loss
2.2
3,8
KE (ii
Σ
me [6]
Σ
E
[t)
8,2
E
2
B d
E
(1
logit
5.52
(1)
[(x1) = E
+
B
3
(1 logit (1)
3.52
(2(1)) = s
nyd
m==
Lemma H.32. If Induction H.2 holds for all iterations t E [T..2 T1,22s) given E for [Quple.'' P (3,4), s' + T(X), we have
2.2 Loss
Loss 2,2
3,3'
135

---

H.4.3 Attention gap is small
Lemma H.33. If Induction H.2 holds for all iterations t E [T1,2,1,: 21,2,,s), then for any sample 7 we have
(0) Attn ans, ppreed2 Attn (t) 1-s,nns S C1,
where 61 > O is small constant.
The proof follows along the same lines as Lemma H.22, and hence the details omitted for brevity. Lemma H.34. If Induction H.2 holds for all iterations t E (T1,2,1,s: T1,2,,s), then for any sample 2 we have
(4) Attn 1
Attn (6)
1-sped, S Attn (t) ans, 1-pred C2
where
log d 48
0 is small constant.
Proof. Let T denote the first time such that
E Attn (1)
ans, -ppedd2 + Attn (t)
Attn (1)
((x1) = s
log d 4.02B
ans, Ippeed
1-aan-
Since [ S ( [Qiile
forp {3, 43. follows that for any sample Z with T r(x1) s.
at time T we have
Attn (t) ans, 1 + Attn (t)
I Attn 1n, 2
log d 4.02B
+ Ö(1/d).
By Lemma H.24, we also have Attn ans, predd2 the following cases:
Attn ans, -
22(1) >> 0. We then consider
If 2 1 E & then A
(T)
is already activated into the linear regime. Moreover, since d
5.j2.Ta2 uo
(T) Attn ans predd22 Attn ans, I
Attn 1-spe-- < 0, follows that
4.02B
Na cannot be activated. Thus, by Lemma H.11,
(WIT) 432, 232, + +AT))
(TT
2 2
E[
logit 5,j2 T(x1) = 8, =,
while
E
2,1,,, + (A(T) +N77)
2, 2,1,2 )
S S
B
E
logit 3.b
T(x1) = 8, Ei].
dny
If z Eg, then
log
(Attn ans. 1-pdddd + Attn
Attn 1-5,n_+
)2B <
ans, pred,
2.01
Hence logit!)
(1) < Oid-100/22011 (11 logit (1)
(.jil(5.j2) By Lemma H.11,
5.j2
E
(VLTT)
632 + +NT))
+N())
+. 11a) 2 Ω
E logit 5,jo ((x1) = 8, E2],
83,2
dn,,
and
E
8,42,1 6,4,2,5
(V))
+
(
8,422
1 [ss]
S -
dr,
B
E
logit! ( S.jo
T(x1) = =s, Ez]
If 2 E Es, we can apply the crude bounds
(M12) 2, +
(2))
9,3,2,1 - N(T
2.3,2
IEs
2 0,
136

---

and
E [
s,4., + +N
-N2t
3,4,,11 +
242.,0 Les 0.
Combining the above, and noting that the gradient contribution from &g U Eg U Eg is negligible, conclude
V
iy
Loss,
2.2
5.8
2
12(2
1B
logit
5.12
((x1) 8, Ei
(116)
- 52
B
I8
logit
5,22 T(x1) = s, Eg],
dmg
V
Loss
2.2)
$ -22 h)
E logit
(
T(x1) &
(117)
Q
5.32
22
B
E
logit
(T) (((1) = 8, E2]
dn
5.12
Finally, observe that
dQ Q - 0(2/d)
Attn (c)
ans, 2pred,2 + Attn (1)
Attn ans,1-+ans,1
ans, spred,1
Q + Q + 0(1/d)
Hence, (116) and (117) together imply that
Attn (f)
ans.1-pred.2 + Attn
(1) ans
Attn (1)
1-pred,
must decrease time 7 +
H.4.4 At the End of Training
Lemma H.35 (At the end of stage 1.2.2). Induction H.2 holds for all iterations T1221
<
T1,22, =0 polytd))
At the end of training, have
nB
(a) Attention concentration:
E[cattn T(x1) = $ 2.52c, for some small constant 0 < c1
1.005 1o005logd,
2B
(b) Loss convergence: Loss 5.8 < cl +3300lo))
polyd
Proof. Lemma H.24 and Lemma H.31 guarantee the continual growth of [Q]]] and [Q] until
the attention weight cattn reaches 3 Combining Lemma H.31, Lemma H.32, Lemma H.33, and Lemma H.34, we conclude that there exists a stopping time
T,2,2,2 = o poly(d)
BBB
and that Induction H.2 holds for all E [T1,2, T2222,a)
Next, we establish an upper bound for <attn at the end of training. this stage, it suffices focus on the main event Eg. We denote
Attn (9) ans, a 1 Attn (4)
If Attn (
1-rre,2 Attn ans, 1+ then for 2 El have
(t)
(2,11)
> 0(1) logd (Attn un. 1
(tt Att 1-3m000 )B
I +
9 2c
2
log d
B
]2B
where the last inequality follows from
Attn aww.am............. 1ppred22 ]
11 - 2c + c)) = 1ta c.
1 2
137

---

If Attn (t)
p-pred2 S Attn (9)
-ans,1 then for Z E E we have
Dogpr(Zans.22 / (2.1) 6(1) eoogd (Attn (t)
Attm (0)
1-300 )
1macc
log
2c
B
2B
where the last inequality follows from the fact that Attn (0)
ans, prod, S
Putting the above cases together, we obtain
C1
2c
log
2 1
-3.01c1,
2
B
2
which implies
S 2 (3.02c1 + 0.5cl -a1 1.26ct.
I Recursive Learning the Attention Layer: Symmetry Case
In this section, we analyze the recursive learning dynamics of the attention layer in the symmetry case. To this end, recall the definitions of the greedy data annotator, the bootstrapped LEGO data distribution, and the associated self-training loss.
Definition L.1 (Greedy data annotator). The greedy language model Dr induced by a network F is defined as
1, if L'ss = argmax (r ZLIL'
(Za ZL.L'
(118)
0, otherwise.
Definition 1,2 (Bootstrapped LEGO distribution). We define D as the LEGO distribution in Assumption 3.2, except that the answers Zans,e, 1 S l L' generated recursively by sampling Zans,t (-(- ZT.(-1) from the greedy language model pr.
Definition I.3 -Self-training loss). Given (fixed) model F and length L, the self-training next-clause- prediction loss is defined by replacing DL.L' in Definition 3.7 with the bootstrapped distribution DL."' from Definition 1.2:
L. L, (F)) 2 E EEg. log Loss F
log pr (Zans Z Z4,1'-13
(119)
Similarly, the per-token loss defined by
Loss E BZL.L''D..
log PF. log
(Zans. 22,L'-1)], [5].
At stage & > 2, let F(T denote the model obtained from the previous stage, trained on the task
The with L4-1 2k-1, Using the greedy annotator Po PTT 1), construct the bootstrapped LEGO distribution D Dia.he with total length L 2 The corresponding self-training loss L. p)
p(% is then defined as in Definition I.3. In this stage, we focus on training via gradient
Losstu.L
descent on Loss I. Loss p( with sequence length L' = 2 and target token i = 5, initialized from the previous-stage model FT)
For notational simplicity, throughout the discussion drop the subscript of the expectation operator when Z is sampled from the ground-trute LEGO distribution; the subscript will explicitly included only when the expectation is taken over the bootstrapped distribution. Moreover, we
Lx.2 f and the attention gap A Lu.2 A when the abbreviate the wrong attention error cattn context is clear.
138

---

L11
Preliminaries
We first present preliminaries the recursive learning attention layer, including its gradient compu- tations and some useful probability lemmas.
L1.1 Gradient Computations
Fact L1 (Gradients of Q). Given F(Tk 1) from the previous stage The with LE-1 = 2h-1 for
2 2, for (p.
(4. 3), (4, (4,4)). we have
-V Loss 2 p - Qp. Loss L
Lx-1.2 has already been Remark L.1. This fact is straighforward verify. Since the attention error Cattn controlled at small constant level in the previous stage for the model F(Te-1), when transitioning from length 2 to 2h, the incorrect attention produced by p(ln-1) cannot increase significantly. In particular, have
I.k 2f
Cattn
S
which remains small. Moreover, the atlention gap A Inttnans,1--pedl Att.nans, 1--ans non-increasing due to the doubling of irrelevant clauses. As result, the probability assigned by F(Tx-1) to the correct prediction of 2., given Z remains significantly larger than that of any incorrect prediction. Consequently, under the greedy data annotator, the self-training loss Ix.2 the ground-truth LEGO distribution.
Loss Loss 1 p(
coincides with the original loss Loss
Lemma L1 (Gradients of Q1,3). Given F(T 1) from the previous stage and s E T(X), for the diagonal entry QQ,,ee, of the block Q1,3, have
VQ Loss Lx.2 p(7k-1))
AttMaans, 1-
E (Z,1)
Σ
sReLU' (AaL.U(AGsr))
jeld
rE[m]
(Wa.. (Ws.jor, s.jors Zpred,2) Asgii + (s.ser
ll Bazr(or))
Moreover, for the off-diagonal entries QQ with s s', we have
VQ,, Loss 1 p(2-1),55 53,5
E Att
Σ (Z2.1)
Σ
sReLU' (Nogr)-
BEId
rE/m]
(Wo.irr Zpred,1) As.ir + bs...)
Il =
Lemma 1.2 (Gradients of Q4.a). Given (T from the previous stage and T(2)), for the diagonal entry Q4,ale, of the block Q4,4, we have
VQ.. Loss 1x.2 1) = E Attnanss1-tans..
Lo.
Σ (Z,
Σ
sReLU' sReLU'(AG.rr)
SEld
rE(m)
(W Zans.1)
Aosir + bs.d.r) L
Moreover, for the off-diagonal entries [Q4,4]a,r with s + s', we have
| VQu. Loss LE.2 p( 3,
Attnans.1--nn..0
Σ
Es, (Z2.1)
Σ
sReLU' (As..r))
jEld
rE(m
WWs.er: Zans,o) As..rr + 1os.er) ( (a0=
139

---

Notations for activated neurons. Recall that
21 U jEr())
2, where 2, = (rr 1 y y }.
Given ZL,61 let
((ZLet-)) U {gr)
be the collection of all group elements chosen the predicate clauses, and similarly define
5 U {.
='=0
We then introduce
( g E Fiber g E (( v v
= "'g.v
For simplicity, we omit the dependence on LL,e- in the notation 2, when is clear from context. The above recalls the relevant notations from Appendix H.2. With these preliminaries in place. we now state the following lemma, which also relies on the feature magiitu bounds established therein.
J Kezz then Lemma 1.3 (Characterizations of Lambda). Given Zhe, with {AttMans, l
(a) for j E T()), for activated neuron T A, we have
2
Ly
Aa.j.r =
Σ
Attnans,1--pedd V0, (ge) + am 1 (Ve = Ō(on).
F=1
(b) for j E T(V)). for any non-activated neuron & 2ly we have
Ass.r - 06)).
(c) for & T(V). any E [m], we have
Aassir Ō(oo).
Lemma 1.4 Given j E T(V). E 21, \ Uy we have sReLU' (As.jrr) = 0.
We ready derive the gradients of the attention layer, building Lemmas I.1 and I.2 and the properties established above.
Lemma 1,5 (Refined expression for the gradient of Q4,3). Given p(Tx-1) from the previous stage and (X),, for the diagonal entry [Q,,3ss,s of the block Q4,3, letting f2 ((92(1n)), have
VQu. Loss p(Tk-1)
EE 1ttteas.1--pred.2
(1 logits.jz)
sReLU (AGszrr)) (Virr(02) - As. + Ō(6o) + Ō(8))
r
logits و (+)sEr("))
Σ
sReLU (ReLU'(Asar) (Vir(sz) Aoj.r + 0(60)) + Ō(69)
rea,
Σ logit, 5 Ō(o8) ((g)) 1-(1))=8
jKT(3))
140

---

Moreover. for the off-diagonal entries [Qisleas with 8 # s' we have
Vqi,a Loss 2 1 8,0
2
E Attmans.1-predd
(1- logito.z)
Σ (Re.U((Asssa)) (Virr(or) - As, + =((0)) + Ō(89)
logits,j
sReLU (Vir((r) As.,. + Ō(o0)) - Ō(67)
(#1ZET()
eea,
+ logit,, ,( 11 (r(e))=) ((T(D))
Lemma 1.6 (Refined expression for the gradient of Q1,1)- Given F(Te-1) from the previous stage and 8 E , the diagonal entry (Q4,ass,s of the block Qs letting j2 = ((92(11)), we have
Q Loss Loss 2
p( 17--)),5 = E Attnans,
(1 (1-lowg logits,j2 )
Σ
sReLU (Visr(ui) - As.ia. + Ō(oo)) + Ō(89)
logits,)
sReLU' (As.arr (V5r(08) As... + Ō(0o)) +((8))
j*izEr()
reat
+ logits,, ((ott,, Ar(anza
3(F(D))
Moreover, for the off-diagonal entries (Q4,3lan with 8 + s', we have
VQu. Loss 2u p(aa-1 1),5
Losst
E
AttMans,
(1 logits.j)) )
Σ
sReLU (Asjarr) (Viar(en) - As.jar + Ö(oo)) + Ō(89))
reajz
Σ logits.
Σ
sReLU' (As.grr) (Vir((0)) Aosir + Ö(oo) Ō(69)
H*1eE))))
r
+ logits, (egits,o(oi)) 1 (r(0))=)
((F()))
Following the above calculations, we can further obtain the gradient summation of Q4,3 and Q4,4 as follows:
Lemma 1.7 (Gradient sum of Q4,3 and Q1,1). Given F( from the previous stage and E T =(X).
letting j2 (192(11)), we have
VQu.a
lu.2 2x.2 phrh 11,5
+
Q4,3 Loss 2n.2 Los plt 2) 10,
E
(1 logit s.
sReLU (Ac.LU(As...r)
r
Att --nss000 (22) - Σ Attn (ger
1/
141

---

(1-AAe 1ta- 1-preds As =O(oo)) + +((8))
Σ
logits.
sReLU (As..rr
rEa
(Attnans,1 +ans,, (Bo)( +
Att pped, V2 (gt)
6
-1- 1tt -pped As,5 + ((oo)) + =((8))
logits, (Attmams.1-ans.1 + 2ttdamss1-pree.22200 1
jÆr())
Notations for gradient decompositions. Next we define some useful notations to further simplify the expressions of gradient.
Lemma 188. Given F(Tu-1) from the previous stage and E T(2), we define the following notations for the gradient decompositions:
1. for (Q1,3ass we have [- VQq.a L.. p(7 1) V.0.s,r = -[N,, -.a + -o t. where
Na. = Attnans, 2 (1 - logits..z?
(120)
sReLU' (As.jarr
(Virr As.ja.r 0(oo) + Ō(87) (x)))=
rEali
AAttn ant 1-ped,
Σ
logit
(121)
Σ (RasU r
(V2)(12) As.ir + Ō(Go)) + +((8)) ( (()))==
NA Attn 1-pppe1
logit,, (
jKT(y)
(122)
2. for [Q4,dss,s. we have where
V VQi.a Loss 2n.2
pftk 1 115 [Wa,,,Le.s +
Na,4, Attillass sanns, (1 logits.s))
(123)
Σ (RoLU((Asja.)) (V)+((1)) As.jorr + Ō(Go) Ō(87) 1 (()))=
reil,
Na
AttDans, anns
logits.""
(124)
Σ
sReLU'(AA.r)) (Vr(())) As.ir + Ō((o)) + Ō((7) ()()))=
r
Ne = +Attn 1anas,
logit 5. (
jdr()))
(125)
142

---

3. for the summation of QQ13ls, and [Quale,s, we have - V Q1,, Loss 2:2 2,
+
Vq,, J,, = =B + +,,22 + Na.2,.i)), where
Na,La,i (1 logits..
Σ
CReLU"(As.ja)
rea
Att (vo)
Σ
attmans. Varr((e)
8f2
+ (1 Attmans,1-ans) Att 1pprd 2 = =0(oo)) = =((89) 11 (
Σ logits.
Σ
sReLU (As
)tiET(())
r
Att
Vir((U)) + Σ att Vg.r(ge)
Attmans.1--ass.11 Attmans. As.j.r + Ō(oo)) x((8)) ((1))=s
Σ
logits, (AttMans.1-ans.) + Attd
JET(D)
L1.2 Probabilistic Events
We introduce the following events:
((ge(v)) 92(11)} s Uk Uz)
ELp11 Lk sE(La]M2)
Σ
ELTT2 = {oo = 3/1 h:
EL
max yEy
Σ
[(((()) = ga(u)} $ Wh
tE|2
We set
log ny
0
L S ny log ny:
L 8
log
4n Le
log nyy
Uk =
1
= @(1),
We =
Lx
Lb
L 2 ny log ny.
n2g
By standard balls-into- -nnn--itt--issst tail bounds and maximum-load estimates (c.g., [109]), obtain
An UxLa+1), exp( (n, lyg
L Ry:
Pr(ZZ tai $ ELL.))
L Ongy).
exp( ((L log n.)),
La ngy
and
Pr(z ET, = 1 1- n-n, 9(1)
I.2 Reducing the Wrong Attention
As discussed in Remark 1.1, both ett and the attention gap remain bounded by a small constant at the beginning of stage k. Hence, we can directly proceed to stage 1.2.2 of the convergence analysis as T Moreover, by the symmetry between [Q4,as.s and [Q4.4]s,a. we may, without loss generality, restrict attention to particular s E T(X') and analyze the corresponding loss Loss in what follows.
Lu.2
Induction L.1. Given F(Te-1) from the previous stage Tha-1 with Lk-1 2 for k 2, let Th
denote the first time that Loss Ia.2 decreases below el +3.001c) For all iterations' S Te,
have the following holds
143

---

1. a,, + [Q
monotonically increases;
a,8
2. for any sample 2 we have Attn 1nsp1-2 Attn small constant C1
ans, lans,1 S foror some sufficiently
1.005 log d
2/3
3. for any Z we have Attn ans, -ans,
Attn ans, S min
{ Lx-1 Ln.2 T4
{
Cz2 00.
where 02 log > 0 is some sufficiently small constant.
(4) 4. for p E (3,4), for s' =(x) * s. [1 =0
otherwise
1Q
1.2.1 Attention and Lambda Preliminaries
Lemma 1.9. If Induction 1.1 holds for all iterations [TR-1, 1, t)th for any sample Z we have
1. Lu [ c1 2
2c Ix-1
2. Attn ans,l-k
(+)
Attn ans,1-k [ any k, k E I \\ {(ans, 1), (pred, 2)};
3. Attn (1)
Attn (4) 1 s (1)) for any k. k' Ihael \ [(ans,
a-s,1-k
Proof. The highest loss occurs when the sample ' makes the wrong answer highly confused with the correct answer. Thus, we consider the case that Dr" Y0 = 92 Yo for all E [Tx]] \ (2) and yo * ø1-
log PF (Zans. (,55.)) (((.))))
(Attn (1)
-x-1 ki
-Attn (1) .m.. )B (Aten (1)
1n ))22)) 2B -prdd,2 Tx
s ((1) max
3ne,1 -
log d- (Attn (0) pond 2s Lx
e 2
1 1igi (er)+ +c1]2B 1-2---0
(Acen'e) 1onddd I Sann )22
2x.
< ①
hei? -ca -c1)28
< 0
Thus, if L.t < fc,, must have
e 201 +0)22 (
Loss
222
which contradicts the definition of 2,,,,, Therefore, it must hold that Lt 2 f/ for all / S Th.
Lemma 1.10. If Induction 1.1 holds for all iterations [T-1, t), then given Loo & Etq 2. we have
1. for the prediction j2, we have
(
Σ 1 (0((1))))1())) Attn ans, I I-ppred Attn (t)
ans. ) 2B 20 ( Fly
B
+ +00)).
RE(T..)
2. for the prediction J/2 T (92(1))). we have
(0)
Σ
(t) (2200))=)2(10) Attn ans, a --pped Attn (1)
1-aan)))))))) 2B 0
B
+ 0(0).
5.3g Tog
my
BE(T..)
3. for other j (2(1) E T T))), noticing that for this case y + yo. Ur then we have
(1)
Σ
1 1v(())) Attnn Attn(t) (t)
Attn ant. tans,1
Attn 1 2B
[E[La)
144

---

( + 0(0). Bl
B
Lemma 1.11. If Induction 1.1 holds for all iterations [Tk-1, b), then given Z bull ELU, we have
1. for the prediction j2, we have
A (1)
Σ (E(Lu)
Attn (4) (2(()))=22(1) Attn I-n pred, .2
B 0 +(())
9..2:
n2y
2. for other j = T(g2(1)) T()) with y # yr. then we have
(1)
=
-[E
E
9(()))=a()))
Attn (t)
Attn ans, 1ans,
Attn
ans, --ans,)) 2B
-
[E[La]
B
0 + O()).
I.2.2 Gradient Lemma
Lemma 1.12. Induction 1.1 holds for all iterations 1 [Tk-1,Te). given E T(X), have
V i0
Loss
2.2
+
V (t) Loss 2,2
8,8
8,8
2 BBBdd
Js attn
(1 -
logit
5.j2 l
TT(1)) s, E O E E
Proof. By Induction I.1, Lemma 1.10, and Lemma I11, we have
2 (0) S.Le.
E (1 5.
Attn (t)
ans, pred, Visire x1 (ge) Attn attmmms,1--nns.u Viaara (31)
(1
Attn ans, l I-nans,
Attn a pred, 22 + Ō((o)) + =(080)
(8)
(126)
m(())== 1 A00 603
SB"
Similarly,
(1) E
[
S.Le.
Σ
(1) l (.r()))))
(A
(0)
Attn (t) I-peed Vi(ga(o))>oo (ge) + Attn ans, Vr(ga(w))) (30)
2x
(1 Attn (4) ans 1-ans,1
Attn (4)
A (9)
Ō((n) +((8)) (()))=)
ans, p (o2
Consider the event EL,,1. For L E EL,,1 n NEL we have A (1)
S B, since
5.j2.
(t)
Σ
Attn 1 Vis ( Attn ans, Jans,0 Viarozm (30)
2n
2 Attn (0)
ans, pped,2 + (Uk- - 1/LLk)
Lx Cattn
B
litto 2
< 2
+ (Us- - 1/)x)
Eattn
B
{ + 1/Tx Uk)
Iou Cattn
B.
Thus, provided Lit 2 4c//3 and choosing Us = [1 + L6
Lg. we indeed have A
G.22 SB B.
145

---

For j = (()((v)) in N
y yo # V1.
Attn
ans, xped, Vi(ma(u)).r (0x) + Attn (t)
1-anns V2
Attn (
sprad, V (o(2((l)) 22. (ge) (v)
2
2
(1- Attn (t) I-s,l Attn (t)
ans, 1-ppeed)))
(127)
* 91- V1- cancellation,
If there exists at least one & + 2 such that ge(v) = 92(v), then by
Attn ans,
ans, pppddd (o(oa(0)))roz (gr) + Attn 1 V=(oa(o))roo (y0)
2
(1 Attn (c)
1-ans, Attn (t)
-ppeed))))
(128)
Else, pick 3/ # y. so that 0
Then
l (o2(0)) S 50
(0(,)) ( (1 logits
which implics
logits.)
(1)
Σ
Attn (2)
1- Vrll:(u)).Tag (gt) - ams. Velaa(w)))ra (v0)
(8)
5.j
Ex2
-0
B Lxnry
(o-ooiits.)) (1 -Attn 1-aass Attn (t)
1-pedd)))))
(129)
For Arll) 3
If Lu.1 E EL n 2 then
Σ Attn 1-ppee= Vi (0x) Attn (t)
(1)
ans, 1-ans0 Visrru in (30)
1x2
(1 Attn (1)
1-s,an Attn(t) (t)
ans, -spred, 2) 5..22
(t)
Σ Attn (t)
ams, 1 p-edd I gi=jo
Viniorn (gt) - Attn ans. 1-ans0 Visrroe (20)
1+2
(1 Attn (8)
ans, Attn
ans, ppred,)) S.j..oo2m
2 (Uk 2.) 2B + A.derraam
(t)
Catin 2
(A
4) -itin (130)
Else, for Zhu.l UEL note that having more identical predictions ge (vy1) j2 or less distraction (i.e., yo = 91) increases (tt)
A.ja. Hence
[11 logits,5)) EL,,1 U 0(1) [[ (oo-ottts,,)) EE EL
and
Σ Attn (t) ans, 1-pred Visrog vs (gr) Attn (t) ans.1- Viarogn (20)
Ex2
+ (1 - Attn (0)
ans, ans, I- - Attn (1)
ans, 2) (t()) (t)
S 0(1) Lo. B.
(131)
146

---

Putting the above together, we have
V
Loss
+
V
Loss
2.2
5,8
8,8
ailiel I E2
+ EE
<B Le ",,
<B
E| 2 O(08)
E (W + +Nt() 2) E i 2
E [V s.he.l A i SB Me UF2
+ E selin
s 2ee ee, Ō(a)).
(
First, by (127), (128), (130),
E (Ar10) s,br,1 + + s,L,,2) ) 1
E (1 logit 5.52) (A
B3
Eatti)
Lx
(1 logit
(1)
5.52
(0())) Ry
je
+ max A (0) v
(,2())()))
1=(m)== I E
2((B/d) E Catin -1 logit r(x1) E1,11 NET
La
Moreover, by (131),
(((1)
8,L2,1 A00
9.22 BB/2 1s ,
2 Pr(Zun. 1 (B/d))
La
(1 logit
[x(x)) S UELAR
attn
5,22
2 (0((1/100) + 1 1/mg) O(B/d) E
Eattn (1 logit
[(xx)) = s, E1r, n EE
5.52
and, by (127), (128), (129),
E B (1)
2. E 2
2 ny
O(B/d) E c
Li logit
5.j0 1 8, EL.11 U U
(1)
Therefore,
V 2,2
Loss
+
V (!))
Loss
2.2
8,8
12 B/dd
La Cattn
logit
(t)
(Tx1)= s, E n n
5,12
Lemma I.13, If Induction 1.1 holds for all iterations E (TR-1, Ta), given & E T(X') for Q4.pla p E (3,4), & T(X),, we have
[[-
V o
Loss?
2,2
s00 V (
Loss 2.2
I,2.3 Attention gap is small
Lemma I.14. If Induction 1.1 holds for all iterations E [Tk-1, Tk), then for any sample Z we have
Attn (8)
Attn (t) a-s,s-s 5 C₁,
ang, 1-pred,2
where C1 > O is small constant.
Proof. The proof follows the same high-level idea as Lemma H.22. The key observation is that for each ha. there always exists at least one y 전 y1 such that, for some index j appearing in the
147

---

gradient term Ar0e)
we have
S,Lu,
2 1-ppedd22 Attn (t)
(+)
aans, B.
Let J denote the set of indices j achieving arg yf (m1)
(1)
A and pick one such index
= 92(1/). Define T' be the first time when Attn (1) time, have
ans 1-pedd Attn (0)
1 2 c1. At this
Σ
logit 5.
(1-0(d-R())) logit 5,32
JEZ
Plugging this into the gradient expressions yields
5,3. + Ar(T) 2,3,
(1-- (ogio (7) (A v'
Attn ans, ppred22
1.522 x
+0dd ((0))(0)( - logit
(1-ogit ( (Va.s. 2(92) - (7)
Saaioon ) (1((1))
and similarly
ACT)
s,,, + +N(7)
((1-0(-11)) logit (
(2)) (- V Vo (21) + A o2
Attn 1-pee_
+ O (1)- logit 5.12 - (Voserre 2(11) A.Jiiioam 11
(T)
Since (T)
o on
S 0 and
Vinioon (11)
(T) A..
2 (BB),
obtain
2.2 Loss
+
Qit Loss 2,2
(( B logit (T) 6.32
7(x1) s).
Therefore, the quantity Attn (t) that
Attn ans, 1-nns_ cannot further increase, and we conclude
pred
Attn
(4) I
(t) Attn
1-nns S Ct.
Lemma I.15. If Induction 1.1 holds for all iterations E (T TR), then for any sample Z bu. we have
Attn (0)
Attn (1)
La - La
1ns.
S
Eattn
C2:
Lx
where 02
log d > small constant. 2B
Proof. The argument parallels Lemma H.34, but in the recursive setting we first note that for every Z E EL.22
La 1 Lb
5, (9.((22)))) va
Attn 1-pped2
Attn ans. aans,1 +
Lk
Catin.
Let T be the first time such that
La 1 Li T(x1) = <
log
E
Attn ans. 1-ppee2
Attn ans, +ans,
Lsk
4.02 B
148

---

Then, for any
A (7)
5,7(02(6)), T
2(Attn 1-peed
(T)
Attn ans,1--pred, Taaeedde--aie 9x)())))
Attn ans, ans,0) v
log d
2.01
2
Hence the corresponding wrong class logit mass is small:
wrong-class
logit
(1) 5,32
I (d-1101/200) logit
Plugging this into the gradient decomposition, we obtain
M(T) 2e, +
M L,
Attn ans, ppodd (1 (1-logit) (T)
5.52 Vo 2(92) (T)
Odd-1.01/20) (1 logit (T)
Σ
(o - A (T)
( )) I1 (((x)==
5,32
y#
Attn ans, p-pred2 -((1)= -1-logit (T)
(Vo. 2(92)
- (T)
3.52. vn (2))))
and
= Attn ans,1 (1 logit,
(T)
) (V5.a ATT)
5.12
Old 1(1-10112.11) (1 logit
Σ
2(11)) - A I (
5.22
y/
Attn ans, p 21(1 - logit logit! ( 3.220 (
Therefore,
B
Loss
2,2
2
E logit
(2)
T(x1) = EL.2).
(132)
d
5,j0
B
V
Loss
S
-
02
E 1
logits
((x1) s, EL, 5,jo
(133)
d
Finally, recall that
19 Qiilhi 5(2/d)
Attn ans, ppred22 Attn
(c)
ans, 1predd Attn Attm ans.1aas--
(1)
(8)
Q + Q + (((1/a)
Combining the two gradient inequalities shows that the quantity
Attn (t) ans, 1-pped2 + Ts -
Le
Attn (0)
Ls attn
I-aan
must strictly decrease at time T + 1.
1,2.4 At the end of stage
Lemma I.16 (At the end of stage). Induction 1.1 holds for all iterations' T3-1 <t T O ppll(d))
1/.3
At the end of stage k, we have
(a) Attention concentration: Eattn I 5.04c, for some small constant C₁
Ls
1.005 log d
0;
(b) Loss convergence: Loss 12,2 Lx,2
p( S ell +3.01c)
polyd
149

---

Proof. By Lemmas 1.9 and 1.12, the diagonal entries [Q and [Q keep increasing until
the total wrong- aanettionn mass satisfies Cattn stopping time
4c1 Combining Lemmas 1.12 I.15, there exists 3
Th o poly(d)
nB
such that Induction 1.1 holds for all t E [Ti. Tk).
Ls To obtain upper bound atnn at the end of stage k, it suffices work the main event EL,,3-
Case 1: L3 (1). For each y E V, most one index l * 2 can satisfy ge(y) 92(y). and thus the number of collisions per y is uniformly bounded by Ws. Hence,
log PF 2,3 Z(2.1)
+1) exp
log d - (Attn ans, 2 +
Σ E+2
Attn (1)
ans,1- (=(()))=9(())) Attn ans, -anss, 2B
0(1) exp log d - (Attn ans, pred, + (1/L4- 1/bx) Lx
Sattn) 2B
0 ex(pp(( ( f
log
2B)
where the last line uses Attn (1)
ans, pred S [1 Sattn + +11 = 1ta
Le
Sam
2
Case 2: L = w(1. If La S ny log ny then WE/L2 T.g = (x)/Tx)): otherwise W LL = O(1 nn))
In either regime Ws/ L = (1). and we similarly obtain
log Po Pp 2,5 2 (2)11) 2 011) exp log d - (Attn (t)
ans, pped, + (o(1) - 1/Lx) Sain) 2B
θ expl
( +
9
Sua
log
(2BBBB)))))
2
B
At convergence we therefore have
+ 2
Fum
log.d
2 33001
2
which rearranges to
<atin
S
2(3.02c + 0.5c, c1) 5.04cI.
In particular, since B/ log d can be taken sufficiently large constant, obtain the clean stage-end bound
Le attn
O(e)).
I,3 Proof of Main Theorem
Theorem I.1 (Recursive self-training (Restatement)). Assume the distribution D- induced from LEGO(X,G,) satisfies Assumption 3.1, 3.2 and 5.1, and assume the transformer network satisfies Assumption 3.3, C.I, C.2, and By < & logg d. Then for any 1 S k < |oogg [X), the transformer F(Te) trained via Algorithm up to length L = 2 and Tw 0 poly(d)) satisfies: F(Tx) is able to solve Thes with I&+11 2k-1, Le..
ACCLESS (F(Da)) = (0110(()))l
Proof. By Lemma 1.16, at the time Tk. we have Lin < 5.04c1, combining with Induction 1.1, non- diagonal entry of Qp.q remains close to 0, thus moving to the next stage, we have fattn 20.16 16c (notice that 10.16G, could be sufficiently small e.g., 0.01 since we can set a reasonably large B) for all & < L Moreover, A Lat1.e < A 1 Hence, we have
Lay1 < Ji <
E EgtRi Bg B (a (12/ (2m. +
150

---

0 B logits, (Zmre (41,5)
0 ee e- racn
Lx+1 Lxtt
(Famn
-oa)/2 22)
Lx+11 1)
2B
(0 - + O (polyr))
1
Thus ACCLEnI (F 2 1 - 0 (poly))
151


Reviews:
[{"review_id": "oP5urIZiHu", "reviewer": "Reviewer_X31B", "summary": "This paper presents a theoretical analysis of a one-layer transformer's ability to learn Chain-of-Thought (CoT) reasoning via gradient descent. Focusing on the synthetic \"state tracking\" task, the authors prove that the model can learn to solve instances of this problem that belong to the complexity class NC1. The work further investigates length generalization, showing that for simple algebraic structures (cyclic groups), the learned solution generalizes to longer sequences. For more complex structures (symmetric groups), where direct generalization is limited, the authors propose a \"recursive self-training\" curriculum and prove that it allows the model to achieve similar polynomial-length generalization.", "questions": "1. Could you elaborate on the role of the smoothed ReLU in your theoretical analysis? Are there specific properties of this function (e.g., its behavior for negative inputs or its specific polynomial form) that are essential for the proof to go through, as opposed to other common smooth activations?\n2. In the recursive self-training curriculum (Algorithm 2), new data is generated using the model's greedy output (Definition 4.3). From a theoretical standpoint, what challenges would arise if stochastic sampling (e.g., with a non-zero temperature) were used instead of greedy decoding to generate the self-training data? Would the convergence guarantees for length generalization still hold?\n3. Theorem 4.2 provides an accuracy of 1−O(2^{−k/poly(d)}) for the bootstrapped model. Could you provide more intuition on how errors might propagate or accumulate as the reasoning length $L$ grows to be a large polynomial in $d$? Does the model become more susceptible to single-step failures as the chain lengthens, and does the theory provide a way to bound this cumulative error?", "limitations": "I do not think that equaling limitations with assumptions makes sense. I suggest the reader to position their work in a more objective way.", "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\n- **Addresses a Fundamental Question:** The paper tackles a central and challenging question in modern AI theory: whether the reasoning capabilities observed in large language models can be provably learned via optimization, moving beyond expressivity results alone.\n- **Detailed Technical Analysis:** The paper provides a detailed and technically involved proof for its specific setting, demonstrating a mechanism by which attention concentration can emerge from a staged training process to solve a sequential task.\n- **Formalization of Self-Training:** The theoretical formalization and proof for the \"recursive self-training\" curriculum is a novel contribution. It provides a potential mathematical explanation for the empirical success of self-improvement methods that have been observed in practice.\n\n### Weaknesses\n\n- While I completely acknowledge the necessity of studying why CoT is helpful from an optimization perspective, I think that the single special task is not representative enough for demonstrating the underlying mechanism why models can learn to leverage CoT. Specifically, the authors use the state tracking task in NC1 for demonstration, but they do not explicitly point out that whether state tracking is hard and representative enough to distinguish  TC0 and NC1. For instance, is it NC1-complete? In your citation [24], the paper prove their theorems on differnt tasks including a whole class of dynamic programming problems, making their results universal and convincing. Without showing that the theorem holds for a general class of tasks (rather than hold specifically for this task because of its special structure), I do not think that \"explaining why CoT can be learn from an optimization perspective\" can be claimed. \n- The focus on one-layer transformer makes the optimization perspective more or less similar to a construction (expressivity perspective), as Alg 2 constructes optimization in a manual and step-by-step way for different layers. Since the authors try to go beyond expressivity, and since, in theory, a two layer transformer is more expressive, I think the authors should also prove, at least more generally, that a two / general layers of transformers can also be optimized to learn CoT well. Otherwise, the whole proof for optimization seems to be not generalizable, making the theorems less convincing.\n- To the best of my knowledge, your citation [24] proposes \"transformers could simulate polynomial-size circuits using poly(L) steps of CoT\" earlier than [25]. What is the reason that you choose to base on a later work?", "was_revised": true, "final_justification": "After the discussion with the authors, most of my concerns have been addressed. I will update my score to weak-accept to reflect this.", "rebuttal": "We thank the reviewer for the time and thoughtful comments on our work. Below we address your concerns.\n\n**Q1**: The single special task is not representative enough\n\n**A1**:\n1. We first emphasize that state tracking over group operations has been widely used as a benchmark for separating architectural expressivity in seminal works [1–4]. More importantly, state tracking problems that distinguish TC0 from NC1 or beyond have been used to provably separate not only different neural architectures, but also the expressivity of transformers with v.s. without CoT reasoning [5]. This makes state tracking a meaningful and theoretically grounded testbed for studying the capabilities of transformers, either using CoT or not. The task is also widely adopted in empirical studies of neural reasoning [6–7].\n2. Importantly, it is a well-established theoretical fact that state tracking over non-solvable groups (e.g., $S_5$) is NC1-complete [1]. We highlight this in direct response to the reviewer’s question, as it is a standard and widely known result in circuit complexity.\n3. Prior theoretical studies on CoT [8–10] necessarily focus on simple, hand-crafted tasks within TC0 (e.g., parity, linear regression), due to the inherent difficulty of analyzing optimization dynamics. In contrast, our work advances this line of research by studying both TC0- and NC1-level CoT tasks, and by providing the first convergence guarantees in these more expressive settings without overclaiming generality beyond this well-motivated domain.\n\n**Q2.** The focus on one-layer transformers makes the optimization perspective more or less similar to a construction\n\n**A2. We believe there may be a misunderstanding regarding our optimization analysis.**  Our work is fundamentally different from and significantly more challenging than the expressivity argument the reviewer suggests.\n1. This confusion likely stems from prior optimization-based CoT theory [8–10], which focuses on simple tasks (mostly the parity problem in TC0) where a single gradient step closely mimics a hand-crafted construction. In contrast, our setting requires studying the **multi-phase training dynamics** of both MLP and attention weights across many gradient steps, and there is no straight path from initialized weights to the optimal weights. The features learned in the network go through sophisticated trajectories that we cannot control and characterize completely but can only guarantee its convergence on a high level, **which is very different from construction-style arguments**. Thus this highly nontrivial learning process is nothing like expressivity type construction. \n2. The reason we introduce a curriculum-style schedule (Alg 2) that separates the learning of MLP (focused on one-step group actions) from the learning of attention (enabling length generalization) is to manage the complexity of the intertwined dynamics between different parts of the network. **This separation does not simplify the learning process into a staged construction**. Both stages still undergo complex multi-phase optimization, and our analysis must carefully track how intermediate solutions influence subsequent stages. This level of coupling introduces deep technical challenges far beyond what construction-based proofs address.\n3. The implication that one-layer transformers are “simple” to analyze is also misleading. The optimization landscape remains highly nontrivial even in this setting [11–13]. In fact, one-layer models likely represent the **current theoretical frontier** for understanding transformer training dynamics. Extending such results to deeper models, without strong layer-wise assumptions, is out of reach for current research.\n\nPrior CoT learning theory has only addressed very simple tasks such as parity with no results on even cyclic group reasoning or length generalization. To our knowledge, we are the first to bridge the gap between the theoretical expressivity and learnability for transformers with CoT, and to provide a length generalization guarantees that go significantly beyond what previous work achieves. Especially since most prior expressivity results rely on absolute positional encoding and thus cannot length generalize.\n\n**Q3.** What is the reason that you choose to base on a later work [25] rather than [24]?\n\n**A3.** We clarify that prior works on expressivity mainly provide background context rather than form the main foundation of our study.  While  [24] was the first to show that transformers with CoT can go beyond TC0 in expressivity, their results do not establish a separation beyond NC1, as the CoT expressivity lower bound is only NC1-hard, even though the tasks may include P-complete problems. We refer more often to [25] because it offers a more comprehensive and aligned characterization of the complexity classes for transformers with and without CoT. That said, we note [25] is not the only work that proves such results, for example,  [1] also establishes similar findings, albeit with a slightly weaker expressivity upper bound.\n\n**Q4.** The role of the smoothed ReLU in your theoretical analysis\n\n**A4.** The smoothed ReLU is chosen due to two key properties that are essential for our optimization analysis:\n1. Its polynomial smoothness around zero is necessary for applying the tensor power method, which plays a central role in proving MLP feature emergence, as shown in prior works [14–16]. \n2. The symmetry around zero serves as a soft barrier that prevents some pre-activations from drifting too far into the negative region. This helps ensure convergence to a clean and interpretable set of features. We adopt this design specifically to enable our length-generalization result. \n\nIn summary, the use of smoothed ReLU makes the optimization analysis more tractable, allowing us to show that the transformer learns clean features for one-step state tracking. This enables us to focus on analyzing length extension and extrapolation, which is the core contribution of this work.\n\n**Q5.** Challenges of using stochastic sampling (e.g., non-zero temperature) instead of greedy decoding for generating self-training data?\n\n**A5.** The main challenge lies more in the clarity and tractability of the proof rather than in any fundamental technical obstacle. In fact, even with non-zero temperature, at the end of each training length $2^k$ for $k \\in N^+$, the probability of generating an incorrect prediction has already fallen below 1/d with high probability over the dataset. A more involved analysis would show that the noise introduced by occasional incorrect predictions due to stochastic sampling does not significantly affect the learning of the attention weights. Sequences with slightly higher error rates occur with exponentially small probability, and thus have a negligible impact on the overall attention gradients. However, to analyze such gradients in detail would require extensive case-by-case discussion, which would make the proof significantly more cumbersome without offering fundamentally new insights.\n\n**Q6.** Intuition on how errors might propagate or accumulate as the reasoning length grows to be a large polynomial in d\n\n**A6.** This is a great question, and we agree that a clearer explanation of the limits of length extrapolation should be included in our paper. At a very high level, as context length grows beyond the training length, attention to important clauses becomes diluted, and the influence of incorrect group operations can accumulate for CoT prediction. If the total attention to misleading tokens becomes significant, the model’s prediction can deteriorate. In the simply transitive action case, the correct operation typically receives $1 - 1/d^c$ attention, while all others get $1/d^c$ in total. Thus, one would need roughly $L \\cdot d^c$ incorrect operations in the context to confuse the model, allowing extrapolation to polynomially longer lengths. Beyond length $\\gg |Y| \\cdot L \\cdot d^c$, predictions approach uniform. In the symmetric group action case, the learned representations are less balanced, and error arises for length $> C \\cdot L$ with a small constant $C < 10$. Beyond that limit, the prediction of transformers for values $y \\in \\mathcal{Y}$ can quickly become uniformly random. \n\n---\n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.\n\n---\nReferences:\n\n[1] Merrill et al. The expressive power of transformers with chain of thought\n\n[2] Merrill et al.,  A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers\n\n[3] Merrill et al. The Illusion of State in State-Space Models\n\n[4] Liu et al. Transformers learn shortcuts to automata\n\n[5] Li et al. Chain of thought empowers transformers to solve inherently serial problems\n\n[6] Kim et al. Entity tracking in language models\n\n[7] Toshniwal et al. Chess as a Testbed for Language Model State Tracking\n\n[8]  Huang et al. Transformers learn to implement multi-step gradient descent with chain of thought\n\n[9]  Kim et al.  Transformers provably solve parity efficiently with a chain of thought\n\n[10] Wen et al. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency\n\n[11] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot\n\n[12] Zhang et al. Trained transformers learn linear models in-context\n\n[13]Li et al.  A Theoretical Understanding of shallow Vision Transformers: Learning, Generalization, and Sample Complexity\n\n[14] Cao et al. Benign Overfitting in Two-layer Convolutional Neural Networks\n\n[15] Allen-Zhu et al., Towards understanding ensemble, knowledge distillation and self-distillation in deep learning\n\n[16] Chen et al. Towards understanding the mixture-of-experts layer in deep learning"}, {"review_id": "xKgpudJkDQ", "reviewer": "Reviewer_kcpH", "summary": "The paper provides a theoretical analysis of the ability of transformers to learn state-tracking problems belonging to $\\text{NC}^1$ by learning to produce the intermediate reasoning steps, as in the popularly used Chain of Thought setting. This means that sequential prediction methods such as CoT provably extend the capabilities of transformers, unless $\\text{NC}^1 = \\text{TC}^0$, which is an open problem but is popularly believed to be false. The proofs do make some assumptions on the dimensionality and structure of the transformer, which means they do not immediately apply to the popular transformer architectures as used by practitioners. The scope of the result nonetheless appears to be significant, since it brings seminal work on the expressivity of transformers closer to the CoT paradigm typically of use today.", "questions": "Suggestions:\n- Please enhance the legibility of Figure 1 substantially. The text is impossible to read when printed out.\n- Please add a non-CoT baseline to your experiment to give some credence to the impact on downstream performance implied by your theory.\n- Some typos:\n  - P.3, Theorem 1.2, second bullet point: \"from its own outputs**, .**\"\n  - P.4, example of LEGO sentence: in answers, $x_L$ should likely be $x_{L'}$\n  - P.4, in Eqn (2), $L$ should be $L'$ (per the preceding sentence)\n  - P.9, \"Note that **in by**\"\n  - P.9, Experiments, \"transformer **with attention with two atetnion heads**\"\n\nQuestions:\n- If I understood your construction on P.5-6 correctly, your transformer in effect uses the same projection matrix $Q$  for the queries and keys, while not using a projection at all for the values. Is that correct? If so, why is this not explicitly listed as an assumption in the paper, since it does not correspond to the more general transformer architecture?\n- Assumption 3.4 appears to be very strong. Could you help me build some intuition for why this is necessary? In particular, why $(p, q) \\notin \\{(4, 3), (4, 4)\\}$?\n- My understanding is that your clause-wise embedding structure is what allows the model to disentangle the differing meanings of the tokens without positional encodings. Is that correct? If so, how well do you expect that your results would carry over to a RoPE (or other PE) setting in which the clauses are **not** hard-coded to be separated, as is the typical case in language modeling?", "limitations": "Please see questions above.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is fairly readable. There are a few minor typos, but nothing major.\n- As far as I can tell, the main results are quite exciting.\n- By basing their study off a task that has also been used to study similar aspects of transformers empirically in the past, the paper does provide some linkage to the broader literature.\n\nWeaknesses:\n- While the limitations and assumptions are acknowledged well in the paper, their downstream impact on the generalizability of the results is not clear to me.\n- Because of the clause-embedding mechanism, and lack of other positional encoding, it is unclear to me whether these results actually correspond as well to CoT as the authors claim, in which the structure of the sequence (and thus the attention mechanism) is substantially different.\n- The experiment could benefit from a non-CoT baseline to display the difference in length generalization empirically. (I realize empirical performance is not the subject being studied here, but it would give some sense of the practical implications of the very theoretical results.)", "was_revised": true, "final_justification": "My primary concerns revolved around the original paper being a bit unclear about (1) how the assumptions affected the applicability of the theoretical results to real-world models, and (2) why those assumptions were needed in the first place. I think the authors' have substantially improved the paper in both of those respects during the rebuttal and am confident that the final paper will be clearer.\n\nI also had some concerns about the empirical results being presented. These have been extended and improved upon (in scope) during the rebuttal; I still think they are a bit too small and naive to really say much about how the theory here extends to practice, but I do not assign much weight to this as I do not believe it to be a realistic goal that every theoretical contribution to the literature should also have to include experiments with state-of-the-art models (as long as the *theoretical* connection between the two is clear). In particular, I do not share vE9B's apparent sentiment that the experiments are grounds to reject the paper on.\n\nI therefore update my score to 5 - Accept. However, I note to the AC that this branch of theoretical modeling of transformers is not my main research area, and so I have limited my confidence score to 3.", "rebuttal": "We thank the reviewer for providing the valuable feedback! We will enhance the readability of the figure and fix the typo in the revision. Below we address your concerns.\n\n**Q1.** While the limitations and assumptions are acknowledged, their impact on generalizability is unclear\n\n**A1**. We thank the reviewer for raising this general yet important point. Indeed, as in most theoretical works, certain simplifying assumptions are necessary to keep the analysis tractable and focused. \nWithin the clearly stated scope of our theoretical model and assumptions, we have provided rigorous explanations of length generalization. While extending these results beyond our current setting would be valuable future work, the primary goal of this paper is to provide some novel insights from a theoretical perspective. We hope that our answers to your later questions can explain to you why our setting shares the structural similarities with real-world CoT performed by actual language models, and how to interpret our idealized results to shed light on practical applications.\n\n**Q2.** It is unclear whether the results truly apply to CoT, given its different sequence structure and attention patterns.\n\n**A2.**  Through our technical construction, we aim to create a setting where the structure of step-by-step reasoning is preserved and is amenable to theoretical analysis. While it may appear different from natural CoT traces on the surface, we emphasize that our design abstracts two essential components of real-world CoT reasoning:\n1. Each CoT step typically focuses on a localized subproblem, even when the overall task involves long-range dependencies. We represent this structure using a **clause-based data format**, which allows us to formalize how reasoning steps are represented, sampled, and generated during training.\n2. Long-range dependencies in CoT are often handled by the retrieval function of attention. This is an important topic of study in both theoretical analyses [3,6,7] and practical transformer components such as *induction heads* [8]. In our framework, we build on the LEGO formulation and adopt the \"variables\" in LEGO as retrieval queries, thereby enabling transformers to attend to relevant context.\n\nIn summary, while our problems are abstract examples of chain-of-thought reasoning, they preserve key structural patterns observed in real-world reasoning traces. We hope these abstractions can offer insights into how attention mechanisms may function in practical CoT models.\n\n**Q3.** The model ties query/key projections and omits value projections. Why isn't this assumption explicitly stated, given it deviates from standard transformer architectures?\n\n**A3**. Thanks for the question. We adopted a quite common simplification in existing theoretical works of transformer training dynamics due to technical tractability [1-5]. Our assumptions delegated all the functionalities of attention to a single attention weight Q, which leaves room to analyze the attention pattern between different tokens while drastically simplifying the analysis because we dropped many interplays between different weight updates. In fact, one can view our setting as we adopted a slightly stronger attention layer, since now the $Q^TK$ is full rank and the value matrix can be thought of as absorbed into the MLP. Thank you for highlighting this point, and we will explicitly include additional references and clarify this assumption to better justify our architectural choices in the revised manuscript.\n\n**Q4.**  Assumption 3.4 appears to be very strong\n\n**A4.** As we pointed out in line 226 in the paper, setting certain parts of the weight matrix to zero aligns with common practices in existing theoretical analyses of transformer training dynamics [1-5]. Technically, for our task, only the blocks $Q_{4,3}$ and $Q_{4,4}$  play a meaningful role in learning: $Q_{4,3}$ is important for in-context retrieval, while $Q_{4,4}$ puts weights on the current state. The remaining components either receive no gradient or capture task-irrelevant biases that complicate the analysis. The delicate balance between these two blocks, and their interaction with irrelevant tokens in the context, is the major technical challenge in our attention dynamics analysis. In fact, some existing works [1-2] assume even stronger simplifications, such as being a scalar multiple of the identity matrix, with only a single learnable parameter, for similar reasons. Even though we have made some compromises in the setting to allow tractable analysis, we hope that our methods can offer some insights on the reasoning ability of general language models. In fact, our result gives a clear example of how transformers can learn to implement the induction head [9], showing that our setting can reproduce certain results on the inner workings of transformers studied in prior works.\n\n**Q5.** How well do the results extend to RoPE or other positional encodings where clause boundaries are not explicitly separated\n\n**A5**. As discussed in Q2, we believe that achieving length generalization beyond the training context requires the model to overcome certain proximity biases. However, we do not aim to discredit the importance of locality that is better captured by other PE such as RoPE, as our clause-based structure actually absorbed all the local correlations to make the analysis possible.  Regarding the extension, we conjecture that currently used positional encodings like RoPE or ALiBi are inherently limited in their length generalizability, due to their position-dependent nature (either absolute or relative), such as shown in some empirical study of architectural design [9]. Prior theoretical works on CoT expressivity [6, 7] also rely on fixed positional encoding schemes and cannot generalize beyond a fixed problem length. We hope our work could motivate further exploration of architectural designs that can enable transformers to extrapolate beyond training context lengths, whether for CoT reasoning or to solve other problems.\n\n**Q6.**  The experiment could benefit from a non-CoT baseline to display the difference in length generalization empirically. \n\n**A6.** Following the reviewer’s suggestion, we conducted additional experiments with a non-CoT baseline in both the simply transitive (cyclic) and symmetric cases to evaluate length generalization. Specifically, we trained the same one-layer transformer under the identical task setup described in the paper, but without intermediate answer clauses, i.e., directly predicting the final answer. For non-CoT models, we explored eight learning rates: $[5e-6, 1e-5, 2e-5, 5e-5, 1e-4, 2e-4, 5e-4]$, whereas our CoT model was trained with a single learning rate ($1e-4$). Additionally, the CoT model was trained exclusively on LEGO sequences of length 5, while the non-CoT baseline was trained on sequences of lengths 5 and 10. \n\n| Cyclic | 5 | 10 | 20 | 50 | 100 | 150 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| CoT (trained on L=5) | 1 | 1 | 1 | 1 | 0.997 | 0.981 |\n| Best (trained on L=5) | 0.164 | 0.169 | 0.151 | 0.181 | 0.192 | 0.19 |\n| Best (trained on L=10) | 0.168 | 0.184 | 0.166 | 0.179 | 0.174 | 0.179 |\n\n| Symmetric | 5 | 10 | 20 | 50 | 100 | 150 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| CoT  (trained on L=5) |0.999 | 0.998 | 0.991 | 0.989 | 0.965 | 0.922 |\n| Best (trained on L=5) | 0.009 | 0.009 | 0.01 | 0.009 | 0.013 | 0.011 |\n| Best (trained on L=10) | 0.01 | 0.01 | 0.013 | 0.005 | 0.008 | 0.011 |\n\nResults show that the transformer with CoT achieves length generalization in our setting, while the non-CoT baseline does not. We also plan to include additional multi-layer transformer experiments to further investigate this setting.\n\n---\n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.\n\n---\n\n[1] Wang et al. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot\n\n[2] Nichani et al. How Transformers Learn Causal Structure with Gradient Descent\n\n[3] Kim et al. Transformers provably solve parity efficiently with a chain of thought\n\n[4] Zhang et al. Trained transformers learn linear models in-context\n\n[5] Kim et al. Transformers learn nonlinear features in context: nonconvex mean-field dynamics on the attention landscape\n\n[6] Feng et al. Towards revealing the mystery behind chain of thought: a theoretical perspective\n\n[7] Li et al. Chain of thought empowers transformers to solve inherently serial problems\n\n[8] Olsson et al. In-context Learning and Induction Heads\n\n[9] Ren et al. Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling"}, {"review_id": "6lurka1nP8", "reviewer": "Reviewer_vE9B", "summary": "The paper proves the learning/generalization of CoT for a 1-layer transformer for LEGO, a reasoning task defined by the authors which involves performing a limited set of group operation on state spaces. The paper also proposes a recursive self-training scheme that enables CoT length extrapolation on the LEGO task.", "questions": "None", "limitations": "none", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n1. The paper is in general written in a clean manner with clear math annotation usage.\n2. There is some novelty in proofing CoT for a 1-layer transformer for LEGO, a reasoning task defined by the authors which involves performing a limited set of group operation on state spaces. The recursive self-training also seems an interesting idea that can be applied to some other reasoning problem.\n\nWeakness:\n1. The paper's title seems to suggest it proves for all reasoning tasks but in reality the paper only tackles a limited subset of the reasoning tasks which only include a limited set of group transformations on a state space. \n2. The evaluation section is weak. There is no detailed description of the datasets (how big is the training set? some example data), the training setup (how many GPUs? trained for how long). There is also little analysis of the results.\n3. Sometimes assumptions are stated but without enough explanation. For example, the authors stated that \"we use no positional encoding (NoPE) in our attention mechanism\" on line 202, but never explained in details why this is the case. Modern LLMs all use positional embeddings which makes the proof not applicable to them.", "was_revised": true, "final_justification": "I will maintain my rating as my concerns are not fully addressed (overclaim in title, weaker evaluation section and the issue of no position embeddings).", "rebuttal": "We thank the reviewer for the time and thoughtful comments on our work. Below we address your concerns.\n\n**Q1.**  The paper's title seems to suggest it proves for all reasoning tasks but in reality the paper only tackles a limited subset of the reasoning tasks which only include a limited set of group transformations on a state space.\n\n**A1.**  We would like to respectfully clarify that our title is not an overclaim in the context of existing theoretical work on reasoning with transformers.  \n1. State tracking has been widely adopted as a standard synthetic testbed for evaluating the reasoning abilities of language models. The seminal work [15] established foundational expressivity results for transformers in algorithmic reasoning and provided comprehensive empirical evidence across architectures, from RNNs to transformers, on tasks ranging from Dyck languages to non-solvable groups. It showed that sufficiently deep transformers can solve most group word problems within the training length. Subsequent studies [12–13] further demonstrated the expressivity gap: constant-depth transformers without CoT cannot solve problems beyond TC0, whereas adding a scratchpad (analogous to CoT) enables transformers to solve NC1-complete problems such as S5 [1,2,15,16]. Our work bridges the gap between theory and empirical observations. We go beyond prior results by providing a formal proof that transformers can learn a length-generalizable CoT solution, overcoming the expressivity limitations of the base architecture. This confirms that transformers can learn solutions that are provably more powerful than what the architecture can express without CoT.\n2. To our knowledge, no prior theoretical paper on CoT has established guarantees for **all** reasoning tasks [1-7]. From the perspective of optimization analysis, existing theoretical studies on CoT [8-10] have exclusively addressed specific, simple tasks within TC0, such as parity or linear regression.  In contrast, our work represents significant advances: 1). We provide the first rigorous analysis of transformer training dynamics for a broader and more challenging class of CoT-style reasoning tasks. Specifically, we study tasks beyond parity within TC0 (simply transitive setting) as well as tasks within NC1 (symmetry group setting); 2). We are the first to establish convergence guarantees that demonstrate not only the transformers' ability to solve these tasks but, crucially, their capacity for **length generalization**. \n\n**Q2.** The evaluation section is weak\n\n**A2**. We would like to emphasize that the empirical part primarily serves as numerical verification of our theoretical results, specifically illustrating length generalization in the cyclic case. Due to space constraints, we omitted detailed dataset descriptions; however, the synthetic data and loss function strictly follow the LEGO structure analyzed theoretically in our paper. Basic setup about network architectures and optimizers is briefly summarized in the paper. Other detailed training setups mentioned by the reviewer, which are less directly related to our theoretical focus, have been intentionally omitted to maintain clarity and conciseness.\n\n**Q3.** Sometimes assumptions are stated, but without enough explanation, For example, the authors stated that \"we use no positional encoding (NoPE) in our attention mechanism\" \n\n**A3**. We thank the reviewer for raising the concern regarding assumptions stated without sufficient explanation. Specifically,  for the use of \"no positional encoding (NoPE),\" we note that it has been widely observed empirically that standard positional embeddings inherently fail to achieve length generalization. This limitation arises because traditional positional embeddings implicitly introduce approximate biases that typically favor local positions and thus fundamentally hinder extrapolation beyond the training sequence length. In contrast, our theoretical analysis supports an alternative approach and suggests an alternate transformer design with no positional embedding that provably achieves length generalization, thereby offering insights and guidance for current empirical research. Furthermore, our theory implicitly reveals that the inability of existing networks to generalize to unseen lengths may stem from the presence of positional embeddings, thereby identifying a potential underlying reason why current settings fail in this regard. In the revision, we will carefully include additional discussion and remarks to explicitly clarify these assumptions and the rationale behind our design choices.\n\n---\n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.\n\n---\n\nReferences:\n\n[1] Feng et al. Towards revealing the mystery behind chain of thought: a theoretical perspective. \n\n[2] Li et al. Chain of thought empowers transformers to solve inherently serial problems. \n\n[3] Merrill et al. The expressive power of transformers with chain of thought.  \n\n[4] Abbe., et al. How far can transformers reason? the globality barrier and inductive scratchpad. \n\n[5] Kim et al. Metastable dynamics of chain-of-thought reasoning: Provable benefits of search, rl and distillation. \n\n[6] Huang et al. A formal framework for understanding length generalization in transformers. \n\n[7] Golowich et al. The role of sparsity for length generalization in transformers.\n\n[8] Huang et al. Transformers learn to implement multi-step  gradient descent with chain of thought.\n\n[9] Kim et al. Transformers provably solve parity efficiently with a chain of thought. \n\n[10] Wen et al. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency. \n\n[11] Kim et al. Entity tracking in language models. \n\n[12] Merrill et al. A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers. \n\n[13] Merrill et al. The Illusion of State in State-Space Models. \n\n[14] Zhang et al. Unveiling Transformers with LEGO: a synthetic reasoning task.\n\n[15] Liu et al. Transformers learn shortcuts to automata.\n\n[16] Wen et al. RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval."}, {"review_id": "k7kpNEqRP1", "reviewer": "Reviewer_2b2M", "summary": "The paper theoretically studies length generalization in trained transformers for LEGO tasks (i.e., state tracking under $G$-actions). The studied transformer model consists of a constrained softmax attention and smoothed ReLU feedforward layer, without positional encoding. The model is trained on CoT loss via a curriculum of one-step and two-step reasoning. It is shown that for simply transitive actions, the CoT model trained for $poly(d)$ time can generalize up to length $poly(d)$ where $d$ is problem size. Moreover, when $G$ is the symmetric group on the state space, a bootstrap self-training curriculum of progressively doubling length can generalize to longer tasks. This also constitutes a novel optimization guarantee for constant depth transformers beyond the $TC^0$ class.", "questions": "See Weaknesses", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe paper rigorously studies a problem of crucial importance, namely the acquisition mechanism of length generalization in CoT transformers. The considered LEGO task is an empirically well-studied synthetic task, and the one-layer transformer model is also very natural (e.g., incorporates a smooth ReLU feedforward layer, compared to other works which fix or discard the feedforward layer). While there are some structural assumptions and a curriculum learning setup is required to separate the parameters dynamics, I think these are still quite weak for optimization guarantees. The paper is overall easy to read and obtains solid generalization guarantees which displays the potential of length generalization from a theoretical perspective.\n\n**Weaknesses/Questions**\n\nThe settings of the two results are restricted to simply transitive actions and action of the full symmetric group for Theorems 4.1 and 4.2, respectively. While the difficulty of analyzing even these settings is evident from the laborious treatment in the appendix, I am curious as to where the technical difficulty arises when attempting to study more general actions, and what we may achieve with future works.\n\n* Why is the assumption of simple transitivity necessary for Theorem 4.1?\n* The symmetric group assumption of Theorem 4.2 is rather restrictive in that it necessitates $|\\mathcal{Y}| < \\log\\log d$ so that $|G|$ is logarithmic. With Cayley's theorem in mind, can this be relaxed/generalized to (certain) subgroups of symmetric groups?\n* The heading of Section 4.2 mentions non-solvable groups. While it is true that large symmetric groups are non-solvable, unless the theory actually leverages the group-theoretic notion of solvability, it would be better to just say \"symmetric group\" for clarity.\n* Is there a setting in the LEGO setup where length generalization provably fails? (e.g., lower bounds without CoT which have been demonstrated in e.g. the parity papers)\n\nSome typos:\n\"action group action\" in line 242, \"atetntion head\" in line 301", "was_revised": true, "final_justification": "The paper is a solid contribution to understanding length generalization on the synthetic but illuminating LEGO task. The setup and optimization analysis, while necessarily stylized due to the complexity of the transformer model, also seems to be novel compared to existing works.", "rebuttal": "We thank the reviewer for the positive comments and thoughtful questions. Below, we provide point-by-point response to the main comments.\n\n**Q1.** Why is the assumption of simple transitivity? \n\n**A1.** The assumption of simple transitivity can be thought of as a natural generalization of cyclic group actions. Our analysis only requires that for each pair of $y_1, y_2$, there exists a unique group operation $g$ such that $g⋅y_1=y_2$​. In our context, this assumption allows us to study group-based reasoning tasks while keeping the structure simple for analysis. \n\n\n**Q2.** Can the symmetric group assumption in Theorem 4.2 be relaxed/generalized?\n\n**A2**. This is a great question. In fact, our proof can be directly generalized to certain subsets of the symmetric group that satisfy the following condition: for any pair of elements $y, y^{'}$ in the domain, the number of group elements mapping $y$ to $y’$ is a fixed constant. This includes subsets of symmetric groups such as the groups of all k-swaps in $S_n$, for $k \\leq n$. These subsets are not necessarily groups (they can be conjugacy classes) but they provide exactly the regularity we need for successful feature learning. If this structure is removed entirely, the distribution of $(g,y)$ pairs becomes imbalanced, which can harm the learning of both single-step reasoning and context-length generalization. While we agree it would be very interesting to extend the analysis to other groups, it is beyond the current scope of our work. \n\n**Q3.** It would be better to just say \"symmetric group\" for clarity.\n\n**A3**. Thank you for your helpful suggestions! We will revise the wording in the revision.\n\n**Q4.** Is there a setting in the LEGO setup where length generalization provably fails?\n\n**A4.** Our theory links length generalization to the model’s ability to perform in-context retrieval: as context grows, the attention mechanism must still assign enough mass to every relevant clause. Two key factors determine whether generalization can succeed:\n1. The peak attention attained by the relevant clauses within training context length;\n2. The rate at which attention is diluted by irrelevant clauses as context length increases.\n\nIn the case of a simply transitive group action, the learned features are sufficiently “pure,” enabling sharply focused attention and yielding a polynomial bound on generalization length. In contrast, in the symmetric group action setting, the learned features do not have a similarly strong guarantee, limiting the length generalization to only a constant multiple of the training length.  While our exact theoretical predictions may not transfer directly to all LEGO settings due to differences in task structure and objective, we hope our analysis offers useful insight into when and why length generalization can succeed or fail, especially in realistic transformer-based reasoning scenarios.\n\n---\nWe thank the reviewer again for your comments. We are more than happy to answer your further questions."}]

Author rebuttal:

Meta-review: This paper provides a theoretical analysis of how a one-layer transformer learns to perform CoT reasoning via gradient descent. The analysis focuses on a synthetic "state tracking" task over group operations. 

Overall, reviewers agree that the paper is well-written; the problem being addressed is a fundamental problem; and the technical analysis is rigorous and solid. Reviewers questioned about the representativeness of the state tracking task, the simplicity of the one-layer transformer, the lack of non-CoT baseline for generation, and assumptions of no PEs. I think the authors did a very good job in addressing these concerns. 

I think this is a nice piece of work with solid theoretical contributions and therefore I recommend to accept this paper.
Decision: Accept (poster)