On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability
Authors: Chenyu Zheng
Wei Huang
Rongzhen Wang
Guoqiang Wu
Jun Zhu
Chongxuan Li
Venue: NeurIPS 2024 poster (2024)
Keywords: Mesa-Optimization
In-context learning
Autoregressive pretraining
Non-convex optimization
Learning theory
Transformers
Primary area: learning_theory
TL;DR: Towards understanding the mechanisms underlying the in-context learning from autoregressive pretraining by rigorously verifying the mesa-optimization hypothesis.

Abstract: Autoregressively trained transformers have brought a profound revolution to the world, especially with their in-context learning (ICL) ability to address downstream tasks. 
Recently, several studies suggest that transformers learn a mesa-optimizer during autoregressive (AR) pretraining to implement ICL. Namely, the forward pass of the trained transformer is equivalent to optimizing an inner objective function in-context.
However, whether the practical non-convex training dynamics will converge to the ideal mesa-optimizer is still unclear.
Towards filling this gap, we investigate the non-convex dynamics of a one-layer linear causal self-attention model autoregressively trained by gradient flow, where the sequences are generated by an AR process $x_{t+1} = W x_t$. First, under a certain condition of data distribution, we prove that an autoregressively trained transformer learns $W$ by implementing one step of gradient descent to minimize an ordinary least squares (OLS) problem in-context. It then applies the learned $\widehat{W}$ for next-token prediction, thereby verifying the mesa-optimization hypothesis. Next, under the same data conditions, we explore the capability limitations of the obtained mesa-optimizer. We show that a stronger assumption related to the moments of data is the sufficient and necessary condition that the learned mesa-optimizer recovers the distribution. Besides, we conduct exploratory analyses beyond the first data condition 
and prove that generally, the trained transformer will not perform vanilla gradient descent for the OLS problem. Finally, our simulation results verify the theoretical results.

