On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability

Authors: chenyu  zheng
wei  huang
rongzhen  wang
guoqiang  wu
jun  zhu
chongxuan  li
data  management
analysis  methods
riken  aip
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:31:40

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.

Main paper: 
 introduction: 
 Foundation models based on transformers [1] have revolutionized the AI community in lots of fields, such as language modeling [2; 3; 4; 5; 6], computer vision [7; 8; 9; 10] and multi-modal learning [11; 12; 13; 14; 15] . The crux behind these large models is a very simple yet profound strategy named autoregressive (AR) pretraining, which encourages transformers to predict the next token when a 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 to adapt their computation and perform downstream tasks based on the information (e.g. examples) provided in the context without any updates to their parameters. However, the reason underlying the emergence of ICL ability is still poorly understood.ICL for linear regression, where they find the ICL performance of trained transformers is close to the OLS. On the theoretical side, considering a 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; 34] . Unfortunately, the pretraining goal of these studies is different from the AR training. Therefore, it is still unclear whether these findings can be transferred to transformers autoregressively trained on sequential data. 
other explanations for icl: 
 In addition to the mesa-optimization hypothesis, there are other explanations for the emergence of ICL. [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 a 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 to ours because we study the ICL of a one-layer linear transformer with AR pretraining on the first-order AR process. 
problem setup: 
 Elementary notations. We define [n] = {1, 2, . . . , n}. We use lowercase, lowercase boldface, and uppercase boldface letters to denote scalars, vectors, and matrices, respectively. For a vector a, we denote its i-th element as a i . For a matrix A, we use A k: , A :k and A ij to denote its k-th row, k-th column and (i, j)-th element, respectively. For a vector a (matrix A), we use a * (A * ) to denote its conjugate transpose. Similarly, we use a and A to denote their element-wise conjugate. We denote the n-dimensional identity matrix by I n . We denote the one vector of size n by 1 n . In addition, we denote the zero vector of size n and the zero matrix of size m × n by 0 n and 0 m×n , respectively. We use ⊗ and to denote the Kronecker product and the Hadamard product, respectively. Besides, we denote Vec(•) the vectorization operator in column-wise order.We consider a first-order AR process as the underlying data distribution, similar to recent works on AR pretraining [16; 17] . Concretely, to generate a sequence (x 1 , . . . , x T ) ∈ C d×T , we first randomly sample a unitary matrix W ∈ C d×d uniformly from a candidate set P W = {diag(λ 1 , . . . , λ d ) | |λ i | = 1, ∀i ∈ [d]} and the initial data point x 1 from a controllable distribution D x1 to be specified later, then the subsequent elements are generated according to the rulex t+1 = W x t for t ∈ [T -1].For convenience, we denote the vector (λ 1 , . . . , λ d ) by λ. We note that the structural assumption on W is standard in the literature on learning problems involving matrices [46; 17] . In addition, it is a natural extension of the recent studies on ICL for linear problems [24; 26; 28; 29; 30] , where they focus on y = w x = 1 d diag(w)x. Furthermore, adopting this new data model is a necessary approach to investigate AR pretraining because the regression dataset in the previous ICL theory is not suitable since each token x i does not have a relation with the context. In this paper, we mainly investigate the impact of the initial distribution D x1 on the convergence result of the AR pretraining.Linear causal self-attention layer. Before introducing the specified transformer module we will analyze in this paper, we first recall the definition of the standard causally-masked self-attention layer [1] , whose parameters θ includes a query matrix W Q ∈ R d k ×de , a key matrix W K ∈ R d k ×de , a value matrix W V ∈ R dv×de , a projection matrix W P ∈ R de×dv and a normalization factor ρ t > 0. At time step t, let E t = (e 1 , . . . , e t ) ∈ R de×t be the tokens embedded from the prompt sequence P t = (x 1 , . . . , x t ) ∈ C d×t , the causal self-attention layer will outputf t (E t ; θ) = e t + W P W V E t • softmax (W K E t ) * W Q e t ρ t ∈ C de ,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; 25; 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 W KQ = W K * W Q and W P V = W P W V , we have θ = (W KQ , W P V ), and the output can be rewritten as:f t (E t ; θ) = e t + W P V E t • E * t W KQ e t ρ t .Though it is called linear attention, we note that the output f t (E t ; θ) is non-linear w.r.t. the input tokens E t due to the existence of the E t E * t . In terms of the normalization factor ρ t , like existing works [24; 26] , we take it to be t -1 because each element in E t E * t is a Hermitian inner product of two vectors of size t.Embeddings. We adopt the natural embedding strategy used in recent studies on AR learning [16; 17] . Given a sequence P t = (x 1 , . . . , x t ), the i-th token is defined as e i = (0 d , x i , x i-1 ) ∈ C 3d , thus the corresponding embedding matrix E t can be formally written as:E t = (e 1 , . . . , e t ) =   0 d 0 d • • • 0 d x 1 x 2 • • • x t x 0 x 1 • • • x t-1   ∈ C 3d×t ,where x 0 = 0 d as a complement. This embedding strategy is a natural extension of the existing theoretical works about ICL for linear regression [30; 24; 25; 26; 28; 29] . The main difference is that the latter only focus on predicting the last query token while we need to 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 t = (x 1 , . . . , x t ), the network's prediction for the next element x t+1 will be the first d coordinates of the output f t (E t ; θ), aligning with the setup adopted in [16; 17] . Namely, we havey t (E t ; θ) = [f t (E t ; θ)] 1:d ∈ C d .Henceforth, we will omit the dependence on E t and θ, and use y t if it is not ambiguous. Since only the first d rows are extracted from the output by the attention layer, the prediction y t just depends on some parts of W P V and W KQ . Concretely, we denote thatW P V =   W P V 11 W P V 12 W P V 13 W P V 21 W P V 22 W P V 23 W P V 31 W P V 32 W P V 33   , W KQ =    W KQ 11 W KQ 12 W KQ 13 W KQ 21 W KQ 22 W KQ 23 W KQ 31 W KQ 32 W KQ 33    , where W P V ij , W KQ ij ∈ R d×d for all i, j ∈ [3]. Then the y t can be written asEQUATIONHere E x t = (e x 1 , . . . , e x t ) ∈ C 2d×t denotes the last 2d rows of the E t , where e x i = (x i , x i-1 ) . Therefore, we only need to analyze the selected parameters in Eq. 1 during the training dynamics. The derivation of Eq. 1 can be found in Appendix A.1.Loss function. To train the transformer model over the next-token prediction task, we focus on minimizing the following population loss:EQUATIONwhere the expectation is taken with respect to the start point x 1 and the transition matrix W . Henceforth, we will suppress the subscripts of the expectation for simplicity. The population loss is a standard objective in the optimization studies [24; 47] , and this objective has been used in recent works on AR modeling [16; 17] . The summation starts from t = 2 because we do not have any information to predict x 2 given only x 1 .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 τ = 0, we assume thatW KQ (0) =   0 d×d 0 d×d 0 d×d 0 d×d 0 d×d 0 d×d 0 d×d a 0 I d 0 d×d   , W P V (0) =   0 d×d b 0 I d 0 d×d 0 d×d 0 d×d 0 d×d 0 d×d 0 d×d 0 d×d   ,where the red submatrices are related to the y t 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 6, we also conduct experiments under standard initialization, which further supports the rationality of Assumption 3.1.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) d dτ θ = -∇L(θ). 
main results: 
 In this section, we present the main theoretical results of this paper. First, in Section 4.1, we prove that when D x1 satisfies some certain condition (Assumption 4.1), the trained transformer implements one step of gradient descent for the minimization of an OLS problem, which validates the rationality of the mesa-optimization hypothesis [16] . Next, in Section 4.2, we 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 4.1, where the exploratory analysis proves that the trained transformer will generally not perform vanilla gradient descent for the OLS problem.In this subsection, we show that under a certain assumption of D x1 , the trained one-layer linear transformer will converge to the mesa-optimizer [16; 17] . Namely, it will perform one step of gradient descent for the minimization of an OLS problem about the received prompt. The sufficient condition of the distribution D x1 can be summarized as follows.Assumption 4.1 (Sufficient condition for the emergence of mesa-optimizer). We assume that the distribution D x1 of the initial tokenx 1 ∈ R d satisfies E x1∼Dx 1 [x 1i1 x r2 1i2 • • • x rn 1in ] = 0 for any subset {i 1 , . . . , i n | n ≤ 4} of [d],and r 2 , . . . r n ∈ N. In addition, we assume thatκ 1 = E[x 4 1j ], κ 2 = E[x 6 1j ] and κ 3 = r =j E[x 2 1j x 4 1r ] are finite constant for any j ∈ [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-diagonal elements of W KQ 32 and W P Vas zero, thus they can keep diagonal structure during the training. We note that any random vectors x 1 whose coordinates x 1i are i.i.d. random variables with zero mean and finite moments of order 2, 4, and 6 satisfy this assumption. For example, it includes the normal distribution N (0 d , σ 2 I d ), 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. a b = κ 1 κ 2 + κ3 T -2 T -1 t=2 1 t-1 .As far as we know, Theorem 4.1 is the first theoretical result for the training dynamics and the mesaoptimization hypothesis of autoregressive transformers. The technical challenge compared to existing ICL theory for regression [24] mainly has two parts. First, our data model breaks the independence between data at 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) to 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 a non-trivial extension of recent work [17] , which characterizes the global minima of the AR modeling loss when imposing the diagonal structure on all parameter matrices during the training and fixing x 1 as 1 d . In comparison, Theorem 4.1 does not depend on 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 mesa-optimization 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 Theorem 4.1 holds. When predicting the (t+1)-th token, the trained transformer obtains W by implementing one step of gradient descent for the OLS problem L OLS,t (W ) =1 2 t-1 i=1 x i+1 -W x i2 , starting from the initialization W = 0 d×d with a step size a b t-1 .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 T te is EQUATIONE x1,W [ ( y Tte ) j (W x Tte ) j ] → 1 5 .Therefore, the prediction y Tte will not converges to the true next token W x Tte .Proposition 4.1 suggests that ICL by AR pretraining [16; 17] is different from ICL by few-shot pretraining [24; 25; 26; 27; 28; 29; 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 N (0 d , σ 2 I d ). However, in the AR learning setting, the trained linear transformer fails.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), [16] 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] 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 Assumption 4.1 does not hold. In this subsection, we conduct exploratory analyses by adopting the setting in [17] , where the initial token x 1 is fixed as1 d .First, sharing the similar but weaker assumption of [17] , we impose W KQ 32 and W P V 12 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 as 1 d , and we clip non-diagonal gradients of W KQ 32 and W P V 12 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, witha b = 1 1 + d-1 T -2 T -1 t=2 1 t-1 .Therefore, the obtained transformer performs one step of gradient descent in this case. Theorem 4.3 can be seen as a complement and an extension of Proposition 2 in [17] from the perspective of optimization. We note that [17] assumes all the parameter matrices to be diagonal and only analyzes the global minima without considering the practical non-convex optimization process.Next, we 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 is 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 32 and W P V 12 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 Appendix A.8). The limiting 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 problem1 2 t-1 i=1 x i+1 -W x i 2 .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 KQ 32 and W P V 12 anymore. When W KQ 32 and W P V 12 are general dense matrices, computation of the gradient will be much more difficult than that in Proposition 4.2. Therefore, we 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 a similar problem. [24; 26; 23] find that when the input's distribution does not satisfy Assumption 4.1 (e.g., N (0 d , Σ)), 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 x 1 is a full one vector, in Section 6. 
