Title: On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability
Venue: NeurIPS 2024
Accepted: yes
Keywords: Mesa-Optimization

In-context learning

Autoregressive pretraining

Non-convex optimization

Learning theory

Transformers
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

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 selfattention model autoregressively trained by gradient flow, where the sequences are generated by an AR process x t +1 = Wx 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 ̂ 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, and the code is available at https://github.com/ML-GSAI/MesaOpt-AR-Transformer .

1 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.

Recently, we are aware that some preliminary studies [16; 17] have attempted to understand the ICL ability from the AR training and connected its mechanisms to a popular hypothesis named mesaoptimization [18], which suggests that transformers learn some algorithms during the AR pretraining. In other words, the inference process of the trained transformers is equivalent to optimizing some inner objective functions on the in-context data.

Concretely, the seminal work [16] constructs a theoretical example where a single linear causallymasked self-attention layer with manually set parameters can predict the next token using onestep gradient-descent 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-based mesa-optimization algorithms. Built upon the setting of [16], recent work [17] precisely characterizes the pretraining loss landscape of the onelayer linear transformer trained on a simple first-order AR process with a fixed full-one initial token. As a result, they find that the optimally-trained 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 [19; 20; 21; 22] and theory [23].

In this paper, we take a further step toward understanding the mesa-optimization in autoregressively trained transformers. Specially, without an explicit diagonal structure assumption, we analyze the non-convex dynamics of a one-layer linear transformer trained by gradient flow on a controllable first-order AR process, and try to answer the following questions rigorously:

When do mesa-optimization algorithms emerge in autoregressively trained transformers?

What is the capability limitation of the mesa-optimizer if it does emerge?

Our first main contribution is to characterize a 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 x 1 whose coordinates x 1 i are i.i.d. random variables with zero mean and finite moments satisfy this condition, including normal distribution N ( 0 d , σ 2 I d ) . Under this assumption, the non-convex dynamics will exactly converge to the theoretical construction in [16] 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 mesaoptimizer under Assumption 4.1 in Section 4.2. We characterize a stronger assumption (Assumption 4.2) related to the moment of data distribution as the necessary and sufficient condition that the mesa-optimizer recovers the true distribution (a.k.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 ICL by few-shot learning pretraining [24; 25; 26; 27; 28; 29; 30; 23] (see details in Section 2), a setup attracting attention from many theorists recently. We think ICL in the setting of AR pretraining needs more attention from the theoretical community.

In Section 4.3, 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 a 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 our theoretical findings in Section 6.

2 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; 29; 34; 30] 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 [31; 35] considers

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.

2.2 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.

3 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 glyph[circledot] to denote the Kronecker product and the Hadamard product, respectively. Besides, we denote Vec( · ) the vectorization operator in column-wise order.

3.1 Data distribution

We consider a first-order AR process as the underlying data distribution, similar to recent works on ARpretraining [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 x 1 to be specified later, then the subsequent elements are generated according to the rule x t +1 = Wx t for t ∈ [ T -1] . For convenience, we denote the vector ( λ 1 , . . . , λ d ) glyph[latticetop] 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 glyph[latticetop] x = 1 glyph[latticetop] 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 x 1 on the convergence result of the AR pretraining.

3.2 Model details

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 × d e , a key matrix W K ∈ R d k × d e , a value matrix W V ∈ R d v × d e , a projection matrix W P ∈ R d e × d v and a normalization factor ρ t > 0 . At time step t , let E t = ( e 1 , . . . , e t ) ∈ R d e × t be the tokens embedded from the prompt sequence

P t = ( x 1 , . . . , x t ) ∈ C d × t , the causal self-attention layer will output

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 PV = W P W V , we have θ = ( W KQ , W PV ) , and the output can be rewritten as:

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 glyph[latticetop] d , x glyph[latticetop] i , x glyph[latticetop] i -1 ) glyph[latticetop] ∈ C 3 d , thus the corresponding embedding matrix E t can be formally written as:

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 have

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 PV and W KQ . Concretely, we denote that

where W PV ij , W KQ ij ∈ R d × d for all i, j ∈ [3] . Then the ̂ y t can be written as

Here E x t = ( e x 1 , . . . , e x t ) ∈ C 2 d × t denotes the last 2 d rows of the E t , where e x i = ( x glyph[latticetop] i , x glyph[latticetop] i -1 ) glyph[latticetop] . 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.

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:

where 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 that

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 ( θ ) .

4 Main results

In this section, we present the main theoretical results of this paper. First, in Section 4.1, we prove that when D x 1 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.