On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability
Chenyu Zheng 1.2, Wei Huang", Rongzhen Wangh? Guoqiang Wu',
Jun Zhu", Chongxuan Li
1 Gaoling School of Artificial Intelligence, Renmin University of China
2 Beijing Key Laboratory Big Data Management and Analysis Methods
RIKEN AIP School of Software, Shandong University
5 Dept. of Comp. Sei. & Tech., BNRist Center, THU-Bosch ML Center, Tsinghua University {cyzheng, wangrz, chongxuaml /ruc. edu. en; wei huang. vr@riken.jp; guoqiangvudsdu. edu. cn; dcszj @mail. tsinghua. edu. cn
Abstract
Autoregressively trained transformers have brought a profound revolution to the world, especially with their in-context learning (ICL) ability to address downstream tasks. Recently, several studies suggest that transformers learn a mesa-optimizer during autoregressive (AR) pretraining implement ICL. Namely, the forward pass of the trained transformer is equivalent to optimizing an inner objective function in-context. However, whether the practical non-convex training dynamics will converge to the ideal mesa-optimizer is still unclear. Towards filling this gap, investigate the non-convex dynamics of one-layer linear causal self- attention model autoregressively trained by gradient flow, where the sequences are generated by AR process #+++1 = W at, First, under certain condition of data distribution, we prove that an autoregressively trained transformer learns W by implementing one step of gradient descent to minimize an ordinary least squares (OLS) problem in-contexi. It then applies the learned W for next-token prediction, thereby verifying the vresa-optimization hypothesis. Next, under the same data conditions, we explore the capability limitations of the obtained mesa-optimizer. We show that stronger assumption related to the moments of data is the sufficient and necessary condition that the learned mesa-optimizer recovers the distribution. Besides, conduct exploratory analyses beyond the first data condition and prove that generally, the trained transformer will not perform vanilla gradient descent for the OLS problem. Finally, simulation results verify the theoretical results, and the code is available at
1 Introduction
1 have revolutionized the Al community in lots of 4 5 6) computer vision 7 g 9 10 and multi-modal learning [1] 12 13 14 15). The crux behind these large models is a very simple yet profound strategy named autoregressive (AR) pretraining, which encourages transformers predict the next token when context is given. In terms of the trained transformers, one of their most intriguing properties is the in-context learning (ICL) ability (2 which allows them adapt their computation and perform downstream tasks based on the information (e.g. examples) provided the context without any updates their parameters. However, the reason underlying the emergence of ICL ability is still poorly understood.
Foundation models based on transformers
fields, such language modeling 2
B
Correspondeecc Chongxuan Li.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Recently, we are aware that some preliminary studies [6; 17 have attempted to understand the ICL ability from the AR training and connected its mechanisms to a popular hypothesis named mesa- optimization [8]. which suggests that transformers learn some algorithms during the AR pretraining. In other words, the inference process of the trained transformers is equívalent to optimizing some inner objective functions on the in-context data.
Concretely, the seminal work [16] constructs a theoretical example where single linear causally- masked self-attention layer with manually set parameters can predict the next token using one- step gradient- learning for an ordinary least squares (OLS) problem over the historical context. Moreover, they conduct numerous empirical studies to establish a close connection between autoregressively trained transformers and gradient- the setting of [16], recent work [17] precisely characterizes the pretraining loss landscape of the one- layer linear transformer trained on a simple first-order AR process with a fixed full-one initial token. As result, they find that the optimally- -ttainned transformer recovers the theoretical construction in [16]. However, their results rely on imposing the diagonal structure on the parameter matrices of the transformer and do not discuss whether the practical non-convex dynamics can converge to the ideal global minima. Besides, it is still unclear about the impact of data distribution on the trained transformer, which has been proven to be important in practice [9] 20 211 [2] and theory 23
oestiioptiiiitiinn algorithms. Built upon
In this paper, we take further step toward understanding the mesa-optimizatiin trained transformers. Specially, without an explicit diagonal structure assumption, we analyze the non-convex dynamics of one-layer linear transformer trained by gradient flow on a controllable first-order AR process, and try to answer the following questions rigorously:
autoregressively
1. When do pesaaoptimizaatin algorithms emerge in autoregressively trained transformers?
2. Whut is the capability limitation of the mexa-optimizer if it does emerge?
Our first main contribution is characterize sufficient condition (Assumption 4.1) on the data distribution for a mesa-optimizer to emerge in the autoregressively trained transformer, in Section 4.1 We note that any initial token 37 whose coordinates I1; are i.i.d. random variables with zero mean and finite moments satisfy this condition, including normal distribution M(a, の² La). Under this assumption, the non-convex dynamics will exactly converge to the theoretical construction in [[6] without any explicit structural assumption in [17]. resulting in the trained transformer implementing one step of gradient descent for the minimization of an OLS problem in-context.
Our second main contribution is to characterize the capability limitation of the obtained mesa- optimizer under Assumption 4.1 in Section 4.2 We characterize a stronger assumption (Assump- tion 4.2) related to the moment of data distribution as the necessary and sufficient condition that the mesa-optimizer recovers the true distribution (a. predict the next token correctly). Unfortunately, we find that the mesa-optimizer can not recover the data distribution when the initial token is sampled
from the standard normal distribution, which suggests that ICL by AR pretraining 16 17 is different
from ICI. by few-shot learning pretraining [24] 25 26
27 28 29, 30 23 (see details in Section 2
setup attracting attention from many thcorists recently. We think ICL in the setting of AR pretraining needs more attention from the theoretical community.
we further study the convergence of the training dynamics when Assumption 4.1 does not hold anymore by adopting the setting in [17]. In this case, as a complement of [17]. we prove that under similar but weaker structural assumption, training dynamics will converge to the theoretical construction in [16] and the trained transformer implements exact gradient-based mesa-optimization. However, we prove that without any structural assumption, the trained transformer will not perform vanilla gradient descent for the OLS problem in general. Finally, we conduct simulations to validate
In Section 4.3
our theoretical findings in Section
2
A Additional related work
2.1
Mesa-optimization in ICL for few-shot linear regression
In addition to AR pretraining, much more empirical 31 25 32 33 and theoretical studies [24]
26 27 28
have given evidence to the mesa-optimization hypothesis when transformers are trained to solve few-shot linear regression problem with the labeled training instances in the context. On the experimental side, for example, the seminal empirical works by [11 considers
29
34 30

---

ICL for linear regression, where they find the ICL performance of trained transformers is close to the OLS. On the theoretical side, considering one-layer linear transformer, 26 24 prove that the global minima of the population training loss is equivalent to one step of preconditioned gradient descent. Notably, [24] further proves that the training dynamics do converge to the global minima, and the obtained mesa-optimizer solves the linear problem. For multi-layer attention models, recent works suggest that they can perform efficient high-order optimization algorithms such as Newton's ' method [36] 37 [4]. Unfortunately, the pretraining goal of these studies is different from the AR training. Therefore, still unclear whether these findings can be transferred to transformers autoregressively trained sequential data.
2.2 Other explanations for ICL
In addition to the mesa-optimization hypothesis, there are other explanations for the emergence of
ICI. (38 (39 40 41 explain ICL. as inference over an implicitly learned topic model. 42 connects ICL to multi-task learning and establishes generalization bounds using the algorithmic stability technique. 43 and 44 study the implicit bias of the next(last)-token classification loss when each token is sampled from finite vocabulary. Specially, [44] proves that self-attention with gradient descent learns an automaton that generates the next token by hard retrieval and soft composition. 45] explains ICL as kernel regression. These results are not directly comparable ours because we study the ICL of a one-layer linear transformer with AR pretraining on the first-order AR process.
3 Problem setup
Elementary notations. We define [n]] = (1.2...., We use lowercase, lowercase boldface, and uppercase boldface letters to denote scalars, vectors, and matrices, respectively. For vector a, denote its i-th element as a, For a matrix A, we use Ag:, A, and Aij to denote its k-th row, k-th column and (i, (ith element, respectívely. For vector (matrix A), use a" (A") to denote its conjugate transpose. Similarly, we a and A to denote their element- wise conjugate. We denote the n--dimensional identity matrix by I,, We denote the one vector of size n by 1,- In addition, we denote the zero vector of size n and the zero matrix of size m X n by O,, and 0mxn. respectively. We
use & and @ to denote the Kronecker product and the Hadamard product, respectively, Besides, we denote Vec(-) the vectorization operator in column-wis order.
3.1 Data distribution
We consider a first-order AR process as the underlying data distribution, similar to recent works on AR pretraining [7]. Concretely, to generate a sequence (z)..... ØT) E CdXT, we first randomly sample unitary matrix W E Cdxd uniformly from a candidate set Pw =
> 1, Vi E [d} and the initial data point Φ from controllable distribution Da, to be specified later, then the subsequent elements are generated according to the rule x = W w; for t E [ 1]. For convenience, we denote the vector (A,...., )..) by X. We note that the structural assumption on W is standard in the literature is a natural extension of the recent studies on ICI. for linear problems 24 26 they focus on y w approach to investigate AR pretraining because the regression dataset in the previous ICL theory is not suitable since each token æ; does have a relation with the context. In this paper, we mainly investigate the impact of the initial distribution Da, on the convergence result of the AR pretraining.
learning problems involving matrices 46
17
In addition, it
28
29
30), where
I diag( (w).. a.. Furthermore, adopting this new data model is necessary
3.2 Model details
Linear causal self-attention layer. Before introducing the specified transformer module we will analyze this paper, we first recall the definition of the standard causally-masker self-attention layer D1. whose parameters 0 includes a query matrix Wº E R xd. a key matrix W E Rd a value matrix WV E Xd. a projection matrix WP e xdd and a normalization factor Pt > O. At time step t. let E, (ei....,e:) dd be the tokens embedded from the prompt sequence
3

---

P. = (#1....,.I.) E Cdxt, the causal self-attention layer will output
(W Ex) (
- WV E, softmax
E ECO,
Pt.
where the softmax operator is applied to each column of the input matrix. Similar to recent theoretical
works on ICL with few-shot pretraining [24 225 16. 17: 26
27
28 [29], we consider the linear
attention module in this work, which modifies the standard causal self-attention by dropping the softmax operator. Reparameterizing WKQ = W K WQ and WPV WFWV we have
(WKQ, WPV), and the output can be rewritten as:
E, W
f. (,; 0) e + W E.
Pt
Though it is called linear attention, we note that the output ft EE; 0) is non-linear w.r.t the input tokens E, due to the existence of the E, In terms of the normalization factor Pt, like existing works [2 26), we take it to be 1 because each element in E, is Hermitian inner product of two vectors of size f.
Embeddings. We adopt the natural embedding strategy used in recent studies AR learning [66 17]. Given sequence P. = (at,...., 222. the i-th token is defined as e; = (0 20, a0 _1 E C3d thus the corresponding embedding matrix E, can be formally written as:
0a
0d
(el, e()
21
32
æt E C3axe
Zp
x
It-1
where To Od complement. This embedding strategy is natural extension of the existing
theoretical works about ICL for linear regression [30 20) 24 26
28 29] The main difference is that
the latter only foeus on predicting the last query token while need predict each historical token. We note that practical transformers do learn similar token construction in the first softmax attention layer (e.g., see Fig. 4B in [16)..
Next-token prediction. Receiving the prompt P: = (....... Zt), the network's s prediction for the next element 2 will be the first d coordinates of the output fi EE 0), aligning with the setup adopted in 16 17). Namely, we have
(e 0) (E; Cd.
Henceforth, we will omit the dependence on E, and 0 and use D. if is not ambiguous. Since only the first d rows are extracted from the output by the attention layer, the prediction ÿt just depends on some parts of PV and W KQ, Concretely, we denote that
WPV
WPV
W
PV 13
KQ w
KQ
w
kQ
13
wrv
W PV
W PV
21
WPV PV
W PV
W
W
23
PV
PV
WKQ
=
w
w w
w
Q
kk
31
32
33
33
where PV W KQ
E Rdxd for all i,j E [3].
Then the ÿt can be written as
ETE. ET
W
KQ
W
KQ Q
(WPPV WPV
W WRQ
KQ
WKQ
E
(1)
Pt.
Here (ef e E C2dxt denotes the last 2d rows of the E,, where e. = (क. e==
,I T..
Therefore, only need to analyze the selected parameters in Eq. T during the training dynamics. The derivation of Eq. can be found Appendix A.I
3.3 Training procedure
Loss function. To train the transformer model over the next-token prediction task, we focus on minimizing the following population loss:
T-1
T-1
L( = Σ
L.
Em
W
2.+ 112
(2)

---

where the expectation is taken with respect to the start point 2j and the transition matrix W. Henceforth, we will suppress the subscripts of the expectation for simplicity. The population loss is standard objective in the optimization studies [24] 47 and this objective has been used in recent works on AR modeling [66 17]. The summation starts from & = because we do not have any information to predict 22 given only 11
Initialization strategy. We adopt the following diagonal initialization strategy, and similar settings have been used in recent works on ICL for linear problem (24 30 17..
Assumption 3.1 (Diagonal initialization). At the initial time T O, we assume that
Oaxd
Oaxd Oaxd
Oaxd Oaxe
Oaxd
bo
Oaxd
w w
Oaxd
W PV (0) = Oaxd
Oaxd
Oaxd
Odxd
Ta
Oaxe)
Oaxd
Oaxd
Oaxd
where the red submatrices are related to the De and changed during the training process.
The most related paper [17] considers a stronger diagonal structure than ours, and it only investigates the loss landscape. Our results deepened the understanding of AR transformers by considering practical training dynamics, We think this assumption might be inevitable for a tractable analysis and leave theory for standard (Gaussian) initialization to future work. In Section experiments under standard initialization, which further supports the rationality of Assumption 3.1
also conduct
Optimization algorithm. We utilize the gradient flow to minimize the learning objective in Eq. 2 which is equivalent to the gradient descent with infinitesimal step size and governed by the ordinary differential equation (ODE) dr (O).
Main results
In this section, we present the main theoretical results of this paper. First, in Section 4.1 we prove that when Da, satisfies some certain condition (Assumption 4.1), the trained transformer implements one step of gradient descent for the minimization an OLS problem, which validates the rationality of the mesa-optimizatioon hypothesis [16]. Next, in Section 4.2 further explore the capability limitation of the obtained mesa-optimizer under Assumption 4.1 where we characterize a stronger assumption (Assumption 4.2 as the necessary and sufficient condition that the mesa-optimizer recovers the true distribution. Finally, we go beyond Assumption [.1] where the exploratory analysis proves that the trained transformer will generally not perform vanilla gradient descent for the OLS problem.
4.1 Trained transformer is a mesa-optimizer
In this subsection, show that under certain assumption of Day, the trained one-layer linear transformer will converge to the mesa-optimizer [6 17. Namely, will perform one step of gradient descent for the minimization of OLS problem about the received prompt. The sufficient condition of the distribution Da, can be summarized as follows.
Assumption 4.1 (Sufficient condition for the emergence of mesa-optimizer). We assume that the distribution Da, of the initial token æi E Rdd satisfies ERi-Da, [ ...xi,, for any subset {iz, [i1,...,44 in n S 4) of [d], and r2.... E In addition, we assume that K1 E[r4 N2 = E[29,] and Kg = Σ イ E Eizs,at are finite constant for any j E [d].
Finding Assumption 4.1 is non-trivial since we need to derive the training dynamics first. The key intuition of this assumption is to keep the gradient of the non non-diagonal elements of W and WPPV as zero, thus they can keep diagonal structure during the training. We note that any random vectors 2 whose coordinates 2 are i.i.d. random variables with zero mean and finite moments of order 2, 4, and satisfy this assumption. For example, includes the normal distribution N ((0a,, Ia), which is a common setting in the learning theory field 47 48 49 50 51]. Under this assumption, the final fixed point found by the gradient flow can be characterized as the following theorem.
Theorem 4.1 (Convergence of the gradient flow, proof in Section 5).
Consider the gradient flow of the layer linear transformer (see Eq. 7 over the population AR pretraining loss (see Eq. 2
5

---

Suppose the initialization satisfies Assumption 3.1
and the initial token's 's distribution Dz, satisfies then the gradient flow converges
Assumption 4.7
W QQ W
WKQ W
=
Oaxx āJa
Oaxd Oaxd
W.P
W = bla
Oaxa)
Though different initialization (4o; bo) lead to different (ū, b), the solutions' product äb satisfies
ab
K1
Σ
K2 K2
As far as we know, Theorem 4.1 is the first theoretical result for the training dynamics and the mesa- optimization hypothesis of autoregressive transformers. The technical challenge compared existing ICL theory for regression [24] mainly has two parts. First, our data model breaks the independence between data different times, which causes difficulty in decomposing and estimating the gradient terms. Second, we modify the embedding strategy (more dimensions), scale the attention model (much more parameters), and change the loss function (more terms) perform the full AR pertaining. All these parts are not well studied in the literature and make the gradients more complicated.
Theorem 4.1 is also non-trivial extension of recent work 177. which characterizes the global minima of the AR modeling loss when imposing the diagonal structure on all parameter matrices during the training and fixing 2 as 1a- In comparison, Theorem 4.1 does not depend the special structure, and further investigates when the mesa-optimizer emerges in practical non-convex optimization.
We highlight that the limiting solution found by the gradient flow shares the same structure with the careful construction in [16), though the pretraining loss is non-convex. Therefore, our result theoretically validates the rationality of the opstiippimiziiionn hypothesis [16] in the AR pretraining setting, which can be formally presented as the following corollary.
Corollary 4.1 (Trained transformer as a mesa-optimizer proof in Appendix(A.3) We suppose that the same precondition of T'Thoormm hollds. holds. When predicting the (L+ 1)th token, the trained transformer
obtains W by implementing one step of gradient descent for the OLS problem LOLs.((W) =
Wx; starting from the initialization W Oaxd with a step size
ūb
4.2 Capability limitation of the mesa-optimizer
Built upon the findings in Theorem 4.1, a simple calculation (details in Appendix A.3 shows that the prediction of the obtained mesa-optimizer given a new test prompt of length Tte is
E. Tie
DT. W
ab
TT.
(3)
It is natural to ask the question: where is the capability limitation of the obtained mesa-optimizer, and what data distribution can the trained transformer learn? Therefore, in this subsection, we study under what assumption of the initial token's distribution Dz, the one step of gradient descent performed by the trained transformer can exactly recover the underlying data distribution. First, leveraging the result from Eq. 3 present negative result, which proves that not all Dz, satisfies Assumption 4. can be recovered by the trained linear transformer.
Proposition 4.1 (AR process with normal distributed initial token can not be learned, proof in Appendix A.4. Let Da, be the multivariate normal distribution (O2,0 o² Ia) with any 0² 0, then the "simple" AR process can not be recovered by the trained transformer even in the ideal case with long training context. Formally, when the training sequence length T., large enough, for any test context length Tua and dimension j E [d), the prediction from the trained transformer satisfies
(TTi))ss
Fay,W W
(W Dj
Therefore, the prediction yr. will not converges to the true next token W &T.

---

Proposition 4. suggests that ICL by AR pretraining [16 17 is different from ICL by few-shot 30 23), which attracts much more attention from the theoretical community. In the latter setting, recent works [24] 26 proves that one step of gradient descent implemented by the trained transformer can in-context learn the linear regression problem with input sampled from (02, Ia).. However, in the AR learning setting, the trained linear transformer fails.
pretraiming 24 25
26
27
28 29
This negative result shows that one-step GD learned by the AR transformer can not recover the distribution, but this can be solved by more complex models. Even for more complex data (W is not diagonal), [6] has empirically verified that multi-layer linear attention can perform multi-step gradient descent to learn the data distribution. Therefore, to address the issue, future works are suggested to study more complex architecture such as softmax attention [30]. multi-head attention [52] [4]. deeper attention layers (55] [6], transformer block [56] [57; [8) and so on. Future theory considering more complex AR transformers can adopt the same data model and token embeddings in
this paper, and try to use similar proof technique to derive the training dynamics.
Proposition 4.1 implies that if we want the trained transformer to recover the data distribution by performing one step of gradient descent, a stronger condition of Dæ, is needed. Under Assumption 4.1 the following sufficient and necessary condition related the moment of Da, is derived from Eq. by letting Dr. converges to W Wr.. when context length Tür and The are large enough.
Assumption 4.2 (Condition for success of mesa-optimizzer), Based on Assumption
4.1 we further
suppose that Ra ETa m
TTe - TT.. for any I and W, when Tte is large enough.
62 Tee
Assumption strong and shows the poor capability of the trained one-layer linear transformer because common distribution (e.g. Gaussian distribution, Gamma distribution, Poisson distribution, etc) always fails to satisfy this condition. Besides, it is a sufficient and necessary condition for the mesa-optimizer to succeed when the distribution Da, has satisfied Assumption 4.1 thus can not be improved in this case. We construct the following example that satisfies Assumption 4.2]
Example 4.1 (sparse vector). If the random vector æ1 R2 is uniformly sampled from the candidate set of size 2d { O, o,.... 0)? (0, C, 0) (O,...... (0..... 0, c) } for any fixed E R, then the
distribution Pas satisfies Assumption 4.2 The derivation can be found in Appendix A.S
For completeness, formally summarize the following distribution learning guarantee for the trained transformer under Assumption 3.1 and 4.1
Theorem 4.2 (Trained transformer succeed learn the distribution satisfies Assumption 4.2 proof
in Appendix [.6].. Suppose that Assumption 3.7 (and4.///hlll. then Assumption 4.2 is the sufficient and necessary condition for the trained transformer to learn the AR process. Formally, when the training sequence length Ter and test context length Tie are large enough, the prediction from the trained transformer satisfies
yT.
W WITW
Ter, Tre - +00.
4.3
Go beyond the Assumption 4.1
The behavior of the gradient flow under Assumption 4.1 has been clearly understood in Theorem 4.1 The follow-up natural question is what solution will be found by the gradient flow when Assump- tion 4.1 does not hold. In this subsection, we conduct exploratory analyses by adopting the setting in where the initial token #1 is fixed as 1
First, sharing the similar but weaker assumption of [17], we impose W and w to stay diagonal during training by masking the non-diagonal gradients, then the trained transformer will perform one step of gradient descent, as suggested by [17] Formally, it can be written as follows.
Theorem 4.3 (Trained transformer as mesa-optimizer with non-diagonal gradient masking, proof in Appendix A.7.. Suppose the initialization satisfies Assumption 3.1] the initial token is fixed La- and we clip non-diagonal gradients of W WKQ and W.PV during the training, then the gradient flow of the one-layer linear transformer over the population AR loss converges to the same structure as the result in Theorem 4.1 with
1
+
E
Therefore, the obtained transformer performs one step of gradient descent in this case.
7

---

Theorem can be seen as complement and an extension of Proposition in [17] from the perspective of optimization. We note that [T7] assumes all the parameter matrices to be diagonal and only analyzes the global minima without considering the practical non-convex optimization process.
Next, adopt some exploratory analyses for the gradient flow without additional non-diagonal gradient masking. The convergence result of the gradient flow can be asserted as the following proposition. The key intuition of its proof that when the parameters matrices share the same structure as the result Theorem 4.1 the non-zero gradients of the non-diagonal elements of W Q and WiP will occur. In addition, we note the result does not depend on Assumption 3.1
Proposition 4.2 (Trained transformer does not perform on step of gradient descent, proof in Ap- pendix A.8 The limitinn point found by the gradient does not share the same structure as that in Theorem 4.1 thus the trained transformer will not implement one step of vanilla gradient descent for minimizing the OLS problem 2 >>
im1 1Pi41-Was =
To fully solve the problem and find the limiting point of the gradient flow in this case (or more generally, any case beyond Assumption 3.1 and 4.1. one can not enjoy the diagonal structure of W and W PV anymore. When W and W 12
PV are general dense matrices, computation of the gradient will be much more difficult than that in Proposition 4.2 Therefore, leave the general rigorous result of convergence without Assumption 3.1 and 4.1 for future work.
We are aware that recent theoretical studies on ICL for linear regression have faced similar problem. (24 26 23] find that when the input's distribution does not satisfy Assumption [41 (e.g.. (0,, Σ),, the trained transformer will implement one step of preconditioned gradient descent on for some inner objective function. We conjecture similar results will hold in the case of in-context AR learning. We will empirically verify this conjecture when is a full one vector, in Section
5 Proof skeleton
In this section, we outline the proof ideas of Theorem 4.1 paper, and also a theoretical base of the more complex proofs of Theorem 4.3 and Proposition 4.2 The full proof of this Theorem placed in Appendix A.2
which is one of the core findings of this
The first key step is to observe that each coordinate of prediction y (Eq. can be written as the output of quadratic function, which will greatly simplify the follow-up gradient operation.
Lemma 5.1 (Simplification of Tt... proof in Appendix A.2.1 Each element of the network's
prediction Øe.j (j [d]) can be expressed as the following.
E Ex
E E
y1.j = BJ
T
Vec(A) Vee (A) E
VecT(A)
B,,
Pr
Pt
where the A and Bj are defined
WkQ
KQ W-RQ 32
W
WkQ R
bii
W PVT
A (an a2d)
B
bi2
WPP 13,j:
33
with a; E 2d and رزl و E Rdd
Next, We calculate the gradient for the parameter matrices of the linear transformer and present the dynamical system result, which is the most complex part the proof of Theorem 4.1
Lemma 5.2 (dynamical system of gradient flow under Assumption 4.1 proof in Appendix A.2.2). Suppose that Assumption 4.1 holds, then the dynamical process of the parameters in the diagonal of W Q and WPP. satisfies
d
T
ab²
(T
2)x²
+
K3
+ b(T 2)
dr
d
T-1
-0b
(T
2)
Σ
K3
+ a(T 2)
dr
while the gradients for all other parameters were kept at zero during the training process.
8

---

Similar ODEs have occurred in existing studies, such as the deep linear networks [59] and recent ICL for linear regression (24]. Notably, these dynamics are the same as those of gradient flow on a non-convex objective function with clear global minima, which is summarized as the following.
Lemma 5.3 (Surrogate objective function, proof in Appendix ).2.3). Suppose that Assumption 4.1 E i-1Kis Kis and (T 2) by 0i and cz. respectively. Then, the holds and denote (T k2 +
dynamics in Lemma 5.2 are the same as those of gradient flow on the following objective function:
1 (c2 -cab))
E(a, b) =
201
whose global minimums satisfy ab c2/01.
Furthermore, We show that although the objective P(a, b) is non-convex, the Polyak-t.ojasiewicz (PL. inequality [60]6]] holds, which implies that gradient flow converges to the global minimum.
Lemma 5.4 (Global convergence of gradient flow, proof in Appendix A.2.4). Suppose that Assump-
tion 4.1 holds, then E(a, b) is a non-convex funcctttion and satisfies the PL inequality as follows.
i E(a, b) da
1
((aabb)
b) 2 2c (o2)++)) [((a,b) (((((((.)))
(((a,.)-mmnn((.. min E(a, b) a.b
+
0b
Therefore, the gradient flow in Lemma
5.2 converges to the global minimum of E(a b).
Finally, Theorem 4.1 can be proved by directly applying the above lemmas.
6 Simulation results
In this section, we conduct simulations to verify and generalize theoretical results. In terms of the train set, we generate 10k: sequences with Ter = 100 and d = 5. In addition, we generate another test set with 10k sequences of the same shape. We train for 200 epochs with vanilla gradient descent, with different diagonal initialization of (ag, bo) by (0.1, 0.1). (0.5, 1.5). (2, 2). The detailed configurations (c.g., step size) and results of different experiments can be found in Appendix B
Initial token sampled from (O,, La). We conduct simulations with o = 0.5, 1, 2 respectively. With any initialization of (ag, bu), simulations show that ab converges to K1 / 1 1/5a², and ÿr. converges to xT../5 /5 in expectation, which verifies Theorem 4.1 and Proposition 4.1 respectively. In the main paper, present the convergence results with our theory in the malll-context scenarios (T₁, = 5), which is placed in Fig. in Appendix BB.
0.5 in Fig. Ia and 1b We also verify
Initial token sampled from Example 4.1. We conduct simulations with scale c= 0.5, 1, 2 respec- tively. With any initialization of (ao: bo), simulations show that ab converges to K1/K2 /K2 = 1 1/e2 (see
details in Appendix A.5), and yr. converges to the truth TT.,, which verifies Theorem 4. and Theorem respectively. In the main paper, we present the results with c = 0.5 in Fig. Ic and ld
The results Fig. 7 in Appendix B.5 show that KQ and WPV converge to dense matrices with strong diagonals, and other matrices converge to 0axds which means that the trained transformer performs somewhat preconditioned gradient descent. The detailed derivation is placed in Appendix B.5)
Initial token fixed as 1., We conduct experiment with
Go beyond the diagonal initialization. Finally, in order to extend our theory, we repeat experiments under Gaussian initialization with different variance (a = 0.001, 0.01, 0.1). The results of Gaussian start points and sparse start points (Example 4.1 can be found in Fig. of Appendix Bn3 Fig. of Appendix B.4 respectively. As result, though the convergence results of parameters are not the same as those under diagonal initialization, they keep the same diagonal structure, which can be understood as GD with adaptive learning rate in different dimensions. In addition, the test results (ratio or MSE loss) under the standard Gaussian initialization are the same as those under diagonal initialization, which further verifies the capability limitation of the trained transformers. To sum up. these experimental results demonstrate that our theoretical results have a certain representativeness. which further supports the rationality of the diagonal initialization.

---

4.0
(a), (0.1,))) 0.11
&
(ac.bg) (..1,0.)
3.5
(ag (0.5, 1.5)
(ag.by) (0.5, 1.5)
Dynamics of
3.0
1/5a²
(ac, by) (2.0,200)
Ratio of pred/true
0.6
(Bc. by) (2.0.
Rario 0,2
2.0
0.4
1.0
0.5
0.0
0.0
0
25
50
75
100
125
150
175
200
0
25
50
75
100
125
150
175
200
Epoch
Epoch
Gaussian with
0.5, dynamics of ab.
(b) Gaussian with
0.5, ratio of yr 1 //T..
12
(dp. (0.1. 0.1)
Dynamics of ab
3.0
3.5
MSE loss at tte
10
(do. (0.5, 1.5)
(ap. by) (2.0, 2.0)
2.5
8
Theoretical MSB
2.0
1.5
(ag. (0.1,
1.0
(g) (0.5, 1.5)
0.5
(ac, by) (..0.
2
1/c² 4.0
0.0
50
125 Epoch
100
175 200
100 Epoch
150 175 200
(c) Example 4.1 with 0.5, dynamics of ab.
(d) Example 4.1 with 0.5, 1 TT.
Figure 1: Simulations results on Gaussian and Example show that the convergence of ab satisfies Theorem 4.1 In addition, the trained transformer can recover the sequence with the initial token from Example|4.] but fails to recover the Gaussian initial token, which verifies Theorem 4.2 and Proposition 4.1 respectively.
7 Conclusion and Discussion
In this paper, we towards understanding the the mechanisms underlying the ICL by analyzing the mesa-optimization hypothesis. To achieve this goal, investigate the non-convex dynamies of one-layer linear transformer autoregressively trained by gradient flow on controllable AR process. First, we find a sufficient condition (Assumption 4.1 for the emergence of mesa-optimizer. Second, we explore the capability of the mesa-optimizer, where we find a sufficient and necessary condition (Assumption 4.2) that the trained transformer recovers the true distribution. Third, we analyze the case where Assumption 4.1 does not hold, and find that the trained transformer will not perform vanilla gradient descent in general. Finally, our simulation results verify the theoretical results.
Limitations and social impact. First, our theory only focuses on the one-layer linear transformer, thus whether the results hold when more complex models are adopted is still unclear. We believe that our analysis can give insight to those cases. Second, the general case where Assumption 3.1l and does not hold is not fully addressed in this paper due technical difficulties. Future work can consider that setting based on our theoretical and empirical findings. Finally, this is mainly theoretical work and we do not direct social impact of our theory.
8 Acknowledgement
This work was supported by Beijing Natural Science Foundation (L247030); NSF of China (Nos. 62076145, 62206159); Beijing Nova Program (No. 20230484416); Major Innovation & Planning Interdisciplinary Platform for the "Double- First Class" Initiative, Renmin University of China; the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China (22XNKJ13); 7R2022QF117), the Fundamental Research Funds of Shandong University. The work was partially done the Engineering Research Center of Text-Generation Intelligent Search and Recommendation, Ministry of Education. G. Wu was also sponsored by the TaiShan Scholars Program.
Natural Science Foundation Shandong Province (Nos.
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pages 5998 -0888 2017.
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurlPS, 2020.
[3] OpenAL. GPT-4 technical report. CoRR, abs/2303.08774 2023.
[4]
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiao- qing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CORR, abs/2307.00 092888 2023.
[5] 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.
[6] Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- Bean-BBptistte Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 11005, 2023.
[7] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICMI., pages 8821-8831, 2021.
[8] Doyup Lee, Chiheon Kim, Saehoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In CVPR, pages 11523-11532, 2022.
[9] Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive mod- eling: Scalable image generation via next-scale prediction. arXiv preprint arXiv:2404.02905. 2024.
[10] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In ICML, volume 119, pages 1691-1703, 2020.
[11] Zhengyuan Yang, Linjie Li, Kevin Lin, Jianfeng Wang, Chung-Ching Lin, Zicheng Liu, and Lijuan Wang. The dawn of Imms: Preliminary explorations with gpt-4v(ision). CoRR, abs/2309. 17421, 2023.
[12] Wenhai Wang, Zhe Chen. Xiaokang Chen, Jiannan Wu, Xizhou Zhu, Gang Zeng, Ping Luo, Tong Lu, Jie Zhou, Yu Qiao, and Jifeng Dai. Visionllm: Large language model is also an open-ended decoder for vision-centric tasks. In NeurIPS, 2023.
[13] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. In NewrIPS, 2023.
11

---

[14] Jun Zhan, Junqi Dai, Jiasheng Ye, Yunhua Zhou, Dong Zhang, Zhigeng Liu, Xin Zhang, Ruibin Yuan, Ge Zhang, Linyang Li, Hang Yan, Jie Fu, Tao Gui, Tianxiang Sun, Yugang Jiang, and Xipeng Qiu. Anygpt: Unified multimodal LLM with discrete sequence modeling. CoRR, abs/2402. 12226, 2024.
[15] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy P. Lillicrap, Jean- Baptiste Alayrac, Radu Soricut, Angeliki Lazaridou. Orhan Firat. Julian Schrittwieser, loannis Antonoglou, Rohan Anil, Sebastian Borgeaud, Andrew M. Dai, Katie Millican, Ethan Dyer, Mia Glaese, Thibault Sottiaux, Benjamin Lee, Fabio Viola, Malcolm Reynolds, Yuanzhong Xu, James Molloy, Jilin Chen, Michael Isard, Paul Barham, Tom Hennigan, Ross Mcllroy, Melvin Johnson, Johan Schalkwyk, Eli Collins, Eliza Rutherford, Erica Moreira, Kareem Ayoub, Megha Goel, Clemens Meyer, Gregory Thornton, Zhen Yang, Henryk Michalewski, Zaheer Abbas, Nathan Schucher, Ankesh Anand, Richard Ives, James Keeling, Karel Lenc, Salem Haykal, Siamak Shakeri, Pranav Shyam, Aakanksha Chowdhery, Roman Ring, Stephen Spencer, Eren Sezener, and al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. CoRR, abs/2403.05530, 2024.
[16] Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymy- rov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimization algorithms in transformers. CoRR, abs/2309.000 005558, 2023.
[17] Michael Eli Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do transformers perform in-context autoregressive learning ? In ICML, 2024.
[18] Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant. Risks from learned optimization in advanced machine learning systems. CoRR, abs/1906.01820, 2019.
[19) Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya K. Singh, Piere H. Richemond, James L. McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In NeurlPS, 2022.
[20] Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and lims by learning learn discrete functions. In ICLR, 2024.
[21] Kartik Ahuja and David Lopez-Paz, A closer look at in-context learning under distribution shifts. CoRR, 1675/ 2023.
[22] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In EMNLP, pages 11048-11064, 2022.
[23] Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Revisiting the equivalence of in-context learning and gradient descent: The impact of data distribution. In ICASSP, pages 7410-7414. IEEE, 2024.
[24] Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1 2024.
[25] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICML, volume 202, pages 35151-351744 2023.
[26] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In NewrIPS, 2023.
[27] Kwangjun Ahn, Xiang Cheng, Minhak Song. Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). In ICLR, 2024.
(28) Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent provably the optimal in in-context learner with one layer of linear self-attention. In ICLR, 2024.
12

---

[29] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In ICLR, 2024.
[30] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In ICML, 2024.
[31] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. In NeurIPS, 2022.
[32] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of ACL, pages 4005-4019. Association for Computational Linguistics, 2023.
[33] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In ICLR, 2023.
[34] Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge. Linear transformers are versatile in-context learners. CoRR, abs/2402.14110, 2024.
[35] Samuel Müller, Noah Hollmann, Sebastian Pineda-Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In ICLR, 2022.
[36] Deqing Fu, Tian-Qu Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order opti- mization methods for in-context learning: A study with linear models. CoRR, 17.06 2023.
[37] Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How well can transformers emulate in-context newton's 's method? CoRR, abs/2403.03183, 2024.
[38] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In ICLR, 2022.
[39] Xinyi Wang, Wanrong Zhu, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. CoRR, abs/2301. 11966, 2023.
[40] Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. In NeurlPS, 2023.
[41] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In ICML, volume 202, pages 19689-19729 19729, 2023.
[42] Yingcong Li. Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in-context learning. In ICML, volume 202, pages 19565-19594, 2023.
[43] Christos Thrampoulidis. Implicit bias of next-token prediction. CORR, abs/2402.18551, 2024.
[44] Yingeong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pages 685-693, 2024.
[45] Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. CoRR, abs/2305. 12766, 2023.
[46] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In NeurlPS, pages 7411-7422, 2019.
[47] Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the ddpm objective. NeurIPS, 36: 1 19649999649,, 2023.
[48] Haiyun He, Hanshu Yan, and Vincent YF Tan. Information-theorrtic characterization of the generalization error for iterative semi-supervised learning. Journal of Machine Learning Research, 23(287):1-52, 2022.
13

---

[49] Ludwig Schmidt, Shibani Santurkar, Dimitris Tsipras. Kunal Talwar, and Aleksander Madry. Adversarially robust generalization requires more data. NeurIPS, 31, 2018.
[50] Chenyu Zheng, Guoqiang Wu, and Chongxuan Li. Toward understanding generative data augmentation. NeurIPS, 36, 2024.
[51] Chenyu Zheng, Guoqiang Wu, Fan Bao, Yue Cao, Chongxuan Li, and Jun Zhu. Revisiting discriminative vs. generative classifiers: Theory and implications. In JCML, volume 202, pages PMLR, 2023.
[52] Yingqian Cui, Jie Ren. Pengfei He, Jiliang Tang, and Yue Xing. Superiority of multi-head attention in in-context linear regression. CoRR, abs/2401.17426 2024.
[53] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality (extended abstract). In Shipra Agrawal and Aaron Roth, editors, COLT, volume 247, page 4573, 2024.
[54] Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. CoRR, abs/2310. 1266800, 2023.
[55] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. CoRR, abs/2402. 14735, 2024.
[56] Ruiqi Zhang, Jingfeng Wu, and Peter Bartlett. In-context learning of linear transformer block: Benefits of the MLP component and one-step GD initialization. CoRR, 1bs/24 2024.
[57] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. CoRR, abs/2402.115077 2024.
[58] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics the attention landscape. In ICML, 2024.
[59] Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions the nonlinear dynamics of learning in deep linear neural networks. In ICLR, 2014.
[60] Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal Vychisli- tel noi Matematiki Matematicheskoi Fiziki, 3(4):643-66 653, 1963.
[61] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewicz condition. In ECML PKDD, volume 9851, pages 795-811, 2016.
[62] Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
14

---

Contents of Appendix
Proofs
16
A.1
Proof of eq. (1)
16
A.2 Proof of Theorem
16
A.2.1
Proof of Lemma 5,
16
A.2.2
Proof of Lemma
17
A.2.3
Proof of Lemma
27
A.2.4
Proof of Lemma
28
A.3
Proof of Corollary|4.1
28
A.4
Proof of Proposition 4.1
29
A.5
Derivation of Example
30
A.6
Proof of Theorem 4.2
30
Proof of Theorem
A.8
Proof of Proposition 4.2
36
Experimental details and additional results
36
GPU and random seed
36
B.2
Step size in simulations
36
Additional results for Gaussian initial token
36
B.4
Additional results for Sparse initial token (Example 4.1)
36
B B.5
full-one Additional results for full one initial token
37
15

---

Appendix A Proofs
A.1 Proof of eq.
Proof. We first calculate the output by causal linear attention layer as the following: O)
=et + P:
WV E,E,WROe
el + WV
1
Σese w
Pr
1
W es (wrkete)) ei) ei)"
Et
Pt
WPY
WPV
WPy PV
WKQ
WKO
W W.O
KQ
Ct
Pt
W
W
W W
12
W PV
W
e;
W KQ
W
W
W
RQ
Re
W kQ
W KQ
Ci
et
31
32
33
WVV H; W x;-1
X
0s
Tt
+
Pt
X
W QT
KT
22
C;
+
W
W 32
ROT
KQT
2;-1
Ft-1
23
+
33
at-1
Fz-1
where xs are the elements that will not contribute to the final ÿr. A further simple computation shows that
Ū. = 02 +
W PV 12
+
WPP
1 KQT + KQT x;-1)" Et
Pt
1
W PV
12
x₁ + WPP ( W.KTT KQT
2,
W KQT
x1-1) x4-1)"x111 II1-
Pt
1
(WPY x; + W x111)
W KTT RoT x; +
x;
W KQT
W 32
0;-1
P:
W
RQT
0;-1
Øt-1
1
(WRY WPY)
E
KQ 20 W
W
W KQ
T
D;
Dt.
P1
(1-11)
W KQ
81-1
0t-1
KQ
W
KQ 23
T
(WPP
WPY) c.
W.Q
QQ
P1
33
1
W.QQ W RQ 32
W-QQ 23
(WiP) PV
2 e:e
W
WKQ
E
Pr.
33
ET E
W
KQ 22
W.k
(W
PPV
e E Cd
13
W KQ
W.QQ KQ
Pt.
32
33
which completes the proof.
A.2
Proof of Theorem
A.2.1
Proof of Lemma 5.1
For the reader's convenience, restate the lemma the following.
16

---

Lemma A.1. Each element of the network's prediction De. (j E [d] can be expressed as the following.
E. E.
E E
Vi.1 = B
T
Vec( () Voe (A) Et 8
B,y
Pr.
Pt
where the A and Bj are defined as
W
bbi
WPPT
A (a1 a2d)
W
W
23
B₁ =
b,2
WPPT
12.44
32
33
13,5
with a E 2d and b 2,2 Rd.
Proof. Based on the result Eq. I we can write
E Er
w.kQ
KQ
w.k
Dt.
(WiP
WP )
W2 k Q
et
Pt
33
E E
B
(a)
a2d) et
Pt
E E
B
Pt
i=1
2d
E E
2 tr
B
a1
Pt
(=1
2d
E E,a
Σ
tr
B;
1=1
E E
tr B
Ct Pt
t=1
a BJ
E E
E E Pr.
tr
G7, a
112d
a2d B
E. E
tr
Vec A)B
Pt
E E
E. Ex*
tr
BJ T
Vec(A)
B. E X
Vec(A)
Pt
Pt
E. E
E
=tr Vec
(A)
E
X
B,
= Vec T
(A)
Bj.
Pt.
Pt
which finishes the proof.
A.2.2 Proof of Lemma 5.2
For the reader's 's convenience, restate the lemma the following. Lemma A.2. Suppose that Assumption 4.7 holds, then the dynamical process of the parameters in the diagonal of W.Q and PP satisfies
2b
(T - 2) +
+ b(T - 2)
dr
1=2
T-1
= Pb
(T 22
K3
+ a(T 2)K1,
dr
2
17

---

while the gradients for all other parameters were kept at zero during the training process.
Proof. The population loss I(0)
in Eq.
2
can be rewritten as
T-1
T-1
d
T
1(0)- =
Σ
(0) = Σ
E
Mh.
set1.l
2
2
T-1
Σ
E
Di, 2Re (sitI-0te))
2
2 j=1
T-1
= Σ
Σ
R
De.
Dt. Re (fit1.0th)) + - F
2
Then, we can calculate the derivatives (0) with respect to By and Vec(A) as
VB
Σ
1 VB
Re (zitl.0tns)
j=1
B, Fi.j - Re Re (zit1,00.s)
= E
V
2
E
V B, Ne. Dt.
E
BB, Re ( (2
and
d Lc()) =
> E
V Voo(A))i - Vvoc(a) (
j=1
d
d
Vieo(
Vioc(a) (2it1)
=
5=1
BB Re (B
Step one: calculate E
Based on Lemma
we have
E. E
Wt BJ.
eạTT &
(A).
Pt
VB, Re (#it)
Then, the EE
can be derived the following.
]
Es E Vec(A)
E
V BB,
Re
(zit1gBes))
=
V B, Re
BJ eee
@
Pr
E E
Vec(A))
E
V B, BJ.
Re
TT
Pt
E.P F
Re
ritid
ex
Vec(A)
E. E
Vecl )
Re
E
Titij
Ex
Pt.
Ee Ee
=
Re
E
Aj
mi
Veol
Act)
(use generating process)
E.d E AcT))
= Re
E
X 02i
18

---

where the penultimate equality uses the property of Kronecker and Vec operator Vec(AXB) B Vx)) we refer Section 10.2 in [62 for details.
For Ex Et"
we can simplify it as
E
E
ej
ei
Pt
Pt
1
Pt.
2i [
01-10
W W - a x W
1
Σ
Pt
Based on the diagonal property of W, we can simplify the x; as the following.
T. ( (W)(W-1) =1)"
M, @ E
where we define Σ = 2f and M, X x Therefore, have
E E
Σ 2
M Σ
WE 1= M + Σ M O
(4)
Pt
Ei=1
M.
Ew
E. A.
Then, leveraging the sparse property of A, we can derive
AcT follows.
E E
E E.
0a
AeT
Pt
Pr.
W (
M, E)a,
1-1
M; 2 E)xx
Pt
Therefore, for any 1 E [d], we have
E E.
Ae
21
a
N
Σ (M. 3 E)
P
A Xf'aa11 x11 a1d
Xj x x
> 4-
I1E
Pt
1-1 d
a
x i
Pt
d
x
Pt
1=
Similarly, for any E [2d] - [d], we have
EP E
AcT
Pt
Pt
To sum up, for any 1 E [2d]. we have
E E Ae; =
dd
LE [d],
Pt
Pt
Pr
2r=1
21 Tm
1 [2d] [d]
19

---

Next, we calculate the E
1j
21
E,
AcE
For any E [d]. we have
Pi
A Ø15"
E.
E
Aet
=
E
} 1 Eij
a
t-1
d
a
B
==1 r=1
discuss them in the following categories,
1. l + j. In this case, E[ = O by Assumption 4.1 thus
E. Er Aci)
)
1-
E xj
015
E 1 0.
Pt
Pt.
2. = j.r = j. Because d,s are i..d and [EA,]*
8(k = 0). we have
E m
(E
E
Er Act
}
a d
Σ
E[1 MA
Pt
Pt
(t 1)
Pt.
Similarly, for any E [2d] - [d],
we have
E. E
Ae x ]
x x1j
= 0.
Pt.
Therefore, for any [2d], we have
E
17'
Es ET
Ex
Act
{
a (t (t-1))E[r]])
l=
1 1
Pt
and the 1-th element of
VB,Re
Re (3it2,50hs)
is
se?)) 1)
E v B,
(Iite)
Re
E
A Tij
E. ET
Pt
(t
1)E(r4gh,
1 (= j,
pr
0,
1 1 j.
Step two: calculate
V Vi,
Based on Lemma 5.1
we have
E E
Vi.5 Vec (A)
V(A)
c
Bj.
Pt
then simplify the E V B, Dr. follows.
E
V B, Ut.
E F
E. E
E VB, Vec T T(A) E 8
Vec ((A)
&
B;
20

---

E
ET!
E E.
E
B,
B
et
Vec(A). Vec (A)e,
Bj
Pr.
Pt.
E
D
D E
E
Vec (A) Aee (A)e,
B
Pt
Pt
E E
E E.
E
Vec(A) Vec (A)ez
Pt
Pt
ES D
E E
E
2Re
et
Vec(A) Vec eA
B)
Pt
Pt
E E
E E
2Re
E
Ct
Vec(A) Vec (A)e) &
Bj
Pt
Pt
We further derive that
ES E
E. E
E
C
Vec(A).
Vec (A)ef (AcT(A)ez %
B₁
Pt
Pt
Es Er
Es Eee
E
Vec
(A)
&
et Bj
T
Vec(A) (A)
Pt
Pt
E
Er
Es Eie Vec(A)
E
B
x*
8
Vec(A)
Pt
E
ET
E. E
E
B
Vecl
Act)
Veel
Acet)
Pt
Pt
EP
E
ET T
E
B
Ae
Aet
Pt.
Pt.
2d
E. T
E E
Bijk
Aet
AeT
Pr
k=1
Es E
E E
Ae
AeT xt
(sparsity of B)
For any and / [2d]. recall that
E E
le [d],
Aet
Pt
lE [2d] [d].
Pt
and for any E [d].
we have
E
E
E. E
1-1
AeT
Aee
11
Pt
Pr.
Pr.
For any / [d], with careful computing, we have
E
E
E E
E
Ae 2
Aet
Pt
Pt.
abb
ΣΣΣ
E
p
21

---

We discuss it in the following categories,
1. j. In this case, Bliiaaai ri,, by Assumption 4.1
thus it becomes 0.
2. | = = j. It becomes
12
a²b 1-1
ab
1-1
a²b (t
Σ Σ
ΣΣ Pi
p 1=
p
i==1
3. 1 J.n
=2 デ becomes
abb 1-1 1-1
a²b Σ
Σ
B
Pi Tサ}
a²b (1 1) Σ
B B
r
4. = j,r # 72. In this case, E[1 1a becomes 0.
B[x Aoم - [
0. thus it
Similarly, for any E [2d] - [d].
we have
E E
E E
E
b
A
Ae,
0.
Pt
Pt.
To sum up, we have
E.
ET
T
E.
E
Ae
AEE
Pt.
Pr.
a²b (t
1>
+ (t 1) E,,, B 1
1=j,
p?
0,
otherwise,
which implies that the 1-th element of E V V, Mi .r. Wr. is
E V. -
a't p7
(t 1)--1) || + (t- 1) Blaj
(=j. otherwise.
2
0,
Step three: calculate VB, (0)) and V,, L(O). Based on steps one and two, the I-th element of VB, (.) can be derived as follows.
V LO):: 2
[ Tr. Dr. -E
V Re (1i))))es)
abb (t 1) + (t - Erzz
Ela mt.1
(t
()=
I= j. otherwise,
pf
a.
0,
Furthermore, the (-th element of V VB,((0) is
T-1 V L(O)) = Σ
V
(0))
1=2
Et=2
aDb pf
(t-
+ (t -
- a (t--t115 --1))Eat)) otherwise,
[=
0,
22

---

Ei=2
<26 [|a,,]
-
Enzj 1
aB(aty)): 1=
0,
otherwise,
d2b (T 2)1[2f,]
Σイー 2-1
a(T
1= otherwise,
0,
@bb(T 2) E
(-
a(T
2),
1 = otherwise,
0,
where the last equality is from Assumption 4.1
Ve ( (wit),)
and Ej= E Vrecla) (sit),
Step four: calculate
E
Based on Lemma
5.1 have
E. E. Bj.
Ti. = Vec V(A)
(A).E,
Pr.
VeocAa) (rs)
Then, the E
can be derived as the following.
E
Veoca) (DttL
E
VeclA) siti Vec (A) V())
Vec
E E
B3
Pt
E ET
E
V Vec(A) Vec (() Re
Tit1s Et X
Bj
Pr.
E E
E
Re
c 8
Bj
Pt
E. E
= Re
E
Ii-1j -e 2
B,
Pt.
E. ET E
B,e
= Re
E
Veel
Pt
In terms of ET
B, BeeeT, based on the sparsity of and Eq. 4 we can derive that
E. ET
E E
E. E.
B,e TT
):sbee"
bee
Pt
=b (EL M
WE M,
Then, for any 8, E [2d], we have
8 [d], r E [d].
E E
s [2d] - [d], E [d].
B c
TT
Pt
s s [d], r E [2d] [d]
8 s [2d] [d],r E [2d] - [d].
23

---

Next, we calculate EE
1j l Ø1
E ET
B
For any E [2d]
[d].r [d]. we have
Es
E
b
A 1
= E
Pt
1=1
t-1
-
141r
We discuss in the following categories,
1. + In this case, El 0 by Assumption 4.1 thus it becomes O.
2. d = r It becomes
0
b
Σ Bpe
B[ 1 [E(1))
Pt
Pt
b
1-1
b (1 - 1)
Pt
3. 8 d = r + j. In this case, -| =
0, thus it becomes 0.
Similarly, for any other s, r, we can calculate that
E E
A
B
0.
Pt
To sum up, have
E.ST
ET
b (t 1)E[ot,),
s=d+j otherwise,
E
x
B
Pe
Pr
0,
(2it (2it1,)bes)
thus the (s, r)th element of E
VARe
ARRe (citi,)rs)
b ( - (1-1)Bba
8 s ==dr = j,
0,
otherwise.
b (1 1)m1,
== d j. otherwise.
0,
Σラ= E ARRe (iitio@ee) )
Finally, calculate
d Σ j=1
E Re (2itt)) = {
b (t - 1)K1,
s s-d=r, otherwise.
=1
pi
0,
Step five: calculate E Vooo( and Dj E
a) A) Dno]
Based on Lemma 5.1
we have
E E
Dt.5 B c c
Vec(A)
Pt
24

---

then we can simplity the E
V A) He.s follows
E
Veal Veo(A)
Wee.
E Ex"
E E
E
V A))(A)
B eT
VVo((A)
Bj
Veo(A)
Pt
Pt
E E.
Ey Er
E
A) Vec( (A)
B, BJ
Et
Voo(A))
Pt
Pt
E E
E E Vco(A)
E
Bj
B
C
Pt
Pt
E
E
E E
E
Bj
BJ
C₁
Veo(AA
Pt
Pt
E
E
E. E
E
2Re
C
Bj-
BJ
et
Voo(A)
Pt
Pt
E
E. E
voo(Aa)
2Re
E
B
B
Gt
P1
Pt
We further derive that
D
E
ET E
E
Bj
BJ_
Et @
Veo(A)
Pt
EC
ES
ET
E ET
B
Vec(
B3
Pt
Pt
E.E.
ET
E ET
E
Vec(A)
By
B3
Pt
Pt.
E. E
E E
E
Vec(A)
Veel
B,ett) Veel
Pt B
Pt.
2d
E.Sh E
E E
E
Akl
Bie
Veel
Bye,T
Pt.
Pt
k,1=1
2d
E E
E E Pt
E
B
Vecl
BjeeT)
(sparsity of A)
k=d+1
Pt
k.k-d
Recall that for any 8, [2d], we have
x - Mant-111jjas Fisti
8 [d], r E [d].
E ETT
21,11
s [2d] [d], r E [d].
BBee
Pt
8 s [d], r E [2d] - [d]
x - 1
s [2d] [d] r E [2d] - [d].
Furthermore, for any ls E 22d - [d]. we can calculate that
E.N E B,e,"
E. E
TT B
Pt.
k.k-d
Pt.
k.k-d
25

---

b t-1
1-1
k- -
Σ
Mj
P
Pt
With careful computing, for any 8 E [2d] - |d], r E [d]. we have
2d
E
3
E E
E
B e
B
TT
k=d+1
k
Pt
k.
Pt
ab²
2d
t-1
Σ
B[A
x R-d
X
k=d+1
We discuss it in the following categories.
8 + r. In this case, E|a Ef 2 01
by Assumption 4.1 thus it becomes
2. 8 d r d == j. It becomes
ab²
ab? P
B
(t - 1)2
pf
3.
k d + j. It becomes
ab²
1-1 1-1
ab² t-1
B
B
Pt 11=
11=1 12=1
ab²
- (1 - 1) 1 1
Pi
4. 8 - = x In these case, E[A thus it becomes 0.
4 A 2 a x =
ス =
Similarly, for any other 8, r, we can calculate that
2d
E Ex
E Ex
E
B₁ee²
BeeeT
-0.
b=d+1
Pt
k.k-d
Pt
To sum up, have
E. E
E E
10
B,ee
B
k=d+1
Pt.
k.k-d
Pr.
ab?
(t 1) 1
== d t = j, = d # j. otherwise,
(t 1)
1
0,
which implies that the (s, -hh element 1EZ
A Dr. is
(t 1
8 s s=d+r,==) =
E
Dt.s
(t 1)E 1
s= - j.
2
0,
otherwise,
at? (t 1)²
s T, j,
db
of (t 1) 1)
8 s= - j. otherwise,
0.
26

---

Based on these results, we can derive
Ut.. sr -{
ab
(t 1) + (t 1)s3], d r,
E
A Vr.
0,
otherwise.
Step six: calculate V Vee A) (0) and V) L(O). Based on steps four and five, the (s,r) r)-th element of V (0) can be derived as follows.
d (LL((0) br Σ
E
Σ
V VaRe (sēti.0re))
1=1
j=1
ab² 1)2k2
[t
(t - 1)so}
_t 1)m1,
s s-d=r otherwise.
pc
0,
Furthermore, the r-t element of V VAL(0) is
T-1
Val(O) or = (O)
2m
T-1
ab²
[t 1)2ki2 + (t 1)ma]
b (t 1)81)
s d dr otherwise,
21=2
P9
Pr
0,
21=2
ab?
1i2 +
1-1 kis
1
bwg
- d dr
0,
otherwise,
ab² (T 2)k² + ΣTz
TT3
b(T - -)) 8 -
0,
otherwise.
Step seven: summarize the result by induction. From the gradient of V (() and VB, L(O),
observe that non-zero gradients only emerge in the diagonal of W. and W.P. and they same. Therefore, the parameter matrices keep the same structure as the initial time. Thus we can summarize the dynamic system as the following.
d -ab²
(T K)
+
+ b(T 2)N1,
dr
d
T-1
b = - bb
(T 2) +
+ +a(T 2)
dr
1
2
which completes the proof.
A.2.3 Proof of Lemma 5.3
For the reader's convenience, we restate the lemma as the following.
Lemma A.3. Suppose that Assumption 4.1 holds and denote (T - -2)22 Ei= I-I Kg and (T-2)k1 by ej and C2 respectively. Then, the dynamics in Lemma 5.2 are the same as those of gradient on the following objective function:
1
ē(a,b)
(c2 c1ab)3
201
whose global minimums satisfy ab 02/e1.
27

---

Proof. Basic calculus shows that
0
d
f(a, b)
2( (jab( -b(c2 - cjab)
dr a
da
2c1
1
[a, b)
-2 c-ab)(-eaa)) = -a(e2 clab) = dr
b.
0b
2c1
Therefore, the dynamics in Lemma 5.2 are the same those of gradient flow on E(a, b), whose global minimums satisfy ab C2/c1.
A.2.4 Proof of Lemma 5.4
For the reader's convenience, restate the lemma the following.
Lemma A.4. Suppose that Assumption 4.1 holds, then Ela, b) is non-convex function and satisfies the PL inequality as follows.
fa, b) da
+ f(a,b) 2c1 2a b2) E(a,b)
+
-lii a.b
min E(a, E(a, 1))
Therefore, the gradient flow in Lemma 5.2 converges to the global minimum of E(a, b).
Proof. First, we prove that E(a, b) is non-convex. The Hessian matrix of E(a, can derived as follows.
V2((a,b) ( c1b²
2cjab - c2)
[2caa - C2
qaa
Its determinant 2t (2crab 02)² (e2 c1ab) 02) when ab
ab a c2 Se,
Thus, E(a, b) is non-convex. Besides, the PL inequality holds because
2 E(a, b) aa
2 ((a,b)) b)
+ 0b
= =c²(cz cjab) + c c2ab)²
(a² (b) - -ciab)²
-2ci(a² + +b²). 2c1 (02 - cjb)2
-2c1la² + +b²) (((ab)) a.
min min ((a
2c +b²) ?(a,b) min {(a, b)
a.b
A.3 Proof of Corollary 4.1
For the reader's convenience, we restate the corollary as the following.
Corollary A.1. We suppose that the same precondition of Theorem 4.7 holds. When predicting the t-th token, the trained transformer implements one step of gradient descent for the minimization of the OLS problem LOLs(W) = Σ!=l ||1+1 Wx,l starting from the initialization W Oaxd with step size
Proof. The proof is stem from the theoretical construction in 16 D. as follows.
First, we simplify the prediction
E
E
WQ 22
KQ W 23
ÿ =
(W.P PV
12
WPP PV 13
kKQ 32
WRQ QQ 33
e
28

---

-bbla
(Ia
E E
Oaxd
Oaxd Oaxd
Odxd
aLa
ei
Pt
t
bId
Oaxd Σ
ex
exe
Oaxd ald
Oaxd Odxd/
et
Pt
bx:
(ám)
Odxd) et
ab
bar ;
2;0
0t
t-1
ab
2
Pt-
Then, connect it to the one step of gradient descent for the OLS problem LoLs(W) Wa a;; 1|2
- x++1
W
V
Øi+1
1
t 1
ab
W
[1+11 (xx)(-x;))
t
1=1
ab
0
[1:41 0z,)(-x())
Øb
Thus, the proof is completed.
A.4 Proof of Proposition
For the reader's 's convenience, we restate the proposition as the following.
Proposition A.1. Let Da,; be the multivariate normal distribution (Os,, a La) with any a2 > then the "simple" AR process can not be recovered by the trained transformer even in the ideal case with long training context. Formally, when the training sequence length Tür is large enough, for any test context length Tie and dimension j E [d], the prediction from the trained transformer satisfies
(TTT..)
Ex w
(WoT.))
5
Therefore, the prediction yr.. will not converges to the true next token W TT..
Proof. First, built upon the results in Theorem when Te, is large enough, have
K1
K1 E E[r
30 1
R2
15g6
5
Second, by directly calculating, we have
(Wan.)) = (W = X
and
āb To
(Dr..); The -
X f
k=1
29

---

0. Therefore, we have
(Frr.) 1i Wx (W (WIT.e)
ER,
ab
Ti--1 d
i Tea x
-
21k
Tte
==1
k=1
ab
Em
Fisl
abo
Tte
Since ab
and converges to 562 when Tr, is large enough, the proof is completed.
502
A.5
Derivation of Example 4.1
Proof. We first prove that the example satisfies Assumption 4.1 Because only one element of æ sampled from Example 4.1 will be non-zero, we have Ex, ~D, Tia "M ...2 = Ew,,D,,, [0] = for any subset {i........ derive that
I 4 of [d]. and r2, E N. In addition, for any j E [d]. we can
d
d-1
61 E[
1j
d
c
+
d
0 =
d
1
d
1-1
6
R2
E
c
0
Σ
Second, prove that it satisfies Assumption 4.2 follows. Without loss of general, we assume that the first coordinate of æ1 is c.
K-
1=
xx0t
diag(c 0, 0)(, (1,A77--,,,0....... ,),
K2
T
WT.
12
(ffic-1 e,o....,0)) TTr..
The proof is finished.
A.6 Proof of Theorem
For the reader's 's convenience, restate the theorem as the following.
Theorem A.1. Suppose that Assumption 4.7 holds, then Assumption 4.2 is the sufficient and necessary condition for the trained transformer to learn the AR process. Formally, when the training sequence length Tur and test context length The are large enough, the prediction from the trained transformer satisfies
UT.
WaT
Ter The
Proof. First, built upon the results in Theorem
when Ter is large enough, we have
Ri
12
T
T 12
t-1
12
Second, when Tie is large enough, by Assumption 4.2
Tea
11 i=
TTo
WT.:
K2 The
Therefore, we have
Tie
> Tio
TT. W
2;00
TT.
W
-
TT. -T.
Wa WTT
T
162
The
which finishes the proof.
30

---

A.7 Proof of Theorem
For the reader's convenience, we restate the theorem as the following.
Theorem A.2. Suppose the initialization satisfies Assumption 3.1 the initial token fixed la, and we clip non-diagonal gradients of W and W.P during the training, then the gradient flow of the one-layer linear transformer over the population AR loss converges to the same structure the result in Theorem 4.1 with
ab
d-1 1 1+ T-2 21=2 t-1
Therefore, the obtained transformer performs one step of gradient descent in this case.
The proof is similar to the proof of Theorem 4.1 in Appendix A.2, But the calculating for the gradients more difficult than that Theorem 4.1 Similarly, we first present and prove the following lemmas.
Lemma A.5 (dynamical system of gradient flow). Under the same assumption as Theorem 4.3 the dynamical process of the parameters in the diagonal of W.k and W12 PV satisfies
KQ
32
T-1
dr - ab
(T 2) +
Σ
+ b(T - 2),
T-1
a²b
(T - 2) + Σ
+ +a(T - 2),
dr
(=2
while the gradients for all other parameters were kept at zero during the training process.
Proof. Recall that in Appendix A.2.2 can be rewritten as
have already known that the population loss L(O) in Eq.
T-1
T-1
L(O) Σ L(
Σ
De - Re (ait))00s) + ,
j=1
Besides, the derivatives of (.) with respect to Vec(A) and Bj are
(.)) = | v B,
Do.
E
V Re (cit))
and
d Σ
( (ait1150s))
V (eo((A)(x(0)
Dt.
Re
j=1
Step one: calculate E
V BB Re (2it1)
Similarly to the step one in Appendix A.2.2 the
V B, Re
(t.))
can be derived as the following.
V B, Re
(xēt1:0hs))
Re
E
Tit
Ct
Esd E
Voc(A)
Pt
E E
Re
E
Veel
AC
(use generating process)
Pt
ENd
E
Re
E
Ac
Pr.
31

---

We note that for any 1 E [2d], based on the sparsity of A, we have
E
E
1 E [d].
ACt)
Pt
LE [2d] |d].
which implies that
E E
a a (1--1) 1); 1= j, Pt 0,
I
17
Aee
=
1
Pt
(2it1. )
and the /-th element of
VB Re
is
(= ]] =
(1
- 1),
1= j, [
E V Re
Pr 0,
Step two: calculate
V B, M7 D0.)
Similarly to the step two in Appendix A.2.2. we can simplify
the E V B, Mi De.5
as follows.
ES
ET"
E ET"
V B, Me. Ut.
2Re
E.
et
Veo(A) Vec T (A)ez
Bi)
Pt
Pt
E
E
ET E
AeT )
=2Re
E
b
AeT
(sparsity of B)
Pt
Pt
For any j E [d] and I E [2d], we can calculate that
E EF
1-1
Aet
Pt
j
and recall that
E E
XPA- LE [d]
Aet
Pt
M LE 12d - -d].
Pt
With careful computing, we have
d'h
#F
[1-- 2 + (d - 1)((1))))
- 1)].
1= j.
E E
(
E E AcT
Act
2²b (t - 1), p
1 LE [d] - otherwise,
Pt
0,
which implies that the /-th element of 1E V WB, Mi. Dr.5 is
[(1 a²th of
1)2 + (d (d-1)(() -)(1-1)]] 1)],
1= LE [d] j j. otherwise.
E
VB
a't (t - 1),
of
0,
Step three: calculate V,, (.) and V L(O). Based steps one and two, the I-th element of V B, (.( can be derived as follows.
VB, [()) 1B
1/8 VB, Dh.
E.
BB, Re (sit)aDiss)
32

---

a (t 1)2 + (d - 1)( - 1)] - 2 (t-1) -1],
1= j
b²b (t - 1),
1 E [d] - j. otherwise.
pi
0,
Furthermore, the I-th element of V VB, is
T-1 VB, = Σ VB, (0):
2m
[[t 1)2 + (d-1)(t- 1)]
-g 1)) 1= j
E ab (t - 1),
I E [d] - j, otherwise.
(5)
0,
we clip the non-diagonal gradient of W PV
12 we have
(
Σ
dbh 1) +
(11
(d - 1)(1 - 1)] (1 (1)1
1=
VB, L(0),
0,
otherwise,
dbb (T - 2) + -
a(T - -2), 1=
otherwise.
Step four: calculate Vvect Re (th)),)0.))) and Σj E Vecc()) Re ( (Dit))Diss) Sim-
ilarly to the step four in Appendix A.2.2 the E. following.
V Re (witlsths) can be derived as the
E Veola) (kitin0hs)
E. E
= Re
E
Fitli
20 63
Bj
Pt
Re
E
Aj 1.
Veel
Er E
BBeg TT
Pt
For any 8, E [2d], we have
s [d], E [d],
E ET
b
S [2d] [d], r E [d].
B,e
TT
Pr.
s [d], r E [2d] - [d],
8 s [2d] [d] r E [2d] - [d],
which implies that
b 1),
8 s= j.
ER
E
E
A
B T
+ j.
Pt
PC¹ 0,
otherwise.
and the r)-th element of E VARe (Bit1)0ts)) is
b b - 1),
s=d+jr=
E
v
Aee (xitL)D(s))
6 Pi
s s j. otherwise.
33