proof skeleton: 
 In this section, we outline the proof ideas of Theorem 4.1, which is one of the core findings of this 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 is placed in Appendix A.2.The first key step is to observe that each coordinate of prediction y t (Eq. 1) can be written as the output of a quadratic function, which will greatly simplify the follow-up gradient operation. Lemma 5.1 (Simplification of y t,j , proof in Appendix A.2.1). Each element of the network's prediction y t,j (j ∈ [d]) can be expressed as the following.y t,j = B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = Vec (A) • e x t ⊗ E x t E x t ρ t • B j ,where the A and B j are defined asA = a 1 . . . a 2d = W KQ 22 W KQ 23 W KQ 32 W KQ 33 , B j = b j1 b j2 = W P V12,j:W P V 13,j:,with a i ∈ R 2d and b j1 , b j2 ∈ R d .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 in 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 KQ 32 and W P V 12 satisfiesd dτ a = -ab 2   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + b(T -2)κ 1 , d dτ b = -a 2 b   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + a(T -2)κ 1 ,while the gradients for all other parameters were kept at zero during the training process.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. (a, b) = 1 2c 1 (c 2 -c 1 ab) 2 ,whose global minimums satisfy ab = c 2 /c 1 .Furthermore, We show that although the objective (a, b) is non-convex, the Polyak-Łojasiewicz (PL) inequality [60; 61] 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 Assumption 4.1 holds, then (a, b) is a non-convex function and satisfies the PL inequality as follows.∂ ∂a (a, b) 2 + ∂ ∂b (a, b) 2 ≥ 2c 1 (a 2 + b 2 ) (a, b) -min a,b (a, b) .Therefore, the gradient flow in Lemma 5.2 converges to the global minimum of (a, b).Finally, Theorem 4.1 can be proved by directly applying the above lemmas. 