4.1 Trained transformer is a mesa-optimizer

In this subsection, we show that under a certain assumption of D x 1 , 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 x 1 can be summarized as follows.

glyph[negationslash]

Assumption 4.1 (Sufficient condition for the emergence of mesa-optimizer) . We assume that the distribution D x 1 of the initial token x 1 ∈ R d satisfies E x 1 ∼D x 1 [ x 1 i 1 x r 2 1 i 2 · · · x r n 1 i n ] = 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 1 j ] , κ 2 = E [ x 6 1 j ] and κ 3 = ∑ r = j E [ x 2 1 j x 4 1 r ] 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 PV 12 as zero, thus they can keep diagonal structure during the training. We note that any random vectors x 1 whose coordinates x 1 i 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.

Theorem 4.1 (Convergence of the gradient flow, proof in Section 5) . Consider the gradient flow of the one-layer linear transformer (see Eq. 1) over the population AR pretraining loss (see Eq. 2).

Suppose the initialization satisfies Assumption 3.1, and the initial token's distribution D x 1 satisfies Assumption 4.1, then the gradient flow converges to

Though different initialization ( a 0 , b 0 ) lead to different ( ˜ a, ˜ b ) , the solutions' product ˜ a ˜ b satisfies

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 -Wx i ‖ 2 , starting from the initialization W = 0 d × d with a step size ˜ a ˜ b t -1 .

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 T te is

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 D x 1 , 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, we present a negative result, which proves that not all D x 1 satisfies Assumption 4.1 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 D x 1 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 satisfies

Therefore, the prediction ̂ y T te will not converges to the true next token Wx T te .

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], multi-head attention [52; 53; 54], deeper attention layers [55; 16], transformer block [56; 57; 58], 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 a 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 x 1 is needed. Under Assumption 4.1, the following sufficient and necessary condition related to the moment of D x 1 is derived from Eq. 3 by letting ̂ y T te converges to Wx T te when context length T tr and T te are large enough.

Assumption 4.2 (Condition for success of mesa-optimizer) . Based on Assumption 4.1, we further suppose that κ 1 κ 2 ∑ Tte -1 i =1 x i x ∗ i T te -1 x T te → x T te for any x 1 and W , when T te is large enough.

Assumption 4.2 is 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 D x 1 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 x 1 ∈ R d is uniformly sampled from the candidate set of size 2 d {± ( c, 0 , . . . , 0) glyph[latticetop] , ± (0 , c, . . . , 0) glyph[latticetop] , ± (0 , . . . , 0 , c ) glyph[latticetop] } for any fixed c ∈ R , then the distribution D x 1 satisfies Assumption 4.2. The derivation can be found in Appendix A.5.

For completeness, we formally summarize the following distribution learning guarantee for the trained transformer under Assumption 3.1 and 4.1.

Theorem 4.2 (Trained transformer succeed to learn the distribution satisfies Assumption 4.2, proof in Appendix A.6) . Suppose that Assumption 3.1 and 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

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 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 as 1 d .

First, sharing the similar but weaker assumption of [17], we impose W KQ 32 and W PV 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 PV 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, with

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 PV 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 problem 1 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 PV 12 anymore. When W KQ 32 and W PV 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.

5 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.

where the A and B j are defined as

with a i ∈ R 2 d and b j 1 , b j 2 ∈ 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 PV 12 satisfies

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.

Lemma 5.3 (Surrogate objective function, proof in Appendix A.2.3) . Suppose that Assumption 4.1 holds and denote ( T -2) κ 2 + ∑ T -1 t =2 1 t -1 κ 3 and ( T -2) κ 1 by c 1 and c 2 , respectively. Then, the dynamics in Lemma 5.2 are the same as those of gradient flow on the following objective function:

whose global minimums satisfy ab = c 2 /c 1 .