---

Finally, we can calculate Σÿ= VAVR VARe (oitis
) as
(t
1),
d= r, Agr
W WKQ
E V ARe (s
Pr
8 d rr A.r E
WKQ KQ 32
0,
otherwise.
Step five: calculate EE V Ve, and Σラ
E
V
Vecl A)
Similarly to the step five
in Appendix A.2.2 the E
Veol A) Dr Wi. is simplified as follows.
E
V Vocl A) Wi,
2Re
E
et
E E
Bj
B
Et
E E.
Vec(a))
Pt
Pt
B !
E
E
E E
-2ee
E
a
Bye
Veo
k=d+1
Pt
k
Pt
For any ls E [2d] [d], we can calculate that
E E
1-1
Beee"
Pt D
Pr
k.k-d
Pr
and recall that for any 8, E [2d]
have
s s [d],
E E T
8 8 [2d] [d] r E [d],
e
B,
TT
[ E [2d] [d].
Pr
Pr
8 s [2d] [d] r E [2d - [d]
With careful computing, we have
2d
E. E
E. ET
E
B,e,¹ 20
Bye T
b=d+1
k.k-d
ab?
of (t 1)2,
= j,r
ab² (t 1), Pi
s s d+ j,
ab²
of (t 1),
d
ab
1),
d d
remains in WKQ otherwise,
0,
which implies that the I-th element of 1 V ATT. Dr. is
ab? (t 1)2,
8= d + j,- j,
ab"
pg (t 1),
8 d - + j,r=j,
(t 1),
s= s=d-r,r
E V De.
(t 1),
= j,r ドj remains in W KQ 32
at
0.
otherwise.
34

---

Based on these results, we can derive
[t - 1)2 + (d - 1)(t 1)]
8 d = r, Agr E kQ
32
ab - 1)t-1))dd2)
- d + r, Agr E W W.k
KQ
P
32
otherwise.
Step six: calculate V Vec(A) and V A) (0).
Based on steps four and five, the (s, r-th
element of V AA can be derived as follows.
V VAL, or 2 j=1
Σ
Σ
E
V
ARe
5=1
aba
(tt 1)2 + (d 1)(
1)]
(t 1),
s-d= r, Aer E WRR
[2 d 2]
d + r, Agr E W 32
pr:
0,
otherwise.
Furthermore, the r-t element of V ( is
VAL(0) (O)(
ΣT-1
ab? - 11- (d- -
L1 1))
d = r, Aor E W.ke
ΣT=
a6
1
d r, Asr WKQ
(6)
O,
otherwise.
If we clip the non-diagonal gradient of W W.K we have
T-1
VAT(0)srr sT
Σ V (.((())s
TT-1 ab² [tt 1)2 + (d 1)(1 - -1)] - (1- 8 s-d=r, W.k
2
0,
otherwise,
{ ab? 22) 2) E b(T - 2), d- Ast E W.K
0,
otherwise
Step seven: summarize the result by induction. From the gradient of V (() and Vn, (()). observe that non-zero gradients only emerge in the diagonal of w and W.P.. and they same. Therefore, the parameter matrices keep the same structure as the initial time. Thus we can summarize the dynamic system as the following.
T-1
- ab?
(T - 2) +
Σ
+ b(T - 2),,
dr
t-1
1=2
T-1 d
= - a²b
(T 2) + Σ
+ aa(T - 2),
dr
t t-
2=2
which completes the proof.
Lemma A.6. Suppose that the precondtions of Theorem 4.3 hold, and denote (T-
- 2) - 전
flow and T 2 by 0i and cz2 respectively. Then, the dynamics are the same as those of gradient flow on the following objective function:
E(a, b) = 201
(c2 -1ab)²
whose global minimums satisfy ab 02/e1.
35

---

Table 1: Step size in different simulations.
o/c
step size
0.5
0.001
Gaussian
1
0.0001
2
0.000002
0.5
0.03
Example 4.1
1
0.001
2
0.0001
1
0.0005
Proof.
The proof the same as that of Lemma 5.3 in Appendix A.2.3
Using the results from the above lemmas and Lemma 5.4 conclude Theorem 4.3
A.8 Proof of Proposition 4.2
For the reader's s convenience, we restate the proposition as the following.
Proposition A.2.
The limiting point found by the gradient does not share the same structure as that thus the trained transformer will not implement one step of gradient descent for | W
in Theorem 4.1 7
minimizing
Proof. From Eq. and Eq 5 we know that when the parameters matrices share the same structure as the result in Theorem 4.1 the non-zero gradients of the non-diagonal elements of W KQ and W will occur, which implies the result.
Appendix B
Experimental details and additional results
B.1
GPU and random seed
The random seed in the experiments is fixed as 1. All experiments done single GeForee RTX 3090 GPU in one hour.
B.2 Step size in simulations
The step size of the gradient descent in different simulations is summarized in Table I
B.3 Additional results for Gaussian initial token
In practice we estimate the ratio of two vectors by calculating the mean of the element divide between two vectors. The results for
1,2 and T 100 with diagonal initialization are presented in Figure 2 The results for a = 1 and = (small-context scenarios) with diagonal initialization are presented in Figure The results for σ 1 and T 100 with Gaussian initialization (a 0.001, 0.01, 0.1) are presented Figure
B.4
Additional results for Sparse initial token (Example 4.1
The results for 1,2 and T = 100 with diagonal initialization are presented in Figure 3 The results for c = 1, 2 and T 100 with Gaussian initialization (a., = 0.001, 0.01, 0.1) are presented
in Figure 6
36

---

4.0
(0p. by) (0.1. (, 0.1)
Ratio of pred/true te
0.7
(ac.bg = (..1,0.)
3.5
(eb. (0.5, 1.5)
0.6
(ag. (0.5, 1.5)
Dynamics of
3.0
1/5g²
Da) 2.00.
0.5
(Bc. by) (2.0,
Ratio 0.2
0.4
2.0
1.5
0.3
1.0
0.5-
0.0
20
40
60
80
100
120
140
0
20
40
60
80
100
120
140
Epoch
Epoch
(a) Gaussian with dynamics of ab.
(b) Gaussian with ratio of ÿr, 1/ /TT..
1.75
0.4
(e) 0.1)
Ratio of pred/truc te
(an. by) - (0.1, 0.1)
(eg. ba)
1.50
(da. by) - (0.5, 1.5)
Dynamics ab
0.3
(ag. by) (2.0, 2.0)
1/5a* - 0.05
1.25
(sp. (2.0,
Ratio
1.00
0.2
0.75
0.50
0.1
0.25
25
100 Epoch
150
175
200
25
50
75
100
125
150
175 200
Epoch
(c) Gaussian with
2. dynamics of ab.
(d) Gaussian with a 2, ratio of yr. //
Figure Simulations results in Gaussian initial token with ab satisfies Theorem I. addition, the trained transformer can not recover the Gaussian initial token, which verifies Proposition
2. results show that convergence
B.5 Additional results for full-one initial token
wKQ and PV with different diagonal initializations are presented in Figure 7 From the results, we know that the gradient descent converges to
WKQ W
KQ W
04xd W1
Oaxd Oaxd
22
PV W.PV
(W2 Oaxa)
W 32
KQ
W
KQ
where W1 and W2 some dense matrices. Similarly to the proof of Corollary 4.1 in Appendix 4.1 we have
E Exp
W.Q KQ whk
KQ W R
W WKQ
9:
(W
W
et
Pt.
33
E. E
(Oaxd W
O Oxx
(W₂
Oaxd
O Oaxd
Pt
Pt
(W2
Oaxd) Σ
(Oaxd
W1
Oaxd
Oaxe
ex
Wy (x W Oaxa) e
Pt
W2 Wi..
Pt
which can be seen as a somewhat preconditioned gradient descent on the OLS problem.
37

---

3,5
(#p. Dy) (0.1, 0.1)
(0.1, (do, bp) m0.1 0.1)
3.0
(ep. by) (0.5, 1.5)
(0s, bu) (0.5. 1.5)
Dynamies ab
(0g, by) (2.0, 2.0)
MSE loss at Tte
bp) 2.0)
2.5
1/e 1.0
80
Thcoretical MSE -
2.0
60
1.5
40
1.0
0.5
20
0.0
25
50
75
100 Epoch
125
150
175
200
75
175
Epoch
(a) Example 4.1 with
1, dynamics of ab.
(b) Example 4.1 with
1 Tr 11.
(#p, by) (0.1, 0.1)
350
(da, bu) (0.1, 0.1)
1.0
(ep, by) (0.5, 1.5)
(#t, bg) (0.5. 1.5)
Dynamics ab
(gg (2.0, 2.00
(ay, by) (2.0, 2.0)
0.8
1/22
Thcoretical MSE
0.6
MSE loss
200
150
0.4
20
40
60
80
100
120
20
40
60
80
100
120
Epoch
Epoch
(c) Example with dynamics ab.
(d) Example
4.1 with c 2, 117
11
Figure 3: Simulations results Example Results show that the convergence of ab satisfies Theorem 4.1 addition. trained transformer can recover the sequence with the initial token from Example 4.1 which verifies Theorem 4.2
38

---

0.35
0.30
0.30
0.25
0.25
0.20
0.20
0.15
0.15
0.10
0.05
0.00
(a) WKQ b 0.1.
(b) WPV a b= 0.1.
1.4
1.2
1.0
0.04
0.6
0.82
0.4
0
0.00
0.0
(c) WKQ 0.5, b 1.5.
(d) W 0.5, 1.5.
0.35
0.25
0.15
0.15
0.10
0.05
0.05
0.00
0.00
(e) kQ
(l) PV a b
(eg. = =0 (0.10....) 0.1)
(0g, by) = (0.5, 1.5)
Dynamics ab
(dp, by) = (2.0, 2.0)
Theory - 0.1343
0.5
0.0
0
50
100
150
200 Epoch
250
300
350
(g) dynamics of ab, o a=1
Figure 4: Results of small-context scenarios (d T 5) with Gaussian start point (a 1) and diagonal initialization, and the theoretical limit is 1/(5a² + -=1 Er= 2 by Theorem 4.1. The read blocks in Assumption 3.1 presented, which related the final prediction. The product ab does converge the results in Theorem
39

---

04
0.1
02
0.0
0.0
-B1
0.2
-03
-0.3
(a) W KK oe 0.001.
(b) WPP PV
o0 0.001.
0.4
0.3
02
0.1
0.1
-.1
-0.1
-0.2
-02
-0.4
0
-0.4
(c) WKQ Tw 0.01.
(d) W
0w 0.01.
0.4
0.3
0.1
0.0
0.0
-0.
-03
-03
-0.4
(e) WKO oi 0.1.
(I) PV ou = 0.1.
0,20
Ratio of pred/true at Tte
0.15
0.10
0.05
d, 0.1
0.00
0, 0.01
0, - 0.001
-0.05
0
50
100
150
200 Epoch
250
300
350
400
(g) Ratio fr..-1/X1.
Figure 5: Results of Gaussian start point (o 1) and standard Gaussian initialization with different variance Tw- The read blocks in Assumption 3.1 are presented, which related to the final prediction. The parameter matrices retain the same strong diagonal structure and test performance as those of the diagonal initialization.
40

---

06
0.6
-02
-0.2
-0.4
(a) KQ Twe 0.001.
(b) W WVV
00 0.001.
02
0.2
0.0
-0.2
-B.4
0.6
-0666
(c) W KQ ou. 0.01.
(d) W PV
0 0.01.
06
0.4
0.
02
0.0
002
-0.4
-0.6
-066
KQQ oe 0.1.
(f) WVV
o 0.1.
Oy - 0.]
Oy
MSE loss at Tte
0.001
Theorctical MSE 0
01
100
200
300
400
500
600
700
800
Epoch
(g) MSE || Drg lB
Figure 6: Results of Example 4.1 1) and standard Gaussian initialization with different variance aw. The read blocks Assumption 3.1 presented. which are related the final prediction. The parameter matrices retain the same strong diagonal structure and test performance as those of the diagonal initialization.
41

---

0.8
0.4
0.2
(a) WKQ 0.1, 0.1.
(b) PV a ..1.0.1b= = 0.1.
1.4
1.0
0.8
0.2
0.1
0.4
(c) KQ
0.5, b = 1.5.
(d) W PV
a 005,, = 1.5.
0.6
0.4
0.0
(e) W KQ a 2, 2.
( W PV = 2,6 2.
Figure 7: WRQ and PV full-one start points with different diagonal initialization. read blocks Assumption 3.1 are presented, which are related to the final prediction.
42

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims in the abstract and introduction accurately match theoretical results and reflect the paper's contribution and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Please see Section for details.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
43

---

Justification:
All assumptions are presented clearly in the main paper (Assumption 3.1 4.
and (.2). Complete proof can be found in Section [ and Appendix The correctness of our proofs can be verified by simulations in Section 6 and Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Ys)]
Justification: We provide complete configurations Section 6 and Appendix B We also upload the code at http:///.tbbb.ooooMLLcMcsAAA//// AR- Transformer for reproduction as well.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- spen-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
cases, in which
44

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer:
[Yes]
Justification: We upload the complete code and sufficient instructions in the supplemental material. It is available at httppp//gtttlaa.ooo////SSS////at Tranfformee
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (httpss:/nnis..cc)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/nips. wttpcc/pcccccpccccccceoooo httpa//gggudeee/deuuuuubbniiiiii sessSuummssiins. seriissssinpiiiiiee
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Please see details in Section 6 Appendix B and attached code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We do not present eтror bar due to insufficient computing resources. However, we have run all experiments with three different initializations to verify the correctness of our theory, see details in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
45

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yss]
Justification: Please see details in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: fetttp:://neurigge www.rrp...cpubuiiieee.eooc cccpuuliic/thiicccudde.........
Answer: [Yes]
Justification:
We conform research conducted in the paper satisfies the NeurlPS Code of
Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Please details in Section 7
Guidelines:
The answer NA means that there is no societal impact of the work performed.
46

---

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

---

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

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
49


Reviews:
[{"review_id": "bhQ5UZqIlY", "reviewer": "Reviewer_ykMD", "summary": "This paper studies the autoregressive training of a one-layer linear attention model for in-context learning of first-order AR processes. It is shown that under certain distribution of the initial point, the gradient flow on the population next-token prediction loss will converge to a model that makes the prediction based on the estimate given by one step of gradient descent over the in-context loss. Furthermore, a sufficient and necessary condition is provided to characterize learning capability of the trained model. Besides, it is shown that for more general data distribution, the trained model does not perform one step of gradient descent over the in-context OLS loss. Numerical simulations are conducted to support the theory.", "questions": "What's the value of the normalization factor $\\rho_t$? This seems to be a non-standard design as $\\rho_t$ depends on $t$. From my reading of line 634, it requires that $\\rho_t=t-1$. Please clarify this and explain why this is reasonable.", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and easy to follow. The explanation of the results is very clear, and logic flow of the organization of the paper is smooth. The current results present a solid contribution to the theoretical understanding of the in-context learning by transformers, especially for AR training.", "weaknesses": "It would be helpful to further highlight the technical contribution. In particular, how does the technique and analysis differ from those existing papers on training dynamics of transformers for in-context learning of regression problems?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer ykMD\nWe thank Reviewer ykMD for the positive score and valuable comments, which can definitely improve the quality of this paper.\n## Weakness 1: further highlight the technical contribution\nThanks for the nice suggestion. We consider the novel AR pertaining which makes our setting more complex. **First**, compared to ICL for regression, our data model breaks the independence between data at different times, which causes difficulty in decomposing and estimating the gradient terms. **Second**, we additionally modify the embedding strategy (more dimensions), scale the attention model (much more parameters), and change the loss function (more terms) to perform the full AR pertaining. All these parts are not well studied in the literature and make the gradients more complicated.\n\nWe will discuss this in detail in the final version.\n\n## Question 1: normalization factor $\\rho_t$\nYes, $\\rho_t = t-1$ as we defined in line 134. **This setting follows the convention in ICL theory for regression problems**. We need this length-aware normalization factor to balance the outputs along different times because each element in $E_tE_t^*$ is an inner product of two vectors of size $t$, which scales as the $t$ increases. In ICL theory for regression problems, we only need to predict the last token, where the width of the embedding matrix $E$ is $T+1$ and the $\\rho$ is defined as $T$. In this paper, the width of $E_t$ is $t$ thus we define $\\rho_t = t-1$.\n\nIn practice, this length-aware normalization factor is a non-standard design mainly due to the existence of the softmax operator, which naturally normalizes the outputs. The theory for softmax attention in the AR setting is a meaningful direction.\n\nWe will definitely discuss this in detail in the final version.", "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": "g9te21qKM5", "reviewer": "Reviewer_dsNs", "summary": "- The authors study the problem of in-context learning an autoregressive (AR) process (defined with a uniformly drawn diagonal unitary transformation) with a one-layer linear causal self-attention model, trained by gradient flow on square loss.\n- Under a specific parameter initialization scheme and a distributional assumption on the initial token, they show that a trained model can implement one step of gradient descent (=mesa-optimizer) applied to a least-square problem estimating the transformation matrix of the AR process.\n- Unfortunately, they show that the obtained mesa-optimizer can fail to learn the AR process with next-token prediction, even with infinitely many training samples and infinitely long test context length.\n- They further propose a stronger assumption being a necessary and sufficient condition for the success of mesa-optimizer. They also provide an example of the distribution of the initial token satisfying the necessary and sufficient condition.\n- They show that, without the distributional assumption, the trained model may not implement a gradient descent step.\n- Lastly, they verify their theory with simulations.", "questions": "- Q1. In lines 186-187, you mentioned that any random vectors with i.i.d. sampled coordinates satisfy Assumption 4.1. Is it true? I think they only satisfy the first part of the assumption (before “In addition”). As a counterexample of the second part, if the distribution of each coordinate is heavy-tailed (e.g., infinite second moment), $\\kappa_1$ and $\\kappa_2$ are infinite.\n- Q2. In Theorem 4.1, does different initialization $(a_0, b_0)$ “may” lead to different $(\\tilde{a}, \\tilde{b})$, or “always” lead to different point?\n- Q3. Questions on Proposition 4.1\n    - Isn’t $T_{te} \\rightarrow \\infty$ unnecessary? This is because we can show that for any $i\\in [T_{te}-1]$, $\\mathbb{E}\\_{x_1} [x_i x_i^\\ast] = W^{i-1} \\mathbb{E}_{x_1} [x_1 x_1^\\top] (W^\\ast)^{i-1} = \\sigma^2 W^{i-1}(W^\\ast)^{i-1} = \\sigma^2 I_d $ because $W$ is unitary and $x_1 \\in \\mathcal{N}(0_d, \\sigma^2 I_d)$.\n    - Isn’t $T_{tr} \\rightarrow \\infty$ also unnecessary? This is because it seems enough to show that $\\tilde{a} \\tilde{b} < \\tfrac{1}{5}$.\n    - Most importantly, is the last statement **correct**? Shouldn’t you prove that $\\mathbb{E}_{x_1} [\\hat{y}_T - Wx_T]$ converges to a nonzero vector? This is because $x_T = W^{T-1} x_1$ is also a random vector and depends on $T$.\n    - From what I read, it seems like the data distribution $\\mathcal{D}_{x_1}$ is the main cause of the failure of the mesa-optimizer. In my opinion, however, this is also because the “single-layer” linear self-attention model can only simulate a “single step” of GD for the OLS problem. Then isn’t the number of layers the problem? I think there are several reasons why the mesa-optimizer fails, and the authors should mention and analyze these factors as many as possible.\n- Q4. I guess Theorem 4.2 requires Assumption 3.1. Or does it?\n\nIf the authors address all my concerns, I will definitely raise my score.", "limitations": "The limitations are provided in Section 7.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written overall.\n- Most of the theoretical analyses seem correct and rigorous. The proof sketch in Section 5 provides a clear and plausible picture of the proof.\n- The analyses do not require/assume the diagonal structure of submatrices of the parameters. Instead, the authors found a sufficient condition (a distributional assumption on the initial token) so that they can prove the diagonal structure.", "weaknesses": "- W1. The analyses and experiments cannot provide insight into the general cases of autoregressive training of Transformers.\n    - This is mainly because the theoretical analyses greatly depend on the particular initialization scheme (Assumption 3.1). The initializations of weight matrices are too sparse: at initialization (and throughout training), they only focus on the relationship between very particular token embeddings. This seems very unnatural, thinking of the training transformers in practice. Quite ironically, the paper’s result corroborates that the initialization scheme is inappropriate for solving in-context learning with next-token prediction, under most of the plausible data distribution.\n    - This assumption might be inevitable for a tractable analysis. However, the experiments are also confined to the framework under Assumption 3.1. I think the authors should have demonstrated many more simulations with various initialization schemes and even with diverse architectures (linear transformers with non-combined weight matrices, non-linear transformers…) to support their arguments (”the data distribution matters for the emergence of mesa-optimization”, “there is a capability limitation of the mesa-optimizer”…).\n- W2. There are some weaknesses in theoretical analyses.\n    - Even though the authors found a condition for the emergence of mesa-optimization in a very particular problem setting, the analysis does not go beyond the “diagonal structure” of weight matrices. Thus, for me, although the contribution is novel, it does not look significant.\n    - For the comments and questions on Proposition 4.1, see Q3.\n- W3. Minor comments on typos & mistakes.\n    - Line 130: $W^{KQ} = {W^K}^\\ast W^Q$ (rather than transpose)\n    - Equation (2): I guess the sum should be applied over $t=1, …, T-1$.\n    - Line 182: “$x_1 \\in \\mathbb{R}^d$”\n    - Lines 207, 289: ‘$W$’ must be bold.\n    - Figures 1b, 1d, and Line 330: “$T_{te} -1$”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer dsNs\nWe thank Reviewer dsNs for the positive feedback and valuable comments, which can improve the quality of this paper.\n## Weakness1: initialization & negative results\nThanks for the nice suggestion. We discuss two concerns respectively.\n### Initialization\nWe argue the diagonal initialization is reasonable and further conduct simulations with standard initialization. Please see the details in common concern 1.\n### Negative results\nThe results show that one-step GD learned by the transformer can not recover the distribution, but this can be solved by more complex models. Even for more complex data ($W$ is not diagonal), [16] has empirically verified that multi-layer linear attention can perform multi-step gradient descent to learn the data distribution. Our work can be seen as a foundation for the theory of more practical transformers.\n## Weakness2.1: diagonal structure\nThanks for the suggestion. The findings of Assumption 4.1 and analysis are not trivial. Please see the details in the common concern 2.\n## Weakness2.2 & Q3: questions on Proposition 4.1\nThanks for the nice comment, which definitely helps us clarify Proposition 4.1!\n### Refined Proposition 4.1\nFirst, we clarify the original intention of Proposition 4.1. We want to prove that in expectation w.r.t. data ($x_1$ and $W$), when the training sequence length $T_{tr}$ is large, the prediction $\\hat{y}_ {T_ {te}}=W(\\tilde{a}\\tilde{b}\\frac{\\sum_ {i=1}^{T_ {te}-1}x_ ix_ i^*}{T_ {te}-1})x_ {T_ {te}}$ of the trained transformer converges to $1/5 * Wx_ {T_ {te}}$, rather than the true target $Wx_{T_{te}}$. In the submitted version, we prove this by claiming $E_{x_1}[\\tilde{a}\\tilde{b}\\frac{\\sum_{i=1}^{T_{te}-1}x_ix_i^*}{T_{te}-1}] \\to \\frac{1}{5}I_d$, which is not rigorous as the reviewer points. We modify the proposition rigorously as follows.\n\n**Refined Proposition 4.1**: Let $D_{x_1}$ be the normal distribution $N(0_d, \\sigma^2I_d)$ with $\\sigma^2 > 0$, then the AR process can not be recovered by the trained transformer even with long context. Formally, when the training sequence length $T_{tr}$ is large, for any fixed test context length $T_{te}$ and dimension $j\\in [d]$, the prediction satisfies $E_ {x_ 1,W}[\\frac{(\\widehat{y}_ {T_ {te}})_ j}{(Wx_ {T_ {te}})_ j}] \\to \\frac{1}{5}$. Therefore, the prediction will not converge to the true next token.\n\n**Proof skeleton:** By directly calculating, we have\n$(Wx_ {T_ {te}})_ j = (W^{T_ {te}}x_ 1)_  j = \\lambda_ j^{T_ {te}} x_ {1j}$, and $(\\hat{y}_ {T_ {te}})_  j = \\frac{\\tilde{a}\\tilde{b}}{T_ {te}-1} \\sum_ {i=1}^{T_ {te}-1} \\sum_ {k=1}^d \\lambda_ j^i \\lambda_ k^{T_ {te}-i} x_ {1j}x_ {1k}^2$.\nTherefore, we have\n$$\nE_ {x_ 1,W}[\\frac{(\\hat{y}_ {T_ {te}})_ j}{(Wx_ {T_ {te}})_ j}] = E_ {x_ 1,W}[\\frac{\\tilde{a}\\tilde{b}}{T_ {te}-1} \\sum_ {i=1}^{T_ {te}-1} \\sum_ {k=1}^d \\lambda_ j^{i-T_ {te}} \\lambda_ k^{T_ {te}-i} x_ {1k}^2] = E_ {x_ 1}[\\frac{\\widetilde{a}\\widetilde{b}}{T_ {te}-1} \\sum_ {i=1}^{T_ {te}-1}  x_ {1j}^2]= \\tilde{a}\\tilde{b} \\sigma^2.\n$$\nSince $\\tilde{a}\\tilde{b} < \\frac{1}{5\\sigma^2}$ and converges to $\\frac{1}{5\\sigma^2}$ when $T_{tr}$ is large, the proof is completed.\n\n### Answers to Q3\nQ3.1: Yes, though for the new proof, the $T_{te}$ does not need to be large.\n\nQ3.2: Yes, $\\widetilde{a}\\widetilde{b} < \\frac{1}{5\\sigma^2}$ and thus we can prove the prediction will not converge to the true next token. However, we consider the long-context scenarios ($T_{tr}$ is large) to improve the readability and make the experimental verification more convenient. Furthermore, we conduct additional experiments in small-context scenarios and the representative results can be found in the latest uploaded PDF.\n\nQ3.3: This is a very helpful comment that inspires us to refine the proposition. We note that $E_ {x_ 1,W}[\\widehat{y}_ {T_ {te}} - Wx_ {T_ {te}}] = \\frac{1}{5}*0 - 0 = 0$ in this case, which fails to interpret the phenomenon. Thus, we choose to calculate the expectation of the ratio. The experiments in the main paper are also conducted to estimate the ratio with 10$k$ test examples (randomly sampled $x_1, W$), which verify the refined theoretical results here.\n\nQ3.4: Yes, the number of layers influences the effect of the mesa-optimizer largely. However, the main message from our paper is that given a **fixed** number of layers, the obtained mesa-optimizer has limited capability, and we investigate this rigorously. For multi-layer attention, [16] has empirically found it performs multi-step (accelerated) GD to solve the problem well. Theory for this case is a meaningful future work.\n\n## Weakness3 & Q4: typos\nThank you for the careful reading! We will modify these typos in the final version. Here, we only clarify that the summation in Eq. (2) is applied over 2 to $T-1$, which is also adopted in recent AR theory [17]. The reasons are as follows.\n1. Given only $x_1$, we do not have any information to predict $x_2$.\n2. Since we set $\\rho_t = t-1$ according to the convention in ICL theory, $\\rho_1 = 0$ will lead to the exploration of the transformer's output.\n\nWe will discuss these in more detail in the final version.\n\n## Q1: distributions that satisfy the Assumption 3.1\nYes, thank you for the careful reading. We will modify the claim to \"any random vectors $x_1$ whose coordinates are i.i.d. random variables with zero mean and finite moments satisfy this assumption\". This class still includes many common distributions like the normal distribution, thus it is still meaningful.\n\n## Q2: convergence of (a,b)\nDifferent initialization $(a_0, b_0)$ always lead to different limiting points $(\\tilde{a},\\tilde{b})$. For example, in the experiments with $x_1$ from $N(0_d, I_d)$, we approximately have\n1. $(a_0,b_0) = (0.1,0.1) \\to (\\tilde{a},\\tilde{b}) = (0.44,0.44)$\n2. $(a_0,b_0) = (0.5,1.5) \\to (\\tilde{a},\\tilde{b}) = (1.41, 0.14)$\n3. $(a_0,b_0) = (1.5,0.5) \\to (\\tilde{a},\\tilde{b}) = (0.14, 1.41)$\n\nWe will definitely discuss these more clearly in the final version.", "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": "Xq4ZIcwaeI", "reviewer": "Reviewer_WaqB", "summary": "This paper studies the emergence of mesa-optimization/in-context learning capabilities of Transformers. In particular, it attempts to fill the gap of our understanding of training dynamics and specifically on non-convex dynamics where the sequences are autoregressively generated as $\\vec{x}\\_{t+1} = W \\vec{x}\\_{t}$. There are many contributions of this paper. First, it shows, under Assumption 4.1, transformer converges to the gradient-descent conjecture for ICL. Under the same assumption, the authors show moments of the data is necessary & sufficient for ICL. Finally, this papers study data distributions beyond Assumption 4.1, and shows in this case, transformers are *not* doing gradient-descent for ICL.", "questions": "Theories are built upon a one-layer linear causal self-attention model, what happens if we extend to (1) more layers or (2) non-linear activations?", "limitations": "Theories are only built upon a one-layer linear causal self-attention model, where there's is still a big gap compared to what's used in practice -- many-layer non-linear activation self-attention and feed-forward model with skip connections. But for theory, this is good enough.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "Presentation is very clear and easy to follow. \n\nContribution is significant. It bridges the gap between two main stream hypothesis of how Transformers perform mesa-optimization/ICL: (1) Transformer is doing gradient descent [1]; (2) Transformers is doing something beyond vanilla gradient descent, such as preconditioned GD  [2] or higher-order optimization methods [3,4]. The authors concluded the divergence between these two hypotheses are due to the data assumption (Assumption 4.1).\n\n[1] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICML, volume 202, pages 35151–35174, 2023.\n \n[2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In NeurIPS, 2023.\n\n[3] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models, 2023.\n\n[4] Giannou, Angeliki, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How Well Can Transformers Emulate In-context Newton's Method?. arXiv preprint arXiv:2403.03183 (2024).", "weaknesses": "When contrasting with vanilla gradient descent conjecture, the authors only mentions the alternative *preconditioned GD* hypothesis. However, there are many other propositions, such as higher-order optimization proposition, such as a *Newton's method* [3,4,5]. I would recommend the authors have a deeper discussion there. For example, in [5], they argued the GD++ method in [1] is actually a higher-order optimization method with superlinear convergence rate as discussed in [3,4]. \n\n[5] Vladymyrov, Max, Johannes von Oswald, Mark Sandler and Rong Ge. Linear Transformers are Versatile In-Context Learners. ArXiv abs/2402.14180 (2024)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer WaqB\n\nWe thank Reviewer WaqB for the positive support and valuable comments, which can definitely improve the quality of this paper.\n\n## Weakness 1: additional related work\n\nThanks for the helpful suggestion. We will cite and discuss the mentioned papers in the final version. We believe that the higher-order optimization hypothesis is very important for understanding practical ICL, especially in multi-layer cases.\n\n## Question 1: multi-layer and non-linear activations\n\nThanks for the nice suggestion. We discuss two parts respectively.\n\n### Multiple layers\nWe think we can explore this setting from expressive power and optimization/generalization perspectives, respectively.\n\n1. Expressive power. We can try to construct the cases where multi-layer transformers perform multi-layer higher-order optimization (e.g., GD++, Newton's method) in the AR pretraining setting. We think [16] can provide many practical insights, and [a,b] can provide many theoretical insights.\n2. Optimization/generalization. In this case, our work can be seen as a foundation. We can try to investigate the training dynamics and the generalization ability of the trained transformers. Here, we think [c] can also give some theoretical insights, which discuss the multi-layer attention in the regression setting.\n\n### Non-linear activations\nWe only discuss the softmax activation here because it is used most widely in practice.\n\n1. [16] shows that there exists a strong connection between softmax attention and GD in the AR pretraining setting, but the mechanism is still unclear.\n2. [16] also shows that in the AR pretraining, the first softmax layer can learn the copy behavior and obtain the embedding construction in this paper, but the mechanism is also still unclear. To this end, we think [51] can give some insights, which prove similar copy behavior of the softmax attention.\n\nWe will discuss these in detail in the final version.\n\n[a] Transformers learn higher-order optimization methods for in-context learning: A study with linear models, 2023.\n\n[b] How Well Can Transformers Emulate In-context Newton's Method?, 2024\n\n[c] Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?, ICML, 2024", "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": "DawLC1LAdz", "reviewer": "Reviewer_c1Lc", "summary": "This paper studies the behavior of a single unit of linear self-attention trained via gradient flow on an autoregressive (AR) loss on token sequences generated by a particular class of noiseless (deterministic) linear systems. The first theoretical results state that under some assumptions, the result of gradient flow from a particular class of initializations is a model whose in-context learning (ICL) predictions can be interpreted as those of another model that has been updated with one step of gradient descent (GD) using the context examples. The second part of the theoretical contribution is to say that a stronger assumption is necessary and sufficient for the result of this GD to give accurate predictions. These results are supported by numerical simulations of the studied setting.", "questions": "Why is the y-axis in Figure 1-b defined as the ratio of two vectors?", "limitations": "The authors discuss limitations.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- This paper aims to address an important issue in the ICL theory literature: most works study the emergence of ICL in transformers trained on an ICL objective (referred to as few-shot pertaining in the paper), whereas in practice transformers are trained on an autoregressive objective.\n- The theoretical results are correct as far as I can tell from skimming the proofs. The notations are clearly defined and consistent. \n- The proof sketch is well-written.\n- The experiments nicely verify the theoretical results. Sufficient details are provided.", "weaknesses": "1. The results do not suggest that “ICL by AR pretraining is more difficult than ICL by few-shot pretraining”, as the authors claim. Instead, they suggest that ICL to predict the $T+1$-th token given the sequence $\\mathbf{W}^0\\mathbf{x}\\_1, \\dots, \\mathbf{W}^{T-1}\\mathbf{x}\\_1$ is harder than ICL to predict $\\mathbf{w}^\\top \\mathbf{x}_{T+1}$ given $\\{(\\mathbf{x}\\_i, \\mathbf{w}^\\top \\mathbf{x}\\_i)\\}\\_{I=1}^T$ and $\\mathbf{x}\\_{T+1}$. These are two very different types of tasks and it is not fair to make conclusions about pretraining strategies based on their differing performance on these different tasks. Rather, one needs to compare the performance of a transformer trained in an autoregressive fashion and one trained on a few-shot learning objective under the same data/task model. In the data model in this, few-shot pertaining would use the same loss as the AR loss (equation 2) except that only the $t=T-1$ terms would remain in the loss. It is not clear whether this objective leads to better or worse solutions, in fact, this objective contains the same amount of information and therefore intuition suggests that it leads to solutions of the same quality. This is a key point as arguably the main message of the paper is that AR pretraining behaves differently than few-shot pertaining and thus the ICL community should focus on AR pretraining. I would suggest that the authors re-frame their discussion around the inadequacy of the regression data model considered by most of the existing ICL literature, as changing the data model is really the key contribution of this paper.\n\n    2. The setting is very artificial such that the significance of the results seems dubious. The model consists only of one, linear attention unit, and the data is assumed to follow a deterministic linear system, i.e. $\\mathbf{x}\\_{i+1}=\\mathbf{Wx}\\_i$ where $\\mathbf{W}$ is diagonal with diagonal entries having magnitude 1. The learning model and data model both being linear means that the ICL prediction of any learning model with parameters $\\mathbf{W}^{KQ}\\_{32}=a\\mathbf{I}_d$ and $\\mathbf{W}^V\\_{12} = b\\mathbf{I}\\_d$ for $ab>0$ and all other parameters being zero can be interpreted as the prediction of a linear model $\\hat{\\mathbf{W}}$ trained directly to approximate $\\mathbf{W}$ with one step of GD starting from $\\mathbf{0}\\_{d\\times d}$. It is not clear at all how such a conclusion would follow if the softmax activation, shown to be crucial in practice, is reinstated, or if other components of the transformer (MLP layer, multiple attention heads, multiple attention layers, etc.) are added. Also, the token embedding is strange and redundant, with consecutive inputs $\\mathbf{x}\\_i, \\mathbf{x}\\_{i-1}$ concatenated in the same token. Another, though less significant, issue is that only the population loss is considered. The authors repeatedly claim that “We believe that our analysis can give insight to the establishment of theory in settings with more complex models” but never explain why. \n\n    3. The model is already initialized as a model that implements gradient descent during ICL, and under the scalar-times-identity covariance assumption (4.1), or the enforced off-diagonal masking, it is not hard to show that the gradient flow dynamics reduce to the dynamics of two scalars ($a$ and $b$) which parameterize a set of models that always does GD during ICL. It is not clear whether the trained transformer can learn to perform GD during ICL if it is not initialized as a model that does this, and the experiments do not address this question.\n\n    4. “Gradient clipping” is used to refer to the masking procedure that zeros out all off-diagonal elements of the gradient. This undersells the severity of this procedure — clipping is a non-zero-thresholding operation applied equally to all gradient elements and is commonly used in practice for optimization and privacy purposes, and does not use any prior knowledge of how the gradient should be structured. The masking procedure described here sends a particular subset of elements of the gradient to zero, to drive the dynamics towards a particular desired solution based on prior knowledge of the data-generating process, which is a much stronger, and impractical, procedure.\n\n    5. Assumption 4.2 and Theorem 4.2 are tautological and not insightful. They should identify the properties of $\\mathbf{x}\\_1$ and $\\mathbf{W}$ that lead to such behavior.\n\n    6. $\\kappa\\_3$ has $\\Omega(d)$ growth, which means the effective step size of the gradient descent step is $\\tilde{O}(T/d)$, meaning $T$ (number of examples in the context length) must be $\\tilde{\\Omega}(d)$ for accurate predictions, which makes sense for this problem setting, but suggests the proposed setting is far from capturing ICL ability in practice, where the number of in-context examples is far fewer than the token embedding dimension. In the experiments, $T=20d$, which means the experiments do not evaluate whether the results can extend to smaller-context scenarios.\n\n    7. Proposition 4.2 is not meaningful because it does not rule out whether a different initialization could lead to a model that performs GD.\n\n    8. Even if we accept all of the simplifications and assumptions (including very long context) made to obtain a trained model that performs GD during ICL, this trained model does not learn the correct step size for GD, so it can’t perform ICL. To me, this suggests that GD is *not* the mechanism that transformers use to perform ICL — even in such an idealized setting, GD cannot explain the ability of transformers to perform ICL. However, the messaging around the paper is not consistent with this, as it continues to push the narrative that trained transformers are mesa-optimizers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer c1Lc \nWe thank Reviewer c1Lc for the valuable comments.\n## Weakness1: key contribution\nThanks for the insightful suggestion. We agree that we do not have direct evidence that \"ICL by AR pretraining is more difficult than ICL by few-shot pretraining\" because the data model is changed. We will revise it to \"ICL by AR pretraining is quite different from that by few-shot pretraining\" following your suggestion.\n\nHowever, **we emphasize that our key contribution is understanding the ICL ability of AR pretraining, rather than changing the data model**. **First**, **adopting the new data model is not our target but a necessary approach** to investigate AR pretraining because the few-shot regression dataset in the previous ICL theory is not suitable since each input $x_{i+1}$ does not have a relation with the context. **Second**, we additionally need to scale the embedding strategy, the attention model, and the loss function to perform the full AR pretraining. All these parts are not well studied in the literature. **Third**, we are the first to investigate the training dynamics and the mesa-optimization hypothesis in this setting, which deepens people's understanding of autoregressive transformers.\n\nBesides, due to the non-trivial nature of the AR pretraining, we think our results will highlight the role of other components in practical transformers (more attention layers, MLP, etc), and boost researchers to study autoregressive transformers.\n\n## Weakness2: artificial setting\nThanks for the advice. **We note that all settings adopted in this paper are reasonable and have attracted attention from the recent ICL theory, which is also agreed by Reviewer WaqB.** We clarify key points here, and we will discuss these in more detail in the final version.\n1. Single-layer one-head linear attention has been a popular setting in the recent ICL theory for few-shot regression [24,25,28,29,27]. As for the AR pretraining, the model in this paper is more complex and has also been adopted in recent works [16,17]. However, they do not investigate its training dynamics thus the convergence properties are still unclear.\n2. Linear data model is proposed in [16] and simplified in [17] for tractable analysis. We use the same data model as [17], while they only analyze the loss landscape under the diagonal assumption for model parameters.\n3. Token embeddings in this paper have been used in [16,17]. It is a natural extension of the embeddings in existing ICL theory because we need to predict each historical token here. In fact, practical transformers do learn similar token construction in the first softmax attention layer (see Fig. 4B in [16]).\n4. Population loss is a common setting in ICL theory which considers loss landscape and training dynamics [17,24,28,49,51,52].\n\n**Implications on more complex models**.  We will revise the claim to \"We believe that our analysis can give insight into the establishment of theory in settings with autoregressive transformers\". Besides. analogous to the development of ICL theory for few-shot regression (MLP [52], softmax attention [30], etc), future works considering more complex transformers can adopt the same data model, token embeddings in this paper, and try to use the similar proof technique to derive the training dynamics.\n## Weakness3: initialization & keeping diagonal structure\nFirst, the diagonal initialization is reasonable and we further conduct simulations with standard initialization. Second, we note that the findings of Assumption 4.1 and derivation are not trivial. Please see the details in the common concern 1 and 2, respectively.\n## Weakness4: gradient clipping\n**We note that our target in Theorem 4.3 is to complement the theoretical result in [17]** from the perspective of optimization, rather than provide a theory for practical gradient clipping. We will change \"gradient clipping\" to \"gradient masking\" in the final version to avoid confusion. \n## Weakness5: Assumption 4.2 and Theorem 4.2\nWe think the sufficient and necessary condition Assumption 4.2 might not be further simplified and we have discussed some properties of $x_1$ in line 238-245. For Theorem 4.2, we will follow the reviewer's suggestion to move Theorem 4.2 to the Appendix.\n## Weakness6: effect of $d$\nIn the final version, **we will clarify that we fix $d$ as a finite number and focus on the asymptotic property w.r.t. $T$**. Besides, we conduct additional experiments in smaller-context scenarios $(d=5, T=3/5/10, x_1\\sim N(0, I_d))$. As a result, the product $ab$ does converge to the results in Theorem 4.1. **The representative results can be found in the latest uploaded PDF**.\n## Weakness7: Proposition 4.2\n**This is a misunderstanding**. In the final version, we will clarify that **Proposition 4.2 does not depend on diagonal initialization**. As we have mentioned in line 272-274, parameters that perform GD are not the critical points of the loss function. Therefore, the trained transformer will not perform vanilla gradient descent with any initialization.\n## Weakness8: rationality of mesa-optimization\n**We disagree with this with our highest respect**. From the result of this paper, we can only conclude that **one-step GD is not enough to learn the data distribution**. However, even for more complex data ($W$ is not diagonal), **[16] has empirically verified that multi-layer linear attention can perform multi-step GD to learn the data distribution**. The theory for the multi-layer attention in AR pretraining setting is a meaningful future work.\n## Question 1: ratio of two vectors\nAs we have mentioned in line 326, we want to verify the correctness of Proposition 4.1, which claims that the prediction will converge to 1/5 of the true data. Therefore, we present the ratio and find it does converge to 1/5. In practice, we estimate this ratio by calculating the mean of the element-wise divide between two vectors. \n\nWe will clarify this detail more carefully in the final version.", "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:
# Common concerns from reviewers
We thank all reviewers for their valuable and constructive comments. We address the common concerns here and post a point-to-point response to each reviewer as well. We believe the quality of the paper has been improved following the reviewers' suggestions.

## Common concern 1: diagonal initialization (from Reviewer dsNs and c1Lc)
1. We agree with Reviewer dsNs that this assumption might be inevitable for a tractable analysis here. However, we note that **diagonal initialization/structure has been used in recent ICL theory for casual transformers [17,30, 51]**. Especially, as we have mentioned repeatedly in the paper, the most related paper [17] considers a stronger diagonal structure than ours, but they only investigate the loss landscape. Our results definitely deepened the understanding of autoregressive transformers by considering practical training dynamics.
2. **We have added the new experimental results under standard initialization with different scales**. We initialize each parameter using normal distribution $N(0,\sigma_w^2)$, where $\sigma_w$ is chosen from $(10^{-3},10^{-2},10^{-1})$, respectively. The results are summarized as follows and **their representatives can be found in the latest uploaded PDF**.  We think it improves our experimental part a lot.
    1. Though the convergence results of parameters are not exactly the same as those under diagonal initialization, **they keep the same diagonal structure, which can be understood as GD with adaptive learning rate in different dimensions**.
    2. The test results (ratio, MSE loss) of the trained transformers under standard initialization **are the same as** those under diagonal initialization, which further verifies the capability of the trained transformers.

To sum up, the experimental results demonstrate that our theoretical results have a certain representativeness, which further supports the rationality of the diagonal initialization.

For the more complex architecture (concern from Reviewer dsNs), we kindly refer the reviewer to the existing work [16] that includes results of (multi-layer) linear attention, (multi-layer) softmax attention, and full transformer on similar AR process ($W$ is non-diagonal) with standard initialization. They have established a strong connection between practical transformers and multi-step GD, and the theory for more complex architectures is meaningful for future work.

We will discuss these in more detail in the final version.

## Common concern 2: keep diagonal structure (from Review dsNs and c1Lc)
1. Finding the condition for the emergence of mesa-optimization (Assumption 4.1) is non-trivial since we need to derive the training dynamics first. However, the derivation of the dynamics is more difficult than the existing ICL theory for regression [24] because our data, model, token construction, and loss function become more complex.
2. This paper does deepen the understanding of autoregressive transformers compared to previous studies, and the extended experiments show that the theory has a certain representativeness. For the general convergence results, as we have discussed in line 279-283, the computation of the gradient will be much more difficult and might be intractable, which is agreed by Review dsNs. Therefore, we leave the rigorous result of general convergence for future work.

We will discuss these in more detail in the final version.

Meta-review: There is consensus among all reviewers that the paper presents a novel significant contribution and is clearly written. Moreover, several reviewers checked the proofs.

We recommend acceptance as a poster.

To the authors: please make sure to take reviewer comments and to fix remaining typos in the final revision.
Decision: Accept (poster)