simulation results: 
 In this section, we conduct simulations to verify and generalize our theoretical results. In terms of the train set, we generate 10k sequences with T tr = 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 (a 0 , b 0 ) by (0.1, 0.1), (0.5, 1.5), (2, 2) . The detailed configurations (e.g., step size) and results of different experiments can be found in Appendix B.Initial token sampled from N (0 d , σ 2 I d ). We conduct simulations with σ = 0.5, 1, 2 respectively. With any initialization of (a 0 , b 0 ), simulations show that ab converges to κ 1 /κ 2 = 1/5σ 2 , and y Tte-1 converges to x Tte /5 in expectation, which verifies Theorem 4.1 and Proposition 4.1, respectively. In the main paper, we present the convergence results with σ = 0.5 in Fig. 1a and 1b . We also verify our theory in the small-context scenarios (T tr = 5), which is placed in Fig. 4 in Appendix B.3.Initial token sampled from Example 4.1. We conduct simulations with scale c = 0.5, 1, 2 respectively. With any initialization of (a 0 , b 0 ), simulations show that ab converges to κ 1 /κ 2 = 1/c 2 (see details in Appendix A.5), and y Tte-1 converges to the truth x Tte , which verifies Theorem 4.1 and Theorem 4.2, respectively. In the main paper, we present the results with c = 0.5 in Fig. 1c and 1d .Initial token fixed as 1 d . We conduct experiment with x 1 = 1 d .The results Fig. 7 in Appendix B.5 show that W KQ 32 and W P V 12 converge to dense matrices with strong diagonals, and other matrices converge to 0 d×d , which means that the trained transformer performs somewhat preconditioned gradient descent. The detailed derivation is placed in Appendix B.5.Go beyond the diagonal initialization. Finally, in order to extend our theory, we repeat experiments under Gaussian initialization with different variance (σ w = 0.001, 0.01, 0.1). The results of Gaussian start points and sparse start points (Example 4.1) can be found in Fig. 5 of Appendix B.3 and Fig. 6 of Appendix B.4, respectively. As a 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.  
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, we investigate the non-convex dynamics of a one-layer linear transformer autoregressively trained by gradient flow on a 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.1 and 4.1 does not hold is not fully addressed in this paper due to 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 see a direct social impact of our theory. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?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 a 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 a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: 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 in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or 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 nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. 
 
 Appendix: A.1 Proof of eq. ( 1)Proof. We first calculate the output by causal linear attention layer as the following:f t (E t ; θ) = e t + 1 ρ t W P V E t E * t W KQ e t = e t + 1 ρ t W P V   t i=1 e i e * i   W KQ e t = e t + 1 ρ t t i=1 W P V e i • W KQ e i * e t = e t + 1 ρ t t i=1        W P V 11 W P V 12 W P V 13 W P V 21 W P V 22 W P V 23 W P V 31 W P V 32 W P V 33   e i •        W KQ 11 W KQ 12 W KQ 13 W KQ 21 W KQ 22 W KQ 23 W KQ 31 W KQ 32 W KQ 33    e i     *      e t =   0 d x t x t-1   + 1 ρ t t i=1       W P V 12 x i + W P V 13 x i-1 × ×   •    × W KQ 22 x i + W KQ 32 x i-1 W KQ 23 x i + W KQ 33 x i-1    *       0 d x t x t-1   ,where ×s are the elements that will not contribute to the final y t . A further simple computation shows that For the reader's convenience, we restate the lemma as the following.y t = 0 d + 1 ρ t t i=1 (W P V 12 x i + W P V 13 x i-1 )(W KQ 22 x i + W KQ 32 x i-1 ) * x t + 1 ρ t t i=1 (W P V 12 x i + W P V 13 x i-1 )(W KQ 23 x i + W KQ 33 x i-1 ) * x t-1 = 1 ρ t t i=1   W P V 12 x i + W P V 13 x i-1 W KQ 22 x i + W KQ 32 x i-1 W KQ 23 x i + W KQ 33 x i-1 *   x t x t-1 = 1 ρ t t i=1    W P V 12 W P V 13 x i x i-1 •   W KQ 22 W KQ 23 W KQ 32 W KQ 33 x i x i-1   *    x t x t-1 = 1 ρ t t i=1    W P V 12 W P V 13 e x i •   W KQ 22 W KQ 23 W KQ 32 W KQ 33 e x i   *   e x t = 1 ρ t W P V 12 W P V 13   t i=1 e x i e x * i   W KQ 22 W KQ 23 W KQ 32 W KQ 33 e x t = W P V 12 W P V 13 E x t E x * t ρ t W KQ 22 W KQ 23 W KQ 32 W KQLemma A.1. Each element of the network's prediction y t,j (j ∈ [d]) can be expressed as the following.y t,j = B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = Vec (A) • e x t ⊗ E x t E x t ρ t • B j ,where the A and B j are defined asA = a 1 . . . a 2d = W KQ 22 W KQ 23 W KQ 32 W KQ 33 , B j = b j1 b j2 = W P V12,j:W P V 13,j:,with a i ∈ R 2d and b j1 , b j2 ∈ R d .Proof. Based on the result in Eq. 1, we can writeEQUATIONEQUATIONE x t E x * t ρt      = tr Vec(A)B j • e x t ⊗ E x t E x * t ρ t = tr B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = tr Vec (A) • e x t ⊗ E x t E x t ρ t • B j = Vec (A) • e x t ⊗ E x t E x t ρ t • B j ,which finishes the proof.For the reader's convenience, we restate the lemma as the following. Lemma A.2. Suppose that Assumption 4.1 holds, then the dynamical process of the parameters in the diagonal of W KQ 32 and W P V 12 satisfiesd dτ a = -ab 2   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + b(T -2)κ 1 , d dτ b = -a 2 b   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + a(T -2)κ 1 ,while the gradients for all other parameters were kept at zero during the training process.Proof. The population loss L(θ) in Eq. 2 can be rewritten asL(θ) = T -1 t=2 L t (θ) = T -1 t=2 E   1 2 d j=1 y t,j -x t+1,j 2   = T -1 t=2 d j=1 E 1 2 y t,j -x t+1,j 2 = T -1 t=2 d j=1 E 1 2 y * t,j y t,j -2Re x * t+1,j y t,j + x * t+1,j x t+1,j = T -1 t=2 d j=1 E 1 2 y * t,j y t,j -Re x * t+1,j y t,j + 1 2x * t+1,j x t+1,j .Then, we can calculate the derivatives of L t (θ) with respect to B j and Vec(A) asEQUATION∇ Vec(A) L t (θ) = d j=1 E 1 2 ∇ Vec(A) y * t,j y t,j -∇ Vec(A) Re x * t+1,j y t,j = 1 2 d j=1 E ∇ Vec(A) y * t,j y t,j - d j=1E ∇ Vec(A) Re x * t+1,j y t,j .Step one: calculate E ∇ Bj Re x * t+1,j y t,j . Based on Lemma 5.1, we havey t,j = B j • e x t ⊗ E x t E x * t ρ t • Vec(A).Then, the E ∇ Bj Re x * t+1,j y t,j can be derived as the following.E ∇ Bj Re x * t+1,j y t,j = E ∇ Bj Re x * t+1,j B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = E ∇ Bj B j • Re x * t+1,j • e x t ⊗ E x t E x * t ρ t • Vec(A) = E Re x * t+1,j • e x t ⊗ E x t E x * t ρ t • Vec(A) = Re E x * t+1,j • e x t ⊗ E x t E x * t ρ t • Vec(A) = Re E λ -t j x 1j • Vec( E x t E x * t ρ t Ae x t ) (use generating process) = Re E λ -t j x 1j • ( E x t E x * t ρ t Ae x t ) ,where the penultimate equality uses the property of Kronecker and Vec operator Vec(AXB) = (B ⊗ A)Vec(X), we refer Section 10.2 in [62] for details.ForE x t E x * t ρt, we can simplify it asE x t E x * t ρ t = 1 ρ t t i=1 e x i e x * i = 1 ρ t t i=1 x i x * i x i x * i-1 x i-1 x * i x i-1 x * i-1 = 1 ρ t t i=1 x i x * i W t-1 i=1 x i x * i t-1 i=1 x i x * i W * t-1 i=1 x i x * i .Based on the diagonal property of W , we can simplify the x i x * i as the following.x i x * i = (W i-1 x 1 )(W i-1 x 1 ) * = M i Σ, where we define Σ = x 1 x * 1 and M i = λ i-1 λ i-1 *. Therefore, we haveEQUATION)Then, leveraging the sparse property of A, we can deriveE x t E x * t ρtAe x t as follows.E x t E x * t ρ t Ae x t = E x t E x * t ρ t 0 d ax t = a ρ t W ( t-1 i=1 M i Σ)x t ( t-1 i=1 M i Σ)x t .Therefore, for any l ∈ [d], we haveE x t E x * t ρ t Ae x t l = a ρ t λ l t-1 i=1 (M i Σ) l: x t = a ρ t λ l t-1 i=1 λ i-1 l λ 1-i 1 x 1l x 11 • • • λ i-1 l λ 1-i d x 1l x 1d   λ t-1 1 x 11 • • • λ t-1 d x 1d   = a ρ t λ l t-1 i=1 d r=1 λ i-1 l λ t-i r x 1l x 2 1r = a ρ t t-1 i=1 d r=1 λ i l λ t-i r x 1l x 2 1r .Similarly, for any l ∈[2d] -[d],we haveE x t E x * t ρ t Ae x t l = a ρ t t-1 i=1 d r=1 λ i-1 l λ t-i r x 1l x 2 1r .To sum up, for any l ∈ [2d], we haveE x t E x * t ρ t Ae x t l = a ρt t-1 i=1 d r=1 λ i l λ t-i r x 1l x 2 1r , l ∈ [d], a ρt t-1 i=1 d r=1 λ i-1 l-d λ t-i r x 1,l-d x 2 1r , l ∈ [2d] -[d].Next, we calculate the E λ -t j x 1j •E x t E x * t ρt Ae x t . For any l ∈ [d],we haveE λ -t j x 1j • E x t E x * t ρ t Ae x t l = E   λ -t j x 1j • a ρ t t-1 i=1 d r=1 λ i l λ t-i r x 1l x 2 1r   = a ρ t t-1 i=1 d r=1 E[λ -t j λ i l λ t-i r ]E[x 1j x 1l x 2 1r ].We discuss them in the following categories,1. l = j. In this case, E[x 1j x 1l x 2 1r ] = 0 by Assumption 4.1, thus E λ -t j x 1j • E x t E x * t ρ t Ae x t l = a ρ t t-1 i=1 d r=1 E[λ -t j λ i l λ t-i r ]0 = 0.2. l = j, r = j. Because λ i s are i.i.d. andE[λ i ] k = δ(k = 0), we have E λ -t j x 1j • E x t E x * t ρ t Ae x t l = a ρ t d r=1 E[λ -t j λ i j λ t-i j ]E[x 1j x 1j x 2 1j ] = a ρ t (t -1)E[x 4 1j ].Similarly, for any l∈ [2d] -[d],we haveE λ -t j x 1j • E x t E x * t ρ t Ae x t l = 0.Therefore, for any l ∈ [2d], we haveE λ -t j x 1j • E x t E x * t ρ t Ae x t = a ρt (t -1)E[x 4 1j ], l = j, 0, l = j,and the l-th element of E ∇ Bj Re x * t+1,j y t,j isE ∇ Bj Re x * t+1,j y t,j l = Re E λ -t j x 1j • ( E x t E x * t ρ t Ae x t ) l = a ρt (t -1)E[x 4 1j ], l = j, 0, l = j.Step two: calculate E ∇ Bj y * t,j y t,j . Based on Lemma 5.1, we havey t,j = Vec (A) • e x t ⊗ E x t E x t ρ t • B j ,then we can simplify the E ∇ Bj y * t,j y t,j as follows.E ∇ Bj y * t,j y t,j = E   ∇ Bj Vec (A) • e x t ⊗ E x t E x t ρ t • B j * Vec (A) • e x t ⊗ E x t E x t ρ t • B j   = E ∇ Bj B j • e x * t ⊗ E x t E x t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x t ρ t • B j = E e x * t ⊗ E x t E x t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x t ρ t • B j + E e x t ⊗ E x t E x * t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x * t ρ t • B j = E 2Re e x t ⊗ E x t E x * t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x * t ρ t • B j = 2Re E e x t ⊗ E x t E x * t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x * t ρ t • B j .We further derive thatE      e x t ⊗ E x t E x * t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x * t ρ t • B j ∈C      = E Vec (A)e x t ⊗ E x t E x * t ρ t • B j • e x t ⊗ E x t E x * t ρ t Vec(A) = E B j • e x * t ⊗ E x t E x t ρ t Vec(A) • e x t ⊗ E x t E x * t ρ t Vec(A) = E B j • Vec( E x t E x t ρ t Ae x t ) • Vec( E x t E x * t ρ t Ae x t ) = E B j • E x t E x t ρ t Ae x t • E x t E x * t ρ t Ae x t = E   2d k=1 B jk E x t E x t ρ t Ae x t k • E x t E x * t ρ t Ae x t   = E   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t   . (sparsity of B)For any and l ∈ [2d], recall thatE x t E x * t ρ t Ae x t l = a ρt t-1 i=1 d r=1 λ i l λ t-i r x 1l x 2 1r , l ∈ [d], a ρt t-1 i=1 d r=1 λ i-1 l-d λ t-i r x 1,l-d x 2 1r , l ∈ [2d] -[d],and for any j ∈ [d] , we haveE x t E x t ρ t Ae x t j = E x t E x * t ρ t Ae x t j = a ρ t t-1 i=1 d r=1 λ -i j λ i-t r x 1j x 2 1r .For any l ∈ [d], with careful computing, we haveE   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t l   = a 2 b ρ 2 t t-1 i1=1 t-1 i2=1 d r1=1 d r2=1 E[λ -i1 j λ i2 l λ i1-t r1 λ t-i2 r2 ]E[x 1j x il x 2 1r1 x 2 1r2 ].We discuss it in the following categories, 1. l = j. In this case, E[x 1j x il x 2 1r1 x 2 1r2 ] = 0 by Assumption 4.1, thus it becomes 0. 2. l = j, r 1 = r 2 = j. It becomesa 2 b ρ 2 t t-1 i1=1 t-1 i2=1 E[λ -i1 j λ i2 j λ i1-t j λ t-i2 j ]E[x 6 1j ] = a 2 b ρ 2 t t-1 i1=1 t-1 i2=1 E[x 6 1j ] = a 2 b ρ 2 t (t -1) 2 E[x 6 1j ].3. l = j, r 1 = r 2 = r = j. It becomesa 2 b ρ 2 t t-1 i1=1 t-1 i2=1 r =j E[λ i2-i1 j λ i1-i2 r ]E[x 2 1j x 4 1r ] = a 2 b ρ 2 t i1=i2 r =j E[x 2 1j x 4 1r ] = a 2 b ρ 2 t (t -1) r =j E[x 2 1j x 4 1r ]. 4. l = j, r 1 = r 2 . In this case, E[λ -i1 j λ i2 l λ i1-t r1 λ t-i2 r2 ] = E[λ i2-i1 j λ i1-t r1 λ t-i2 r2 ] = 0, thus it becomes 0. Similarly, for any l ∈ [2d] -[d], we have E   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t l   = 0.To sum up, we haveE   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t l   = a 2 b ρ 2 t (t -1) 2 E[x 6 1j ] + (t -1) r =j E[x 2 1j x 4 1r ] , l = j, 0,otherwise, which implies that the l-th element of 1 2 E ∇ Bj y * t,j y t,j is1 2 E ∇ Bj y * t,j y t,j l = a 2 b ρ 2 t (t -1) 2 E[x 6 1j ] + (t -1) r =j E[x 2 1j x 4 1r ] , l = j, 0, otherwise.Step three: calculate ∇ Bj L t (θ) and ∇ Bj L(θ). Based on steps one and two, the l-th element of ∇ Bj L t (θ) can be derived as follows.∇ Bj L t (θ) l = 1 2 E ∇ Bj y * t,j y t,j l -E ∇ Bj Re x * t+1,j y t,j l = a 2 b ρ 2 t (t -1) 2 E[x 6 1j ] + (t -1) r =j E[x 2 1j x 4 1r ] -a ρt (t -1)E[x 4 1j ], l = j, 0, otherwise,.Furthermore, the l-th element of ∇ Bj L(θ) is∇ Bj L(θ) l = T -1 t=2 ∇ Bj L t (θ) l =    T -1 t=2 a 2 b ρ 2 t (t -1) 2 E[x 6 1j ] + (t -1) r =j E[x 2 1j x 4 1r ] -a ρt (t -1)E[x 4 1j ] , l = j, 0, otherwise, =    T -1 t=2 a 2 b E[x 6 1j ] + 1 t-1 r =j E[x 2 1j x 4 1r ] -aE[x 4 1j ] , l = j, 0, otherwise, = a 2 b (T -2)E[x 6 1j ] + T -1 t=2 1 t-1 r =j E[x 2 1j x 4 1r ] -a(T -2)E[x 4 1j ], l = j, 0, otherwise, = a 2 b (T -2)κ 2 + T -1 t=2 1 t-1 κ 3 -a(T -2)κ 1 , l = j, 0, otherwise,where the last equality is from Assumption 4.1.Step four: calculate E ∇ Vec(A) Re x * t+1,j y t,j and d j=1 E ∇ Vec(A) Re x * t+1,j y t,j . Based on Lemma 5.1, we havey t,j = Vec (A) • e x t ⊗ E x t E x t ρ t • B j .Then, the E ∇ Vec(A) Re x * t+1,j y t,j can be derived as the following.E ∇ Vec(A) Re x * t+1,j y t,j = E   ∇ Vec(A) Re x * t+1,j Vec (A) • e x t ⊗ E x t E x t ρ t • B j   = E   ∇ Vec(A) Vec (A) • Re x * t+1,j • e x t ⊗ E x t E x t ρ t • B j   = E   Re x * t+1,j • e x t ⊗ E x t E x t ρ t • B j   = Re   E x * t+1,j • e x t ⊗ E x t E x t ρ t • B j   = Re   E λ -t j x 1j • Vec( E x t E x t ρ t B j e x t )   .In terms ofE x t E x t ρtB j e x t , based on the sparsity of B and Eq. 4, we can derive thatE x t E x t ρ t B j e x t = ( E x t E x t ρ t ) :j be x t = ( E x t E x * t ρ t ) j: be x t = b t i=1 M i Σ W t-1 i=1 M i Σ j: e x t .Then, for any s, r ∈ [2d], we haveE x t E x t ρ t B j e x t sr =            b ρt t i=1 λ i-1 j λ 1-i s λ t-1 r x 1j x 1s x 1r , s ∈ [d], r ∈ [d], b ρt t-1 i=1 λ i j λ 1-i s-d λ t-1 r x 1j x 1,s-d x 1r , s ∈ [2d] -[d], r ∈ [d], b ρt t i=1 λ i-1 j λ 1-i s λ t-2 r-d x 1j x 1s x 1,r-d , s ∈ [d], r ∈ [2d] -[d], b ρt t-1 i=1 λ i j λ 1-i s-d λ t-2 r-d x 1j x 1,s-d x 1,r-d , s ∈ [2d] -[d], r ∈ [2d] -[d]. Next, we calculate E λ -t j x 1j • E x t E x t ρt B j e x t . For any s ∈ [2d] -[d], r ∈ [d], we have E   λ -t j x 1j E x t E x t ρ t B j e x t sr   = E   b ρ t t-1 i=1 λ i-t j λ 1-i s-d λ t-1 r x 2 1j x 1,s-d x 1r   = b ρ t t-1 i=1 E[λ i-t j λ 1-i s-d λ t-1 r ]E[x 2 1j x 1,s-d x 1r ].We discuss it in the following categories,1. s -d = r. In this case, E[x 2 1j x 1,s-d x 1r ] = 0 by Assumption 4.1, thus it becomes 0. 2. s -d = r = j. It becomes b ρ t t-1 i=1 E[λ i-t j λ 1-i s-d λ t-1 r ]E[x 2 1j x 1,s-d x 1r ] = b ρ t t-1 i=1 E[λ i-t j λ 1-i j λ t-1 j ]E[x 4 1j ] = b ρ t t-1 i=1 E[x 4 1j ] = b ρ t (t -1)E[x 4 1j ]. 3. s -d = r = j. In this case, E[λ i-t j λ 1-i s-d λ t-1 r ] = E[λ i-t j λ t-i s-d ] = 0, thus it becomes 0.Similarly, for any other s, r, we can calculate thatE   λ -t j E x t E x t ρ t B j e x t sr   = 0.To sum up, we haveE   λ -t j E x t E x t ρ t B j e x t sr   = b ρt (t -1)E[x 4 1j ], s = d + j, r = j, 0, otherwise, thus the (s, r)-th element of E ∇ A Re x * t+1,j y t,j is E ∇ A Re x * t+1,j y t,j sr = b ρt (t -1)E[x 4 1j ], s = d + j, r = j, 0, otherwise, = b ρt (t -1)κ 1 , s = d + j, r = j, 0, otherwise.Finally, we can calculated j=1 E ∇ A Re x * t+1,j y t,j as d j=1 E ∇ A Re x * t+1,j y t,j sr = b ρt (t -1)κ 1 , s -d = r, 0, otherwise.Step five: calculate E ∇ Vec(A) y * t,j y t,j and d j=1 E ∇ Vec(A) y * t,j y t,j . Based on Lemma 5.1, we havey t,j = B j • e x t ⊗ E x t E x * t ρ t • Vec(A),then we can simplify the E ∇ Vec(A) y * t,j y t,j as followsE ∇ Vec(A) y * t,j y t,j = E ∇ Vec(A) B j • e x t ⊗ E x t E x * t ρ t • Vec(A) * B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = E ∇ Vec(A) Vec(A) • e x t ⊗ E x t E x * t ρ t B j • B j • e x t ⊗ E x t E x * t ρ t • Vec(A) = E e x t ⊗ E x t E x * t ρ t B j • B j • e x t ⊗ E x t E x * t ρ t • Vec(A) + E e x t ⊗ E x t E x t ρ t B j • B j • e x * t ⊗ E x t E x t ρ t • Vec(A) = E   2Re e x t ⊗ E x t E x t ρ t B j • B j • e x * t ⊗ E x t E x t ρ t • Vec(A)   = 2Re   E e x t ⊗ E x t E x t ρ t B j • B j • e x * t ⊗ E x t E x t ρ t • Vec(A)   .We further derive thatE      e x t ⊗ E x t E x t ρ t B j • B j • e x * t ⊗ E x t E x t ρ t • Vec(A) ∈C      = E B j • e x * t ⊗ E x t E x t ρ t • Vec(A) • e x t ⊗ E x t E x t ρ t B j = E Vec(A) • e x t ⊗ E x t E x * t ρ t B j • e x t ⊗ E x t E x t ρ t B j = E Vec(A) • Vec( E x t E x * t ρ t B j e x * t ) • Vec( E x t E x t ρ t B j e x t ) = E   2d k,l=1 A kl E x t E x * t ρ t B j e x * t kl • Vec( E x t E x t ρ t B j e x t )   = E   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • Vec( E x t E x t ρ t B j e x t )   . (sparsity of A)Recall that for any s, r ∈ [2d], we haveE x t E x t ρ t B j e x t sr =            b ρt t i=1 λ i-1 j λ 1-i s λ t-1 r x 1j x 1s x 1r , s ∈ [d], r ∈ [d], b ρt t-1 i=1 λ i j λ 1-i s-d λ t-1 r x 1j x 1,s-d x 1r , s ∈ [2d] -[d], r ∈ [d], b ρt t i=1 λ i-1 j λ 1-i s λ t-2 r-d x 1j x 1s x 1,r-d , s ∈ [d], r ∈ [2d] -[d], b ρt t-1 i=1 λ i j λ 1-i s-d λ t-2 r-d x 1j x 1,s-d x 1,r-d , s ∈ [2d] -[d], r ∈ [2d] -[d].Furthermore, for any k∈ [2d] -[d],we can calculate thatE x t E x * t ρ t B j e x * t k,k-d = E x t E x t ρ t B j e x t k,k-d = b ρ t t-1 i=1 λ -i j λ i-1 k-d λ 1-t k-d x 1j x 2 1,k-d = b ρ t t-1 i=1 λ -i j λ i-t k-d x 1j x 2 1,k-d , With careful computing, for any s ∈ [2d] -[d], r ∈ [d],we haveE   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • E x t E x t ρ t B j e x t sr   = ab 2 ρ 2 t 2d k=d+1 t-1 i1=1 t-1 i2=1 E[λ i2-i1 j λ i1-t k-d λ 1-i2 s-d λ t-1 r ]E[x 2 1j x 2 1,k-d x 1,s-d x 1r ].We discuss it in the following categories,1. s -d = r. In this case, E[x 2 1j x 2 1,k-d x 1,s-d x 1r ] = 0 by Assumption 4.1, thus it becomes 0. 2. s -d = r = k -d = j. It becomes ab 2 ρ 2 t t-1 i1=1 t-1 i2=1 E[λ 0 j ]E[x 6 1j ] = ab 2 ρ 2 t (t -1) 2 E[x 6 1j ]. 3. s -d = r = k -d = j. It becomes ab 2 ρ 2 t t-1 i1=1 t-1 i2=1 E[λ i2-i1 j λ i1-i2 r ]E[x 2 1j x 4 1,r ] = ab 2 ρ 2 t t-1 i1=1 E[λ i1-i1 j λ i1-i1 r ]E[x 2 1j x 4 1,r ] = ab 2 ρ 2 t (t -1)E[x 2 1j x 4 1,r ]. 4. s -d = r = k -d. In these case, E[λ i2-i1 j λ i1-t k-d λ 1-i2 s-d λ t-1 r ] = E[λ i2-i1 j λ i1-t k-d λ t-i2 r ] = 0, thus it becomes 0.Similarly, for any other s, r, we can calculate thatE   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • E x t E x t ρ t B j e x t sr   = 0.To sum up, we haveE   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • E x t E x t ρ t B j e x t sr   =        ab 2 ρ 2 t (t -1) 2 E[x 6 1j ], s = d + r, r = j, ab 2 ρ 2 t (t -1)E[x 2 1j x 4 1r ], s = d + r, r = j, 0,otherwise, which implies that the (s, r)-th element of 1 2 E ∇ A y * t,j y t,j is1 2 E ∇ A y * t,j y t,j sr =        ab 2 ρ 2 t (t -1) 2 E[x 6 1j ], s = d + r, r = j, ab 2 ρ 2 t (t -1)E[x 2 1j x 4 1r ], s = d + r, r = j, 0, otherwise, =        ab 2 ρ 2 t (t -1) 2 κ 2 , s = d + r, r = j, ab 2 ρ 2 t (t -1)E[x 2 1j x 4 1r ], s = d + r, r = j, 0, otherwise, . 1 2 d j=1 E ∇ A y * t,j y t,j sr = ab 2 ρ 2 t (t -1) 2 κ 2 + (t -1)κ 3 , s -d = r, 0, otherwise.Step six: calculate ∇ Vec(A) L t (θ) and ∇ Vec(A) L(θ). Based on steps four and five, the (s, r)-th element of ∇ A L t (θ) can be derived as follows.∇ A L t (θ) sr = 1 2 d j=1 E ∇ A y * t,j y t,j - d j=1 E ∇ A Re x * t+1,j y t,j = ab 2 ρ 2 t (t -1) 2 κ 2 + (t -1)κ 3 -b ρt (t -1)κ 1 , s -d = r, 0, otherwise.Furthermore, the (s, r)-th element of ∇ A L(θ) is∇ A L(θ) sr = T -1 t=2 ∇ A L t (θ) sr = T -1 t=2 ab 2 ρ 2 t (t -1) 2 κ 2 + (t -1)κ 3 -b ρt (t -1)κ 1 , s -d = r, 0, otherwise, =    T -1 t=2 ab 2 κ 2 + 1 t-1 κ 3 -bκ 1 , s -d = r, 0, otherwise, = ab 2 (T -2)κ 2 + T -1 t=2 1 t-1 κ 3 -b(T -2)κ 1 , s -d = r, 0,otherwise.Step seven: summarize the result by induction. From the gradient of ∇ A L(θ) and ∇ Bj L(θ), we observe that non-zero gradients only emerge in the diagonal of W KQ 32 and W P V 12 , and they are same. Therefore, the parameter matrices keep the same structure as the initial time. Thus we can summarize the dynamic system as the following.d dτ a = -ab 2   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + b(T -2)κ 1 , d dτ b = -a 2 b   (T -2)κ 2 + T -1 t=2 1 t -1 κ 3   + a(T -2)κ 1 .which completes the proof.For the reader's convenience, we restate the lemma as the following. (a, b) = 1 2c 1 (c 2 -c 1 ab) 2 ,whose global minimums satisfy ab = c 2 /c 1 .Proof. Basic calculus shows that∂ ∂a (a, b) = 1 2c 1 2(c 2 -c 1 ab)(-c 1 b) = -b(c 2 -c 1 ab) = - d dτ a, ∂ ∂b (a, b) = 1 2c 1 2(c 2 -c 1 ab)(-c 1 a) = -a(c 2 -c 1 ab) = - d dτ b.Therefore, the dynamics in Lemma 5.2 are the same as those of gradient flow on (a, b), whose global minimums satisfy ab = c 2 /c 1 .A.2.4 Proof of Lemma 5.4For the reader's convenience, we restate the lemma as the following.Lemma A.4. Suppose that Assumption 4.1 holds, then (a, b) is a non-convex function and satisfies the PL inequality as follows.EQUATIONTherefore, the gradient flow in Lemma 5.2 converges to the global minimum of (a, b).Proof. First, we prove that (a, b) is non-convex. The Hessian matrix of (a, b) can be derived as follows.∇ 2 (a, b) = c 1 b 2 2c 1 ab -c 2 2c 1 ab -c 2 c 1 a 2 . Its determinant c 1 a 2 b 2 -(2c 1 ab -c 2 ) 2 = (c 2 -c 1 ab)(3c 1 ab -c 2 ) < 0 when ab < c2 3c1 or ab > c2 c1 . Thus, (a, b) is non-convex.Besides, the PL inequality holds becauseEQUATIONFor the reader's convenience, we restate the corollary as the following.Corollary A.1. We suppose that the same precondition of Theorem 4.1 holds. When predicting the t-th token, the trained transformer implements one step of gradient descent for the minimization of the OLS problemL OLS (W ) = 1 2 t-1 i=1 x t+1 -W x t 2, starting from the initialization W = 0 d×d with a step size a b t-1 .Proof. The proof is stem from the theoretical construction in [16] . First, we simplify the prediction y t as follows.y t = W P V 12 W P V 13 E x t E x * t ρ t W KQ 22 W KQ 23 W KQ 32 W KQ 33 e x t = bI d 0 d×d E x t E x * t ρ t 0 d×d 0 d×d aI d 0 d×d e x t = 1 ρ t bI d 0 d×d t i=1 e x i e x * i 0 d×d 0 d×d aI d 0 d×d e x t = 1 ρ t t i=1 bx i ax * i-1 0 d×d e x t = 1 ρ t t i=1 bx i ax * i-1 x t =   a b t -1 t i=1 x i x * i-1   x t =   a b t -1 t-1 i=1 x i+1 x * i   x t .Then, we connect it to the one step of gradient descent for the OLS problem L OLS (W ) =1 2 t-1 i=1 x i+1 -W x i 2 . W - a b t -1 ∇ W 1 2 t-1 i=1 x i+1 -W x i 2 = W - a b t -1 t-1 i=1 (x i+1 -W x i )(-x * i ) = 0 - a b t -1 t-1 i=1 (x i+1 -0x i )(-x * i ) = a b t -1 t-1 i=1 x i+1 x * i .Thus, the proof is completed.For the reader's convenience, we restate the proposition as the following. Proposition A.1. Let D x1 be the multivariate normal distribution N (0 d , σ 2 I d ) with any σ 2 > 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 tr is large enough, for any test context length T te and dimension j ∈ [d], the prediction from the trained transformer satisfiesE x1,W [ ( y Tte ) j (W x Tte ) j ] → 1 5 .Therefore, the prediction y Tte will not converges to the true next token W x Tte .Proof. First, built upon the results in Theorem 4.1, when T tr is large enough, we havea b = κ 1 κ 2 + κ3Ttr-2Ttr-1 t=21 t-1 → κ 1 κ 2 = E[x 4 1j ] E[x 6 1j ] = 3σ 4 15σ 6 = 1 5σ 2 .Second, by directly calculating, we haveEQUATION( y Tte ) j = a b T te -1 Tte-1 i=1 d k=1 λ i j λ Tte-i k x 1j x 2 1k .Therefore, we haveE x1,W [ ( y Tte ) j (W x Tte ) j ] = E x1,W [ a b T te -1 Tte-1 i=1 d k=1 λ i-Tte j λ Tte-i k x 2 1k ] = E x1 [ a b T te -1 Tte-1 i=1 x 2 1j ] = a bσ 2 .Since a b < 1 5σ 2 and converges to 1 5σ 2 when T tr is large enough, the proof is completed.A.5 Derivation of Example 4.1Proof. We first prove that the example satisfies Assumption 4.1. Because only one element of x 1 sampled from Example 4.1 will be non-zero, we haveE x1∼Dx 1 [x 1i1 x r2 1i2 • • • x rn 1in ] = E x1∼Dx 1 [0] = 0 for any subset {i 1 , . . . , i n | n ≤ 4} of [d],and r 2 , . . . r n ∈ N. In addition, for any j ∈ [d], we can derive thatκ 1 = E[x 4 1j ] = 1 d • c 4 + d -1 d • 0 = c 4 d •, κ 2 = E[x 6 1j ] = 1 d • c 6 + d -1 d • 0 = c 6 d •, κ 3 = r =j E[x 2 1j x 4 1r ] = 0.Second, we prove that it satisfies Assumption 4.2 as follows. Without loss of general, we assume that the first coordinate of x 1 is c.κ 1 κ 2 Tte-1 i=1 x i x * i T te -1 x Tte = 1 c 2 diag(c 2 , 0, . . . , 0)(λ Tte-1 1 c, 0, . . . , 0) = (λ Tte-1 1 c, 0, . . . , 0) = x Tte .The proof is finished.A.6 Proof of Theorem 4.2For the reader's convenience, we restate the theorem as the following. Theorem A.1. Suppose that Assumption 4.1 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 T tr and test context length T te are large enough, the prediction from the trained transformer satisfies y Tte → W x Tte , T tr , T te → +∞.Proof. First, built upon the results in Theorem 4.1, when T tr is large enough, we havea b = κ 1 κ 2 + κ3Ttr-2Ttr-1 t=21 t-1 → κ 1 κ 2 .Second, when T te is large enough, by Assumption 4.2κ 1 κ 2 Tte-1 i=1 x i x * i T te -1 x Tte → x Tte .Therefore, we havey Tte = W a b Tte-1 i=1 x i x * i T te -1 x Tte → Ttr W κ 1 κ 2 Tte-1 i=1 x i x * i T te -1 x Tte → Tte W x Ttewhich finishes the proof.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 is fixed as 1 d , and we clip non-diagonal gradients of W KQ 32 and W P V 12 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, witha b = 1 1 + d-1 T -2 T -1 t=2 1 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 is more difficult than that of 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 in Theorem 4.3, the dynamical process of the parameters in the diagonal of W KQ 32 and W P V 12 satisfiesd dτ a = -ab 2   (T -2) + T -1 t=2 d -1 t -1   + b(T -2), d dτ b = -a 2 b   (T -2) + T -1 t=2 d -1 t -1   + a(T -2),while the gradients for all other parameters were kept at zero during the training process.Proof. Recall that in Appendix A.2.2, we have already known that the population loss L(θ) in Eq. 2 can be rewritten asL(θ) = T -1 t=2 L t (θ) = T -1 t=2 d j=1 E 1 2y * t,j y t,j -Re x * t+1,j y t,j + 1 2x * t+1,j x t+1,j .Besides, the derivatives of L t (θ) with respect to Vec(A) and B j areEQUATION∇ Vec(A) L t (θ) = 1 2 d j=1 E ∇ Vec(A) y * t,j y t,j - d j=1 E ∇ Vec(A) Re x * t+1,j y t,j .Step one: calculate E ∇ Bj Re x * t+1,j y t,j . Similarly to the step one in Appendix A.2.2, the E ∇ Bj Re x * t+1,j y t,j can be derived as the following.E ∇ Bj Re x * t+1,j y t,j = Re E x * t+1,j • e x t ⊗ E x t E x * t ρ t • Vec(A) = Re E λ -t j • Vec( E x t E x * t ρ t Ae x t ) (use generating process) = Re E λ -t j • ( E x t E x * t ρ t Ae x t ) .We note that for any l ∈ [2d], based on the sparsity of A, we haveE x t E x * t ρ t Ae x t l = a ρt t-1 i=1 d r=1 λ i l λ t-i r , l ∈ [d], a ρt t-1 i=1 d r=1 λ i-1 l-d λ t-i r , l ∈ [2d] -[d],which implies thatE λ -t j • E x t E x * t ρ t Ae x t = a ρt (t -1), l = j, 0, l = j,and the l-th element of E ∇ Bj Re x * t+1,j y t,j isE ∇ Bj Re x * t+1,j y t,j l = a ρt (t -1), l = j, 0, l = j.Step two: calculate E ∇ Bj y * t,j y t,j . Similarly to the step two in Appendix A.2.2, we can simplify the E ∇ Bj y * t,j y t,j as follows.E ∇ Bj y * t,j y t,j = 2Re E e x t ⊗ E x t E x * t ρ t Vec(A) • Vec (A)e x t ⊗ E x t E x * t ρ t • B j = 2Re   E   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t      . (sparsity of B)For any j ∈ [d] and l ∈ [2d], we can calculate thatE x t E x t ρ t Ae x t j = a ρ t t-1 i=1 d r=1 λ -i j λ i-t r ,and recall thatE x t E x * t ρ t Ae x t l = a ρt t-1 i=1 d r=1 λ i l λ t-i r , l ∈ [d], a ρt t-1 i=1 d r=1 λ i-1 l-d λ t-i r , l ∈ [2d] -[d].With careful computing, we haveE   b E x t E x t ρ t Ae x t j • E x t E x * t ρ t Ae x t l   =        a 2 b ρ 2 t (t -1) 2 + (d -1)(t -1) , l = j, a 2 b ρ 2 t (t -1), l ∈ [d] -j, 0,otherwise, which implies that the l-th element of 1 2 E ∇ Bj y * t,j y t,j is1 2 E ∇ Bj y * t,j y t,j l =        a 2 b ρ 2 t (t -1) 2 + (d -1)(t -1) , l = j, a 2 b ρ 2 t (t -1), l ∈ [d] -j, 0,otherwise.Step three: calculate ∇ Bj L t (θ) and ∇ Bj L(θ). Based on steps one and two, the l-th element of ∇ Bj L t (θ) can be derived as follows.∇ Bj L t (θ) l = 1 2 E ∇ Bj y * t,j y t,j l -E ∇ Bj Re x * t+1,j y t,j l =        a 2 b ρ 2 t (t -1) 2 + (d -1)(t -1) -a ρt (t -1), l = j, a 2 b ρ 2 t (t -1), l ∈ [d] -j, 0, otherwise.Furthermore, the l-th element of ∇ Bj L(θ) is∇ Bj L(θ) l = T -1 t=2 ∇ Bj L t (θ) l =        T -1 t=2 a 2 b ρ 2 t (t -1) 2 + (d -1)(t -1) -a ρt (t -1) , l = j, T -1 t=2 a 2 b ρ 2 t (t -1), l ∈ [d] -j, 0,otherwise.(EQUATIONIf we clip the non-diagonal gradient of W P V 12 , we have∇ Bj L(θ) l = T -1 t=2 a 2 b ρ 2 t (t -1) 2 + (d -1)(t -1) -a ρt (t -1) , l = j, 0, otherwise, = a 2 b (T -2) + T -1 t=2 d-1 t-1 -a(T -2), l = j, 0,otherwise.Step four: calculate E ∇ Vec(A) Re x * t+1,j y t,j and d j=1 E ∇ Vec(A) Re x * t+1,j y t,j . Similarly to the step four in Appendix A.2.2, the E ∇ Vec(A) Re x * t+1,j y t,j can be derived as the following.E ∇ Vec(A) Re x * t+1,j y t,j = Re   E x * t+1,j • e x t ⊗ E x t E x t ρ t • B j   = Re   E λ -t j • Vec( E x t E x t ρ t B j e x t )   .For any s, r ∈ [2d], we haveE x t E x t ρ t B j e x t sr =            b ρt ( t i=1 λ i-1 j λ 1-i s )λ t-1 r , s ∈ [d], r ∈ [d], b ρt ( t-1 i=1 λ i j λ 1-i s-d )λ t-1 r , s ∈ [2d] -[d], r ∈ [d], b ρt ( t i=1 λ i-1 j λ 1-i s )λ t-2 r-d , s ∈ [d], r ∈ [2d] -[d], b ρt ( t-1 i=1 λ i j λ 1-i s-d )λ t-2 r-d , s ∈ [2d] -[d], r ∈ [2d] -[d],which implies that Step five: calculate E ∇ Vec(A) y * t,j y t,j and d j=1 E ∇ Vec(A) y * t,j y t,j . Similarly to the step five in Appendix A.2.2, the E ∇ Vec(A) y * t,j y t,j is simplified as follows.E   λ -t j E x t E x t ρ t B j e xE ∇ Vec(A) y * t,j y t,j = 2Re  E e x t ⊗ E x t E x t ρ t B j • B j • e x * t ⊗ E x t E x t ρ t • Vec(A)   = 2Re   E   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • Vec( E x t E x t ρ t B j e x t )      . For any k ∈ [2d] -[d],we can calculate thatE x t E x * t ρ t B j e x * t k,k-d = b ρ t ( t-1 i=1 λ -i j λ i-1 k-d )λ 1-t k-d = b ρ t t-1 i=1 λ -i j λ i-t k-d ,and recall that for any s, r ∈ [2d], we haveE x t E x t ρ t B j e x t sr =            b ρt ( t i=1 λ i-1 j λ 1-i s )λ t-1 r , s ∈ [d], r ∈ [d], b ρt ( t-1 i=1 λ i j λ 1-i s-d )λ t-1 r , s ∈ [2d] -[d], r ∈ [d], b ρt ( t i=1 λ i-1 j λ 1-i s )λ t-2 r-d , s ∈ [d], r ∈ [2d] -[d], b ρt ( t-1 i=1 λ i j λ 1-i s-d )λ t-2 r-d , s ∈ [2d] -[d], r ∈ [2d] -[d].With careful computing, we haveE   2d k=d+1 a E x t E x * t ρ t B j e x * t k,k-d • E x t E x t ρ t B j e x t sr   =                        ab 2 ρ 2 t (t -1) 2 , s = d + j, r = j, ab 2 ρ 2 t (t -1), s = d + j, r = j, ab 2 ρ 2 t (t -1), s = d + r, r = j, ab 2 ρ 2 t (t -1), s = d + j, r = j, ab 2 ρ 2 t , remains in W KQ 32 , 0,otherwise, which implies that the l-th element of 1 2 E ∇ A y * t,j y t,j is1 2 E ∇ A y * t,j y t,j =                        ab 2 ρ 2 t (t -1) 2 , s = d + j, r = j, ab 2 ρ 2 t (t -1), s = d + j, r = j, ab 2 ρ 2 t (t -1), s = d + r, r = j, ab 2 ρ 2 t (t -1), s = d + j, r = j, ab 2 ρ 2 t , remains in W KQ 32 , 0, otherwise. 1 2 d j=1 E ∇ A y * t,j y t,j =        ab 2 ρ 2 t (t -1) 2 + (d -1)(t -1) , s -d = r, A sr ∈ W KQ 32 , ab 2 ρ 2 t 2(t -1) + d -2 , s -d = r, A sr ∈ W KQ 32 , 0, otherwise.Step six: calculate ∇ Vec(A) L t (θ) and ∇ Vec(A) L(θ). Based on steps four and five, the (s, r)-th element of ∇ A L t (θ) can be derived as follows.∇ A L t (θ) sr = 1 2 d j=1 E ∇ A y * t,j y t,j - d j=1 E ∇ A Re x * t+1,j y t,j =        ab 2 ρ 2 t (t -1) 2 + (d -1)(t -1) -b ρt (t -1), s -d = r, A sr ∈ W KQ 32 , ab 2 ρ 2 t 2(t -1) + d -2 -b ρt , s -d = r, A sr ∈ W KQ 32 , 0, otherwise.Furthermore, the (s, r)-th element of ∇ A L(θ) isEQUATIONIf we clip the non-diagonal gradient of W KQ 32 , we have∇ A L(θ) sr = T -1 t=2 ∇ A L t (θ) sr = T -1 t=2 ab 2 ρ 2 t (t -1) 2 + (d -1)(t -1) -b ρt (t -1) , s -d = r, A sr ∈ W KQ 32 , 0, otherwise, = ab 2 (T -2) + T -1 t=2 d-1 t-1 -b(T -2), s -d = r, A sr ∈ W KQ 32 , 0, otherwise.Step seven: summarize the result by induction. From the gradient of ∇ A L(θ) and ∇ Bj L(θ), we observe that non-zero gradients only emerge in the diagonal of W KQ 32 and W P V 12 , and they are same. Therefore, the parameter matrices keep the same structure as the initial time. Thus we can summarize the dynamic system as the following.d dτ a = -ab 2   (T -2) + T -1 t=2 d -1 t -1   + b(T -2), d dτ b = -a 2 b   (T -2) + T -1 t=2 d -1 t -1   + a(T -2),which completes the proof.Lemma A.6. Suppose that the precondtions of Theorem 4.3 hold, and denote(T -2) + T -1 t=2 d-1 t-1and T -2 by c 1 and c 2 , respectively. Then, the dynamics are the same as those of gradient flow on the following objective function:(a, b) = 1 2c 1 (c 2 -c 1 ab) 2 ,whose global minimums satisfy ab = c 2 /c 1 . For the reader'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 in Theorem 4.1, thus the trained transformer will not implement one step of gradient descent for minimizing1 2 t-1 i=1 x i+1 -W x i 2 .Proof. From Eq. 6 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 32 and W P V W KQ and W P V with different diagonal initializations are presented in Figure 7 . From the results, we know that the gradient descent converges to  W KQ 22 W KQ 23 W KQ 32 W KQ 33   = 0 d×d 0 d×d W 1 0 d×d , W P V 12 W P V 13 = W 2 0 d×d ,where W 1 and W 2 are some dense matrices. Similarly to the proof of Corollary 4.1 in Appendix 4.1, we havey t = W P V 12 W P V 13 E x t E x * t ρ t W KQ 22 W KQ 23 W KQ 32 W KQ 33 e x t = W 2 0 d×d E x t E x * t ρ t 0 d×d 0 d×d W 1 0 d×d e x t = 1 ρ t W 2 0 d×d t i=1 e x i e x * i 0 d×d 0 d×d W 1 0 d×d e x t = 1 ρ t t i=1 W 2 x i x * i-1 W 1 0 d×d e x t = 1 ρ t t i=1 W 2 x i x * i-1 W 1 x t .which can be seen as a somewhat preconditioned gradient descent on the OLS problem. • If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).will occur, which implies the result. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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?
questions: 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: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_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.
- 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.
- 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.
- 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.
- They show that, without the distributional assumption, the trained model may not implement a gradient descent step.
- Lastly, they verify their theory with simulations.
strength_weakness: strength_weakness: - W1. The analyses and experiments cannot provide insight into the general cases of autoregressive training of Transformers.
    - 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.
    - 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”…).