Furthermore, We show that although the objective ˜ glyph[lscript] ( 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 ˜ glyph[lscript] ( a, b ) is a non-convex function and satisfies the PL inequality as follows.

Therefore, the gradient flow in Lemma 5.2 converges to the global minimum of ˜ glyph[lscript] ( 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 our theoretical results. In terms of the train set, we generate 10 k sequences with T tr = 100 and d = 5 . In addition, we generate another test set with 10 k 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 T te -1 converges to x T te / 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 T te -1 converges to the truth x T te , 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 PV 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.

̂

(a) Gaussian with

σ

= 0

.

5

, dynamics of

ab

.

σ

= 0

.

5

, ratio of

̂

y

T

te

-

1

/

x

T

te

4.0

(ao,bo)=(0.1,0.1)

0.8

(ao,bo)=(0.1,0.1)

Dynamics of ab

3.0

3.5

(ao,bo)=(2.0,2.0)

(ao,bo)=(0.5,1.5)

Ratio of pred/true at

T

(ao,bo)=(2.0,2.0)

(ao,bo)= (0.5, 1.5)

2.5-

1/5o2 = 0.8

0.6

Ratio =0.2

2.0

0.4

1.5

1.0

0.2

0.5

0.0

0.0

0

25

50

75

100125150175200

0

25

50

75

100125150175200

Epoch

Epoch

4.0

12

(ao,bo)=(0.1,0.1)

3.5

(ao,bo)=(0.5,1.5)

Dynamics of ab

10

3.0

(ao,bo)=(2.0, 2.0)

2.5

¥8

Theoretical MSE=0

2.0

loss

1.5

(ao,bo)=(0.1, 0.1)

MSE

1.0

(ao,bo)=(0.5,1.5)

0.5

(ao,bo)=(2.0,2.0)

2

1/c2 = 4.0

0.0

0

0

25

50

75

100

125

150

175

200

0

25

50

75

100125150175200

Epoch

Epoch

(c) Example 4.1 with

c

= 0

.

5

, dynamics of ab

.

(d) Example 4.1 with

c

= 0

.

5

,

‖

y

T

te

-

1

-

x

T

te

.

2

2

‖

.

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, 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.

References

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-6008, 2017.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, 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 NeurIPS , 2020.

OpenAI. GPT-4 technical report. CoRR , abs/2303.08774, 2023.

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, Xiaoqing 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.09288, 2023.

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):1113, 2023.

Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICML , pages 8821-8831, 2021.

Doyup Lee, Chiheon Kim, Saehoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In CVPR , pages 11523-11532, 2022.

Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. arXiv preprint arXiv:2404.02905 , 2024.

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.

Zhengyuan Yang, Linjie Li, Kevin Lin, Jianfeng Wang, Chung-Ching Lin, Zicheng Liu, and Lijuan Wang. The dawn of lmms: Preliminary explorations with gpt-4v(ision). CoRR , abs/2309.17421, 2023.

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.

Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. In NeurIPS , 2023.

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.

Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy P. Lillicrap, JeanBaptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, Ioannis 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 McIlroy, 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 et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. CoRR , abs/2403.05530, 2024.

Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimization algorithms in transformers. CoRR , abs/2309.05858, 2023.

Michael Eli Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do transformers perform in-context autoregressive learning ? In ICML , 2024.

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.

Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya K. Singh, Pierre H. Richemond, James L. McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In NeurIPS , 2022.

Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and llms by learning to learn discrete functions. In ICLR , 2024.

Kartik Ahuja and David Lopez-Paz. A closer look at in-context learning under distribution shifts. CoRR , abs/2305.16704, 2023.

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.

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.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55, 2024.

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.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In NeurIPS , 2023.

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.

Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In ICLR , 2024.

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.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In ICML , 2024.

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.

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.

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.

Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge. Linear transformers are versatile in-context learners. CoRR , abs/2402.14180, 2024.

Samuel Müller, Noah Hollmann, Sebastian Pineda-Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In ICLR , 2022.

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. CoRR , abs/2310.17086, 2023.

Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How well can transformers emulate in-context newton's method? CoRR , abs/2403.03183, 2024.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In ICLR , 2022.

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.11916, 2023.

Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. In NeurIPS , 2023.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In ICML , volume 202, pages 19689-19729, 2023.

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In ICML , volume 202, pages 19565-19594, 2023.

Christos Thrampoulidis. Implicit bias of next-token prediction. CoRR , abs/2402.18551, 2024.

Yingcong 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.

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.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In NeurIPS , pages 7411-7422, 2019.

Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the ddpm objective. NeurIPS , 36:19636-19649, 2023.

Haiyun He, Hanshu Yan, and Vincent YF Tan. Information-theoretic characterization of the generalization error for iterative semi-supervised learning. Journal of Machine Learning Research , 23(287):1-52, 2022.

Ludwig Schmidt, Shibani Santurkar, Dimitris Tsipras, Kunal Talwar, and Aleksander Madry. Adversarially robust generalization requires more data. NeurIPS , 31, 2018.

Chenyu Zheng, Guoqiang Wu, and Chongxuan Li. Toward understanding generative data augmentation. NeurIPS , 36, 2024.

Chenyu Zheng, Guoqiang Wu, Fan Bao, Yue Cao, Chongxuan Li, and Jun Zhu. Revisiting discriminative vs. generative classifiers: Theory and implications. In ICML , volume 202, pages 42420-42477. PMLR, 2023.

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.

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.

Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. CoRR , abs/2310.12680, 2023.

Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. CoRR , abs/2402.14735, 2024.

Ruiqi Zhang, Jingfeng Wu, and Peter L. Bartlett. In-context learning of a linear transformer block: Benefits of the MLP component and one-step GD initialization. CoRR , abs/2402.14951, 2024.

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.15607, 2024.

Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In ICML , 2024.

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In ICLR , 2014.

Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal Vychislitel'noi Matematiki i Matematicheskoi Fiziki , 3(4):643-653, 1963.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition. In ECML PKDD , volume 9851, pages 795-811, 2016.

Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark , 7(15):510, 2008.

Contents of Appendix

Appendix A Proofs

A.1 Proof of eq. (1)

Proof. We first calculate the output by causal linear attention layer as the following: f t ( E t ; θ )

where × s are the elements that will not contribute to the final ̂ y t . Afurther simple computation shows that

which completes the proof.

A.2 Proof of Theorem 4.1

A.2.1 Proof of Lemma 5.1

For the reader's convenience, we restate the lemma as the following.

Lemma A.1. Each element of the network's prediction ̂ y t,j ( j ∈ [ d ] ) can be expressed as the following.

where the A and B j are defined as

with a i ∈ R 2 d and b j 1 , b j 2 ∈ R d .

Proof. Based on the result in Eq. 1, we can write

which finishes the proof.

A.2.2 Proof of Lemma 5.2

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 PV 12 satisfies

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 as

Then, we can calculate the derivatives of L t ( θ ) with respect to B j and Vec( A ) as

and

Step one: calculate E [ ∇ B j Re ( x ∗ t +1 ,j ̂ y t,j ) ] . Based on Lemma 5.1, we have ̂ y t,j = B glyph[latticetop] j · e x glyph[latticetop] t ⊗ E x t E x ∗ t ρ t · Vec( A ) .

where the penultimate equality uses the property of Kronecker and Vec operator Vec( AXB ) = ( B glyph[latticetop] ⊗ A )Vec( X ) , we refer Section 10.2 in [62] for details.

For E x t E x ∗ t ρ t , we can simplify it as

Based on the diagonal property of W , we can simplify the x i x ∗ i as the following.

where we define ̂ Σ = x 1 x ∗ 1 and M i = λ i -1 λ i -1 ∗ . Therefore, we have

Then, leveraging the sparse property of A , we can derive E x t E x ∗ t ρ t Ae x t as follows.

Therefore, for any l ∈ [ d ] , we have

Similarly, for any l ∈ [2 d ] -[ d ] , we have

To sum up, for any l ∈ [2 d ] , we have

We discuss them in the following categories, glyph[negationslash]

1. l = j . In this case, E [ x 1 j x 1 l x 2 1 r ] = 0 by Assumption 4.1, thus

Similarly, for any l ∈ [2 d ] -[ d ] , we have

Therefore, for any l ∈ [2 d ] , we have

and the l -th element of E [ ∇ B j Re ( x ∗ t +1 ,j ̂ y t,j ) ] is

glyph[negationslash]

Step two: calculate E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] . Based on Lemma 5.1, we have

