Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought
Authors: Jianhao Huang
Zixuan Wang
Jason D. Lee
Venue: ICLR 2025 Spotlight (2025)
Keywords: Chain of Thought
Transformer optimization
Training dynamics
Primary area: learning theory

Abstract: Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs), particularly in arithmetic and reasoning tasks, by instructing the model to produce intermediate reasoning steps. Despite the remarkable empirical success of CoT and its theoretical advantages in enhancing expressivity, the mechanisms underlying CoT training remain largely unexplored. In this paper, we study the training dynamics of transformers over a CoT objective on a in-context weight prediction task for linear regression. We prove that while a one-layer linear transformer without CoT can only implement a single step of gradient descent (GD) and fails to recover the ground-truth weight vector, a transformer with CoT prompting can learn to perform multi-step GD autoregressively, achieving near-exact recovery. Furthermore, we show that the trained transformer effectively generalizes on the unseen data. Empirically, we demonstrate that CoT prompting yields substantial performance improvements.

Published as a conference paper at ICLR 2025
TRANSFORMERS LEARN TO IMPLEMENT MULTI-STEP GRADIENT DESCENT WITH CHAIN OF THOUGHT
Jianhao Huang', Zixuan Wang"* Jason D. Leea Shanghai Jiaotong University, PPrineeoon University huang.jh20216s juu. edu. en, {wangzx,
ABSTRACT
Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs)) particularly in arithmetic and rea- soning tasks, by instructing the model to produce intermediate reasoning steps. Despite the remarkable empirical success of CoT and its theoretical advantages enhancing expressivity, the mechanisms underlying CoT training remain largely unexplored. In this paper, study the training dynamics of transformers over a CoT objective on an in-context weight prediction task for linear regression. We prove that while a one-layer linear transformer without CoT can only implement single step of gradient descent (GD) and fails to recover the ground-truth weight vector, a transformer with CoT prompting can learn to perform multi-step GD autoregressively, achieving near-exact recovery. Furthermore, we show that the trained transformer effectively generalizes on the unseen data. Empirically, we demonstrate that CoT prompting yields substantial performance improvements.
INTRODUCTION
Transformer-bases Large Language Models (LLMs) have demonstrated significant success across various language modeling tasks, achieving state-of-the-art performance in numerous domains OOpenA 2023). Remarkably, these models have also unlocked complex reasoning abilitics, par- ticularly mathematical problem-solving and coding tasks (Chowdhery et al. 2023; Anil et al.
2022 Achiam et al. 2023). A key method driving this advancement the Chain of Thought (Co'T), which enables LLMs generate intermediate reasoning steps autoregressively rather than providing a direct answer. This process effectively improves the model's capacity to solve complex problems. In practice, CoT reasoning can be elicited either by providing few-shot CoT examples or by append-
ing prompts like "let's think step by step" to bootstrap the model's 's response (Kojima et al. 2022
Wei al. 2022 Suzgun et al.) 2022: Nye et al. (021).
Theoretically, CoT enables LLMs to perform multi-step sequential computations by generating in- termediate results, thereby significantly improving the expressive power of transformers (Li et al.
2024b Feng et al. 2024; Merrill & Sabharwal 2023a) compared to standard decoder transformers that generate direct outputs without intermediate reasoning (Liu al., 2022, Merrill & Sabharwal 2023b). Despite these theoretical insights, it remains unclear how transformers are trained on CoT data to effectively execute multi-step reasoning. Furthermore, is unknown whether a transformer trained specifically with an auto-regressive objective with multi-step CoT can substantially outper- form one trained to directly output answers without COT
This paper takes initial step beyond expressiveness to study the training dynamics of transformers when trained on CoT data. Specifically, following the modified in-context learning (ICL) setting on linear regression proposed Ahn et al. (2023); Zhang al. (223)), we use the training process with the CoT framework implemented. We name the task in-context weight prediction where the goal is to predict the linear weight vector from the sequence of input prompts. Instead of performing direct ICL and outputting a prediction, the transformer with CoT prompting is allowed generate multiple intermediate steps before arriving at the final answer. We theoretically investigate the transformer's training trajectory on the CoT objective and show the expressiveness
testbed to analyze
"Equal Contribution (alphabetical order).

---

Published as a conference paper at ICLR 2025
gap between transformers trained with CoT and those without. Our main results show this separation is learnable: gradient- graddeent-based algorithm can learn the constructed transformer with COT.
We summarize contributions follows:
Expressiveness Gap. We characterize the global optimum of the population loss for the in- context weight prediction task on linear regression using a one-layer transformer without CoT prompting. Our results show that, without CoT, the transformer at the global minimizer effec- tively performs a single step of gradient descent (Thcorrrr in predicting the d-dimensional weight vector w" Rd when the number of examples for ICL is n (dd (Corollary 3.1 In contrast, we demonstrate that a one-layer transformer with CoT prompting can achieve near-exact recovery by executing multi-step GD (Theorem 3.2)
leading to significant errors
Convergence. We prove the convergence results of rumning gradient flow the population CoT loss under mild assumptions (Theorem 4.1). Our analysis uses a novel stage-wise approach combining dynamics analysis and landscape properties: the parameters initially approach the global minimizer, followed by local convergence toward the final solution. Our proof technique involves a novel characterization of the complicated population gradient. Furthermore, we prove that the trained transformer can exhibit both in-distribution and out-of of-distributtion generalization (Theorem (.2) at inference time. We are the first to establish the learnable separation between transformers with and without CoT under the in-context linear regression setting. We empirically validate that the trained transformer converges to the minimizer predicted by our theory, with distinct performance gap between models trained with and without CoT prompting.
Outline. In Section 2 formalize the problem setting including the data model, the one-layer transformer architecture, and the CoT prompting format. In Section 3 we theoretically show the performance gap between the transformer with and without COT. Section results, including our dynamics analysis and out-of-distribution (OOD) generalization result. Sec- tion cempiriiri validates the advantage of COT.
4 consists of our main
1.1 RELATED WORKS
Training dynamics of transformers. Several works have studied the training process of spe- cific transformer architectures. Jelassi et al. (2022): Li et al. (2023) examined the training process and sample complexity of Vision Transformer (Dosovitskiy et al.) 2020). [Tarzanagh et al. (2023); Ataee Tarzanagh et af. (202)); Li et al. (002a) explored the connection between the optimiza- tion landscape of self-attention mechanisms and the Support Vector Machine problem. (Tian et al.)
2023a c) provided insights into the training dynamics of the self-attention and MLP layers during the training process respectively.
related line of research focuses on Markov-like data models. Bietti et al. (024) studied the in- duction head mechanism from the perspective of associative memory. Nichani al. 2024) demon- strated that a simplified two-layer transformer provably learns a generalized induction head on latent causal graphs. Chen et al. (2024b) further proved that modified two-layer multi-heal transformer can learn in-context generalized n-gram. [Edelman et al. (2024) investigated the multi-stage phase transitions during training on bigram and n-gram 2 3). Additionally, Makkuva al. (2024) studied the loss landscape of transformers trained on sequences from Markov Chain.
Another growing body of literature aims understand the training dynamics of in-context learning (ICL). Garg et al. (2022) first empirically studied the ICL capabilities of transformers over a variety of function classes. Akyürek et al. (2022): Von Oswald et al. (2023) investigated the behavior of transformers on random ICL instances of linear regression. Several works have also established the existence of deep transformers capable of implementing multi-step gradient descent (GD) across dif- ferent domains (Fu et al. 2023 Bai et al. 22033] Giannou et al., 2023). Mahankali et al. (2023); (Ahn] ct al. (2024) analyzed the loss landscape of the linear regression ICL task and Zhang et al. (2023) proved global convergence on one-layer linear self-attention layer using gradient flow. Gatmiry et al. (2024) demonstrated that a linear looped transformer with specific update procedures can learn to implement mult GD for linear regression. Further analyses of training dynamics under more realistic assumptions about data models and architectures have been conducted by HHang et al. 2023): Kim & Suzuki (2024); Chen al. (202a)). For a detailed discussion see Appendix A.
2

---

Published as a conference paper at ICLR 2025
Compared to prior works, our study and Huang et al. (2023) Ahn et al. (2024); Zhang et al. (0023)
Tarzanagh et al. (2023); Nichani et al. (2024): Kim Suzuki (2024); Wang et al. (2024 } Chen et al.
(2024b): Ren et al. (024) all use similar reparameterizations that combine key and query matrices
to simplity the training dynamics. Moreover, many previous studies (Tian et al. 2023a Zhang et al.,
2023 Huang al.) 2023 Nichani al. 2024 Kim & Suzuki 2024 Chen et al. 2024a Gatmiry
et al.
2024) adopted the population loss to facilitate the analysis of these dynamics.
A closely related work is Gatmiry et al. (2024), which shows that looped transformer can imple- ment multi-step GD on the ICL linear regression task to directly predict the query answer in context. In comparison, the goal of our setting is to predict the weight vector from the input examples using a realistic CoT autoregressive generation process. Theoretically, we also establish a performance gap between transformers with CoT and those without. See Appendix A.2 for a more detailed discussion.
Chain of Thought and Scratchpad The CoT prompting method was first introduced by Wei et al. (2022) to enhance the multi-step reasoning capability of LLMss Before the formalization of CoT, Nye et al. (2021) demonstrated that allowing language models to generate intermediate results on "scratchpuds" dramatically boosts the multi-step computation ability of LLMs. [Wag et al. (202bb); Yao et al. (2024); CCreswell et al. (2022) ; (hou et al. (2022) further proposed variants of the CoT/scratchpad method to improve the efficiency and reliability of generation.
Recently, several works have attempted to understand CoT from both experimental and theoreti- cal perspectives. Wang et al. (2022a): Saparov & He (2022); Shi et al. (2022); Paul et al. (2023)
empirically studied the capability of Co',, providing valuable insights
its reasoning processes.
Meanwhile, Wu et (202)) ) TTtunov et al. (2023); Hou et al. (2023); Cabannes et al. (0024) inves-
tigated CoT through the lens of mechanistic interpretability. On the theoretical side, [Lu al. (2022);
Merrill & Sabharwal (223a); Li et al. 2024b); Feng et al. (2024) explored the expressive power of transformers with CoT, showing that CoT can significantly extend the expressivity of transformers in the context of circuit complexity. Hu et al. (2024) investigated the statistical foundations of COT. However, the training dynamics of CoT remain largely unexplored. To the best of our knowledge, this work is among the first theoretical analyses of training dynamics on CoT//cratchpad objectives.
PRELIMINARIES
In this section, describe the modified in-context learning linear regression task, i.e. in-context weight prediction, the one-layer linear self-attention architecture, and the Chain of Thought (CoT) prompting formulation.
Notation We use [T'] to denote the set 2, ...,,,,,, T'.. Scalars are in lower-case unbolded letters (y,a, a, etc.). Matrices and vectors are denoted in upper-case bold letters (W, V. etc.) and lower-case bold letters (x, w, etc.), respectively. W (ij) Wk.), Wi-J] respectively denotes the (i. j-th entry, i-th row, and j-th column of the matrix W. W!.-i] means the last column of the matrix W. The notation W.j denotes block matrices/vectors on the i-th row and j-th column according to context. For norm, ||-II denotes & norm and ||
r denotes the Frobenius norm. We use 1(-) to denote the indicator function. We use Ö(-) to hide logarithmic factors in the asymptotic notations.
2.1 IN-CONTEXT WEIGHT PREDICTION
Previous works (Zhang al. 2023 Ahn et al. 2023 2024 Akyürek et al. 2022 Mahankali et al.) 2023) focus on the in-context learning (ICL) task on linear regression. We suppose the data sequence sampled from linear regression task where the ground-trutt truth
" ~ Ia) x ~ N( yi w for all i i [n].
(1)
The goal of in-context learning to predict the correct label w Aquerry given a query Buuery and the previous example pairs (æ,, y:). Most previous works (Zhang et al. Mahankali et al.) (023) show the transformer predicts the query label Mquery by implicitly doing a one-step gradient descent without predicting the linear classifier w=.
2023 Ahn et al., 2024
In this work, we go one step further: instead of directly outputting the query label, we require the transformers to implement gradient descent to learn the ground-truth weight vector w". We call this task in-context weight prediction for linear regression. Specifically, the data sequence is in the following format:
3

---