- W2. There are some weaknesses in theoretical analyses.
    - 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.
    - For the comments and questions on Proposition 4.1, see Q3.
- W3. Minor comments on typos & mistakes.
    - Line 130: $W^{KQ} = {W^K}^\ast W^Q$ (rather than transpose)
    - Equation (2): I guess the sum should be applied over $t=1, …, T-1$.
    - Line 182: “$x_1 \in \mathbb{R}^d$”
    - Lines 207, 289: ‘$W$’ must be bold.
    - Figures 1b, 1d, and Line 330: “$T_{te} -1$”
questions: 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.
- 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?
- Q3. Questions on Proposition 4.1
    - 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)$.
    - Isn’t $T_{tr} \rightarrow \infty$ also unnecessary? This is because it seems enough to show that $\tilde{a} \tilde{b} < \tfrac{1}{5}$.
    - 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$.
    - 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.
- Q4. I guess Theorem 4.2 requires Assumption 3.1. Or does it?

If the authors address all my concerns, I will definitely raise my score.
limitations: limitations: The limitations are provided in Section 7.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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]. 

[5] Vladymyrov, Max, Johannes von Oswald, Mark Sandler and Rong Ge. Linear Transformers are Versatile In-Context Learners. ArXiv abs/2402.14180 (2024)
questions: 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

    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. 

    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.

    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.

    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.

    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.

    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.

    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.