then we can simplify the E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] as follows.

glyph[negationslash]

We further derive that

For any and l ∈ [2 d ] , recall that

and for any j ∈ [ d ] , we have

For any l ∈ [ d ] , with careful computing, we have

We discuss it in the following categories, glyph[negationslash]

l = j . In this case, E [ x 1 j x il x 2 1 r 1 x 2 1 r 2 ] = 0 by Assumption 4.1, thus it becomes 0 .

l = j, r 1 = r 2 = j . It becomes

glyph[negationslash]

l = j, r 1 = r 2 = r = j . It becomes

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

l = j, r 1 = r 2 . In this case, E [ λ -i 1 j λ i 2 l λ i 1 -t r 1 λ t -i 2 r 2 ] = E [ λ i 2 -i 1 j λ i 1 -t r 1 λ t -i 2 r 2 ] = 0 , thus it becomes 0 .

Similarly, for any l ∈ [2 d ] -[ d ] , we have

glyph[negationslash]

To sum up, we have

which implies that the l -th element of 1 2 E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] is

glyph[negationslash]

Step three: calculate ∇ B j L t ( θ ) and ∇ B j L ( θ ) . Based on steps one and two, the l -th element of ∇ B j L t ( θ ) can be derived as follows.

glyph[negationslash]

Furthermore, the l -th element of ∇ B j L ( θ ) is