Published as a conference paper at ICLR 2025
0
X
Zo-
y1
0
:
Odxn
y
wy
0
E e x(++))
(2)
0
0
1
xm
where X : [ , [n is the data matrix and wo the initialization of the linear parameter iù. We assume Wg 02 for simplicity, and define de 2d 2. Our setting similar the setting in Bai et al. (2023) where multi-layer transformers are constructed to do explicit multi-step GD on the weight vector wu. We separate the input example space and the weight vector space as Bai et al. 2023) (the {Pilie[N ()) in order to facilitate training. Moreover, we add a dummy token (an extra I) the end of each token similar what Bai et al. (2023) did in their input sequence format.
2.2 LINEAR SELF-ATTENTION LAYER
We consider one-layer linear self aelf-attention (LSA) module with residual connection, following the setting in Zhang et al. (202)); Ahn et al. (2023); Gatmiry al. (2024): we remove the softmax(-) non-linearity, consolidate the projection and value matrix into a single matrix V E RR, and merge the key and query matrices into W Rd.xdd, We denote
Z Zzz
LLS Z: V. W) Z+VZ
n
The prediction of the transformer will be the last token of the output sequence, namely
Z T W Z
FLSS (Z; V, W [.--1] 2 +Z
11
Since the first (d+1) 1) entries of the full weight tokens (0, 0, 1) zero, only part of the W and V affect the prediction. We can rewrite the parameter V. W into block matrices
V1
V12
V13
Vi4
W1
W12
W13
W14
V21
022
V23
1/4
W
W
1F22
W23
U/44
ER (222+2)x
V31
V22
33
V31
W
31
W32
33
W31
V41
042
V43
C44
W41
1642
W43
W44
where the block matrices are the following shape (i.j E {1,2]):
; W E gdxd, V2i-125 W2:-1,2j V2l.2j-1 W E gdx1, U2.22 W2r,2 E R. 12--1,55-1,
In the following sections, will show only V51, W1s, and u224 affects the prediction. We will further prove that all other entries are always zero along the training trajectory initialized at zero.
2.3 CHAIN-O--HHOME
PROMPTING
In language modeling tasks, transformers have been proven to be versatile in various downstream tasks. However, transformers struggle to solve mathematical or scientific problems with one single generation, where several reasoning steps are required. CoT was then proposed make transformers learn to generate intermediate results auto-regressively before reaching the answer.
With CoT, allow the transformer generate i steps before it outputs the final prediction thş for the ground-truth w". Specifically, given the generated input sequence Z; at the i-th step of generation, have fLSA (Z:)).-) as the prediction of the next token ((++1)-th token), and append it to the end of the current sequence s.t. 2.+1 = Žú. FLSA (2.) After generation steps, the CoT process induces & intermediate sequences (Z: 1 in the following form:
Æ1
*
2.
/1
Yn
R ( i
(Inference)
0
0
We
w.
0
Here, we define un; = JLSA (2)-1) as the i-th step prediction for the weight vector. The other entries in the same column are irrelevant and denote them as Finally, the transformer inputs the last generated sequence 2, back the transformer once again generate the final output aik-1 fLSA (x))(42:22411-1 the prediction of the weight vector w².

---

Published as a conference paper at ICLR 2025
Different from the inference time generation, the training process is similar to pre- pre-training the ground-trutt sequence to predict the next token. Specifically, we input the transformer with CoT ground-truth sequences Z;
nn
0
Z
V1
Yn
0
0
E g ((+++1))
i [k]
(Training)
0
WD
W1
wi
0
1
1
where w; Wi-1 1 X(X W is the ground-truth intermediate weight vector after i gradient steps on the linear regression objective. Each gradient step adopts a fixed learning rate for all possible training instances {X, w) when generating the ground-truth sequence Z₁. Note that Zi is the corresponding ground-truth sequence of 2,.
In the training objective for the i-th step, the transformer is required to predict the next token Zi+1|.-1] (Oz, 0, w;+1; 1) given the i-th ground-truth intermediate sequence Z, Finally, we predict the final ground-truth weight vector w" with the final intermediate sequence Zy. The CoT training objective given sample prompt X, y then becomes:
CoTT X w V. W)
= Σ 1
(LLA (Oa. O, w;,1,
(5)
Here we denote W w = w for clarity. Following Zhang et al. (2023); Nichani et al. (2024); Kim & Suzuki (2024); Tian et al. (2023b); [Chen et al. (2024a); Gatmiry et al. (2024), we consider the gradient flow dynamics over the population loss of the CoT objective:
C (V, W)
[X, "x, W)]
(6)
The following differential equation gives the For clarity, write the expectation Ex wo []- gradient flow dynamics of the parameters:
d0
VL CoT CoT (0),
0: (V. W).
dt
When measuring the performance after training, we apply the CoT inference procedure to generate k intermediate sequences { [ and consider the final output token f( |: by inputting the last generated sequence 2g. The performance evaluation measured on the error between the final output f( f(2x) and the ground-truth w²
||EESA (L (0,, O, wi, 1) 1)]²]
Evall (V W)
x
When CoT prompting is not used (A: = 0), the evaluation loss Eval equivalent to CCoT
EXPRESSIVENESS IMPROVEMENT WITH CHAIN OF THOUGHT
In this section, we theoretically explore the performance gap on our data model between transform- ers with CoT and those without. We first prove that one-layer transformer without CoT can only implement a one-step GD and cannot recover the ground-trute truth, while it can near-exactly predict the ground-truth parameter with CoT by implementing multi-step GD.
3.1 ONE-LAYER TRANSFORMER CANNOT RECOVER GROUND- TRUTH
For the ICL linear regression task, the optimal prediction given by a one-layer linear transformer is equivalent single step of GD the MSE objective of linear regression (Mahankali al. 2023).
What about our task on predicting the ground-truth weight vector w" in context? The following theorem proves that the optimal solution is still one-step GD solution.
Theorem 3.1 (Lower bound without CoT). If the global minimizer of C (V. W) is (V*, W*), the corresponding one-layer transformer fLSA (Zo);--1 implements one step GD on linear model with some learning rate 1)" n404I and the transformer outputs (0a, 0, X y ,).
5

---

Published as a conference paper at ICLR 2025
We briefly present the high-level intuitions in the proof and the detailed proof is deferred to Ap- pendix B. We use a similar technique in |Mahankali et al. (2023) when proving the optimality of step GD in the ICL task. The key strategy of the proof Is to replace (0a, 0, w", 1) in the evaluation loss Eval (V, W) (Equation 7 with (0a, O, Xy ,1) the following form.
Evall V W)
I 2
fLSSA (I.SAZZ) (Zo) Od 0, T X
|1 +C
In order to prove this equation above, we show the gradient of the original loss Equation 7 and this formula are identical. We first obtain the closed-form formula of the expected gradient for both sides with regard to X, w". Then simplify the gradient expressions, and eventually prove them equal.
the symmetric property of the distribution of X, w to
The equivalent form of loss indicates that the evaluation loss only depends the lp distance between the output of the linear self-attention module and (Oa, O, " Xy",1) Therefore, any (V, W) is a global minimizer of this loss function if and only the output of FLSA (2x) is (Oa, O, x' Xy 1). Meanwhile, one can assign
0
0
0
1
0
0
-1
V
W
(8)
17
0 0
0
and the one-layer transformer achieves the optimal solution, which concludes the proof.
Is one-step gradient solution good enough? Most of the previous ICL work Zhang et al. (202 ; Ahn et al. (2023); Gatmiry al. (2024) consider the number of examples n - to00 when d is fixed. In this case, the one-step GD solution can perfectly find the ground-truth weight vector w". However, a simple corollary of this theorem indicates that the one-step solution has a non-negligible error when there are limited samples, e.g. n @(d). This number examples n is required to guarantee the reconstruction of w² Rd
Corollary 3.1. For any parameters (V, W) in the transforrr transformer, Evall (V, W)
+ L
Moreover, if = (d)) EEal W) ((d)
d +00.
Proof. By Theorem 3.1 we directly calculate the evaluation loss on the global optimum:
2
Eval (V, W
2 Ex,
17" xx w w²
= Ex tr
17" xx
2
n
since Eng. u=0
I. Apply E.
XXX nI and (XX )²]
nnn + d + 1)I,
1 Ex tr
xxxx 2 1 - (I-""xxT²
I- " n
x2
(n + d 1)d = d2
(
2nj"d
2
n
and we finish the proof by substituting with @dd).
3.2 ONE-LAYER TRANSFORMER WITH COT CAN IMPLEMENT MULTI-STEP GD
The previous subsection shows that the one-step solution by the one-layer transformer without CoT is not the endgame. Nevertheless, CoT can become the savior for this simple transformer because enables the transformer to generate several intermediate computation steps improve the final per- formance. The following theorem shows that with the reinforcement of CoT, there exists a one-layer transformer that can perform multi-step GD using intermediate generations. We show that Ollg d) steps of CoT can remarkably improve the performance, reducing the error from O poly a to O(1/ poly d). With constant learning rate, O(log d) steps of GD is also necessary to reconstruct w accurately. The proof is deferred Appendix B.2

---

Published as a conference paper at ICLR 2025
Theorem 3.2 (Informal). There exists V and W s.t. fLSA (Zx)), output (O Wk: 1) where wy: := -1 2 x )²)=" is the k-step GD solution with learning rate n on a linear regres- sion model. Moreover, if = ((d), k = Dilog d), E (0.1, 1), then the evaluation loss
Evall (V", W-) = xx wo ( xxxxxxx
7
k+1
w m01
1 poly(d)
S
2
With the one-step GD solution in Theorem 3.1 the proof is straightforward: we assign the param- eters (V, W) in the same form of Equation 8). with the 1f" replaced by n. However, the transformer is allowed to generate & steps before reaching the final output. We can inductively calculate the i-th step of generation, showing that the output is exactly the i-th gradient step:
FLSA ZZ 1)) (0d, O, wt, 1), = .... +1
After k+1 steps, we have the final output (I - (I " XX )(+1) n by induction and the evalu- ation loss becomes Equation 9). By Lemma D.4 the final loss is upper bounded by poly(ai This is strictly better than a one-step GD solution by comparing with Corollary 3.1
Now we theoretically display the expressivity improvement of transformers brought by CoT. In the following sections, we will further prove that this separation is learnable simply by gradient flow.
GRADIENT DYNAMICS OVER CHAIN OF THOUGHT
In this section, we go beyond the construction and prove our convergence result on the CoT objec- tive. We show that the final solution found by gradient flow is approximately our construction in Theorem 3.2 which is significantly better than the one-step gradient descent solution without COT.
4.1 MAIN RESULTS
According to our construction in Theorem 3.2 we use the following specific initialization to zero the irrelevant blocks while keeping the essential blocks W13, V31, and 4244- Assumption 4.1 (Initialization), Let 0 be parameter. We assume the initialization of the
parameters satisfies that
0
0
0
0
W
13
(0)
0
0
0
W/44
V
W
Vn1
0
0
0
0
0
0
Here W1 >d f== AW u and (01 Σd X u symmetric and simultaneously
diagonalizable, AY - la, 11. Further, fix W244 for all t >> 0.
This initialization follows Chen et al. (024a) by assuming V51 and W13 share the same set of eigenvectors. It is close to the particular symmetric random initialization schemes discussed in Zhang al. (2023) with scaling factor a. We use this specific initialization to zero out the irrelevant blocks, and we fix W21 = -1 to break the homogeneity of the model and avoid multiple global minimizers. Those simplification facilitates the analysis of the complex dynamical system.
Now we prove that under appropriate initialization, gradient flow will nearly converge to the global minimizer. We provide proof sketch in the next subsection. See Appendix C.3 for details.
Theorem 4.1 (Informal, Global Convergence). Suppose = Š2(d), with Under Assumption with o = ((1), ie run gradient flow on the population loss in Equation (5
(0.1, = O(log d)
then after time = O d + log } we have CCoT(t (t) S efor any e (poly(z):
4.2
PROOF IDEAS
In this subsection, we briefly outline the proof of Theorem
first prove that under Assumption 4.1 the gradient dy- namics only depend on the parameter blocks Wia(t), Vai W24. while other blocks stay zero (Lemma C This is because the Gaussian data assumption makes sure the gradients of these
Before analyzing the dynamics,

---

Published as a conference paper at ICLR 2025
blocks are zero once initialized at zero, except for Wts(t), V11(t), W21. By this lemma, we can simplify the linear self-attention formula and consider the following equivalent yet simplified loss (we denote W == W13, V = V31, and W24 is fixed as 1.):
1 -(V x w"
[ (0) Ex ~
(Vxx W + Jw,
n
1=0
(1+ wk - -O- (LVXXT +1 w"
+ Ex
For ease of presentation, we denote S == XXx. To analyze the gradient dynamics, we first need to compute the exact closed-form gradient instead of keeping the expectation. However, the formula involves the i-th step weight vector w; = (T- -I -))))))))))))))))))) w", involving the higher order
moments of the Wishart matrix $D whose closed form is hard to obtain. In our paper, we provide a tighter estimate compared to previous work (Gatmiry et al., 2024) using the concentration of the Wishart matrix S (Vershynin, 2018) when n = O(d log d) to estimate the expectation. In particular, we use the exponential decaying tail probability bound for the operator norm of the error For example, when estimating the expectation [[[ - n)S)",, we can decompose the
ós = S expectation into two cases: when | is small, (I - n5)² = (1 - n)* I; when ||ōS|| is larger than a threshold, the rest part of the expectation can be controlled by integrating the exponential decaying tail probability The concentration lemmas are provided in Appendix D
The motivation behind a better concentration estimation is to ensure nearly independent dynamics along different eigenspaces {u:] of W and V. As an extreme case, we consider -> oc and S converges to I almost surely. Now the gradient component on the u.u.? subspace is only dependent on X and A without any other Xy. A.W .j o involved. That means there no interaction be- tween two different subspaces, i.e. the dynamics are independent. However, some interactions are introduced since the concentration error 85 * 0 when n is finite. Therefore, the improved charac- terization of the expected gradient is essential upper bound the interaction between the dynamics of different eigenspaces {u, )d leading to a nearly independent evolution at initialization. This independence property motivates us to conduct the following stage- analysis:
Stage 1: W, V converges near-optimal. In this stage, the dynamics along each direction u, stay nearly independent. Specifically, we can expand the gradient flow dynamics for V. W and project them into the eigenspaces u.u. to get the dynamics of the eigenvalues X! := u Vu.. \ := u, u, The dynamics of eigenvalues are characterized by the following Lemma 4.1 where we can prove that the interaction terms between different subspaces are bounded by O log 2 d).
Lemma 4.1 (Informal version of Lemma
C.6
The dynamics of XY and AW are given by the
following equations with /
S
0
S O
log d
log" d
A W
drp
(k + 1)( 1/1 )
2
2 W
15 1- 11 w - 1 + +06
=
+
-
(1-
+
dt
11
(1(2 m)
2
dxW
k +111- 1 m 25² ( 1- 1/(2 7)
1-1 v 6/
w
+ 2 2- 11
di
This nearly independent evolution along each eigenvector u: enables us to analyze the individual dynamics of X and A. at the beginning of training. Under Assumption 4.1 X .WW are initialized
(11) By Lemma 4.1 we prove by induction that the eigenvalues will go through two phases: (1) Xÿ
1 To deal with the similar problem, Gatmiry al. (2024) proposed simple combinatorial method to estimate the expectation. We use the same technique to get a certain form of the expectation (see Appendix []. but the hound is not tight enough to get the desired results. See discussion in Appendix A.2
"his method can keep the (1 (1-1) n1' factor to prevent introducing unwanted estimation errors when is large.
8

---

Published as a conference paper at ICLR 2025
increases yet stay smaller than -0
k(1-AW)) while ir increases to 1-o(1). ) (2) My stays ((1)-
close 1, and \ also converges to -cl)-sl -7. Here all ((1) terms some O(1/ llg" d) terms for some constant с > 0. That indicates that the distance between the eigenvalues and the target Ay + / / - converge to ( log" d) for all j E [d] at the end of Stage I.
Stage 2: Local convergence. One may expect that after Stage 1, the transformer can approximate gradient steps quite accurately since the parameter V. W are both o(1) to ground-truth along each direction u, Unfortunately, the sum of error in directions can still be (dd) since we can only reduce the error (1/ poly log d) in each direction. So for now, the solution cannot recover the weight vector wn at this stage. To address this issue, we further consider the exact form of the interaction terms 6,4 .5 and analyze the local convergence. By fine fine-grained expansion of the error terms, notice that 6/4 and 6y are always coupled with some individual residual like (1 - AW)
(n + some weighted average those individual residuals. Meanwhile, the coefficient of the residual in the interaction terms is still upper bounded by O(1/ poly log d). That enables to derive some gradient lower bound similar to PL-conditions (Lemma C.12) when V, W are close to the ground-truth, leading to local convergence to near-optimal at a linear rate.
The final training error is some O poly )) which depends on the inference step k and ground-truth Note that the optimal loss value is also at least polynomially small in d given (log d) CoT steps. Therefore, now we can conclude that the transformer can learn to implement multi-step GD when given intermediate ground-truth states after optimizing the CoT loss with gradient flow.
(olyya)
4.3 OUT-OF-DISTRIBUTION GENERALIZATION
INFERENCE
In this section, we prove that after training. the transformer not only correctly predicts the weight vector context with CoT generation. but also can generalize outtof-distribution (OOD). The fol- lowing theorem shows that the trained transformer obtained from Theorem 4.1 with CoT generalizes over other problem instances when the input example sequence has an OOD covariance, as long as the covariance is not too ill-conditioned. Here E is defined the OOD evaluation loss in eq. with the in-context examples æ; ~ (0, Σ)
and weight vector w ~N0
Eval W)
2 E. No.
[1
(0a, 0, w".
Theorem 4.2 (Informal, Theorem C.2 Suppose = (dd). E (0.1,0.9). k' = (log d). Assume
the out-of-diss ofdiistiibttion covariance is well-conditioned: 5
" < Amin (E) < Amaax (E) 2-6 for some
constant 8 > 0. Then after training in Theorem 4.1 we have C (t) S é for any e polyyay
Note that this theorem covers both in-distribution (when 77 = 8) and OOD tasks at evaluation, indicating that the transformer trained to implement a general iterative optimization algorithm. Moreover, the inference step number 1 in this theorem can go beyond the training CoT steps k achieving better estimation for w=.
One may think once the next-toked preedicineiiiion training loss _CoT converges to zero based on ground- truth CoT data, the transformer naturally learns to do multi-step reasoning at inference, i.e. Eval small. However, at the i-th generation step, the transformer is predicting the next weight token w; based on the previous generation W; instead the ground-truth intermediate step w, It is possible that prediction error for each step accumulates we expand the sum of all the prediction errors at each step and show a converging series of errors throughout the inference process. That ensures can achieve any O( (pllytaa) )-small evaluation loss when we have k' (loo d) reasoning steps. The detailed proof is provided in Appendix C.4
even increases exponentially. In this theorem,
EXPERIMENTS
In this section, we introduce our experimental setup on our in-context weight vector prediction task to numerically validate theoretical results. Specifically, show that parameters of the

---

Published as a conference paper at ICLR 2025
transformer match the prediction of our theory when optimized the CoT loss. Furthermore, present the gap of evaluation loss Eval in Equation between transformers with and without CoT.
Experimental Setup We train the transformer architecture in Equation 3 the synthetic data. The data distribution follows our in-context weight prediction task Equation ① In particular, we choose the token dimensions d 10, number of in-context examples
20, and GD learning rate = 0.4 for generating the ground-trun truth intermediate states. We use a batch size B = 1000 and run Adam with learning rate 0.001 for 750 iterations. More details refer Appendix E
Global convergence Our experiments show that the structure that weights of the full model exhibit consistent with Theorem 3.2 At final convergence, all of the entries of W converge except the elements on the diagonal in the top-right comner block (the red box in the heatmap of W, Figure D. while all the entries of V are near zero except elements on the diagonal in the bottom-left corner (the red box in the heatmap of V. Figure 1 Also, the pattern shows W13 = aI, W244 -Q, and V11
" with some scaling factor a which equivalent the construction stated in Theorem 3.2 and Theorem That means the transformer implements one step of gradient descent (0,,
(w: w w²), 0) before the residual connection, and autoregressive CoT process enables model to perform multi-step GD.
Performance improvement We empirically verify the evaluation loss gap between transformers with and without CoT shown by Theorem 3.1 and Theorem 3.2 Our experiments in Figure 2 demonstrate that the evaluation loss of transformers with CoT converges to near zero even when achieve
k 10. In comparison, the optimal expected loss that the one-layer linear transformer (the pink dashed line, from Corollary 3.I is much larger than any of the model that applies multiple steps of computation. We also observe that evaluation loss convergence keeps decreasing when the number of reasoning steps increases from 10 to 40, which is consistent with Theorem where larger allows for smaller error c.
k-10
k-20
k-30
k-40
1-
100
200
600
700
Healmap V
Heatmap of
ieeation
Model weights: We present the heatmap of the weights of the trained transformer. We initialize V, W randomly at = O, where 20, d 10 and k 20. After training, all entries of and W converge to zero except the two blocks highlighted in the red box. More- over, the pattern matches the theoretical results.
Figure 1:
Figure 2: k-step v.s. 1-step: We plot the evaluation loss [vval when 10. We randomly initialize the transformer. For transformers with CoT, loss converges to near zero while transformers without CoT cannot. Moreover, the loss at convergence decreases when increases.
20,
CONCLUSION
This paper investigates the training dynamics of transformers when the Chain of Thought (CoT) prompting introduced. By focusing on the in-context weight prediction task, our theoretical results demonstrate that transformers learn implement iterative algorithms like multi-step GD with the enhancement of CoT, highlighting the essential role of CoT in multi-step reasoning tasks. Our empirical findings corroborate these theoretical insights, indicating that CoT prompting provides significant performance benefits.
There are still many open problems. Can we move beyond population loss on the in-context weight prediction task and show a sample complexity guarantee? Can CoT empower the transformer to acquire compositional reasoning capability instead of doing the same iterative steps?
IIn Figure 0 while all 0 works for the construction. Empirically, the sign of a depends the random initializatio and both positive and negative solutions exist.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
JDL acknowledges support of the NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877.. 2023.
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Lin- ear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082 2023.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211. 2022.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Am- brose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546-38556, 2022.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token se- lection in attention mechanism. Advances in Neural Information Processing Systems, 36:48314- 48362, 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint arXiv:230.04633, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Vivien Cabannes, Charles Arnal Wassim Bouaziz, Alice Yang, Francois Charton, and Julia Kempe. Iteration head: mechanistic study of Sfain-of-thought arXiv preprint arXiv:2446.0221 2024.
Siyu Chen, Heejune Sheen, Tianhao Wang. and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. /9442, 2024a.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.10559 2024b.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Antonia Creswell, Murray Shanahan, and Irina Higgins. Celection-inferencee Exploiting large lan- guage models for interpretable logical reasoning. arXiv preprint arXiv: 2022.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning. arXiv preprint 16111:2308.06612, 2023.

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Benjamin L Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint 20Xi::2402./10044 2024.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order op- timization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310. 11100/7088 17086, 2023.
Shivam Garg. Dimitris Tsipras, Perey S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neurul Information Processing Systems, 35:30583-30598, 2022.
Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learn- ing? In Forty-first International Conference on Machine Learning, 2024. URL https:
77 //openzeview.o net www.ooonnndddddddaaaa-
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dim- itris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301. 13 13196, 2023.
Yifan Hou, Jiaoda Li, Yu Fei, Alessandro Stolfo, Wangchunshu Zhou, Guangtao Zeng, Antoine Bosselut, and Mrinmaya Sachan. Towards mechanistic interpretation of multi-step reasoning capabilities of language models. arXiv preprint arXiv:2310. 14491, 2023.
Xinyang Hu, Fengzhuo Zhang, Siyu Chen, and Zhuoran Yang. Unveiling the statistical foundations of -fahnoof- prompting methods. arXiv preprini (ttl:2400./4511// 2024.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37888 2022.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances neural information processing systems, 35:22199-295 222133 2022.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shal- low vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06011 2023.
Yingcong Li, Yixiao Huang, Muhammed Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention In International Conference Artificial Intelligence and Statistics, pp. 693. PMLR, 2024a.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint arXiv:2402. /2875 2024b.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata. arXiv preprint arXiv:2210. 10749, 2022.
12

---

Published as a conference paper at ICLR 2025
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576. 2023.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.0416/ 2024.
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:2310.07922, 2023a.
William Merrill and Ashish Sabharwal. The parallelism tradeoff:: Limitations of log-precision trans- formers. Transactions of the Association for Computational Linguistics, 11:511-545, 2023b.
Eshaan Nichani, Alex Damian, and Jason D) Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.1 1 14735, 2024.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021.
OpenAI. Gpt-4 technical report, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Ed- ward Yang. Zach Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019.
Debjit Paul, Mete Ismayilzada, Maxime Peyrard, Beatriz Borges, Antoine Bosselut, Robert West, and Boi Faltings. Refiner: Reasoning feedback on intermediate representations. arXiv preprint arXiv:2304.00 0/04, 2023.
Yunwei Ren, Zixuan Wang, and Jason D Lee. Learning and transferring sparse contextual bigrams with linear transformers. In The Thirty-eighth Annual Conference Neural Information Pro- cessing Systems, 2024.
Abulhair Saparov and He He. Language models are greedy reasoners: systematic formal analysis of chain-of-thought, arXiv preprint arXiv:2210.012400 2022.
Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang. Suraj Srivats, Soroush Vosoughi, Hyung Won Chung, Yi Tay, Sebastian Ruder, Denny Zhou, et al. Language models are multi- lingual tffhuutf- reasoners. arXiv preprint arXiv:2210.03057 2022.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe Le, Ed Chi, Denny Zhou, et al. Challenging big-bench tasks and whether thain-of-thought can solve them. arXiv preprint arXiv:2210.09261 2022.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Sean and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint 1.XI::23005.63880 2023a.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer, July 2023b. URL http::/apxxxii. org/abs// 2305. 16380 arXiv:2305.163. [cs].
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying mul- tilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00533 2023c.
13

---

Published as a conference paper at ICLR 2025
Rasul Tutunov, Antoine Grosnit, Juliusz Ziomek, Jun Wang, and Haitham Bou-Ammar. Why can large language models generate correct chain-of-thous th-tugght? arXiv preprint arXiv:2310.13577 2023.
Roman Vershynin. Hight-dimessional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordv- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174, PMLR, 2023.
Boshi Wang. Sewon Min, Xiang Deng, Jiaming Shen, You Wu, Luke Zettlemoyer, and Huan Sun. prompting: An empirical study of what matters. arXiv
Towards understanding preprint arXiv:2212. 2022a.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le. Ed Chi, Sharan Narang, Aakanksha Chowdh- cry, and Denny Zhou, Self-consistency improves chain of thought reasoning in language models. arXiv preprint vrXhv:2203.1177/ 2022b.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference Machine Learning. 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.
Skyler Wu, Eric Meng Shen, Charumathi Badrinath, Jiaqi Ma. and Himabindu Lakkaraju. Analyz- ing thain-of-thought prompting large language models via gradient-based feature attributions. arXiv preprint arXiv:2307. /0.33077.33,,.... 13339, 2023.
Shunyu Yao. Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in Neural Information Processing Systems, 36, 2024.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett, Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.009927 2023.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc Le, et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205. 10625, 2022.
14

---

Published as a conference paper at ICLR 2025
A DISCUSSION AND LIMITATION
A.1 RELATED WORKS ON EXPRESSIVENESS
Our work closely related the previous works in multi-step GD using multi-layer attention layers, including [Bii et al. (2023): ; Fu et al. )023); Ding et al. (2023): ; Ahn et al. (2024); Giannou et al. (2023 Gatmiry et al. (2024). These works guarantee that transformers are expressive enough to do in-context learning by implementing gradient descent, and they serve as the foundation of our work which focuses on optimization. Most of them focus on the in-context learning setup as the testbed so we naturally follow the setup to understand the advantage of CoT.
Most of the above works on expressiveness focus on those iterative algorithms, e.g. (pre- conditioned) gradient descent various objectives (Bai et al. 2023 Ahn al., 2024, Ding et al.) 2023), Newton methods/matrix inverse (Giannou et al. 2023), etc. Those papers have similar con- structive proof techniques using multi-layer transtormers: they construct basic block(s) to rep- resent one step of some iterative algorithm and stack them up to do multi-steps of that algorithm. Sometimes the blocks can be even the same, which means "looped" transformer, i.e. implementing the same transformer blocks several times as a loop, can express those al algorithms. In our warm-up construction for better understanding of the setup, we similar techniques to construct the linear transformer that allows auto-regressive generation to iteratively implement the block. However, we require the practical auto-regressive setting, which is novel in the literature.
Most importantly, despite the close relation between our work and those previous expressiveness papers, our work mainly focuses on the optimization perspective. siveness because there is no guarantee that one can algorithmically find the constructed solutions in the previous work. Ahh et al. (2024); Gatmiry et al. (2024) are the only two papers related to optimization of multi-layer transformers over in-context linear regression setup. [Ahn et al. (2024) analyzed the global optimizer/critical points for multi-layer transformers, but they didn't 't prove that any gradient-based algorithm can reach those solutions. Compared to all the works above, proof' techniques for the main theorems are completely orthogonal and not straightforward extensions of the previous papers like Bai et al. (2023).
big step beyond expres-
Gatmiry et al. (2024) is the most related work to us. They also proved some results on learning to implement multi-step GD by looped transformer. We will highlight the differences and our novel
contributions of our work in the next section.
A.2 DISCUSSION ON GATMIRY ET AI. 2024
In this section, compare work with Gatmiry et al. (024). We begin by outlining the simi- larities and connections between the two works before highlighting our theoretical contributions in contrast to Gatmiry al. (2024).
Both Gatmiry et al. 2024) and our study analyze the dynamics of a one-layer linear transformer in the context of a linear regression task, demonstrating that transformers implement multi-step gradient descent. We adopt similar architectural frameworks to those in /Zhang et al. ): Ahhn
al. (2024) 2023); Mahankali et al. (2023). as well as several other works. The key connection between our work and Gatmiry et al. (2024) lies in the observation that both looped transformers and transformers with CoT prompting through autoregressive generation are capable of naturally implementing iterative algorithms like gradient descent.
However, our data model and training objective are intrinsically different from those in Gatmiry al. (2224), leading to distinct insights. While Gatmiry et al. (2024) focuses on an ICL setting for linear regression tasks involving examples and query, our task is centered on predicting the ground weight vector w" within context, i.e. in-context weight prediction. The final converging solutions are totally different, even though they both are equivalent to some type of GD. From the perspective of the training objective, (aammiry et al. (2024) uses standard squared loss over the ICL objective. In contrast, we use sum of squared losses across all intermediate steps, corresponding to the CoT loss highlight the effectiveness in improving the performance
defined in Equation (6). Therefore, of the CoT prompting on a shallow transformer, while Gattmiyy et al. (2024) stress a multi-layer transformer with shared weights (looped transformer) do multi-step GD through the layers.
15

---

Published as a conference paper at ICLR 2025
From technical perspective, Gatmiry et al. 20244) fix the outer layer and train only the matrix A, which is analogous to our matrix W. In contrast, our work allows for training both layers of the transformer, providing stronger analysis of training dynamiics. Our proof strategy is also novel, given that our training dynamics are more complicated: obtaining our final solution requires solving challenging d-dimensional dynamical system, whereas prior work in ICL. reduces the outer layer to scalar.
As a more profound theoretical contribution, we rigorously establish clear performance gap between the one-layer transformer without CoT and the ones with CoT. Specifically, the one-layer transformer without CoT restricted a single step of GD, with the final error (dd poly logd), while one-layer transformer with CoT can achieve O(1/ poly d) loss with only (log d) steps. the other hand, Gatmiry al. (2024) do not show their transformer implementing the multi-step GD can outperform the transformer with one-step GD. According to Theorem 4.2, the looped transformer in their setting can only provably get the final loss down to 1/ where is v/
the number of loops. However, a one-layer transformer achieve (d /n) loss by implementing one-step of GD, asymptotically better than the multi-step solution in Gatmiry et al. (2024).
A.3 LOOPED TRANSFORMER LEARNS TO IMPLEMENT MULTI-STEP GRADIENT DESCENT (IMPROVED VERSION)
In this section, we also discuss how our estimation technique improves the result in Gatmiry et al. (2024) for looped transformers. The gap between our analysis lies in our different methods of calcu- lating the terms in the gradient concerning Wishart matrices. For intuition, we introduce the novel expectation calculation method Section 4 which asymptotically improves the estimation of higher moments of Wishart matrices in Gatmiry et al. (2024). We adopt the combinatorial technique
in Gatmiry et al. 2024) to compute the form of E SASTSS* but when we calculate the expected gradient we use the concentration tail bound technique to calculate the expectation.
We will demonstrate that applying our techniques proves looped transformers outperform their coun- terparts without looping the ICL setting. Here, we follow the notations, models and loss defined in Gatmiry et al. (2024.. We focus on the simple setting with Σ = I, but the same technique should also be applied to cases with other covariance matrices.
Definition A.1 (Corollary A.4 in Gatmiry et al. (2024)).
The loss for loop Transformer is
[(A,0) 0) = E tr (I A
26
By using our techniques, we strengthen the conclusion and obtain the following theorem. Theorem A.1. Suppose M(dL2 log² d) and ||(())] O(). Consider the gradient flow with respect to the loss L(A,0):
A(t) -VVA((A()).0)) 0).
Then, for any E> o [Lidloz' In particular, given any polynomially small & > ( (pply(a)) ) there exists L O(log d). the final loss C(A(t), 0) E.
>((( d
)
L.
11' :) we have L(A(I), S.
after time 122 Ω 1
Remark. This result also gets arbitrary polynomially small loss the case in the CoT setting, establishing the separation between looped transformer and the transformer without loop.
Proof. Denote X to be the i-th eigenvalue of A and u, to be the corresponding eigenvector. Let
k arg max
Asl. X = dp-
Suppose u E arg max leld |1 - X is the corresponding eigenvector. Apply Lemma A.3 we have
dA
df = 2L(I - A) 20-1 + I | 2 AZI.
16

---

Published as a conference paper at ICLR 2025
Multiply on both sides (note that is fixed direction, so the time-differential is zero), obtain
d(Au)
di = 2L(1 - 2) + A1| u.
2
Note that the gradient of A has the same eigenvectors u, as A. Therefore, A2L have the same eigenvector A and have
d(1 2)
du (I A)u 2 2)22-1 (1 -1)22-1- A.Lu.
dt
dt
By Lemma A.3
log d ) 1 The dynamics of (1 - X)² become: we have | $ (
d(1 - x)2
4L( x)ah 2(1 2) DELu S 3L A)2L,
We therefore upper bound the difference between À and 1 by solving the ODE:
(1 -(t)²² - ((1 ((1-MM) 1(0))2-2 3L 1)t) - -
4
When the training time 2 2
1
d
the largest eigenvalue of - A, i.e.
17
I - A = o (
)
To satisfy the condition on Lemma A.3 the & should be lower bounded by @ @ (((Laload))
Now consider the loss expression:
A:SA:) 24 E tr ( (II- 2
L(A.0) 0)
SA) =I
tr
(II AI- 22 || A1
(By Lemma A.2p
< 2d I 2 < O(E).
By the computation above, ((A. 0) ( In particular, when L = clog d and 2(L'd d) with e > 1. the loss is smaller than ((d d' Hence, only O d) steps of looping can achieve arbitrary polynomially small loss O (oly
Remark. If we have O(log d) steps of GD using this looped transformer, we can get an arbitrary polynomially small Joss. is a huge improvement compared to Gatmiry al. (2024). and this result successfully establishes separation between the looped transformer and the ones without the loop.
Lemma A.1 (Equation (24) in Gatmiry et al. (024))).
The derivative of the loss can be written as
2L-1
VAC(A,0) = Σ
E
SAY'S SA) 2L 21-1-11
SA)2
1=0
Lemma A.2. Assume ((dL?), || A 2 0
1.'dlog2
12
B (I SI) - 24
26-1 A1- | A 26-11
where A has 0 operator norm.
17

---

Published as a conference paper at ICLR 2025
Proof. Denote 85 S - I. Then expand the term in the expectation:
(I S( - -5)) = (I - A - is A) + i5) - A - ös A) 2)26-1-i
Now take expectation to both sides. Note that [E|ss] 0, so all terms containing first order 8s vanish. We denote
I
All
Al
2- A = -
A) + 68) A)2L 2)2
- -i(I -) ös A(I - A) - (2L i) - A)) 2) öS A
to be the sum of all higher order terms (the degree of 6$ 2 2). We can estimate the expectation using similar technique in Lemma D.I
A) 21 (OS A)
2L
1)
(Term 1)
k
11 A1
1-i 21
(I-A)'85( A)'8S(I A) 2) (65 1
A)
2L 1 1--1
(Term 2)
k
|-- 21
k=1
21-1-i
2L - 1 1-i [I (-) - - 16S* A) 88(1 - A) 2t (6S A)
|| A|- 22-1
1=0
(Term 3)
To get an estimate of the operator norm, we bound each term (Term I to Term 3) respectively.
Term 1:
[1 A) 2L. 1k (8S AA
2L - 1)
| (A51|-|1A41 ||Al| | |-
(2L 1)
k
| Ali 21
2=2
k=2
Note that || of order 0( 0(v5) ( with high probability, the term in the middle is less than 1 and
the dominating term of the error is ((2I. 1) ||6S11-14!)
O
1
log
Term 2:
21 1-i 1-i
[I 85 2) (8S A)*
2L - 1 1-1
k
|| ||-A|| 21
k=1
I
All
2L-1-i
||:5||| ||A||
(2L-1-2
|
|1 - All
k=1
Term 3:
21-1-1
[I A) A)* *6S *)= A)* AA) A)2 2 (8S AA)
2L - 1 -1
Σ
Σ
1
|| A1-
k=1
1=0
1
All
|8S|| |A||
21-1-1
|
Σ
(2L 1
|
All
| All
I=U
Now we upper bound the operator norm of the error term A A A=EE[A]
||||
B
A
1E
=
= E
|لقا|
1
| S
C'
+ 11112 |
log d
log d
C"
5
+
Pr |
>
ds.
log
d
c² Toga
18

---

Published as a conference paper at ICLR 2025
When || 2 where C" some constant C1 S.t.
logd can first upper bound the || with ||6S]| using there exists
||ه|| S max (CCI)))SS)) (C: (C1))))))22))
Therefore, when ||ر s, ||oS|| 2 min 8//2 -1(22+
}. To apply the tail bound, we need to make
CTI
sure we pick some s' such that max 02) & CiI probability, where ā C(V / + ) Now since
CiL } to upper bound the integral of
log min 2//2 11/22144) 2 Co
C"
CT
log
for some constant Ca. Therefore, we just need max min C2 2/ Va C3 " s1/10440), vi Cs L C3
5 V
S min 8/12 21e44
Ci C.t. b, i.e. s' S
n
8'/4 VL
VI
Applying the tail bound (18) with s' min C2
Va
C3
2 Vo C. y I CB
8/14 8/AVA Vn
VL where Cz, Cs, Ca. Os are some constant, we have the error term for the tail expectation,
vL
8//2 s1/(2444)
}} ds
Pr | 2
2 ds S
Pr | min
CL
CLL
s 2
exp -o
2 dds.
Now we estimate the upper bound of error with
52/(L+2)
82/(2++4)
C Vsn
8²2 min C
n, C.
n, C²
L
1
Gān g1(((+2);
For the first term, let x
L
{ -C3.
s1/(L+2)
2 exp
nyds
L
L
L+2
exp( 1ap(-aJahtld dx
Ch
1.2
L+2
Chn n
C²
ziz
1++1
Chn
C
$ 2()
exp
C2
1
log d
L2
log d
log d'
The second term, let x C3.
81/120+4)
b
m
81/(2(+1)
f exp
CG-
n
nds
or logd
L
2L+1
L
4[L + 2)
exp( (xp(-z)a
2. 2L+3 dx
cJn
C3n
(lord)
2L+4
C
C²
2L14
21+3
Cān
/
C"
24.4
1
44
exp
Cyn
L
log d
I
log d
log d'
Cjen. 1.
For the third term, let x
Cs sn
000
L
2
exp
ds =
expl d=
1:2
oin
Cin
Tog
2 2
C"
Cin
1
exp
<
C
logd
I
log d'
19

---

Published as a conference paper at ICLR 2025
The fourth term, let
C
In
1/2
82/2
2
exp
CG-
ds
4L2
exp( Jdd
n n²Ch
00 b loge
1//2
4L2
C²
1/2
C²
1/2
C3 T.
-02
n²
logd
exp
log
log d
Therefore, we plug this error back into the upper bound of ||||
C"
||| -
Pr
|
ds
log
C²
log
C"
1122 21((2+41)
} ds s c0 (
Pr
|ōs| 2 min
(
log
d
Cik'
Cik
lage
Lemma A.3. Assume n = (d(d))), for any A that | - All 2 LRd d
we have the
following gradient estimate:
V VAC(A,0) - 21 A) || - All 2- 1..-111 A2r.
where A22/l S
log
Proof. We use our technique to estimate the derivative of the loss. By Lemma A.1
we have
21
V VA[(A,0) - E -A)""
1=0
Apply Lemma A.2 to each term in the summation, we have
(T SA) S(T -A)"-1111 (I A) || All 21
where A, has O
21 -operator norm. Denote A2L = -i=0
A; and add all terms together, we
log
obtain
VAC(A,0) = - 4)2-1--22-L 2
2 AZL.
where Aet has ( Tog
operator norm.
A.4 LIMITATION AND FUTURE DIRECTIONS
Architecture and parameterization In this work, we use the single-layer linear transformer to analyze the training dynamics. Moreover, we adopt the same reparameterizztioo and similar initial- ization in previous works (Zhang al. 2023] Tian al. 2023a Chen et al., 2024aa] Mahankali al.) 2023 Ahn et al., 2024 It deviates from the practical softmax attention with Q, K. V parameteri- zation and random initialization, which is limitation this work.
However, analyzing the linear counterpart of the model before targeting the more difficult practical models is common the development of learning theory. As for linear attention, the connection between linear attention and softmax attention is also partially justified by the empirical observations in Ahnn et al. (2023). Analyzing the dynamics using more practical architectures will be a very important and fundamental future direction.
20

---

Published as a conference paper at ICLR 2025
Population loss and sample complexity Following most of the previous work, we population loss when analyzing the training trajectory instead of using finite sample loss. This modification is to simplify the analysis and focus on the population dynamics without noise. A possible future step to generalize this analysis to finite sample setting and train the model with online SGD.
CoT on iterative tasks In this work, mainly focus iterative tasks, of the simplest forms where multi-step CoT can help yield better performance. That serves as the imitial step towards understanding why CoT helps reasoning following the first principle. As limitation, though CoT can empower the transformer to acquire compositional reasoning capability instead of doing the much harder question beyond our paper's scope. very important
same iterative step. future direction and definitely worth further exploring.
B PROOFS OF THEOREMS IN SECTION 3
In this section, we prove the expressiveness results of the linear transformers with and without CoT. In Appendix B.I prove that one-layer linear transformer without CoT can only obtain the one-step gradient descent solution. Appendix B.2 we prove that there exists a one-layer linear transformer that implements multi-step gradient descent with the CoT prompting. As corollaries, there exists separation between the one-step and multi-step solutions.
B. PROOF OF THEOREM 3.1
We first restate the theorem:
Theorem B.1 (Lower bound without CoT), If the global minimizer of Evall (V. W) W*), the corresponding one-layer transformer fLSA [Zo implements step GD linear model with some learning rate 1) ntd+T and the transformer outputs 1 XXy
Proof. Recall the loss expression in Equation when k ==0
1
C(V,W) = Ex..
|LLA (Zo) -1-]] -1 -O0a, O, ,,,,))1 1)1l
2
Zo W Zo -00,, 0, ".ww)
Ex x
VZn
(since wh = Oz.))
2
n
The key insight of the proof is to replace the w" with the one-step GD solution xXyT.
C(V, W)
1 E -
V
Z W 20
(0.s
2 _ .0)
+
After proving this property, we can conclude that the optimal solution without CoT is exactly the one-step solution 2 Xyy We prove this result by showing the gradient of those two loss functions the same.
First, before calculating the gradient, we extract the identical parts of the loss. Notice that the ground-truth entries are all zero in i 1, ,d, d Therefore, that part of error the norm of the output fLSA Z [ [--1] in those corresponding entries:
1, 2d + 2 positions in both expressions.
Z W Zop-111-11
Z W Z [2d+2,-1]
E 2
V
E 2
V V
which is the same for both expressions. Therefore, we just need to consider
Z W Zol.-11
fLSA Zo [d+ Vi42224111 Zo
which corresponds to the ground- signals. Here V(442.241.1 need to prove that
[Vs1, Vs2, V3s, Vsul. We only
E fLSA Zo) (Z0))422241111 w0 w² 2=12=BB (iS (00)((22))1222241111 " xx " +C
21

---

Published as a conference paper at ICLR 2025
for constant C.
We show the gradients of both sides are the same, and equivalently the differential of both sides should be the same. The differential of L.H.S. is
E (40 - =0(l3)
2E (0)
( (00
and the differential of R.H.S. is
- Xx w w0'(²)
JisSa (4-22411111 - E
2E
(Ji.A (Zo) -4124111111-111------ x x w0) dfisA ( 042:24+1--11
Therefore, we only need to prove that
d/i.sA - (xxxx x dfi.sA (20)4422441111
(10)
We expand this expression FLSA Zo) (++)) (Note that now we don't have the assumption of initialization):
Zo W 20%--11
(142:20+1.1 Zo
n
X
0
[V
1/2
V/3
V34
y
X
y
T
Onxd
O 1
W
0
Oaxn
W0
0 xx
0
wo
wh 1
01xn
XXT
Xy
Odxd
0₂
W14
VV
V/2
Va3
V34]
Y
yy
01xd
W22
(since we = 0g)
Oaxd
0.
0axe
0
W34
01xd
0
0 01xd
W144
XX
W 14
+ w24 y
V31
V32
Va3
Vs4]
y
W 14
+
W24yy'
Oa
W44
Xx (W14 + W24w")) V34w44
Vai + V32w"
(y X w)
and the differential of fLSA (Zo))a+:22411-11
dfLsA (Z)
XX (W14 + w2)")) + V11w1
Val + Vs2w
n
(dVs) dV
"2w"" ) - (dV31 w44 + V34dw4) x (W14 XXT(W14
n
+ Vai View"
XX (dW1. dw2rww)
Now, prove Equation (10), we compare the differential for each parameter both sides. For all parameter, we start from the left side and prove it equal to the right.
VaI: The Var term of differential in d/I.sAA Zo (142:224411111 is LdV XX
(W14 XXW w2w""),
E
w
dV
XX (W14 + "2ww0)
n
tr
w -d
-V
XX
(1 + w2)w0))
(It is scalar in the trace.)
22

---

Published as a conference paper at ICLR 2025
L
- -VV3 X (W14 + "22 "2ow")w""
Ba
(d ((Viw22))
(E [XXXX = nl B.w.aaB................... B(w") = 0, Elw0' " [ Ia.)
E
tr
dVs X Xxx]
(E X [XXX))) n(n Id, n+d+1*
22
E tr
XX dV X.X (W14 + W22 W22 ww"
(E [ 0, E/w0" "0" J Ia.)
E
x
wO XV X w2sl0))
n
So those two d 3V terms are identical.
V32:
The V22 term of differential in dfi.A (Zo) [d):22414+11111 is dV32 w² x (W14 + u24t""),
dV32
EI
w
w
xx
(W w2222)
dV w
E
tr
w
XX T (W + w2ww")
(It is a scalar in the trace.)
dV32 w"
E
tr
xx T (W11 + w21w0
dV32
II
tr
w"
xx "inw"*
[E[w"] 0 and w T XX T w""w is odd)
3V32
E
tr
W 14
XxX
"n "
(W X w is scalar.)
E
tr
(dV32 W
(E | nLa, Elw0" u T Ia.)
E
Ir
32 XX xx
(B (B((XXT)) n(n + 1d xtd+1
IE
tr
XX
V22 XX (W14 + u24 ywo )
(R (B(ur"] = 0, Flue u0> T] = Ic.))
1 w w²)T dV xx (W14 + w2sw")
= E
x
n
So those two d 22 terms are identical.
d Was Vsa The Vaa term of differential dfLSA (z)) (2)) is n
E
dV W44
0 E
xx "T aVsawss (
n
n
since E[w"] = 0g. Therefore those two are equal.
Wi4: The W14 term of differential in dfLSA
Vsi (Vss
Vi2w* XXT XXTdW1.
E
Vi1 + Vs2w
X
E
tr
w
T 1
n (Vs + Vizw* T
XXTdWi dW11
(It is scalar in the trace.)
E
tr
uD Vszw"*
XX xdWi
(E[w>) = 0a-)
E
tr
V- un "0
(V 32 w is scalar.)
E
(V22
W14)]
(E X X T = nLa Elwe" un T Ia.))
tr
XX
V Xx (xx_dw1))
(B x 1 ( 1 = n
?2
23

---

Published as a conference paper at ICLR 2025
1 (V Vi2w" x XdW111 Thus the two dW14 14 terms are the same.
W247 The W24 term in dfL.SA (20) Vai + Vs2w*T X doo2ar""
E w
Vsi + V32w*
xx " du22
1 V/s + V32w" 1)
R
tr w=
xx " dwzss
(It is scalar in the trace.)
E tr 1
(ww= Tw xxX " uddza
(B(w=") = 0g.)
Bir (V3icWw24)
(E[X X 1 nla, Elwo* w [ Ia.)
" 1 xx 231 Xxx du22))
(E (xI((X]]) n(n + d 1e -+d+1
11
2 [ w²)T - (Vol Va2wm x Txuwwar dw24
Therefore the differential for w24 are the same.
W44: The was term of differential in dfi.sA (022221-2+ is Visdwas,
E
w
V34dw44
0 E xx w²) Vaadwas
since E[ww]] = Og. Therefore those two are also equal.
In conclusion, Equation I0 holds since all the differential terms are equal. Therefore, 3C
E fLsA ( (2) w"
(4 xxx 20 +
which finishes our proof.
B.2 PROOF OF THEOREM 3.2
Here we restate the Theorem 3.2 and provide the detailed proof.
Theorem B.2. Suppose n = (d log d), 2 C'log d, r E (0.1,0.9). There exists V and W s.t. XXXT) f "'w"" the k-step GD
f [-1) outputs (0a, O, 1x-1,1) where w; == (I- 2 n
solution with learning rate a linear regression model. Moreover the evaluation loss
Evall W*)
2 Ex
-
" X
k+1
w=
(11)
72
d log
Proof. We construct V and W in the following way,
0
0
0
0
0
I
0
0
-1
V
W
(12)
I
0
0
0
0
Now the transformer is allowed generate l steps before reaching the final output. We can induc- tively calculate the i-th step of generation, showing that the output is exactly the parameter after i-th gradient step (i = 2,
1):
Z, W Z
fLSA [Z.) (Os, O, wi 1) + VZ
(Os, O, wc 1) + n
(Oz, O, Va X (Wistt)u;-aa -
24

---

Published as a conference paper at ICLR 2025
(0,, 0, wi, 1) + (Oa, 0, -2 x
(wn; 20 w²), ),0)
n
(Os, 0, w 1)
After k+1 steps, we have the final output (I - (I- XXXT)++1 jee" by induction and the evalua- tion loss becomes Equation ( By Lemma D.4 the final loss is
k+1
E ,,,0
n
Ex, m
tr (I
! xx1
2k++2
WB(w" w T] I.)
2
n
Fxx -0 (1
!1 X x
2k+2
n
1
tr((1 11)* (1 8)
(By Lemma D.4
2
d(1 n)* s d
C log -
C PROOF OF THEOREM 4.1
C.I GRADIENT COMPUTATION OF THE FULL MODEL OVER THE COT OBJECTIVE
In this appendix, we compute the gradient of the full model given the Assumption 4.1 and prove the equivalence between the dynamics of the full model and a simplified model. Throughout the appendix, we denote the S XX for simplicity. And recall the i-th step of the linear classifier un; (I (I "
In Section 2.2 we have the full attention model
Z' T W Z
fL.SA (Z; V, W ).--1] - VZ
n
and the Chain of Thought (CoT) objective
k
CCOT (V, W) = Ex
[LLAA - (0,, 0, w1+1, 1)11
2
We define the error for the i-th step
C. E. w [LSA (Z)) (0.,, u 1)[l
By linearity of expectation, know the gradient of the CoT objective is the sum of gradients of all CoT steps: V.CoT Σ VL, Now we can calculate the gradients of V, W based on the loss
of each CoT step:
Lemma C.1 (Gradients of the full model). The gradient of V, W are given by the following equa- tions:
1 VvL Ex
Z, W Z [ip (0s, O. w:+1 w,,
VZ.
Z ZiT-
W Z,Z,
n
0m0
z. W [1-1] (02z O. W.+1 - w. O)
- VWL
Ex m
Z.2, V
VZ.
2 ZiT.-11
n
n
0=
Proof. The loss is given by eq. @
k
k
Cool (V. (V.W)
-xx
fLSA Zz) -O0a. O, w1+1 1)ll²
= = c
2 1=0
25

---

Published as a conference paper at ICLR 2025
Take differential of the loss for the i-th step C and we have
dLi Ex so (fLSA Zi) [.-[]]]] (0a, O, wi+1, 1)) dfiis (f)ss(ZZ)):.-)
2, W Zz [[:-1]]
Ex s0 (fLSA Zi) (O,, wt+1, 1)) T d VZ.
n
Z: W Z[[:.-1
-x (FLSA [i)) (Os, wi+1, 1)) d(V)
Z, dW [[:--1]
+ Ex,= m (JLSA p. (Oa. 20)+1, VZ,
"
Then the gradients W, V the L, are:
VvL;
Ex
(O,, O, wi+1, Zi WT Z:E:
7
W/WZ (1:-1) (0,, 0, w1+1 - w,, 0)
Exx n
VZ.
Zij.-1) 2.2,
n
W
E w- Z,Z. (fLSA Zi) [..1]] (0a. 0, ws+1: 1)) Zi(i--1]
Z, WZ ((.-1
Ex w- V. n
VZ
(0,, 0, w1+1 wi 0)
Z Zi
n
Take the sum of the two equations above from = 0 to k; and we finish the proof.
Now consider the gradient flow (GF) trajectory (note that w2s is fixed under Assumption
d0
CoT (0),
= (V, W {w224).
Recall block matrix form of V, W:
V11
V12
V13
V11
W
W1
12
13
W14
14
V
V21
U22
V23
1/44
W
W
W22
W
U211
V1
V22
V33
V3a
W
W32
W
W31
VA1
U42
V43
0/4
WAI
U42
W43
1744
According the construction Theorem rameter blocks, while the others should be zeroed out. Next, we prove that if we initialize those irrelevant blocks to 0, then they will stay at along the gradient descent trajectory.
the blocks Wis, Va1 W24 are the only relevant pa-
Lemma C.2. Under the Assumption 4.1
when the linear transformer is trained under GF, have have the following form:
for all t 0, the parameters (t), ((t)
0
0
0
0
0
W (t)
0
0
(tt)=
(t)
V (t)
0
0
0
0
0
0
0
0
0
Proof. To prove this lemma, prove that when the irrelevant blocks are 0, the gradients L, L, for those blocks are always 0 and they never update the corresponding parameter block. Also, note that W24 -1 for all t >
Z, W W. Z.p.-1).
First, we calculate the output of the linear self-attention VZ,
2,
W
[t
V Z,
0
0
0
X
0
0
0
0
0
0
y
0
0
Z,
0
71
(1)
0
0
0
0.xo
Wo
W1
W.
0
0
0
0
1
1
1
26

---

Published as a conference paper at ICLR 2025
dxn
0d
0,
X
0
0
0
(t)
01xn
0
0
y/
0
71
(1)
0g
0,
Oaxn
W0g
W1
W0;
02
01xn
0
0
01xm
1
Odxn
0g
0.
0g
1
O1 Xnxn
0
0
X Wis (fa(w;; yn
0
n
(t)
0g
Od
0++1
Vai ( X
(W13 (f) ww w*)
01xn
0
0
The last line is because y
X T "w". Now, we consider the gradient for V:
1 Ex .0
( (VZ.
Z:
W Z
VvL; =
(0,s 0, W:41 - w)
Zi Zil.-11 W Z
0g
Z T 11 WT ZiZg
n
Va1(1XXX T (Wis(!)w; w") -nnnww -n(ww+1 - w;;
0₂
w (
-1
112
Ex.w
Vai(1) (Wia(t)ws (iii)x) w0) - no( n(wi+1 wi)
0g
Z:Z;
0
0₂
w
W
13
XX
y
1
Ex.w' w
W
W
via(t)
Xy
yy
112
(xx (Wis (f)tt))sss w. w0) no ((W++1 w;)
0a2 0
0
0
V. Li(()
Vvs (()
0 0
0
Therefore, know all blocks of the gradient are except the positions of V and V32.
Now look at V Vo2 fe::
1
V f.
Ex _
[VVs (Wis(!)us - -)²) n(wi+1 - -w:))
n2
(w,Wii(t)xx - yy
Ex
[VVx (Wis(1)us - w²) n(wi+1 ww,)
n
(w, ( " - w T- w0))
Note that w; (I- (I- S) w" for all [k), and Wk+11 w" Therefore, for all i {0, 1.... k+1} the formula inside the expectation is an odd function of w Since w0 N(0, Ia), the expectation should be Og-
Similarly, calculate the gradient of the W:
Z
W
2 [1-1
VWL. Exx m
Z,Z, T
VZ.
(0., 0, W:41 wis T)
Z Z.T.-11
0₂
1
E xx
V.
VT
Zip.-11)
n
Xxx (Wia (1) ww w²) -n(w111 wi)
0
Odxd
X Vt)
0,
2
1
Ex
w
Odxd
Odxd
yX Va(
Oaxd
Vxt(t)XXX (Wis(t)ws (WW))XX((Ws(()uur) w") n1(w -w:)
0
Z ZTT-11
Oaxd
Oaxd
27

---

Published as a conference paper at ICLR 2025
XX (t) Vsi(t) (s((t)
( (Wis(t)w: =-) -n(w;+1 wi)
[0
yx (ta(t) T Van (W), -w") n(We+1 - w;
E
Ex. '
12.
0g
un;
0
0
0
V W1 (it)
Pwi,C((t) Vanza (t)
0
0
V
W
Ci(t)
0
0
0
0
Since we fix W24, we only consider the remaining three blocks. First, we consider the gradient of the vector block W14:
Vw (i)) = 1E
Ex.. XX Val (() (Wis())u;; - =²) - - -n(w;+1 w:)]
Notice that the (t) (X) (Wis(t) w; w -1w;;+1-- ) odd in wt. There-
fore the expectation is Og- Similarly, we consider the other block W23:
V C.tt)
Ex [F (yxTVa((t)) Vaa Vx ( w") - -n(wi+1 - -w,)
n12
[ww' xx V Xxx (Wis(t)w; w²) -n(wi+1 - w.: w
=, Ex w-
01xd-
In conclusion, all the blocks have zero gradient except VI,, Wis given that they are all zero matrices. Under Assumption 4.1 all the irrelevant blocks remain zero matrices for all t 2 0.
By Lemma C.2 prove that along the gradient flow trajectory under Assumption 4. the objective of the linear self-attention model with residual connection can be equivalently transform to the following simplified form.
Lemma C.3. Under Assumption
4. we have the training objective
[CoT (V, W)
Ex um Σ (SW -ww) -Awill
2
where S = = XX and Aw. == w w; i = O, I..., k is the residual for each step i.
Proof. Given the following CoT objective,
CCOT (V, W) Ex m
Σ JiSSA (2a)(ZZ))) Z.) (0a, O, We+1, 1)
By Lemma plug in the V, W expressions and get:
Z Z.WZ Z Z. (Oa, W,+1 -w;; 0)T
(Zs) 1.--1 |.
(0,, O, w:+1, 1) = VZ.
Og. 0, Val [XX T W13m, - Xyy) Ann,, 0)
Since y X "", we have
fLSA (Zi)).-11 - (O,, O, w:+1, 1) (0,, 0, Vai (S Wiaw; - Sw") - Aww, 0) T
Put back to the loss expression and we complete the proof.
Now the chain of thought loss can be rewritten into the form by Lemma C33 we can directly calculate the gradient update using the simplified loss for clarity. We denote the only relevant blocks W :=
W13 13 and V V11- Moreover, further expand the CoT loss with Aw; -1)
xx (w,;
28

---

Published as a conference paper at ICLR 2025
w²) for i {0, 10,1,... ,k - 1}. and Awy, = "² CoT loss:
Wk. That leads the following expression of the
k-1
CoT (0) (0)
Ex w
wi + VS W wi w0)
wi+1
1=
(13)
+ Ex, 2
wk + VS Wwy w" w
Observe that the final loss only depends the (d 2) to (2d 2) entries of the transformer's output, indicating we can simplify the model bit and prune out the irrelevant part. We can define simplified one-layer transformer get the loss form above, where the dynamics of the equivalent model is exactly the same with the original dynamics of W13 and V31. Accordingly, the last token input of the transformer for i-th step becomes w; and the label becomes w++1 since the other entries in the original input/label (0, O, Wi. 1) do not affect prediction.
Definition C.1 (Reduced transformer). Let 0 (V, W). Define
fo Z.) w; + VS (Ww.-Wo w"
be the reduced model of the one-layer transformer in Equation denote [o(w:) == Jo(X, Z.).
For ease of presentation,
In the following sections, we will consider the equivalent form of transformer. Here present the gradient with regard to the reduced model. For clarification, throughout this section we will denote
I -- ()s) u" the (k 1)-th update, and w" is the ground-truth.
wk-1
Lemma C.4. The gradient of V and W are given by the following expectations:
Oc
B (folu))) ((lk:)))w::i)(( W w
+ ( - w' (w] W ~ =)
ov
i=
sV (fo(w,) w
SV
(Wk+1 w w')
ow
Proof. Given the equivalent CoT loss in Equation 13 we take the gradient with regard to V of the loss and we have
k-1
OC
E (folw)) we+1) w. W
w= [)))
w
E (fo(ww) w²) (w) W w
ov
k=
Σ
B (folue)) wi+z) w W w s + E [Wwt11 - w0) -wo----=))) w W w+ s
The second step is because we subtract E (fo(w))) wh+1) w W w s
from the second
term and put into the summation. Similarly, the partial derivative of W should be:
ac
k-1 1
Σ
SV T (fo(w:) wi+1) + SV (fo(ws) w w
dw
i=
Σ E
SV (fo(w:) - w00 [
E SV (wx+1 w')
1=0
Therefore we complete the proof.
C.2 GRADIENT CHARACTERIZATION OVER THE COT OBJECTIVE
In this section, we compute the exact gradient for the reduced model parameters facilitate analysis on the dynamics. For clarification, throughout this section we will denote wk++1
[ (I + w 1 as the (I: + 1)-th update, and w" is the ground-truth.
29

---

Published as a conference paper at ICLR 2025
We first compute our gradients for the simplfied model defined in Definition C.1 which is equivalent
to the full model's dynamics. Recall that under assumption 4.1 we have V, W are simultaneously
diagonalizable, with the orthonormal basis {u: V We denote the orthogonal matrix formed by
the basis as U. We will observe that u, are always the eigenvector of V, W, so we denote V
UAYUT, W = UT For clarity, we ignore the timestamp when calculating the gradients
and dynamics.
We present an accurate estimate of the gradient in the following Lemma C.5 We intensively use the concentration lemma in Appendix D to separate the main terms dominating the gradient flow dynamics, and some bounded eπror terms that may complicate the analysis. We also call the error terms as 'interaction terms', since they contain the interactions between two subspaces u and uyu.. The structure of the interaction terms AV, A W are further characterized in this lemma, which is essential for the final local convergence analysis.
Lemma C.5. Suppose n Ogd d). 1l E 0.91, k [clogd] d). Under Assumption 4.1 if we run gradient flow on the population loss in Equation (6). then the gradient of V and W are characterized by the following equations:
OL U
2
1
A AW2
1 A w + (k + 1)1 AP
UT
k +
1
- k+
-2
k-1--
OV
7
(222 7)
11 A
W
+
+I+AV
2-1 17
Oc U
U
h: 2
AW
k+1- 2
1- 17 A V 2 7
A
Ow
7 112
where the error terms (interaction terms)
AV
I o
Tog
W
<
Iog
op
op
Moreover, there exist diagonal matrices
V
B
AW BW with
Ing a
operator norm,
C² VDDV CW DW EEW with O
operator norm and EV ,FW with (((1-n)*)
dloog
operator norm s.l. the error terms A AW can be written as
A V A + I AV
[
B + tr
DD DV + E + tr (I- A W OW + +tr (IAR +nI)A
A AW
C + (I-AW)
W
= A + nl
A w
+ I
A
BW
(I A W A 22
Ew + F
+ tr
Proof. Recall the gradients formula of and W by Lemma C.4
ac
Σ (Fo(uos) - -w:+1= (1:+1)(()) W w w s) + E [Wx+1 w²) w
k
(w) W
av
1=0
DC
Σ R
SV (()((()) wai1)07 + E [VV
(wa-11 - w w))un
ow
1=0
We expand the reduced model (w:)) in Definition C.I and get the residual term
fo(w,) W w. + Vs W ww; w=') Wi+1
Vs W wi - -w") + S w')
(vSW- + n/S w. - (v + np1) Sw"
Substitute fo(,)) w:41 term in the dynamics by the equation above, we have
k ac = Σ E (vSW+ + ns) (I- -n))) W s > (vSW++ + +s) (1-(I- (I- (I-(I-nS))) -nsss'ss
ov
1=0
30

---

Published as a conference paper at ICLR 2025
V + m (1 ( s|
+ Σ B
>Σ>((
+ n 5²
==
(I )ssn 1+1 -I- -- W - [
Σ V + In R SW I -- (I -15))) 2 WT. s
(Term 1)
17
Σ
s 1 w) (I- (I- -nss)) ( WTS s)
(Term 2)
==
Σ
V + lll B Sw I (I - -ns)'ss s)
(Term 3)
t=
17
Σ
E S I W) (1 (n-n)))) ss
(Term 4)
1=
Σ
(V nl E (I ( WT s)
(Term 5)
+
Σ
(V n1) n E[5
(Term 6)
t=
E n) (
W -
(Term 7)
To get an accurate estimate of the gradient, we apply Lemma C.14 Lemma C.15 respectively to each of the terms (Term I to Term 7) and separate the interaction terms introduced by the moments of Wishart matrix, which is bounded by
Iog
Consider Term 7 and the i-th term in the summation of Term to Term 6. By Lemma C.14 and
Lemma C.15 there exist diagonal matrices E₁ E [6] satisfying ||E, llop
I 0
such that
Tog
SW - (I- -IIn))))) W s = U
2 [1- (1- n)* A A
+ E1
U
(1- n)*) 2
[ E S I w) [ I- (I (I-)))))) W s
= U
A AW)
+ E2 UT
E SW [ - 1/8))) s U
(1- m) ) w - Ea UT
E E I w) (I- n))) s) U
(1-(1 m)* [ (I-AW A
+ + UT
E s I - (I -)-)))))) W s) ( n)* b) AA= Es U
B[sa] - U
By Lemma D.4 there exists diagonal matrix E satisfying |1E llop - <0(11- [[[--n)) n)*)
such that
B ( 11S)) +1
+1 ((1 (1)) (I--(-- (1-))) W T -1) U
Moreover, there exist a1, <0 1n
az a4. a5 < 0 Tlog such that
E2 = ( W
+
I AW
W +tr
+tr I (I-A A )
(0sA" w + aaa) +as tr ((1 (I-AW)AW))
and exist 31 5 ( log
1
32 50(
- such that
E4 = B1 (1- - A ) + 32 Lr (I-AW)
31

---

Published as a conference paper at ICLR 2025
+nn((5 (51 - E3- - + +n)+ - x1 We define A as the sum of all the interaction terms (A + m/
for the i-th term in the summation of dynamics of V. From the analysis above, there exist diagonal such that (note every matrix is
matrices AY, BY, c.. D. with their operator norm O logg d
diagonal, so they commute)
A = (A + +nI A + +0
11 2
tr ((1- A W
A W B.
+ (1-AW) W C. + +0
() tr I (I-AW) AW W D.V
We define A is obvious that |
as the interaction term brought by Term 7 since there is no summation in Term 7. It
S <0(11-n)
(1-n)*
ap
Now we denote
AV= Σ A A. 1=0
to be the sum of all interaction term of the dynamics of A From the definition of A and A above, there exist diagonal matrices AV, BV, CV. DV and Eý satisfying A V
C S
o log'dl d Olog d))
)): B |. | s ((alege 1 a) and | s 0(11 -n) such that (because alogg
AV_ = (A + m1) A
+tr (I-AW)AM) B A AW) C +1B (1- (1- W D + EE
+ (1-
Sum up all the seven terms together and we have
2(1 - (1 - m) k+1
1 1 (1 - 11) 2)22
A W (A (AVAW +n UT
U
k +
0v
17
1/2 71)
h ++| 1 1- (1 - m) k+1
A AAVAW A W + +nl) UT
U
11
k+1 1 1- (1 - -m) k+1
A W V +nl) U
U
"
+ U (k+ 1 + U +UAVUT
Denote Ey to be the sum of all 0 -1)*)
terms in the dynamics of V:
2(1 - -n)*+ k+1
(1 11 26+22
A WAW + +nnn (1 m1) 4-1
EV=
A W
AA W +n W +ni)
17
17(2 - 17)
17
Denote E² Eÿ¹ + +E and denote A V
A V
+ EY
we have
T OL U =
[[ k + 1
2
1
A 2
2 k 1
1) A + (I + (++1)1 AF
+
0v
1
712 - 11)
1 17 A
w
+ I + AV
2- 1)
Moreover, A has the form
A V AA V + nI
A AV +tr tr
W
A
BV
(I AW
DV Oo +tr A--AW) ) E
32

---

Published as a conference paper at ICLR 2025
Similar to the calculation of the dynamics of V, we can also have
ac B SV (Fo(w))) 10)+)000
R SV
(Wk+1
200 )00k
ow
E SV ( SW +
Σ cB(SVT (v (V+n))s(I (I- (I- - =
SV (I - - k+1 [ II - ]ns)
Σ B AS ( W - -1) + (V + +n))s) T (T- (1))) }
E SV (v + (11 5 (I ---1 )))))] - E SV T - -1/5) (I-(I- (1-(
-E(SV
k -
SV W - 1) (1- (1 =)s)))
-Σ
SV S +
(1 - (v-n))((-((InSS)
i=0
Σ E SVV (v + +n))s (I (I-n) 1=0
E - (I (I h+2 (I-
We apply Lemma C.16 and Lemma C.17 to each term, similarly define A for i E [K] U {0} as the sum of all interaction terms for the i-th term in the smmation of dynamies of W. There exists diagonal matrics AW, BW C,W. DWW, EW with their operator norm O log' ( such that
A AW=
(A V + nI AW + (I (I-AW) B.W +0
+0(1)
1 d
tr I I-AW C.W
0 (2 Ir (IAP + +DDADD D. + - 1 ir (I ( ( B.W
We define A W as the interaction term brought by the last term
E E[SV (I- m)) ++1 (I- ns)*
]
It is clear that
AW
W
<0((1n))) ))²)
Similarly denote
op
A A
W
AW
1=
then there exist diagonal matrices A W BW C.W
DW EW, FWW satisfying wwl : Bw S
o lop d ).
| | DW Ew s < a (
FWW S so (1- m))) such that
A ww_
(A v + +nI) AW (I IT- A W.
BW +tr [ - A w cw
+1r ((A nAV++DJAN +n DW +ir (U- A 2
EW + FW
Denote F.W to be the sum of all (I--)))) terms in the dynamics of W, FW = F.W + F.
and A AW- Aw + F.W. Thus we have
k Dc = - (1 - A + +nt)OT
k (0(1-(-n))) (AV +
ow
1=0
0=0
33

---

Published as a conference paper at ICLR 2025
2(1 - (1 (1-
-)) (1-11nn)
1 - (1 17) 2k+2
A (W + UJ
= U
k 1
+
17
1722 11)
h + 1- 1- (1 - 17) 1+1
(A + +
U +U
U
=U k 2 +
1
A 2
A W
k+ 1K+1- 12
1V2
1- " A v + A w UT
11
11(2 11)
2 - 11
Moreover, A has the form
A w
A V
+ 1/ ) AW + (I- (I-A AW) BW +
CW
+ tr (AA + nI)AAA) DW + ( (I A
EW + F.
Since ||A + B Bllop I || op
+ ||
[B|]
and ||AB]
op
I
|A| op
B
op it is obvious that
S
0
1
W
S
0
1
op
log² d
op
log
2
After obtaining the estimation of the gradient by lemma C.5 we can decompose the gradient updates into the dynamics along each eigenspace u, which can be characterized by the following lemma.
Lemma C.6. Suppose V -ド= u,J, W Eド AW u,u.. The dynamics of the eigen- values of V and W are given by the following equations:
aay
(ks (1+ +0+11) 1)(1111 11 1
+ 2 2 (1 1W + 1/2 11)
1- 1+ 57
+
dt
2
drw
k: + 1 - _
21² (1 - >
1- 25² 2/m
+ 1- Nÿ 6m
dt
1/2 7)
2 1l
676 S o
Toy ( ), bow! < Tog d
where
Proof. This is directly obtained from Lemma C.5
C.3 PROOF OF THE MAIN THEOREM
In this section, we prove Theorem 4.1 which characterizes the CoT loss of the trained transformer. First, we restate the theorem.
Theorem C.I (Global Convergence). Suppose n O(dlog' d), r (0.1, k [clog d),
clog -
> 2. Under Assumption 4.1 with some constant 3 (2-1) + if we run gradi-
flow on the population loss in Equation 6 then after time L O Olog d log 2), we have
[COT (f) S cor any c26 0
logd d
Proof. According the previous sections, we can reduce the original optimization problem to Equa- tion [3). and consider the equivalent reduced model (Definition C.I.. By Lemma C.5 we fully characterized the gradient expression, which decomposes the gradient of V and W into main signal terms with large norm at initialization (terms before A A AW) and interaction terms (AV. AW) with bounded norm O ( for all it > 0.
The decomposition motivates us to conduct stage-wise analysis. We first analyze the dynamics in Stage 1 when the distance between the parameters V. W and the ground-truth is larger than
34

---

Published as a conference paper at ICLR 2025
O (0ggza) In this stage, the bounded error can be dominated by the signal terms in the gradient, leading to nearly independent dynamics along each direction ut- After this stage, enter Stage as local convergence phase. We describe the dynamics below in detail.
Stage In the first stage, the dynamies are dominated by the main terms, and the interaction terms A V W can be somehow be ignored. Specifically, by Lemma C.6 given the dynamics of Xy W.
dap - (k (k+ (1- 2 +
- -, (1 1- + 17(2 11)
1- 17 A + 67
1
di
2
dip k 16+1- 1
1 (1 (1- ) 1/2 7)
1- 1 w
1 1 1 x/ 8%
+
dt
2 -
we can conclude that the dynamics of the eigenvalue 15 A W mainly depend on themselves when the main term (terms before 6/6 .5 are larger than O log a) which is within the stage 1. That is, the dynamics within the subspace u for V, W almost independent with other subspaces. In this stage, we focus on the analysis of 2 > depending on their own value. The first stage can be further divided into two phases.
Stage 1, Phase 1. At the beginning of training, we have
3(1 11)
3(1 17) 1
(0 2(2 11) (k + 1) (0))
<0
(2 - 11) k: + <0
then by Lemma C.8 we can prove an upper bound of 2I when ih < - (k + 1) 12
1
3(1 111-1) 11) 2(2 - 11) (k + 1) AW)
1
dAw de 2 (+)) There-
according to the dynamics for both sides. With this upper bound, we prove
fore, XW will converge to - (k + 1) in f = Ollog d) time (Lemma C.9
Stage 1, Phase 2.
After time f1, we have 2 very close to the ground-truth value 1. Meanwhile, the lower bound for 2ÿ still holds, and it will further decrease. Specifically,
3(1 -n) 1) AW (4)=1- =1 -1)+=11(+++-- 11 x7
2(2 - n) (Is + 1) (1-1)))((()))
By Lemma C.10 we can prove that 2, will stay close to 1 -
1 - 2(k + 1) E (t) < 1 (k + 1)
for any > t Lemma C.IT
With this condition, converging condition for (2)² + 7) can be deducted from
d (>) + +1) 2n 1) (A) + 17)
< 1
dt
Lemma C.11 shows that 2y + converges to (A: + 1) tt2 in 1a = O(log log d) time.
35

---

Published as a conference paper at ICLR 2025
Stage 2. Now the eigenvalues are already close ground-truth:
(1 (4 + (2) + + - 0 +1)
1 (4 -12)-111 -
2(k + +1 1)-
According to the expansion of the error terms in Lemma C.5 we notice that oy and 6y are always coupled with some individual residual like (AV + 1/d,, (A - I), or some weighted average (A V + +)))Aa) Meanwhile, the coefficient of this kind of residual the interaction terms is still upper bounded by O(1/ log d). That helps us to derive the PL-condition like gradient lower bound (Lemma C.12
A + nn )
dtr (I (
2
dtr
dt
df
1
tr (AV + n11)
21
112
II (I-A A ) (k + 1)
$ 2n - 17)
+a
2
where a 0 11)* 2
By Lemma C.12 we know b + v/ and 1 converge to (e (0g (1) in
to = 0 (log time. At this time, there exist diagonal matrices A and B satisfying || - 9(1)
and B op - (1) such that
AV I+5.A
A W I+6.B. B.
Now we consider the CoT loss given by Lemma
1-1
COT (
Ex,
VSW + nySw, -V +nI)Su
2
Ex w (I+ VSW (VS I)w"
Apply Lemma C.13 we directly obtain that
CoT (0) = 0 (ogd
Since E (0(ct (d log (0g((11g1g)) (1-n))))), .1) the CoT loss is smaller than @(d (1- log d). The
local convergence takes 1/ = O (log }) (log 4) Considering all stages, at time & = 1 12 +
Oloo d) O()), have
CoT (0) S E.
C.3 C.3.1 TECHNICAL LEMMA IN APPENDIX
Lemma C.7. Assume An S (k+ 1) 22 3(1 v)
S SA 0, holds that
2(2 n) (1+1)(1-200)
1 + 3(1-9)
2 (k+1) (1c1))(11))) 1-1W
<0
(14)
dl
Proof. Directly consider the derivative
3(1-)) カ(ーツ nl
(k+1)
(2+1)
1-
dap
3(1 11
di
+
dt
dt
2(k + 1)(2 - n) (1-
dt
36

---

Published as a conference paper at ICLR 2025
Substitute the derivatives with the equations in Lemma C.6 we have
day
3(1 11)
1 diw
dt + 2(k + 1) -1) (1- A ) dt
=- (k+ + -[(*+1) (1-11 )
--[M
2
2 w
(1 - -
1
--" >/ - -(1+07)
1 17
+
+ 1/22 - 11)
2 2- 17
3(1- - 11)
1
- 11²² (1 - 1 )
1 1-1 1/² 3mm
+ 1- 11 " -0p
+ 2(k: + 1)(2 - 7)
+
(1 (1-)
17
17(2 - 17)
2 - 17
Since-
3(1-7) 2(277) (1+ 1) -240 S 2 < 0, we have
di + 3(1 - 17)
dAw
2(k + 1) - 1)
(1 (1- 1/11 2 di
S (k + 1)( - 15²) 2
- ~سt (1 - (1-1,²)
1 2
3(1 - 11)
1
+ 17
+ 12(2 - 11)
2(2 - 11)
(kc + 1)(11 1 )
+ 1 1ー //= - + (l)
2-
2 1 )
) (1-217) )
3(1-1) n1)
3(1 -1)
(k: 11)
2(k- + 1)(2 -n) (1 (1-)
2
2(2 11) (k + 1)(1 13rr)
2
1-n 3(1- -n)
1
1 /
(/22 - 1))
22 - 1)) (k- + 1) 11(111111 - W/)
3(1 - 1)
(1 (-) +
k + 1 (222 11) 1,r
1 3(1 -1 1)
1
2n722 - 17) 1, w²
3(1 - 11)
2(2 11
(k+ )!)
+ 1 1- W - (1+07) + 3((1 - m)) 3
1
2-n n
2(2 - m)
(k: + 1)² (1-11
33(1- - 11)33 11 w
3(1 - 11)
1
+
22( - 1) 1/2 n)
(k + 1)3 (1 - >/) (1-1W)
2(2 - 7)
(kk + 1) [ 1/r)
1-7 17
(1-2!1
1 3(1 - 11) A/
+ k 1+1- 2-7
2(2 - 7)
2 17
1
3(1 11)
//2 -0y+
3((1 - 7) 3
1
1,
+
(A: + 1) (1 - ) 2n n)
22(2 - m)
(k+1)² (1-
33(1 - 1) 3 1 17 1 w
1
3(1 - 7) 2(2 1)
1
26
+
22 - n)] (/2 1)
(k + 1)3/ (1 (1- 1)²(1
wr)
(K+ + 1) (1 2/r)
Put in the assumption on XW that 2h < 1 - (k+ + 1) 1/. we have
aai dt
-1)
1 dx
2(k 1)(2 1) (1 (1- ,,
2 dt
1 1+1 +
1 3(1 - 11)
(k + 1) 2n - 37) 2 + +108
111-1)
[3(1 - 17) 3 1
<
+
+
2(2 - 1))
k 1+1n - 17)
2(2 2((-77) 7)
(k: + 1)² a
37

---

Published as a conference paper at ICLR 2025
3(1 11) 11)33
1-
1
3(1 11)
6/
2(2 77
11(2 11)
(k + 1)
2(2 17)
1 17
1
+0
2(2 11)
log + d
Lemma C.8 (Upper bound of Aÿ), Under Assumption 4.1 if - (k (k+1) (I+1)T+ 1+=1)
in it holds that
3(1 6(1-1) 17)
1
(15)
2(2 - 11) (k: + 1)
Proof. We prove by induction. First, check the initialization 2 (0) s -0, S 1 (0) $ / If
3(1-1) 2-1 K+I then we have
3(1 - 17) 1 (0)
1
< - + 3(1 - 17)
(2 - 17) k + 0
2(2 - 11) (k: + 1)(1 (0)
If the incquality holds until some time 11, that is for any t 11, we have
2 (1) < 2(2 1)) (1e + 1)
3(1- - n))
1
but
15(tt) (ti) 2 3(1-1)
2(2 - 17) (k: + 1) 21 (41)
By Lemma C.7 we have
Ay +
7) 1
2(2-m) (111)(11177)
<0
dt
Therefore, there exists some time t' & such that
3(1 - 17)
x (t') 2 2(2 (Ir + 11 >W (t))
which is contradiction. Hence, the proof is complete.
Lemma C.9 (AWW converges to near optimal). Under Assumption 4.1 it takes Olog d) time for X converge 1 (k: + -
Proof. Recall the gradient of X, W
in Lemma C.6
Mrp W 1:
11+1
x 1
W
1-
v2 Ay
1- +
v
dt
17
1/2 17)
2 11
Substitute xÿ with Lemma C.8 have
2
drw w
2)
3(1 11
(1-177
2
k:
2(2 - 11) (k + 1) (1-
7
38

---

Published as a conference paper at ICLR 2025
1- 7 1222 7)
3(1 - 11)
1
1/8
2(2 7) (k + 1) 11r
1 1/
3(1 11)
1
6.
2
2(2 71) (K + 11 1
4 (k: + 1)
3(1 7)
(1 mir
5
2(2 7) (k= (
1- 1 2
3(1 11)
1
1
2(2 11) (k + 1)(11 11 >
3 (1-n)²
1
10 (2 - n)² (k + 1) 13)
2 1(1 1) ?
5 5 7)² (k+ 1) (1 - )
1 1 - 1)) 1
2 5 (2- n)² k+1 1
In O(log d) time, W can converge to - (I: + 11 1
Lemma C.10. Assume > (41) =1 - (k+ 1) is and A/(t.) <
3(1-)) 2(2- (+1)(1- (11)
for
any 1 2 4 it holds that
1 2(k +1) < 2 (t) 1+ (ks + 1
Proof. First, it is clear that the inequality holds time tj. If the inequality doesn't hold, then there exists t' > t such that
- 2(k 21) < - (+) <1++) to
for any 1ı S t < t
2/ (1)) -1 2(k (( -/
or
1 1- 2(k 1)) < 2mm(1) (1) <1 (1+ 3
for any 41 St < t'
(1) = 1 + (k- 1
In the first case, it suffices to prove
20(0)=
3(1 7) 2(2 n7)
(k- 1) + <
3(1 - 11)
1
2(2 11) (k+ + +1)(111111 2/ (17)
to show
yp 0
dt
t==
which says there exists t S t" < t' such that
(t") 5 2 + +1) :
and leads to contradiction. Recall the gradient of 25 in Lemma C.6
diy = - (k: (1+))))-))))
) 2 +
2 (1 - ) + 1722
+ スー - (1 )5)
dt
39

---

Published as a conference paper at ICLR 2025
5 -
1(k + 1) à +
-
4
i
+
(k + 1) :
+
1
1 2(A + 111 1/ + (k + 1) 8 1g
11
12(2 - 11
1 -
(k: + -2s + 07
+ 2 17
2 17
2
1
< (22 1 1)
X
2(2 11)
and thus we have
X (t)
2y(t)= - Ce by (t- m
4
If C S 0, then
())) S 7 - 3(1 1)
(k= 1) 1)- it2
4
2(2 - 17)
else
S ((0)=1000) = 3(1 17)
2(2 - 17) (k + 1) i0
In the second case,
2 (t)
3(1 - 11)
(k 1) i/3
2(2 17)
still holds for any t St t'. Recall the gradient of W in Lemma C.6
d dt 1==1"
k + 1 - 1
1/ 3 (1 (1-1,) +
1 1-
-- 7
1 1- 11 x/ 6:²
=
+
17
1)(2 - 77)
2 2-17
[k+1= 1 1 1
25 (k: + 1) f22 + (2 - n)
1 1- π
Nj²
+ (k + 1) 1)- :2
1 11 x 6.
+
2 11
(k+ (k+1) 1) 1/² +
1 11
u² + |
2
2(2 17)
9(1 m) 3
(k: + 1) 1 3(1 m)
1/2
(k: + 1) p/ + 1,0%
8(2 - n)
4(2 11)
s (1 1)
(k: + 1) +t2
(2- - 11)
There exists 4 S 1" < l' such that
X 2+
which is a contradiction. Hence, the proof is complete.
Lemma C.11 (2) converges to near optimal). Assume
1 - 2(k 11 < AA (1) < 1 (k + 11)
then it takes Ollog log d) time for |25- +1| to converge to (k + 1) 1)- i._
Proof. From Lemma
C.10 we know
dxp dt
(k + 1) (1 11" )
2 +
2 2 1 )
7/2 - 11) 15021
1
A. + 1 1- 1 ,,="" - (1+
+
- n
4(k+ + ((x+1)=²+ + 17
+ [(k [
-/1
+ +1) 1)- t2 + (k+ + 1) $ +
1
1 1+ + 1) 122 + (k+ + ")=] 28
11(2-1) - m)
1
+ 1- 7
(R: + 1) Th + |%
2- 17 2 17
40

---

Published as a conference paper at ICLR 2025
(1) + +n) + ((((-)))+
1/(2 - n) + +
and
diy V
(k + 1) (1 1 (1-1,²² 2 2 ! (1 1,²) + 7/2 - 1)) 1/1²² 27
1 + ,ー (1+
+
dt
"
2
(k: + +1) -/- + (k + 1)-8) $1
1
[ 1(k + 1) F22 + 4(k + 1) 8
$1 25
+ 1/2 - 11)
1
2(1 - 1)
(k + (k+1)= - -1055
2 2-1 17
2 --1 1l
12 n1) +0 (kk +0((k=
1)- b (2g + nn) + +0 - -0(k-11) =)=
Therefore,
d a(a + +1)
2 (2/ +1) dip
dt
dt
$ 11(2 - 7) + 0 0 (kk ( + n)² + ()( +n)
If A + "
converges to = (I: + 1) ,े, then
d {Ap + n)
211(2 11) ( + :7 )²
S
dt
Thus, there exists @(1) such that
é = (>8 + +n) 2 S e² exp 227(2 17) +
In Olog (?)) Ollog log d) time, |ay + n can converge c.
Lemma C.12 (Local convergence). Suppose k [elog d). Assume
| W (t) -1 1 s 2(k 1)) l
| (+)) (()+n/
then there exists a 0(11 -n)*) N 0 such that A V and A W comply with
dir (AV +n1)² 27 dir (I-AW)²
dl
dl
S 2n 11) tr (AV + =x)]]]
7 (k + 1) tr (I (I-AW) (I-AW) )|
21 + a,
2
thus bay + ml and
1 1 1- can converge to e E
(a & lox (o (+tr)), in O (log !) time.
Proof. Consider the error term in Lemma
C.6 more carefully, we have
dxÿ =
= - (1: + 11++1(11101
2
+ 2 ;
1 1 + 1/2 - -)) 1 2 27
(25 +1)
dt
+ 17( + -11(1-1 ( 2
)²
3 21 x
-1 (1 - )
+
2 -
+ +(28 + =n)o 0 log d
1 (1gz)) )
+ (1 1 0
log"
41

---

Published as a conference paper at ICLR 2025
+ tr ( (1-A") )0
1 (o)ogE)) + tr
(1leedd dlogg d
d dloggd d
a
+ +0(11 (1 -n)*)
and
dvw W dt
152 (1 1/)
1- 11
1 (2/ +1)
=
h: +
+ 1/2 77)
+
17
172 - 7)
+ (25 + m) 0 log² d
+ (1
1 0
1
log² d
tr (1- A W
+
(AV
v + +nl)AA ) 0
dloggd
<dlog" d)
tr (I-A W )AV?
I dlog
+ - m)*)
Now we consider the decay rate of the distance between 1ÿ. A/ and their ground truth.
(>y +1)² 2
dt + d(W -1)
dt
(k + 10(1 (1-1 1,²)
2
+ 2 /// (1-1)7
1
1/²² +0 (1052
log" d )1
1
(25 + +n)²
+
1
1/22 - 17)
3 21 WW
1 1
1/- -1+ +
Tog d ) (1 (1 1,)) (1p +1)
1
2 2- 7
(2 17)
[ (+1 1
2
1
1/² - 7n(k + 1)( +n) + o (102
log" d ) (1- 1,7)² 2
1
+
"
7(2 -1 1)
+ (Ay + +1) tr I - A W
+ (2p + +1) tr (I- -W W
dlog d
dlogg d d'log" d
+ +(1- ) tr (I- (I- A W
)0
+ (1- )}) tr ((A + +
1
dloggd
(1- tr (I-A W
dlog'' d ) + 0(11 - -n)²
+
1
1
(>5 +n)
+
1/22 11)
loge d
2(( -))
1
)) (1) (>8 +n)
o
(2-1) 1)
log 12 d
|)P(k= + 1) + ((+tt) ][
(1 - 1,m) ?
+ (28 + +7) to (I- 1- )o dlog? d
+ (2g + +n) tr (I- ) )o
dllog" a
+ +1- ir (I-AW) A A )
+ (1=) ir ((^² + +
1
dlogg
dllog d
(1- tr (1 A W
dlog2d a) ( + 0
(1))))))) n)*)
Utilizing Mean Inequality, we have
(1-13) +
+n)
2
(2) +
+ 2 1 - 1,m)
Vk
42

---

Published as a conference paper at ICLR 2025
Insert the inequality into the equation and we have
d (Ap
+ +1)
2
d
(AW-1)
+
dt
dt
( 1
)[0r +n)² - [|P(K+ () + (1-15²² 2
+ +0
1/(2- 1)
log² d
+ (Ay + +1) ir [ - A W 0
+ (1)+n) tr (1 (T a of
dlogg d
dlog" d) 1 d dogg d
+ (1-2,²) tr (I-A )o
+ (1- (1- tr (( + 1 )0(
dlogg
+(-- ir (1- A W
A 2
(Hloged)
dloggd d) ) 0 (1 - -))²
There exist (11, 022, 033, Org 0tg = ( dlog"d
1
2 and ogg 0(11 - -))*) 2 0 such that
d ((z² + +1) 2
d d(2 - 1) 2
) (15 +1) 2
1
1
dt + dt
5
+0
17(2 - 11
logg d
[P'(k= + 1) + 0 (1:12 )[ 1 - 1- )²
+ ai A/ + 1 Ir r A Aw) )) + 02 A/ +1
Ir (11 A W )w)
+ 03 - 1 tr ( A w )) + aa
1 - -> tr ((A4= )
+ a5 1 - 1
tr
+ +06-
Notice that for diagonal matrices A and B, we have
tr (AB) S tr (AB) = Σ andis S
Σ
laxi |bal S Σ |a = tr(1A)IB1I
Plug in the inequality and we have
{1 + n) 2
a(2/r -1) 1)²
d
+
dt
dt
1
1
(45 =n)² 2- - [11e(k- + 1) + ( [[ (1- 2
5
+ 0
11(2- 7)
loge t d
b/ + 17
ir [ A W )) az 2/ + ir (r- ^ A²) ||
03
1 tr
A () )) +
+as
1 tr (|AV + ml)) |
- A W/ )) | | | +
a5
1 1
tr
[
1
1
(Ap- + ))²
2 2- |1p + -[p((*+1)+ + (1+t)
)] -((((++)++(ttt]]
(1-1 2
+0
(12 1)
log+ t d
an + 022
A W
1)=
|18
m|
tr A ( 1) + | |1- xpo tr (/AP V + nll)
A AD2
)) 111 |
tr (r- A W 1) + 066
03 as
43

---

Published as a conference paper at ICLR 2025
Take the sum of both sides separately, we have
(A V nI))
dtr (I-AM)
drr
+
dt
logs ) tr [[AP + +n1) +nt)²
1
[np(k
11)+++++++++ 1) 11(k++1++(pp) +0 tr (1-
+
1/(2 11)
+ (a1 + a2
A w
A 1) tr (Arr + tnl) tr (l--aW)
+ a4
+
03 + as
A
)
tr²
(1I (I- A ) + ag
From Jensen's 's Inequality with f(x) = x², we have
Σ de
Ei d
Σi X?
Therefore, it holds for diagonal matrix A E Rdxd that
4ra (A) S dir (A²)
Plug in the inequality and we have
(A v + nI) 27 dtr (I - (I-AØ)²
tr
dt
dt
[(a + +nI) 2
T (k + 1) tr (I- (1- 27 + tag
)
$ 2n1( - 77)
Because k = [clog d), we have 0(((1 n)*)) = d clug (1 So in (l time, |a8 +m
and 1 1 converge to e E (0 d5 log (1-n)) 11).
Lemma C.13. Suppose E (0 di log (1
1) ))
and there exist diagonal matrices A and B
satisfying Al S (1) and [B op S <0(1) such that
A V
n + + A
A W - B,
then it holds that
CCOT =
Proof. Now we consider the CoT loss given by Lemma C.3
1-1
CoT (0)
Ex
SW + +nSw (v + +n[)Sw"
2
I=0
Ex wi
(I+ SSW - (VS+ 15++IJui 1)w"
Plug in the expression of A V and A W
we get
82
(AS SB + 5ASB) (I (I- (I (I-((1)))'') - -AS 1
CoT (0)
E
(16)
2
1
-(1 S) + +A -(I- + 6BB -1--
(17)
2
44

---

Published as a conference paper at ICLR 2025
We first consider the term in the summation:
B
-AS- + SASB) [ (I (I-
(I-n5))) - AS
(-nSB++AASH (1(())))) - -AS(1 - -1/5)"
(-nSB + +ASB) (I- (I (-1BBS + +5BSA)
- 2tr (-1|BB + 6ASB)
(I-1S))SA +tr AS n)S)²* SA
(-n1 + AA)E SB I - - -11)))))' BS (-n) + +OA)
ur
(Term 1)
- ((n) + 5A)E
SB I - (I -nnss)) (I- ( A)
(Term 2)
tr AR S(I 1/ s )
(Term 3)
Apple Lemma C.15 to the expectation in Term 1, we have
E SB I - (I (I-)) 2
BS
= (1 (1- m)) B²
B² + +0
1 log dl ) B² [B² tr (B) (B) tr O0
1
tr (B)I
d2
It is obvious that
E SB I (
(I-(I111)))) BS
S (())
Therefore, for Term we have
tr (-nI +
SB I - (I (n=) BS
(-nI +8A)
S d|| I + 8A
E SB I-
( (1)))
BS
S 0dd).
op
(all matrices in the inequality are diagonal matrices.)
Similarly, for Term 2 and Term have
tr (-n + 8A)E (I-(1-19")u--ss)))) ns)ss A) 0())
tr AE S(I - 2))) A)
3 0()).
Add Term 2, together and we have
E (4S SB SASB) (I I- (I- (I-(I))))) - -AS|2 - sO()..
We then consider the second term in Equation 17
(I n) + A s (I n)) +8B I (I - -n/ss)))))))))))))))) ,/
AVV
SA SAASB
SB 1 - (1-1)))*]|
= tr (E(((I+AA (V)(+Ass)(I-0 1))) S/ (+ [+SAA]]] 1)
(1
28tr [
1 + A V 5) (I - nsn (1 I- (I-nS)" (I((1=)))*) Bs AV)
+ 8r
AVVE
SB - (I -ss)s BS BS AV)
<0(8)))
n1* <0)
45

---

Published as a conference paper at ICLR 2025
Recall the CoT loss in Equation (6) and Equation (77 that
By the analysis above, directly obtain
C (0)
0
Hence, the proof is complete.
Lemma C.14. Suppose $ = 1 5n
i= I n = ldd d). h. Ollog d), 7 = O e
(0.1, 0.9), || S (1), ||P (1).
Then the expectation
S(SA - -I (I-nn)) ns) rs] ) = (1 - n) AT +
where | op ( k'd
S 0
Moreover, the error is in the form
log
A 01 AT 02 (A)T
0z tr( 04 (r(A) tr tr(T)) as tr(AT)
where 01
0 k'd
02 0r3, Os = 0 - 024
nd
Proof. We can directly get the lemma by applying Lemma D.2 [S -
Lemma C.15. Suppose S 1 =
== (0.1, || I (11,, || - (1).
2.a.
O(d d). k. O(log d), 0(1)
Then the expectation
2 rs ( (1-(1-)))) AT+
E
SA -
-I (1-1)s)*
where |
lop
0 Rdd
< 0
Moreover, the error in the form
log
A a + 02 (r( + 03 tr(T) + ar bn
where 01 = 0 R'd
02 0z3 oz = 0 .4
( 0
e nd
Proof. We can directly get the lemma by applying Lemma D.2 to E[SATS] E
SA(I - nSS)
and E SA(I n) 2 rs
Lemma C.16. Suppose S
1 - x,x, n = Od d), k: Ollogd), 9(1)
n
(0.1, || ( [ | op (1). Then the expectation
SAST I I- (I (I-n1)
(I-ns)")) n (1- m)) AT AT+
where
A lop = 0 Rd 0 loga
Moreover, the error is in the form
= AT
tr()) ag tr(T)A a4 tr(A) tr (r(( ag tr(AF)
where 0n = ( k'd
0 12
a0==
02 Cc3, as
Proof. We can directly get the lemma by applying Lemma D.3 to SASTT) SAST -15)) n75)*
Lemma C.17. Suppose S = 1
I
T = Odd d). k O(logd), = O(1)
==
(..1, || I ((1), || 5 (1). Then the expectation
E SAST (I- [- (I
A AT+
where
|A
0 Kdd
< <0
Moreover, the error in the form
7
Tog
A 01 AT
02 tr (A) r 03 Ar(T) + ou (r on
where in = 0 R'd
022 03 05 = +.
044 = O nd
46

---

Published as a conference paper at ICLR 2025
Proof. We can directly get the lemma by applying Lemma D.3 to E[SAST], SAST -1s) and E SAST - n1S)
C.4 OUT-OF-DISTRIBUTIO GENERALIZATION
We restate the formal theorem here. We still denote 5 ;= n X for simplicity. Note that the number of steps & can be different/larger compared to the step number in the previous training theorem.
Theorem C.2. Suppose n = Og d), E (0.1, Ao - C'log d. Assume the out-of- distribution inpul data x; ~ (O,,,, E) E [n] where E S Amia S Amax ( 2- for some constant 8 >> 0.1, and w ~ (0,, I). Then the trained transformer in Theorem 4.1 satisfies that (t) S for any E (a minm 1-)) log" d, 1).
Proof. Recall the definition of the evaluation loss and our reduced transformer (Definition C.1
Eval (V, W) = Ex,w || (2x) (2)))) - (0s, 0, w", . |
F [1.o((u))) w w01
where 2, is the generated sequence after k steps and wxs == So (w)x=1)) is the k-th generated inter- mediate weight vector. Note that each step the transformer is inputted with the last step prediction. We define the prediction error at each step i is Aw; == w; w fe ( - wi We expand the term fe (e(x)) - w" and sum the eтror accumulation as follows:
fo(w)) w" < (Wi++ - w²) + (fe( Wk+1)
(Wk++ 1 un*) + uid VS(Wwe un) uk++1
< (Wx+1 - w²) + (w VS(W w/ w²)
ws+1 w + (I VSW Aw,.
After one step of decomposition, we notice that the error Aw can be decomposed into two parts: (1) The approximation error predicting w with ground-truth input wx. We define it
A pred 1-1
== ur + VS( we w²)) u1
(2) The accumulated error from the last inference step: (I (sw Awy. Therefore, we can inductively calculate the sum of the error:
(w)s) w0 S (Wx+1 - w²) + W. VS(Wun) un w²) - wk+1 w + (I+ VSW A Awk.
-W+++1 - w + k+1
+
I + VSW Awk.
=(w++1 w²) A
pred k+1
+
[ Vsw
A pred + (I+
+(1+++++++ (I+ Aw
= -w+++1 - w²) + Σ (S
pred k-i+1
(Awo = by definition.)
Then we have our evaluation loss upper bounded:
& w²)
2
B (oo1)o(())) - -w011²]
1
E
[wo+1
Σ
I + VSW)
prod k-i+1
2
1=0
k: 2
12) (*)
E E]](w w w01 +
1 (I VV+V
A
pred
x-1+1
2
1=0
47

---

Published as a conference paper at ICLR 2025
We first consider the first term: E [[wk+1 w²)l
E [1++1
w011 [I (I 115))t2)
w" w011== tr (B(I - ()5)2++2)
$ 2d( 8) 5)2k+2 S 2d log +1
(Lemma [D5)
Then we consider the second summation term. Since the parameters of the reduced model V -nI A, W I B, where ||
| A
op S d (r+")) for some constant
want to bound the prediction error given the ground-truth input. By Lemma D.6 have
k
E
+ VSW)
pred k-i+1
1=0
E
Σ)
[I VSW) (wk=: (S(Ww:-:- -ww) wk (K=(411)
1-0
$ 0 (d -los(( k).
Therefore, plug those back to Equation the total evaluation loss should be upper bounded by
C Evaal (0)
( - (dd-Cloogminn F +) 12) ((d--Clog(min .'r. + log" d)
D SUPPLEMENTARY LEMMAS
D.1 CONCENTRATION LEMMAS
In this appendix, we prove some concentration lemmas to estimate the expected gradient more ac-
curately. Throughout the proof, A, T are both symmetric matrices with orthonormal eigenbasis (u:)zz
Lemma D.1. Suppose S = -= E x = Odd d), = O(log d). = O(1) e
(0.1009) || S (1)). Then the expectation
E[SA(I - s] (1 (1-1)* (1-
where |A 5 <0(k'
k'dd 0
Moreover. the error is in the form A a1 + 0g Ir)
log t²
where 0n 0 k'd
02 0
n
Proof. Denote 85 = S J.
Then we expand the term SA(I - -1s))*
SA(I //
-II+ (I+55)A((1 -n) 111--1687 ( 85)
(1 (I (I A 1
17 55
(I + 85)
(1 (1-n)
(1 (I (I+ AS)A
(1 11)
km
8S +
(k)
2
I-n 17
I
85 + Σ
--
1
n7
$
85'
(I+
kn
8S +
k
852 +
7
8
8s
(1 11)
2
j=3
Now take expectation to both sides. Note that E(ō5) = 0, so all the terms only contain first order 8s vanish. We denote
kn
(1 - -- SA(SAAAT-- - n/ (1 x) A + 5SA + A&S
A
1
48

---

Published as a conference paper at ICLR 2025
which denotes all the higher order terms (the degree of öS 2 2.)
Since we have the tail bound for as in Theorem 4.6.1 Vershynin (2018) (In this lemma || is operator norm if without specification):
Pr (1685) > max (6, s ==x where = c
cV
d 8
(18)
vn
We can estimate the expectation using this property. First, given s vd ||S5|| max (5,
C
(since n = @(dlog d)). we can upper bound the operator norm of A:
2
| llpp S
A
7
+
Σ
85'
2
1
j=3
iSA
km
11)
8S
+
(k)
2
1-11 17
k
(
kk
1-
1
ー1
5
85'
f=3
10
Kn
k)
2
k Σ
(k)
+
ASA
I
8S +
S5
85
85
(1 1))
2
1- 1
1
A
kn
5S +
k)
1)
2
SS
55
os
7))
2
1
n
op
Now upper bound all matrices with their operator norm and combine all terms with the same degree of öS. We have
|| A op
S
Σ
k+2
|A
kk)
2 ( k ) (1 1- ( g- - 2) 1- 17
17
+
k
=2)(( (1 11)
|1S
1 1
k+2
<
[A][(()))) + 2(9k) + (9 115$| ||55||
( 9 (*) s $k..
2=
< 4
k+2
|||| (9h))
d C
[155) S /
81
1
k2d C
< ||
S
(#1)
n
9kdl/2 n//2
n
log
Given this upper bound, we can now upper bound the operator norm of the error term A = =E[Ā].
Suppose u = arg max ul|=1
Aull then the operator norm becomes:
|||| u T EE
A r
k2d n
E B 1e Aul
[ |1ı S CH
+
C²
n
S c 1dd
xc
+
Pr
8 ds
n
When || A
we can first upper bound the || with [[8S1| using the second 5 there exists some constant C > 0 s.t.
> where 8 2 C'APd
row of eq.
A S
|||| (9k])SSS)) < max (CCAN1SS)),(CLAA)OSSS)))
j=2
49

---

Published as a conference paper at ICLR 2025
Therefore, when |لق 2 s, ||65|| 2 min 21/2 21/4112)
GG. GG }. To apply the tail bound, we need to make
sure we pick some s' such that max (ó, 82) I min{cr probability, where d =C( C(V2= + (.). Now since 8 C'kdd
21/2 22/0+2) to upper bound the integral of
Gi
min
4//2
z Ca 4 n
Cike
Gik
for some constant Ca. Therefore, we just need max min C2 2100+2)) vn C3
min Cp
i.e. 8 S
C
V1 Vi
C4 Van k: Cs VE
s//4 V
vE
Applying the tail bound
18) with &' = min { (c2
C3
Va
C Len
G Cs
21/4/a vk
&
VA
where C2, C3, Ca, Cg are some constant, we have the error term for the tail expectation.
$2/2 s2/0+22)
Pr [
2 8
ds <
Pr
|8S|| 2
min
ds }
Cik''
Cik
S 2
expl-s""]dd
Now we estimate the upper bound of error with
22((++2)
s s1/(k++)
8² = min C2
n, C3
C2
sn
C3
sn
k2
h
n,
k
For the first term, let = Cān 2/((+2))
=
52/(0+2)
2
exp
-C2.
122 n
000
22
(l++2)/2
= (ks + 2)
expt-
xz) 2/2
Czn
22
(6+2)/2
Chn
1/2
Chn
C'k
2
k2d
< (k 2)
exp
Cz n
2
n
k2
n
n1
11/(++2;
The second term, let x = C3.
k n:
sl/(k+2)
2
exp
C2.
n ds
k
k
k+2
2(kk + 22)
expl -
kt1 dx
Eqp
xt2
k C2n
k+2
C5n
(C'2dd)
k+1
Cin
(212dd)
+
kdd
s 2
exp
<
k
n
k
n
n
Caen. k²
For the third term, let x
:
Cisn 22
0
12
2
exp
ds =
expf dd
C' cte
Cn
< 12
C' d
C₂n
kd
exp
<
07
k2
The fourth term, let x CJ.
1/22
=xa n:
s1/2
2
exp
cg k
ds
CALRA
50

---

Published as a conference paper at ICLR 2025
4h2
0
expl 1oof-xJan xddr
n2 n2
4k2
C3
n
C"k²d
1/2
C3 n2 / CC'
1/2
k2d d
n²
k
exp
k
71
Therefore, we plug this error back to the upper bound of ||||:
|||| S C² kdd
Pr
A 2 s
ds
S C k2d
Pr
||55|| 2
mint
82/2 81/(k+2)
} ds ds==0
(
0
12d
S 0
1
Cik'
C₁k
log
Finally, since by Lemma
D.8
we know the error is in the form A a1 a2 tr (A) for all A. k2
Therefore a1
0
kdd
022 0
Lemma D.2. Suppose S = E.
| op S 11 , n = Odd d), = O(logd). = (1) e
(0.1, ||A I (1),
(1).
Then the expectation
E[SA(I- TTS] (1 (1- 11- (AF + A),
where
A lop = =((22) kdd) < 0
1
Moreover, the error is in the form
log
A ln AT + 02 tr ( + 03 tr(T)A + ar oss
where 01 0 k'd
a22 03 0. 0 ( ( 0 n
Proof. Denote 85 S I.
Then we expand the term SA(I rs:
SA(I - TS
=(1 + (-55S)((111 11)1 T +8)
(1 (1-n)) 7)* (I + A5)A
I
1 85
I(I + 65)
(1-1)
(1-n)*(( (I + 8
I
kn
8S +
k
85
k Σ
(3)
kk)
j
r
(1 11)
2
1 1
[-
j=3
k
I Kn 8S +
(k)
(I
85
r
(1 1)
2
1 1)
j=3
Take expectation to both sides. Note that E[6S] = 0, so all the first order term vanish. We denote
kn
(1
= - = S - (1 - ~) A + öS AT + -
8S
A SST
1- 1
which denotes all the higher order terms (the degree of 8S 2 2.)
We can estimate the expectation using similar technique as in Lemma D.1
First, given vd and (since Od log d)), we upper bound the operator norm of Å:
|| & max 52) C
2
Ä op <
A
7
852
+
Σ
SS
r
2
1
3=3
op
8SA
Km
88
k)
1
8
(k)
-1
8
r
(11
(1-1)
2
7)
1
"
3=3
op
51

---

Published as a conference paper at ICLR 2025
+
ÖSA
I
km
85 +
(1)
kk)
1
2
Σ
()
k
ós'
res
(1
1)
2
1
11
1
+
A
km
85 + 7)
kk)
11
2
85
+
Σ
(
)
85'
ris
2
1-1
j=3
op
Now upper bound all matrices with their operator norm and combine all terms with the same degree of öS. We have
op
k+2
I (0) (1=m))
11
+ 2
(1))(
-
+
6
j-2 2) 1 1-11
k
))
17
|SSS|s
1+2
[T]]] + + (9
1 9, (?) s k).)
k+2
| (99))
CV
(185) s /
n2
j=2
81C
kdd
1
< |||||||||
I
Ca
0
9kd//2 n/7
log d
n
Now upper bound the operator norm of the error term A then the operator norm becomes:
=
E A
Suppose
||alll |u|
arg u|luuu
||||
T A/j
B u Su
(u(IAI S
C
kdd
c²
n
n
S c k2d +
xc
Pr Pr[1
N s ds
n
oryaa
When ||| 2 s where s >
C'kdd there exists some constant C1 > 0 s.t.
81((k+2) S max ((C,A11SS1)* (CIK))5S1))++))
Therefore, when ||ق|| s, || minf 1/2
}. Like Lemma D.I applying the tail bound
C. Cik 1/8440
18 with s' & min { C22 / C3 some constant, we have the error term for the tail expectation
vI
C, k
C3
8//4 v vk
where C2, C2, Ca.y C3 are
*
Vk
11/2 s1/(k+2)
Pr Pr[IA||
s ds <
Pr
|
min
C₁ke }
ds
cr
crezd
C
100 < 2
exp( -8 2 Jds.
Use the exact same argument, 2 Lo exp -2 dss < Rdd Thus, the upper bound of |||| is:
|||| S Cr 12d
Pr
2 s
ds
n2
Crgaa
I C k2dd
g1/2 s1/(k+2)
k2d
Pr
||ST|| 2 min
ds = 0
<
Cik
log
Finally by Lemma D.8. we know the error is in the form A = ajAA + a2 tr [A) + ag tr T +
as (u(r(() (r(T) for all A, P. Therefore an = 0 kdd
022 033 = 0
k2'
04 = 0 md
52

---

Published as a conference paper at ICLR 2025
Lemma D.3. Suppose S = E 2,,; (d d), k: (0.1, || I (1), lop I (1). Then the expectation
(loggd). @(1)
E S ( (1 (1- (AT + A),
where |op I 00(
Moreover, the error is in the form
log"
A = aj AT + + tr(A)r((+)))) + 03 tr (r( + 04
as tr(Ar)
where 0 k'd
022 Oa,, 05 = 0 12 0k4 1 nd
Proof. Denote 85 = -- I. Then we expand the term SIST(I-1SS))
IIST(I-123
&
= =(1- -1- (I+ (5)A (T))(T+
oS
(1 - 7))
= - (I+ AS)AT
I
kn
65 +
kk) ( 17
2
85
Σ
k
-1
85
(1 n)
2 1 1-1)
1
1 -
j=3
kn
())
k)
( 17
2
+(1- (I+
85 +
85 +
(1 17)
2
Take expectation to both sides. Note that E(A5] = 0, so all the first order term vanish. We denote
(1 (1- = - - (1- n)* A + 8S AT + ASST
kn
AT 85
1
which denotes all the higher order terms (the degree of óS 2 2.)
We can estimate the expectation using similar technique as in Lemma D.1 Given s = /d and
||oS||| < max 5d, C. (since n (dd d)), we upper bound the operator norm of A. We directly expand the formula and upper bound all matrices with their operator norm and combine all terms with the same degree of 8S. We have
A op
1++2
(6)
K
1 1-
+ 2
(,),,,)
h: 1111)
1-
+ ( k (
11 1 17
|8S1|
3
k+2
[(]]) + 2(9k))11 + 2 1155|
( S (7) s k)
3=2
k+2
d
k2d n
<
4
|
((K))
CV
S C
<
0
(155) < C a
log
Now upper bound the operator norm of A = E[ Suppose a arg max |||||=
|uu then u|
k2dd
1 11(15 ||قه| C k2d
n })
||||
F. 1
Au
1 |||ا| S C
n
12d
x0
Pr
Pr[1o|| D 8 ds
n
Craza
When |||| 2 8 where 2 C'k'd
there exists some constant C₁ 0 s.t.
n
< max (CCKK))SS))(((CCC)) [S1I)
53

---

Published as a conference paper at ICLR 2025
Therefore, when ||ق|| 2 s, ||ō$||| 2 min
C Gjk }. Like Lemma
1/2 g1/th12)
D.1
applying the tail bound where C22 C3, Ca, C3 are
18 with 8' S min C2 <1/(0+2) /n C3
2/12440) v/ C
Lan
C3
214
vn
vk
some constant, we have
8/2 /2
S exp( - s²2 ds.
Pr [ 2 8
sfss ds I
Pr
||0 2 minf
S Joyas
]ds.
C,k' C2k
Use the exact same argument, Seya exp( --2 v
Ld Thus, the upper bound of |||| is:
|| S C
k2d
Pr 16ST| 2 min!
1/2 //(++22)
ds - 0 k2d
S
0
1
C
Cik
log
Finally by Lemma D.8. know the error is in the form A aI
02 Lr
(A)r
03
tr
A
+
04 tr(A) (r(T) for all A, r. Therefore 0n = 0 k
022 023 0
04
0
m
22
0
Lemma D.4. Suppose = - , n = (d $ d), k = Ollogd). = (1) e
(0.1, 0.9), |oo I (1).. Then there exists = 0 Rdd
S 0 log the expectation is
E[A(E n)))* = (1 n))* (1 + 6)A,
Proof. Denote 85 = S I. Then we expand the term A -n))
A(I 1/5)** = (11 (1-1)*^ A - (1-1)* 69
(1
= (11- n)* k A I (1 n)
kn
55 +
k
17
SS
k
k
-1
55
77
17
kn Take expectation to both sides. Note that E(As] = 0, so all the first order term vanish. We denote
1) A (1)) k - (1 - -1)² (A 1-
A (
65
which denotes all the higher order terms (the degree of 8$ > 2.)
We can estimate the expectation using similar technique as in Lemma D.1 First, given vd and
||6S||| I max (8, / (since O(d d)), we upper bound the operator norm of Ä:
A |
op
<
A
7 1 1-m)
85
Σ
(1-n))
(
1-
-1
SS
3=3
op
Now upper bound all matrices by operator norm and combine all terms with the same degree of oS:
|
(K)
1 7 (
)²)
|
<
k+2
G S 99 (7) < k)))
op
J=2
&1C2k2d
k2d
<
|||||
S C
(1 S C )
9 1/2
n
log d
Now upper bound the operator norm of the error. Suppose = arg maxu:||uu =1
|uul [||
we have
[1113)
S C k2d + +1(131 ||ل
C² k2d
||ه|| = u T A
E
Aul
k2d cr le
Pr | > s sdss ds
n
54

---

Published as a conference paper at ICLR 2025
When ||قق| 2 where
C'k²d there exists some constant C1 0 s.t.
S max ((CIKI1SS1)), (CIK)0SS)))+3
Therefore, when ||ا| 2 8, |||||| 2 minf
C.ki Oik }. Like Lemma D.1 applying the tail bound 2/ vk
1/2 21/(2+2)
18 with s' min { C₂ 3 v
C3
y C4 yan Cs
2'/4 yn
where C2, Cs, Ca, Cs are
&
some constant, we have the error term for the tail expectation
31/2 s
00 / Cruad
Pr
|| 2 s
ds
Pr
|6S|| 2 min!
}
ds 5
exp s22 ds.
Cik'
C,k
Use the exact same argument, Seige exp( -s2
Js ds S kd
Thus, the upper bound of ||||
is:
k2d
// /2/(+22)
Cik ds o h2dd
|||| S C
Pr[
||S||
min
S
0
C
log"
3
Finally by Lemma D.8.
we know the error is in the form A ai A for all A. So a:
O k2d
13
D.2 CONCENTRATION LEMMAS FOR OUT-OF-DISTRIBUTIO, DATA
For non-isotropic covariance Gaussian data input, we also have the concentration around the covari- ance Σ when n (d d) for O. We still denote 5 _ The following lemmas are involved in the calculation for the evaluation process, for in-distribution and out-of-distribution input examples X.
Lemma D.5. Suppose S E x where x; (Oa, Σ), S Amin (m) S Amax (Σ)
2-6 for some constant 8 > 0.1, = (d d), = O(log d). 17 O(1) E (0.1,09), Then the
expectation
tr [B(I - S))) - 2d(1 - 6)*
Proof. Denote 85 = S S- Σ. Then we expand the term A(I - n1S)):
(I (1) -1* 1 - E
is
1 8
[ (I-1ス) (11
1/> Kn
I - 1/
k-1
I
k-j
(1 8)6
8S
85
(1-
8
1-8
1
j=2
Take expectation to both sides. Note that E[(is] = 0, so all the first order term vanish. We denote
(1 1)- (I- 68) (11
I 1 k
Kn
I-
1
h-1
(1 5)
5
which denotes all the higher order terms (the degree of 85 > 2). Note
1ー op
11. 1.
We can estimate the expectation using similar technique as in Lemma
D.1
First, given vd and
||ōSS||| < max (8,59) / (since = (d log d)), we upper bound the operator norm of A:
k
k-j
1 1-6 2
1-8
SS
1=2
Now upper bound all matrices by operator norm and combine all terms with the same degree of öS:
())
(k)
A op S
Σ
(1)
|8SI|
Σ
(9k)
|SS
- S 9, (g) S k))
J=2
1=2
55

---

Published as a conference paper at ICLR 2025
&PC
k2d
S C
0
(|65)) - C CVA
(9kd)/2) 1//
log d
n
Now upper bound the operator norm of the error. Suppose u == arg ina mnx w:|luu =1
|
we have
S C² 12dd
+ + |لقه
> C² k:2d
}
| = u
= u
Àu
}
kdd C
Pr
Pr[1i 2 $ ds
When || 2 8 where 8 2 Cd there exists some constant C > S.t.
max ((CIK))SS1) (CIk]]SS1))
Therefore, when ||ق|| 2 s, ||S||| 2 min
1/2
21(0422)
Like Lemma
D.1 applying the tail bound
C.E G₁ 2/1240
18 with s' S min { C2 -1/(++2) C3 some constant, we have the error term for the tail expectation
v C. 0
23/4 vn C
where C2, C3, Ca, C3 are
&
VA
vk
81/2 Cik'
s 81/(e22)
/ 00 } ds s
P ds <
Pr
165/1 > min!
exp( }
ds.
Cr
Cik
oryad
Use the exact same argument, sers expl -8 S
Rd
Thus, the upper bound of |||| is:
x
|||| < C k2d
81/2 8 81((k+2)
} ds= o
k2d 0o (x2))
1
1
Pr
| 2 min
n
Cik'
Cik
n
log
3 d
2
Finally, the absolute value of the trace should be upper bounded by
tr (1-0) 8)th
(1-0)(((
((T
I- Eミ + 2))
< 2d(1 - 5)
The next lemma deals with the prediction error.
Lemma D.6. Suppose S = 1 -T n.p. where 0; ~ NOaa, E), and the covariance matrix
n
satisfies < Amin ( S Amaxx (Σ) < 2-d for some constant ō >> O. Assume O (ogg d), k
7
Ollgg)) d), 17 (()) (0.1,0.9). Denote that A V nI.B B W I, || B op S (d--)). Then for any < k.
E|| (1 + (WW VS( W Wi-ii - wu +11)
< 0
(1 (1-8)22
8)21
d-=2++1
Proof. We will adopt similar method as we did throughout Lemma D.I to Lemma D.4 First, we expand the left hand side loss:
E I (SW) VS Wi-- w²) - w 3+1 1)) 2
e (I+ VSW) (VSW+ + SI (I -ns) ) - -iV + n1) )l²
(I++VWW) (VSW + SI (I - 175) x -i (V + +nI)s) ||
-E
s <d (1+ (WS + S1 (I 1) k 1) (V +nl)s)
56

---

Published as a conference paper at ICLR 2025
The second equation is due to w, (I - (I (I--1) n)))),, and we arranged to stress the error terms. The
third line is because E(w" "uu""" "o""]] ]]=
I. The last line is || || v vdll 'lopp
Now we expand each term of the expression within the operator norm into A, B, I, and S:
I + VSW - 7/5 + ASB 1,SB + AS.
VSW + +nS = ASB 11SB AS, V + + Therefore the formula becomes (consider each term separately)
(1+ = -1 (1-n5 ASB TJB + AS)
(VVWW +
- ( -1- + +n))s)
= (-1|BB + ASB) -1-)) AS(I -)
We still denote 55 S Σ. We first consider when the concentration holds, a.k.a |AS|| & C/2
Since ||||
B < <(d their enror are dominated by CV 4. reduce this case to the previous
Lemma D.5
Therefore we can upper bound the expression by
|| n ASB SBB + ASI| S 3 2 8)* [[-nBB + ASB) (n-ns)² - - - -1 50(d
I+ VSW
That means this part of the expectation is upper bounded by d (1 O 0
(1-8)2)
d-2c+1
Then we estimate the tail expectation. We first upper bound the above formula by ||55|
I I+ <+VS = =I + ASB - SB + ASIÍ s 00((1 5) min || 13
[- ASB)(I - - AS -7
OOdd""minf | ())
Use the same argument as in Lemma D.I can also be upper bounded by 0 -
to calculate the integral of tail bound, the tail expectation Combine those two part and we finish the proof.
D.3 THE FORM OF EXPECTATION
Lemma D.7. Suppose S = ET x. then the expectation is in the the following form for any k:
E S T 5t unu, 5kr annuu T + aruiy
+ azl for any # t.
E S T S u,u, sk a I.I.
Proof. We notice that by changing the basis to (u.)d=1.
Su S* u. 5*1 = UB sU)e,e (UT (UT)U) e.e, (UT (UTSU) k UT.
(19)
Define e; T Since gaussian is isotropic, have E[æ,] U E/æ,] = 0. After change the basis, the covariance matrix of zi, should also be the same:
Cov(æ.) = T ( I.
Therefore o; has the same distribution as æ; and we have
UE (UT sU) (UT SU) e,e, (U SU) K' UT UUE Se,e e,e s s UT.
57

---

Published as a conference paper at ICLR 2025
Subsequently, only need consider the expectation of Se,e, she s S* Decompose æ; into the sum of basis vectors and we get æ; Ej=1 Eijej.
Plug in the decomposition into the expectation and have
1 k+ E
Se,e, S ese, s
S
=E
Σ
Σ
Rio eo
e,e, II
Σ
Σ
Paiaat Six $2111
Fo
Jo.J Eld]
4=
jarJat41 Eld]
a
Σ
Σ
e,e;
Fisha Richats E/u 32441
jar E(d
=E
Σ
Σ
Figjo Fiahi ... ". Dageerdatesa (aa
jor" So(e+e))+1(d)
T T ejo e Cse
Bj e,
ejzu ein eon Janss eley Bjaiz 324t (i((tt)) e (
Σ for "or- Elul jor )
Σ
E Fiojo
Ih Ii In -
c e. e, eja T2n ete, Cjaus Ja Ciquie (e((=e)))
Note that e,e, e # only need to consider the case where j2g-1 = j2g for any E [1, W.. By symmetry, we know e 5 ee, sk is a diagonal matrix, so we have jo = J2(K+H()++ We denote
only when = b. e e,,eeeee # only when b = Therefore,
Ejo ejo ej e.e, en e ... eee e Te,, exe, ein Je, Bjap ej
to be one of the standard basis in Rddd space. It is non-zero matrix when j1 = and jk+1 = t. By the analysis above, we have
n k+22 Se.e, S c,e, st1
Σ
Σ
E Tinjo Pioh Digte To Jo Ejo-
Eim] jo. Jy Eld
Let P(2k + 2) be the set of all distinct ways partitioning {iojo, i0Ji ... i i jo into k + unordered pairs p = ((1,,22), (P2k+1,2a++2)) From Isserlis' theorem, we have
k+K
E
Dinjo Fioj
Ir we Li Fin
Σ
II
E pm pp2+1)
DEP(2k+2)
i=0
Plug it in the expectation and we have
k+K
n ++2
Se,e, S" She
Σ
Σ
Σ
II
E Tpp Tai 1 Ejo
fnv de E
E(n) jas" dut e El4 pEP(2R++2)
i=0
k+k
Σ Σ
Σ
II
E pa Fras Ejo-
PEP(2k+2) io. E/u jo. Jleld
1=0
To make sure the term in the summation is non-zero, P2 P2g should hold for any S 9 S k+1. Now consider the graph Gp and 9 with vertices {0, ........................ A: + k'}. je, is paired with inzj then we put an edge between in and u2 into Gp and put an edge between on and U2 into Go which means iu; = ioz and j.r = Jug. Therefore, for a cycle C = (M1. u2. .,)) in Gp or Gp we have Jup. Note that
bu, iu,
du- or ju, Ju, of the circle. Here we use C(-) to denote the set of circles the graph and use |C(-)] to denote the
have d choices for the value
number of circles in the graph. Let e* be the cycle in 9% which includes the vertex jn-
58

---

Published as a conference paper at ICLR 2025
Case 1: s # t. For the partition p where ji E c* and je+1 c + cx, there is only choice for and ct to take. So the term in the summation should be n nl " Simi- larly, for the partition p where jk+1 E e and j1 É c + ct, the term in the summation should be ((G) di _ For the partition p where j1 E d' c" and jk+1 E e" cc, there is only one choice for c' and e" to take. Therefore, the expectation should be
n 2 E S St ete s
Σ n a e
E PekliEe Jide
=
Pijic dheige
Σ
n -(.)) - E
+
Σ
n c( e
Σ
I ( e,e,
Piji jetide"
P:je+i Ec Ji dide
Σ
n C(
+
Recall Equation
19 we prove that
E S s u s aiuuu azuut + azI.
For the partition p where ji, 18+1 E ct, there is only one choice for c take. So the term in the summation should be nl C(G))la C(o))l-le,,c For the partition p where ji c and fu+1 E c o c*, there is only one choice for e and c* to take. So the term in the summation should be n C(G,) d (C()))-) ee,e. Similarly, for the partition p where j&+1 E c and j1 E c c", the term the summation should be nl C(S,)) C(c))122 ee,e, For the partition p where ji c' & c* and j&+1 d" + c², there is only one choice for d' and c" to take. Therefore, the expectation should be
Case 2:
n 2 e S s
Σ
n IGGN C e.e. +
Σ
n ICG C e,
PijiseesEe'
Piice" Je+1ge
+
Σ
-
e.
Σ
n [(
P
Pisideti
Σ
nl C(C)))-2
n
Piidet1Ee
Piji Jetige"
Σ
Σ
n C
e,e, +
Pik Jige"
P
Recall Equation 19 we prove that
B Su,u uu, skr
u + ozs.
Hence, the proof is complete.
Lemma D.8. Suppose S 1 D1 I then the expectation is in the following form for any k:
E. SAS'SS =AAT
+ 32 tr()) + (T(T( + B + B3 tr Lr(AT)
where A Ejn Duniuuuuu
59

---

Published as a conference paper at ICLR 2025
Proof. By lemma D.7 have:
B SAS*
d
M (a tauu, + tazf)
Σ
X x (asu: +az1)
j=1 igf
The first term here can be expand into the following form:
M (arnt + a2 + tas1)
an Ar 02 tr())r 0g tr(A) (r(T)) (a1 a2) aiz -o))+roatttt((11) agtr(AA)
Mcanwhile, the second term is directly ar + as tr(AT) We pick B2 a1, B3 a2 B4 03, 31 = 04 On 02 A5 = O5 03, and we complete the proof.
60

---

Published as a conference paper at ICLR 2025
E EXPERIMENTAL DETAILS
For all our experiments, we use pytorch Paszke et al. (2099) and models are trained on an NVIDIA RTX A6000. Each experiment takes about hour.
Setup In all our experiments, we choose d = 10, 20 and 17 = 0.4. The architecture is
Z W [ [t.-1]
FLSA Z; W [.. 2 VZ
π2
and data is drawn from the distribution in Equation (T). The batch size B is 1000 and the learning rate a is 0.001. The total time is r = 750 iterations. In the first experiment, k is chosen as 20 while 10, 20, 30, 40 in the second experiment. The baseline (evaluation loss of transformers without CoT) is given by Corollary 3.1
where 17" = n+d+1
1 2
+2 d - 2n''d d (n + d 1)d
Eval V, W)
In-distribution Generalization We empirically verify the evaluation loss gap between transform- with and without CoT shown by Theorem 3.1 and Theorem 3.2 Our experiments in Figure demonstrate that the evaluation loss of transformers with CoT converges to near zero even when k 10. See 5 details.
nut-of-distribution Generalization In addition, we empirically verify the OOD generalization result shown by Theorem 4.2 We sample different covariance matrices from distribution which complies to
2
Amiu
)Σ) Amax
17
where 77 0.4 and 0.4. 10 experiments taken to show the generality of our results for each set of experiment. Our experiment in Figure exhibits that the OOD loss of transformers with CoT converges near when & 10, 20, 30, 40 the training dos/in-distribution loss converges to zero. The final loss also drops when the number of reasoning steps increases.
k-10
k-20
4
k-30
Loss
k-40
3
OOD
2
0
0
50
100
150
200
250
300
350
iteration
Figure 3: OOD Generalization: We plot the OOD loss L when n = 20, d = 10. Each set of experiments sampled 10 different E. The mean results presented line charts, with variance represented by shaded areas. As shown, OOD loss will converge to near zero.
Given all experiments above, conclude that transformers with CoT can converge construc- tion 4.1)). surpass those without CoT (Corollary 3.1 Theorem 3.2) and generalize well to unseen data (Theorem 4.2
61


Reviews:
[{"review_id": "osYDyDwaIr", "reviewer": "Reviewer_iVV1", "summary": "This paper studies the single-layer linear self-attention (LSA) model in the context of solving in-context weight prediction problems for linear regression tasks. Unlike standard ICL, where given a query feature $x_q$, the model is trained to predict the label $y=x_q^\\top w^*$, this work trains the model to predict the task feature $w^*$ directly. As a result, the loss is evaluated based on the prediction of the task feature, defined as $\\ell=||LSA(X,y)_{[:,-1]}-w^*||^2$.\n\n1. Given input $X,y$, gradient descent (GD) with an appropriately chosen learning rate returns task feature predictions $w_0,w_1,\\cdots,w_k$. The authors introduce CoT prompting by appending the intermediate GD steps $w_0,w_1,\\cdots,w_k$ to the input, demonstrating that this approach reduces the loss compared to scenarios without CoT prompting.\n\n2. Under certain initialization assumptions, the authors present convergence results using gradient flow analysis and further demonstrate that their findings can generalize to out-of-distribution (OOD) settings.", "questions": "1. In some places, the notations are unclear:\n    - The dimensions of model output $f$ and $w$ do not seem to align in Lines 199 and 209.\n    - Is $\\sigma<1/2$ in Assumption 4.1?\n    - Could you clarify whether the loss in Theorem 4.1 corresponds to Eq. (7)?\n\n2. Could the authors explain the reasoning behind setting $V_{31}$ and $W_{13}$ to have the same set of eigenvalues in Assumption 4.1?\n\n3. Could the authors clarify in Theorem 3.1 what the model prediction is: $w^*$ or $w_1$?\n\n4. See Weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized, and the theoretical analysis appears solid.\n2. The paper introduces CoT prompting to enhance the expressivity of single-layer linear attention models in ICL tasks.", "weaknesses": "1. Limited explanation is provided about the training setting. As described, the ground-truth GD steps $w_0,w_1...$ are assumed to be available. To generate this data, the gradient of the data model, such as the linear model in this case, must be known. While, in standard ICL settings, only input-label pairs are required. Additionally, each $w^*$ is randomly sampled for each prompt, which implies that generating $M$ training samples would require $M$ gradient calculations over an $n\\times d$-dimensional dataset, typically where $M\\to\\infty$.\n\n2. The model's performance and loss are highly dependent on the learning rate $\\eta$. In this work, $\\eta$ is fixed, meaning that optimal losses can only be achieved by setting $\\eta=\\frac{n}{n+d+1}$. In standard ICL settings, $\\eta$ is often implicitly learned. As a result, the definition of global minimization in the paper is somewhat ambiguous.\n\n3. By making $n$ and $k$ dependent on $d$, it is unclear how varying values of $n$ and $k$ affect convergence and evaluation losses.\n\n4. The CoT + one-layer approach in this paper appears closely related to ICL + multi-layer methods. The paper could benefit from a discussion on this connection.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QLPV31hFt4", "reviewer": "Reviewer_Kfz9", "summary": "The authors investigate the training dynamics of transformers trained with a Chain of Thought (CoT) objective, specifically in an in-context weight prediction task for linear regression. Under certain assumptions, it proves that a one-layer transformer without CoT training can only perform a single gradient descent step, resulting in suboptimal recovery of the weight vector. However, with CoT training, the transformer can execute multi-step gradient descent, enabling near-exact recovery and some out-of-distribution generalization. The authors provide theoretical results demonstrating the global convergence of the training via gradient flow and empirical evidence showcasing the superior performance of transformers trained with CoT compared to those without.", "questions": "See the weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "This paper provides some theoretical analysis on CoT, which is lacking and deserves much more attention. These seem to be new results on CoT, albeit in the setting of a one-layer linear transformer.", "weaknesses": "-I’m not sure about the significance of these results. The comparison here involves a transformer trained on a CoT objective, which seems to be unusual. It’s entirely possible (and likely) that CoT data is contained in the large pretraining corpuses of today’s LLMs, but it remains a next-token prediction task. There is no difference between the training objectives of an LLM trained on text with or without CoT.\n\n-The results are for one-layer linear transformers. This seems to be a common model choice in the literature for obvious reasons, but it deviates so far from the setup of an LLM that the utility of these results is questionable.\n\n-What is the significance of assuming the number of samples to be bounded by dlog^5d? It seems like a contrived assumption to make the proofs work. For example, couldn’t Corollary 3.1 instead be framed as saying at least a quadratic number of samples are needed to control the evaluation error?\n\n-The main theorem 4.1 relies on seemingly restrictive assumptions. Even though these are settings used in previous works, can you say anything in a more general setting?\n\n-While the theorems may be new, it is hard to tell whether there is any novelty in the methods or whether these are straightforward extensions of techniques used in other works the authors mentioned such as Bai et al.. The manuscript would benefit from clearly outlining where novel ideas are needed.\n\n-Theorem 4.2 is confusing. Can the authors provide an intuitive explanation for why the model exhibits OOD generalization within the provided bounds of the spectrum of the covariance matrix? Also, I don’t think L^{eval}_{\\Sigma} is defined anywhere.\n\n-The manuscript would benefit from empirical analysis for Theorem 4.2.\n\n-In general, the manuscript would benefit from some intuitive explanations of the assumptions in the theorems. Furthermore, while a theoretical analysis of CoT is interesting, I don’t think the results are very impactful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yEJyAjsbnM", "reviewer": "Reviewer_D4BQ", "summary": "This paper introduces a novel in-context learning task, termed in-context weight prediction for linear regression. The authors demonstrate that single-layer transformers exhibit limitations in this task when working with limited number of examples, and they show that incorporating a chain-of-thought mechanism significantly enhances performance. Through analysis of gradient dynamics in single-layer transformers using chain-of-thought, they establish convergence results that align well with their empirical findings.", "questions": "Please refer to the weakness section, the reviewer would like to better understand how the authors reconcile the apparent disconnect between CoT steps in practical applications (where more steps don't necessarily yield better performance) and their theoretical analysis. Additionally, how can the insights from this work inform the design of effective step-by-step instructions for improved CoT performance?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "While existing literature always try to understand the expressive power of CoT from a computational complexity perspective, this work offers a fresh analytical framework, and they provide convergence analysis and sufficient  experiments to validate their findings. The approach is both innovative and provides valuable insights for future research directions.", "weaknesses": "Although this work presents interesting findings, the construction and proof techniques largely adapt existing methodologies. Furthermore, similar results might be achievable using looped transformers as demonstrated in [1], where:\n$$\nZ_0 = \\begin{pmatrix}\nx_0 & x_n& x_{query} \\\\\\\\\ny_0 & y_n& 0\n\\end{pmatrix}, Z_k = f(Z_{k-1})\n$$\nwhere $f$ is a linear transformer with fixed parameters. In other words, while existing work demonstrate looped transformer (inference a (single layer) fixed parameter transformer with $k$ times can get a similar form $w^{k} x_{query}$, where $w^k$  defined like theorem 3.2 in this work, the theoretical contribution that directly analysis $w^k$ seems limited. So a more detailed discussion between this work and analysis in [1], and a meanfull explaination why such a in-context learning task framework is necessary, is desired.\n\n[1] Ding, N., Levinboim, T., Wu, J., Goodman, S., & Soricut, R. (2023). CausalLM is not optimal for in-context learning. arXiv preprint arXiv:2308.06912.\n\nNote that [1] analysis a multilayer transformer with the same parameters as each layer, so I refer this work as looped transformers, other works like [2,3] also analysis similar transformers.\n\n[2] Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 45614-45650.\n\n[3] Gatmiry, K., Saunshi, N., Reddi, S. J., Jegelka, S., & Kumar, S. (2024). Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?. arXiv preprint arXiv:2410.08292.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)