questions: questions: Why is the y-axis in Figure 1-b defined as the ratio of two vectors?
limitations: limitations: The authors discuss limitations.

Comments: title: Summary of the revision
Reply to rebuttal
Thank you for the rebuttal
comment: We thank the reviewer once again for the insightful comments and positive score, which give us much confidence to finish this rebuttal process.
# Summary of the revision

We sincerely thank all reviewers for their valuable comments and positive support, which help us further improve the quality of this paper. We also thank ACs for organizing this nice review process. We have thoroughly addressed the detailed comments and summarized the revision in the final version as follows.

## New experimental results
1. We add complete experimental results under standard initialization with different scales ($\sigma_w = 0.001/0.01/0.1$). Representatives of them have been presented in the latest uploaded PDF (Fig. a-f).
2. We add complete experimental results in small-context scenarios ($T = 3/5/10,d=5$). Representative results have been presented in the latest uploaded PDF (Fig. g-i).

## Refined theoretical result
We will refine the theory and proof of Proposition 4.1 as we presented in the answer to Reviewer dsNs.

## Writing
1. We fix the typos pointed out by Reviewer dsNs.
2. We clarify that Proposition 4.2 does not depend on diagonal initialization
3. We change "gradient clipping" to "gradient masking" in Section 4.3 to avoid confusion.
4. We clarify the ratio calculation between prediction and true vector in Fig. 1b.