.

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

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 have

In terms of E x t E x glyph[latticetop] t ρ t B j e x glyph[latticetop] t , based on the sparsity of B and Eq. 4, we can derive that

Then, for any s, r ∈ [2 d ] , we have

We discuss it in the following categories, glyph[negationslash]

glyph[negationslash]

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 that

To sum up, we have

E   λ -t j ( E x t E x glyph[latticetop] t ρ t B j e x glyph[latticetop] t ) sr   = { b ρ t ( t -1) E [ x 4 1 j ] , 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 1 j ] , s = d + j, r = j, 0 , otherwise , = { b ρ t ( t -1) κ 1 , s = d + j, r = j, 0 , otherwise . Finally, we can calculate ∑ d 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 have

then we can simplify the

We further derive that

Recall that for any s, r ∈ [2 d ] , we have

Furthermore, for any k ∈ [2 d ] -[ d ] , we can calculate that

glyph[negationslash]

With careful computing, for any s ∈ [2 d ] -[ d ] , r ∈ [ d ] , we have

We discuss it in the following categories, glyph[negationslash]

1. s -d = r . In this case, E [ x 2 1 j x 2 1 ,k -d x 1 ,s -d x 1 r ] = 0 by Assumption 4.1, thus it becomes 0 .

2. s -d = r = k -d = j . It becomes

3. s -d = r = k -d = j . It becomes

glyph[negationslash]

4. s -d = r = k -d . In these case, E [ λ i 2 -i 1 j λ i 1 -t k -d λ 1 -i 2 s -d λ t -1 r ] = E [ λ i 2 -i 1 j λ i 1 -t k -d λ t -i 2 r ] = 0 , thus it becomes 0 .

Similarly, for any other s, r , we can calculate that

To sum up, we have

glyph[negationslash]

which implies that the ( s, r ) -th element of 1 2 E [ ∇ A ̂ y ∗ t,j ̂ y t,j ] is glyph[negationslash]

glyph[negationslash]

Based on these results, we can derive

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.

Furthermore, the ( s, r ) -th element of ∇ A L ( θ ) is

Step seven: summarize the result by induction. From the gradient of ∇ A L ( θ ) and ∇ B j L ( θ ) , we observe that non-zero gradients only emerge in the diagonal of W KQ 32 and W PV 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.

which completes the proof.

A.2.3 Proof of Lemma 5.3

For the reader's convenience, we restate the lemma as the following.

LemmaA.3. Suppose that Assumption 4.1 holds and denote ( T -2) κ 2 + ∑ T -1 t =2 1 t -1 κ 3 and ( T -2) κ 1 by c 1 and c 2 , respectively. Then, the dynamics in Lemma 5.2 are the same as those of gradient flow on the following objective function:

whose global minimums satisfy ab = c 2 /c 1 .

Proof. Basic calculus shows that

Therefore, the dynamics in Lemma 5.2 are the same as those of gradient flow on ˜ glyph[lscript] ( a, b ) , whose global minimums satisfy ab = c 2 /c 1 .

A.2.4 Proof of Lemma 5.4

For the reader's convenience, we restate the lemma as the following.

Lemma A.4. Suppose that Assumption 4.1 holds, then ˜ glyph[lscript] ( a, b ) is a non-convex function and satisfies the PL inequality as follows.

Therefore, the gradient flow in Lemma 5.2 converges to the global minimum of ˜ glyph[lscript] ( a, b ) .

Proof. First, we prove that ˜ glyph[lscript] ( a, b ) is non-convex. The Hessian matrix of ˜ glyph[lscript] ( a, b ) can be derived as follows.

Its determinant c 1 a 2 b 2 -(2 c 1 ab -c 2 ) 2 = ( c 2 -c 1 ab )(3 c 1 ab -c 2 ) < 0 when ab < c 2 3 c 1 or ab > c 2 c 1 . Thus, ˜ glyph[lscript] ( a, b ) is non-convex.

Besides, the PL inequality holds because

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.1 holds. When predicting the t -th token, the trained transformer implements one step of gradient descent for the minimization of the OLS problem L 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.

and

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 .

Thus, the proof is completed.

A.4 Proof of Proposition 4.1

For the reader's convenience, we restate the proposition as the following.

Proposition A.1. Let D x 1 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 satisfies

Therefore, the prediction ̂ y T te will not converges to the true next token Wx T te .

Proof. First, built upon the results in Theorem 4.1, when T tr is large enough, we have

Second, by directly calculating, we have

Therefore, we have

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.1

Proof. 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 have E x 1 ∼D x 1 [ x 1 i 1 x r 2 1 i 2 · · · x r n 1 i n ] = E x 1 ∼D x 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 glyph[negationslash]

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 .

The proof is finished.

A.6 Proof of Theorem 4.2

For 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

Proof. First, built upon the results in Theorem 4.1, when T tr is large enough, we have

Second, when T te is large enough, by Assumption 4.2

Therefore, we have

which finishes the proof.

A.7 Proof of Theorem 4.3

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 PV 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, with

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 PV 12 satisfies

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 as

Besides, the derivatives of L t ( θ ) with respect to Vec( A ) and B j are

and

Step one: calculate E [ ∇ B j Re ( x ∗ t +1 ,j ̂ y t,j ) ] . Similarly to the step one in Appendix A.2.2, the E [ ∇ B j Re ( x ∗ t +1 ,j ̂ y t,j ) ] can be derived as the following.

We note that for any l ∈ [2 d ] , based on the sparsity of A , we have

which implies that

and the l -th element of glyph[negationslash]

glyph[negationslash]

Step two: calculate E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] . Similarly to the step two in Appendix A.2.2, we can simplify the E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] as follows.