## Discussion

### Contribution & Impact
1. We add more discussion about the key contribution of this paper in Section 4.
2. We add more discussion about our technical improvement in Sections 4 and 5, compared with existing studies on ICL.
3. We add more discussion about the contribution and the necessity of finding Assumption 4.1 to keep the diagonal structure, compared with existing studies. We also discuss its limitations.
4. We add more discussion about the negative results of one-layer linear attention, which can be solved when considering more complex transformers in the future.
5. We add more discussion about the impact of our work on more complex autoregressive transformers (e.g., multi-layer attention, non-linear activations).

### Setting & Assumptions
1. We add more discussion about the rationality of our setting (model, data distribution, token construction, and normalization factor).
2. We add more discussion about the necessity and limitation of the diagonal initialization near Assumption 3.1 and Theorem 4.1.
3. We discuss how the experimental results under standard initialization extend our theory and support the rationality of the diagonal initialization in Section 6.
4. We add more discussion about the properties of Assumption 4.2.

### Related work
1. We cite and discuss the additional related papers mentioned by Reviewer WaqB.
Thanks for the helpful comments and positive score. We are glad that we have addressed the reviewer's concerns.
I thank the authors for the explanation. I don't have further questions.
Thanks for the helpful comments and positive feedback, as well as for increasing the score. We are glad that we have addressed the reviewer's main concerns.