For any j ∈ [ d ] and l ∈ [2 d ] , we can calculate that

and recall that

With careful computing, we have

which implies that the l -th element of 1 2 E [ ∇ B j ̂ y ∗ t,j ̂ y t,j ] is

Step three: calculate ∇ B j L t ( θ ) and ∇ B j L ( θ ) . Based on steps one and two, the l -th element of ∇ B j L t ( θ ) can be derived as follows.

Furthermore, the l -th element of ∇ B j L ( θ ) is

If we clip the non-diagonal gradient of W PV 12 , we have

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.

For any s, r ∈ [2 d ] , we have

which implies that

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

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.

For any k ∈ [2 d ] -[ d ] , we can calculate that

and recall that for any s, r ∈ [2 d ] , we have

With careful computing, we have

which implies that the l -th element of 1 2 E [ ∇ A ̂ y ∗ t,j ̂ y t,j ] is

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

glyph[negationslash]

,

Based on these results, we can derive

glyph[negationslash]

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.

glyph[negationslash]

Furthermore, the ( s, r ) -th element of ∇ A L ( θ ) is

If we clip the non-diagonal gradient of W KQ 32 , we have

Step seven: summarize the result by induction. From the gradient of ∇ A L ( θ ) and ∇ B j L ( θ ) , we observe that non-zero gradients only emerge in the diagonal of W KQ 32 and W PV 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.

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 -1 and 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:

whose global minimums satisfy ab = c 2 /c 1 .

glyph[negationslash]

Proof. The proof is the same as that of Lemma 5.3 in Appendix A.2.3.

Using the results from the above lemmas and Lemma 5.4, we can conclude Theorem 4.3.

A.8 Proof of Proposition 4.2

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 minimizing 1 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 PV 12 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 are done on a single GeForce 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 1.

B.3 Additional results for Gaussian initial token

In practice, we estimate the ratio of two vectors by calculating the mean of the element-wise divide between two vectors. The results for σ = 1 , 2 and T = 100 with diagonal initialization are presented in Figure 2. The results for σ = 1 and T = 5 (small-context scenarios) with diagonal initialization are presented in Figure 4. The results for σ = 1 and T = 100 with Gaussian initialization ( σ w = 0 . 001 , 0 . 01 , 0 . 1 ) are presented in Figure 5.

B.4 Additional results for Sparse initial token (Example 4.1)

The results for c = 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 ( σ w = 0 . 001 , 0 . 01 , 0 . 1 ) are presented in Figure 6.

(a) Gaussian with

σ

= 1

, dynamics of

ab

.

σ

= 1

, ratio of

̂

y

T

te

-

1

/

x

T

te

.

4.0

(ao,bo)=(0.1,0.1)

0.7

(ao,bo)=(0.1,0.1)

Dynamics of ab

3.5

3.0

(ao,bo)=(2.0,2.0)

(ao,bo)=(0.5,1.5)

Ratio of pred/true at

0.6

(ao,bo)= (0.5, 1.5)

(ao,bo)= (2.0,2.0)

0.5

2.5-

1/5o2 = 0.2

Ratio = 0.2

0.4

2.0-

1.5

0.3

0.2

1.0

0.5

0.11

0.0

0.0

0

20

40

60

80

100

120

140

0

20

40

60

80

100

120

140

Epoch

Epoch

1.75

0.4

(ao,bo)= (0.1,0.1)

(ao,bo)=(0.1,0.1)

1.50

(ao,bo)=(0.5,1.5)

(ao,bo)=(0.5,1.5)

0.3

1/5o2 = 0.05

(ao,bo)= (2.0,2.0)

pred/true a

1.25

(ao,bo)=(2.0,2.0)

Ratio =0.2

1.00

0.2

0.75

0.50-

0.1

Ratio

0.25

0.0

0.00

0

25

50

75

100

125150

175

200

0

25

50

75

100

125

150

175200

Epoch

Epoch

σ

= 2

, dynamics of ab

.

σ

= 2

, ratio of

̂

y

T

te

-

1

/

x

T

te

.

B.5 Additional results for full-one initial token

W KQ and W PV with different diagonal initializations are presented in Figure 7. From the results, we know that the gradient descent converges to

where W 1 and W 2 are some dense matrices. Similarly to the proof of Corollary 4.1 in Appendix 4.1, we have

which can be seen as a somewhat preconditioned gradient descent on the OLS problem.

(a) Example 4.1 with

c

= 1

, dynamics of

ab

.

(b) Example 4.1 with

c

= 1

,

‖

̂

y

T

te

-

1

-

x

T

te

‖

2

2

.

120

3.5

(ao,bo) =(0.1, 0.1)

(ao,bo)=(0.1,0.1)

Dynamics of ab

3.0

(ao,bo)=(0.5,1.5)

100

(ao,bo)=(0.5,1.5)

2.5

(ao,bo)= (2.0, 2.0)

(ao,bo)=(2.0,2.0)

2.0

1/c2 = 1.0

80

Theoretical MSE =0

loss

60

1.5-

MSE

40

1.0

0.5

20

0.0

0

0

25

50

75

100125 150

175

200

0

25

50

75

100125150

175200

Epoch

Epoch

(ao,bo)=(0.1, 0.1)

350

(ao,bo)=(0.1, 0.1)

1.0

(ao,bo)=(0.5,1.5)

(ao,bo)=(0.5,1.5)

Dynamics of ab

(ao,bo)= (2.0, 2.0)

300

(ao,bo)=(2.0,2.0)

0.8

1/c2 = 0.25

250

Theoretical MSE =0

0.6

MSE loss

200

0.4

150

100

0.2

50

0.0

01

0

20

40

60

80

100

120

0

20

40

60

80

100

120

Epoch

Epoch

(c) Example 4.1 with

c

= 2

, dynamics of ab

.

(d) Example 4.1 with

c

= 2

,

‖

̂

y

T

te

-

1

-

x

T

te

‖

2

2

.

(a)

W

KQ

, a

=

b

= 0

.

1

.

(b)

W

PV

, a

=

b

= 0

.

1

.

(c)

W

KQ

, a

= 0

.

5

, b

= 1

.

5

.

(d)

W

PV

, a

= 0

.

5

, b

= 1

.

5

.

(e)

W

KQ

, a

=

b

= 2

.

(f)

W

PV

, a

=

b

= 2

.

0.35

0.35

0.30

0.30

0.25

0

0.25

1

0.20

0.20

2

0.15

6

3

0.15

0.10

4

0.10

8

0.05

Q

b

8

0.05

0.00

0

2

4

6

8

0.00

1.4

0

1.2

0.08

2

0

1.0

0.06

1

0.8

2

0.04

0.6

6

3

0.02

4

0.4

8

0

0.2

0.00

2

4

6

8

0.0

0.35

0.35

0.30

0.30

0.25

-o

0.25

1

0.20

0.20

2

0.15

0.15

3

6

0.10

4

0.10

0.05

8

0

0.05

0.00

6

8

0.00

(ao,bo)=(0.1, 0.1)

2.0

(ao,bo)=(0.5,1.5)

Dynamics of ab

(ao,bo)=(2.0,2.0)

1.5

Theory = 0.1343

1.0

0.5

0.0

0

50

100

150200

250

300

350

Epoch

(g) dynamics of ab

,

σ

= 1

(a)

W

KQ

, σ

w

= 0

.

001

.

(b)

W

PV

, σ

w

= 0

.

001

.

(c)

W

KQ

, σ

w

= 0

.

01

.

(d)

W

PV

, σ

w

= 0

.

01

.

(e)

W

KQ

, σ

w

= 0

.

1

.

(f)

W

PV

, σ

w

= 0

.

1

.