For the claim of theoretical contribution, we agree with the reviewer's suggestion. **We promise that we will put more emphasis on our results contributing a closed-form characterization of the limiting points of AR pre-training**. We will also **discuss the necessity and the limitation of the diagonal initialization more clearly** in the final version.

For the complete experimental results in small-context ($T=3/10$) and standard-initialization ($\sigma_w=0.001/0.1$) scenarios, **the claims in the common response still hold, and we will add them to the final version**. We note that due to the space limitation, we only present the representative result in each scenario during the rebuttal period.

We thank the reviewer once again for the valuable suggestions and insightful comments, which definitely improve the quality of this paper a lot.
I thank the reviewers for their thorough response, which has helped to alleviate my concerns. I especially appreciate the experimental results with non-diagonal initialization, pointers to related work and clarification on Proposition 4.2. I am still not overly impressed by the results, especially as they start with diagonal initialization (Theorem 4.1 and Theorem 4.3), and use gradient masking (Theorem 4.3). Nevertheless, as the authors point out, even in these simplified settings characterizing the limiting points of gradient flow is a novel contribution. I would recommend that in the revision, the authors put more emphasis on their results contributing a closed-form characterization of the convergence points of AR pre-training, rather than showing that AR pre-training results in a mesa-optimizer, since the latter point is weakened by the assumption that the initialization is already a mesa-optimizer. As a minor point, the empirical results with T=3 and T=10 appear to be missing from the attached pdf.
We are glad that we have addressed most of the reviewer's concerns. For Q1, we will follow the reviewer's suggestion and further modify the claim to "any random vectors whose coordinates are i.i.d. random variables with zero mean and finite moments of order 2, 4, and 6 satisfy this assumption".

We thank the reviewer once again for the valuable suggestions and insightful comments, which improve the quality of this paper.
Thank you for your detailed response. 

- The general response and the additional experiments on standard initialization have addressed my concerns (W1 & W2). I also agree that finding Assumption 4.1 is non-trivial.
- Also, I am happy to see that the refined Proposition 4.1 and its proof sketch have satisfactorily addressed my concern on the theoretical contribution.
- Regarding the response to Q1, I think it is important to mention which orders of moments are finite.

Overall, most of my concerns are addressed. Although I am still not absolutely certain, now I think the paper is above the acceptance bar. Hence I will raise my score to 5.