(g) Ratio of

y

T

te

-

1

/

x

T

te

.

0.4

0.4

0.3

0.3

2

0.2

0

0.2

0.1

1

0.1

0.0

2

0.0

6

-0.1

3

-0.1

-0.2

4

-0.2

-0.3

-0.3

-0.4

6

8

-0.4

0.4

0.4

0.3

0.3

0.2

-0

0.2

0.1

1

0.1

4

0.0

2

0.0

6

-0.1

3

-0.1

-0.2

4

-0.2

8

0.3

0

8

-0.3

-0.4

2

6

8

-0.4

C

0.4

0.4

0.3

0.3

0.2

-o

0.2

0.1

1

0.1

t

0.0

2

0.0

6

-0.1

3

-0.1

-0.2

4-

-0.2

-0.3

0

8

-0.3

-0.4

6

8

0.4

0.20

0.15

Ratio of pred/true a

0.10

0.05

Ow =0.1

0.00

Ow =0.01

Ow=0.001

Ratio = 0.2

-0.05

0

50100 150200250 300

350

400

Epoch

̂

(a)

W

KQ

, σ

w

= 0

.

001

.

(b)

W

PV

, σ

w

= 0

.

001

.

(c)

W

KQ

, σ

w

= 0

.

01

.

(d)

W

PV

, σ

w

= 0

.

01

.

(e)

W

KQ

, σ

w

= 0

.

1

.

(f)

W

PV

, σ

w

= 0

.

1

.

0

0.6

0.6

0.4

0.4

0

0.2

0.2

1

0.0

2-

0.0

-0.2

3-

-0.2

-0.4

4

-0.4

8

-0.6

-0.6

2

8

0.6

0.6

0.4

0.4

-0

0.2

0.2

1

4

0.0

2-

0.0

-0.2

6

3-

-0.2

-0.4

4

-0.4

8

-0.6

0

8

-0.6

2

4

6

8

0.6

0.6

0.4

0.4

0

0.2

0.2

1

0.0

2-

0.0

6

-0.2

3

-0.2

-0.4

4

-0.4

8

0

8

-0.6

-0.6

8

50

Ow= 0.1

Ow=0.01

MSE loss at Tte

40

Ow=0.001

301

Theoretical MSE=0

20-

10

0

0

008004009009 00 000 007 001

Epoch

T

(g) MSE

‖

̂

y

te

-

1

-

x

T

te

‖

2

2

.

(a)

W

KQ

, a

= 0

.

1

, b

= 0

.

1

.

(b)

W

PV

, a

= 0

.

1

, b

= 0

.

1

.

(c)

W

KQ

, a

= 0

.

5

, b

= 1

.

5

.

(d)

W

PV

, a

= 0

.

5

, b

= 1

.

5

.

(e)

W

KQ

, a

= 2

, b

= 2

.

(f)

W

PV

, a

= 2

, b

= 2

.

0.8

0.8

2

0.6

0

0.6

1

+

0.4

2

0.4

6

3

0.2

4

0.2

8

0

0.0

0

8

0.0

C

0.5

1.4

1.2

2

0.4

0

1.0

4

0.3

1

0.8

2

6

0.2

3

0.6

0.1

4

0.4

8

0

0.2

0.0

8

0.0

0.8

0.8

2

0.6

0

0.6

1

4

0.4

2

0.4

6

3

0.2

4

0.2

8

0

0.0

8

0.0

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 our 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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Please see Section 7 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 a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions are presented clearly in the main paper (Assumption 3.1, 4.1 and 4.2). Complete proof can be found in Section 5 and Appendix A. The correctness of our proofs can be verified by simulations in Section 6 and Appendix B.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 experimental 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: [Yes]

Justification: We provide complete configurations in Section 6 and Appendix B. We also upload the code at https://github.com/ML-GSAI/MesaOpt-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 be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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) If 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We upload the complete code and sufficient instructions in the supplemental material. It is available at https://github.com/ML-GSAI/MesaOpt-AR-Transformer .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: We do not present an error 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 6 and Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

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 as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We conform research conducted in the paper satisfies the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws 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 see details in Section 7.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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).

11. 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.

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/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.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

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

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

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nWhy is the y-axis in Figure 1-b defined as the ratio of two vectors?\n\nLIMITATIONS:\nThe authors discuss limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nPresentation 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).\n\nWEAKNESSES:\nWhen 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)\n\nQUESTIONS:\nTheories 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?\n\nLIMITATIONS:\nTheories 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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$”\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe limitations are provided in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

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

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?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-15