One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
Authors: Arvind V. Mahankali
Tatsunori Hashimoto
Tengyu Ma
Venue: ICLR 2024 poster (2024)
Keywords: Linear Self-Attention
In-context learning
Gradient Descent
Theoretical Understanding
Primary area: learning theory

Abstract: Recent works have empirically analyzed in-context learning and shown that transformers trained on synthetic linear regression tasks can learn to implement ridge regression, which is the Bayes-optimal predictor, given sufficient capacity (Akyurek et al., 2023), while one-layer transformers with linear self-attention and no MLP layer will learn to implement one step of gradient descent (GD) on a least-squares linear regression objective (von Oswald et al., 2022). However, the theory behind these observations remains poorly understood. We theoretically study transformers with a single layer of linear self-attention, trained on synthetic noisy linear regression data. First, we mathematically show that when the covariates are drawn from a standard Gaussian distribution, the one-layer transformer which minimizes the pre-training loss will implement a single step of GD on the least-squares linear regression objective. Then, we find that changing the distribution of the covariates and weight vector to a non-isotropic Gaussian distribution has a strong impact on the learned algorithm: the global minimizer of the pre-training loss now implements a single step of $\textit{pre-conditioned}$ GD. However, if only the distribution of the responses is changed, then this does not have a large effect on the learned algorithm: even when the response comes from a more general family of $\textit{nonlinear}$ functions, the global minimizer of the pre-training loss still implements a single step of GD on a least-squares linear regression objective.

Under review as a conference paper at ICLR 2024
ONE STEP OF GRADIENT DESCENT IS PROVABLY THE OPTIMAL IN-CONTEXT LEARNER WITH ONE LAYER OF LINEAR SELF-ATTENTION
Anonymous authors Paper under double-blee blind review
ABSTRACT
Recent works have empirically analyzed in-context learning and shown that transformers trained on synthetic linear regression tasks can learn to implement ridge regression, which is the Bayes-optimal predictor, given sufficient capacity (Akyürek et al.] 2023), while one-layer transformers with linear self-attention and no MLP layer will learn to implement one step of gradient descent (GD) on a least-squares linear regression objective (von Oswald et al.] (022). However, the theory behind these observations remains poorly understood. We theoretically study transformers with a single layer of linear self-attention, trained on synthetic noisy linear regression data. First, we mathematically show that when the covari- ates are drawn from standard Gaussian distribution, the one-layer transformer which minimizes the pre-traie pre-training loss will implement single step of GD on the squaresuares linear regression objective. Then, we find that changing the distribu- tion of the covariates and weight vector to non-isotropic Gaussian distribution has a strong impact on the learned algorithm: the global minimizer of the pre- training loss now implements single step of pre-conditioned GD. However, if only the distribution of the responses is changed, then this does not have a large effect the learned algorithm: even when the response comes from more gen- eral family of nonlinear functions, the global minimizer of the pre-training loss still implements a single step of GD least-squares linear regression objective.
INTRODUCTION
Large language models (LLMs) demonstrate the surprising ability of in-context learning, where an LLM "learns" solve task by conditioning on a prompt containing input-output exemplars (Brown)
et al. 2020; Lieber et al.] 2021 Radford et al. 2019 Wang & Komatsuzaki 2021). Recent works
have advanced the understanding of in-context learning via empirical analysis (Min al. 2022 Wei
et al. 2023 Akyürek et al. 2023
von Oswald et al. 2022 Dai et al. 2023), but theoretical amalysis
remains limited (Xie et al. 2022
A recent line of work (Garg et al. 2022 Akyürek et al. 2023 von Oswald et al., 2022 Dai et al.,
2023) empirically finds that transtormers can be trained to implement algorithms that solve linear regression problems in-context. Specifically, in each input sequence the transformer is given a set of in-context examples (I,, yt), where y; = w T T, + e, with a shared and hidden random coefficient vector and random noise é, and test example x: The transformer is then trained to predict y = w T x + e, where denotes random noise from the same distribution as &,. These works find that the transformer outputs prediction ý which is similar to the predictions of existing, interpretable linear regression algorithms, such as gradient descent (GD) or ordinary least squares, applied to the
dataset consisting of the pairs (3t,, U:). In particular, von Oswald al. 2002) empirically show that a one-layer transformer with linear self-attention and no MLP layer will implement single step of gradient descent when trained such distribution.
Several works (e.g. Akyürek et al. 2023); ; Liu et al. 2023); Giannou et al. (203)) theoretically study the expressive power of transformers. In the context of linear regression tasks, Akyürek et al.
In some settings in these works, the noise set O.

---

Under review as a conference paper at ICLR 2024
describe how transformers can represent gradient descent, or Giannou et al. ( describe how transformers can represent Newton's algorithm for matrix inver- sion. However, in addition to the expressive power of transformers, is also of interest to understand the behavior of transformers trained with gradient-based algorithms. Furthermore, it is still useful to understand the behavior of models with restricted -thochh practical L.LMs are very pressive, they need to perform many tasks simultancously, and therefore the capacity per problem may still be relatively limited. Thus, motivated by yon Oswald al. (2022), we theoretically study the global minima of the pre-training loss for one-layer transformers with linear self-attention on the linear regression data distribution described above.
2023)
updates, and
Contributions. In this paper, we study transformers with one linear self-attention layer, and mathe- matically investigate which algorithms the transformers implement for synthetically generated linear regression datasets. We prove that the transformer which implements single step of gradient de- scent on a least squares linear regression objective is the global minimizer of the pre-training loss. This exactly matches the empirical findings of von Oswald et al. 2022
Concretely, we consider a setup similar to von Oswald et al. (2022) ;: Akyürek et al. (2023). The model we study is a transformer with one linear single-head self eelf-attenioon layer, which is the same model as the one empirically studied by yoo Oswald et al.] (2022). The training data for this transformer consist of sequences of the form (I1, y1, ................................ Fn. Br.) where the x, sampled from (0, Taxa) and y/i = w Iy + 6, where w is sampled from (0, Taxa) once per sequence, and the 6; are i.i.d. Gaussian noise with variance o² The pre-training loss is the expected error that the transformer achieves when predicting y w x given the test example x and the context (I yl,,. Nn: Vn), i.e. the pre-training loss is L (x (12.......,....... V...., ). ((y ")) where ý the output of the transformer given (z V/, ................................ : yn) and x: input.
We show in Section that the transformer which is the global minimizer of the pre-training loss L implements one step of gradient descent on a linear regression objective with the dataset consisting of the (I;, y:). More concretely, the transformer implements the prediction algorithm
y. x.
(1)
where 17 is a learning rate independent of the Tis the y; and x. However, one step of GD is also preferred in part due the distribution of the x;, In particular, if the covariance of 24 is no longer the identity matrix, we show (Section 4 that the global minimum of the pre-training loss corresponds to one step of GD, but with condconni........
Interestingly. our theory also suggests that the distribution of y y does not play such significant role in the algorithm learned by the transformer. In Section we study a setting where y: lay is nonlinear, but satisfies some mild assumptions, such as invariance to rotations of the distribution of the xy. As a concrete special case, the target function can be a neural network with any depth/width and i.i.d. random Gaussian weights. We show in Section [5) that a one-layer transformer with linear self-attention, which minimizes the pre-training loss, still implements one step of GD on a linear regression objective. Intuitively, this is likely because of the constraint imposed by the architecture, which prevents the transformer from making use of any more complex structure in the yt-
We discuss the closely related works of Ahn et al. (2023) and Zhang et al. (2023) which are concurrent with and independent with our work and were posted prior to our work on arXiv. Ahn et al. (2023) give theoretical results very similar to ours. They study one-layer transformers with linear self-attention with the same parameterization as von Oswald et al. (2022). and show that with isotropic 24, the global minimizer of the pre-training loss corresponds to one step of gradient descent on a linear model. They also show that for more general covariance matrices, the global minimizer of the pre-training loss corresponds to one step of pre-conditioned gradient descent, where the pre-conditioner matrix can be computed in terms of the covariance of
Concurrent Works.
Different from our work,AAhn (with linear self-attention) with residual connections trained on linear regression data. First, they study a restricted parameterization where in each layer, the product of the projection and value
Ahn et al. (2023) also show additional results for multi-layer transformers
This result is not exactly the same as our result in Section since we assume u (O Σー1) while they assume w ~ (00 Taxa).
2

---

Under review as a conference paper at ICLR 2024
matrices has only one nonzero entry. In this setting, for two-layer transformers with linear self- attention, they show that the global minimizer corresponds to two steps of pre-conditioned GD with diagonal pre-conditioner matrices, when the data is isotropic. For linear transformers with k layers, they show that k steps of pre-conditioned GD corresponds to a critical point of the pre-training loss f where the pre-conditioner matrix is the inverse of the covariance matrix of the x;; Next, they study a less restrictive paramcterization where the product of the projection and value matrices can be almost fully dense, and show that a certain critical point of the pre-training loss for k-layer linear transformers corresponds to k steps of generalized version of the GD++ algorithm, which was empirically observed by von Oswald et al. (0022) be the algorithm learned by k-layer linear transformers.
Zhang et al. 2023) also theoretically study a setting similar to ours. They not only show that the global minimizer of the pre-training loss implements one step of GD (the same result as ours), but also show that a one-layer linear transformer trained with gradient flow will converge minimizer. They also show that the transformer implements step of pre-conditioned GD when the T2 are non-isotropic. They also characterize how the training prompt lengths and test prompt length affect the test-time prediction error of the trained transformer. Additionally, they consider the behavior of the trained transformer under distribution shifts, as well as the training dynamics when the covariance matrices of the x; in different training prompts can be different.
global
One additional contribution of our work is that we also consider the case where the target function in the pre-training data is not linear function (Section 5 This suggests that, compared to the distribution of the covariates, the distribution of the responses at training time does not have as strong of an effect on the algorithm learned by the transformer. We note that our proof in this setting is not too different from our proof in Section 3 Zhang et al. (2023) consider the case where the y, 's in the test time prompt are obtained from a nonlinear target function, and consider the performance on this prompt of the transformer trained on prompts with a linear target function this is different from our setting Section 5 since we consider the case where the training prompts themselves are obtained with nonlinear target function. We discuss additional related works in Appendix A
2
SETUP
Our setup is similar to von Oswald et al. (2022).
One-Layer Transformer with Linear Self-Attention. linear self-attention layer with width s consists of the following parameters: key matrix Wx E Ix* a query matrix Wq E Rax*, and a value matrix Wv E R* Given a sequence of T > tokens (vi, .2,... VT), the output of the linear self-attention layer is defined to be (01, 02 ................................ (T), where for i E [T] with i 1,
6 Σ (c Wo)
(2)
and é = 0. In particular, the output on the 7th token is
T-1
êT =
(3)
1=1
As in the theoretical construction oon Oswald et al. 2022), we do not consider the attention score between a token Y; and itself. Our overall transformer is then defined to be a linear self-attention layer with key matrix WK. query matrix Wq. and value matrix Wy, together with a linear head h E RR which is applied to the last token. Thus, the final output of the transformer is h êT. We will later instantiate this one-layer transformer with 8 = d ++1 1, where d the dimension of the inputs x;, We note that this corresponds to single head of linear self-attention, while one could also consider melti-head self-attention.
One technical point is that they show there exist transformers representing this form of conditiioooo GD having arbitrarily small gradient, but not that there exists transformer with gradient exactly which represents this form of pre-conditioned GD.
1 'eere, they assume that z; N (0, Σ) and u (0, Σート), which is the same assumption our result in Section 4 but different from their result for one-layer transformers where x, ~N E))
3

---

Under review as a conference paper at ICLR 2024
Linear Regression Data Distribution. The pretraining data distribution consists of sequences D (I1, V1...., In+1, Yn+1). Here, the exemplars I; are sampled i.dd from M(0, Iaxd). Then, weight vector E R is sampled from (0,, Laxa), freshly for each sequence. Finally, Ye is computed as = w Ti + 6 where c ~ (0, o²) for some > 0. We consider the vector
33 IR++1 be token - in other words, the sequence (x1, , In+1 Mo+1) is
Ut
considered have n + 1 tokens (rather than 2(n 1) tokens). We use T denote the distribution of sequences defined in this way.
At both training and time, (I1, (11,B1...... Zna yny In+1, Mn+1) 1 is generated according to the pre- training distribution T, i.e. the 2; are sampled i..d. from NomnN(O (0, laxa), a new weight vector w E Rd also sampled from AO, Laxa), and y; = t where the e, are sampled i.i.d from (/,,NO0,,00 o Then, the in-context learner is presented with Tı, ........ Zn Un. In+1, and must predict yn+1 We refer to T1, -i,..., In as the support exemplars and In+1 as the query exemplar. Here, U1 01....e Un are
defined above, but 1a+1
Ip-1
following the notation of von Oswald et al. (2022) We
note that this is not significantly different from the standard in-context learning setting, since even though the final token Un+1 has 0 as any additional information about Vn+1-
coordinate, it does not provide the transformer with
Loss Function. Given a one-layer transformer with linear self-attention and width 1, with key matrix WK E R(d+1)x(d=1), query matrix We E ((d+1)x(4+11 and value matrix Wv g(d+1)x(d+1) and with head h E Rd+1, the loss of this transformer on our linear regression data distribution is formally defined as
L(WK, Wo, Wv, h) ED-T 10+1 - 1)),)
(4)
where defined above, 6441 1 is the output of the linear self-attention layer on the (n + 1) token, which in this case is In+1 0
We now rewrite the loss function and one-layer transformer in a more convenient form. As a convenient shorthand, for any test-time sequence D (#1,81:.... ......... In+1. 0), we write =
(I1, ; Zn: Vn), i.e. the prefix of D that does not include In41 and Un41- We also define
G
(5)
D1
91
With this notation, we can write the prediction obtained from the transformer on the final token as
Ia+1 h T W.WWGGW WQUn-1
(6)
where Un+1 =
In+1 0
Additionally, we also define the matrix X E Raad as the matrix whose
tth row is the row vector x.T. i.e.
32
X
and we define the vector IR" the vector whose tt entry is yis i.e.
y1
1/2
27
8
Yn
'If' we were to treat i;; and yi separate tokens, then we would need to deal with attention scores between and yi for i of j. as well as attention scores between Yi and zú for i +f j. Our current setup simplifies the analysis.

---

Under review as a conference paper at ICLR 2024
Finally, it is worth noting that can write the loss function as
L(W K, Wo, Wv, ED~T((h W WyG 1n++)
(9)
Thus, for w Rd+1 and M E R(d+1)x((+1) we define
L(w, M) Eø-rl(w' Gi,Muu41 1
(10)
then ( h, Wx WQ) ) = L(WK, K Wo: ,v, h). Note that we have a slight abuse of notation, and L has two different meanings depending on the number of arguments. Finally, with the change of variables M We and w W we can write the prediction of the transformer
as UTG. Gp Mu41.. Thus, the output of the transformer only depends on the parameters through w GR.MM
Additional Notation. For a matrix A E Rdxd, we write Anj.: to denote the sub-matrix of that contains the rows of A with indices between i and (inclusive). Similarly, write A denote the sub-matrix of that contains the columns of A with indices between and j (inclusive). We write Auj.Ae to denote the sub-matrix of A containing the entries with row indices between i and (inclusive) and column indices between k and | (inclusive).
3 MAIN RESULT FOR LINEAR MODELS
Theorem 1 (Global Minimum for Linear Regression Data). Suppose (Wi. W WV, is global minimizer of the loss L. Then, the corresponding one-layer transformer with linear self- attention implements one step of gradient descent on a linear model with some learning rate
n O. More concretely, given a query token Un+1
In+1 the transformer outputs
E E x E
0. Here given a prefix D of test-time data sequence D. we let w denote the solution to ridge regression on X and ÿ with regularization strength o?.
Ei Waa, Unth, where
The minimizer (W WQ W h)) is not unique, though the linear predictor implemented by the minimizer unique
the discussion after Lemma 2 One such construction is as follows. yon Oswald et al. 2022)) describe essentially the same construction, but our result shows that it is a global minimum of the loss function, while von Oswald et al. (2022) do not theoretically study the construction aside from showing that it is equivalent to one step of gradient descent. We define
W = [laxd 9) W. = [Taxa
-larre
(Tore
!)
0
: WV = 0 h
(11)
BB 1 of
Here, the unique value of which makes this construction global minimum is E 1 xx D) To see why this construction implements single step of gradient descent on a linear model, note that given test time inputs 31, 1/,, Fng Uny In+1, if we write v; = y/1
2; for S n and Up+1
then the output of the corresponding transformer would be In+1
0
(h*)T T Σ (Win))o) (WR) ( = Σ yeo Ta+1
(12)
On the other hand, consider linear regression with total squared error the loss function, using the Ij and ye. Here, the loss function would be L(w) = "i= (w" x u²)², meaning that the gradient is Lww = Σn= (w x; ,, In particular, if we initialize gradient descent at Wi = 0, then after one step of gradient descent with learning rate 7). the iterate would be at w1 = = Σn= y; observe that the final expression in Equation 122 is exactly w. z.+1-
Now, give overview of the proof of Theorem I By the discussion in Section 2 it suffices to show that L((W Th*, (W;) WG) is the global minimum of L(w, MM. The first step of the proof rewrite the loss in more convenient form, getting rid of the expectation 2n and yot11
5

---

Under review as a conference paper at ICLR 2024
Lemma 1. Let w be the solution to ridge regression with regularization strength o² on the emplars (I1, Yr)..... .... (In, Yn) given in a context D. Then, there exists a constant C. 2 0, which is independent of w, M, such that L(w, M) C ED-TTIMM
du 1l22.
As discussed towards the end of Section 2 the prediction can be written as w ToT G G MeM++1 en+1 where In+1
Pa+1
meaning that the effective linear predictor implemented by the transformer is
the linear function from I to R with weight vector M. Thus, we can interpret Lemma as saying that the loss function encourages the effective linear predictor to match the Bayes-optime predictor wD Note that it is not possible for the effective linear predictor of the transformer to match ü exactly, since the transformer can only implement linear or quadratic function of the Iy, while representing w requires computing (XT X + o²1)= and this is much more complex function of the z4. We prove Lemma 1 manipulations of random variables we give detailed proof in Appendix C
Jsiig the fact that ED.... [Vn+1] v
D xn+1 and standard
Next, the key step is to replace " in the above lemma by n y.
Lemma 2. There exists constant C₁ 2 0 which is independent of w, M, such that
L(w, M) C1 + B B.
n 51
(13)
Lemma 2 says that the loss depends entirely on how far the effective linear predictor is from ny.' y7
It immediately follows from this lemma that (Wx, WQ. Wv, h)) is global minimizer of the loss and only if the effective linear predictor of the corresponding transformer is n J. Thus, Theorem follows almost directly from Lemma 2 and in the rest of this section, we give an outline of the proof of Lemma 2 the detailed proofs of Theorem and Lemma 2
are Appendix C
Proof Strategy for Lemma 2 and L' (w, M). defined as
Our overall proof strategy is to show that the gradients of L(w, M)
Lww, M) == 15
12 L'(w, M) = E M. G x 012. (14)
are equal at every w, M, from which Lemma 2 immediately follows o For simplicity, we write A = M. so without loss of generality, we can instead show that the gradients of the loss functions Ji (A, w) and Jz( A, w)) are identical, where J1 and J2 are defined as
(1, w) = ==bb-TlAAGru - J2(A, w) == EB nx 012.
(15)
In this section, discuss the gradients with respect to that the gradients with respect to A are the same. We have
the proof ideas show
V
21E 2w GBA (AG úp))
(16)
and
V V2,2(A, w) 2E G (AG - GAG T Thus, showing that these two gradients are equal for all A, w reduces to showing that for all A, we have
(17)
E E GBA -nBp.-CGA''X
(18)
12
Recall that Gū is defined as G5
Σ
y:z; vł
Our first key observation that for any
yor
i [n] and any odd positive integer k. Blyy X O, since y, = w x, + 6, and both w T x and 6
have distributions which are symmetric around 0. This observation also extends to any odd-degree
This because L and L' are detined everywhere on [g++1 x (g(++) and for any two differentiable functions f.g defined on an open connected subset S C R* if the gradients of f and are identical on S, then of and are equal on S up to additive constant. This can he shown using the fundamental theorem of calculus.

---

Under review as a conference paper at ICLR 2024
monomial of the y4- Using this observation, can simplify the left-hand side of Equation 18 as follows. We can first write it as
I (A] 2 i + yT uD yiu; (A 11dd B + (A 14+1. ) w.
B Jib = EBTT E
(19)
Then, since T0 (XTX +
TT. each entry of wD has odd degree in the yı, mean- in the above equation vanish after taking the
ing the terms 2 (A)). # and yl (AT)) (AA4))) 2. expectation. Thus, we obtain
y d y:u0 "
X " T7 XA
G On = -E E
Σ
= E
(20)
Since each entry of nx'' has odd degree in the yi, in order to simplify the right-hand side of Equation 18). we can apply the same argument but with ū replaced by n 1xT J7 obtaining
X JAT 4+1 X
(21)
87 XA 11d X
Thus, showing Equation
18) reduces to showing that
TA d+
X JA -
d-1 x 1:d
E B
y XA
BNT
J XAA
(22)
To show Equation 22p. our key tool is Lemma
4 which follows from Lemmal 3
Lemma 3. There exists a scalar C1 such that I E [XT Vy x] = cilaxd, and there exists a scalar C2 such that E ENTT
[XT Tリ D
c2Taxd-
Lemma 4. If B [ x M BB-Tl0 18 X 0. then E B V0
X
gro 51 =0
Overview of Proof of Lemma 3 and Lemma We give overview of how prove Lemma and Lemma 4 here, and defer the full proofs to Appendix C To show that E T yy" X] is scalar multiple of the identity, we first use the fact that even when all of the æ; are rotated by rotation matrix R, the distribution of 7|X remains the same, since the weight vector w is drawn from (0, Laxd) which is rotationally invariant distribution. Thus, if we define MX) E'yg'
X]
as function of X E Raxd then for any rotation matrix R E IR Rdxdd we have
M(XR)) Elyy" |
Elyy" X M(X),
(23)
where the second equality is because multiplying X on the right by RD corresponds to rotating each of the x; by R. Additionally, if rotate the x; by R. then Fğ- vy X] remains the same - this is because the distribution of the x; is unchanged due to the rotational invariance of the Gaussian distribution, and the conditional distribution yi xz is unchanged when rotate x; by R.
This implies that
X Vy X| B M(X)X] B T (ART)TM(XXI M(X ((
(24)
where the second equality is because, as we observed above, ED [XT yy X] remains the same when we rotate each of the x; by R. Since M(XRT) -MX),, we have
E((XRT) ( )XR
(25)
which implies that E[X yy X] RE yy"
X|RT for any rotation matrix R, and therefore E[X vy X is scalar multiple of the identity matrix. To finish the proof of Lemma 3 we show that F [X' ju 5] is a scalar multiple of the identity using essentially the same argument. To
show Lemma trace equal 0.
simply take the trace of E B "g X X ÿũū jū and select v) so that this
Finishing the Proof of Lemma
2 Recall that, to show that the gradients of J1 and Jg (defined in However, this is
Equation (15) with respect to are equal, suffices show Equation direct consequence of Lemma This is because we can rewrite Equation
as
X ju A.141 tr(
X X X X tr(x
E D
E EDTT
(26)
XA! XA.
"V XAT 1.d)

---

Under review as a conference paper at ICLR 2024
This shows that the gradients of J1 and J2 with respect
equal, and similar arguments to show that the gradients of J1 and J2 with respect to A are equal. As mentioned above, this implies
fw w n T is constant that is independent that E BD Il Grw 10 " 112 E D [MM G of M and w. as desired.
RESULTS FOR DIFFERENT DATA COVARIANCE MATRICES
In this section, we consider the setting where the x; 's have a covariance that is different from the identity matrix, and we show that the loss is minimized when the one-layer transformer implements one step of gradient descent with preconditioning. This suggests that the distribution of the x; 's has significant effeet the algorithm that the transformer implements.
Data Distribution. Concretely, the data distribution is the same as before, but the H, are sampled from from.N(O. (0,,
Σ), where Σ E Raxd is positive semi-definiee (PSD) matrix. The outputs are generated according to y/ = w H, + <s. where w ~ (N E- This can equivalently be written as 2; = 2!/2u, where W4 ~ (O, Taxa), and y/2 = (w') Us + where w' ~ NOO Taxa)- We keep all other definitions, such as the loss function, the same before.
Theorem 2 (Global Minimum for 1-Layer 1-Head Linear Self-Attention with Skewed Covari- ance). Suppose (Wk, W Wh"""") ) is a global minimizer of the loss when the data is gen- erated according to the distribution given in this section. Then, the corresponding one-layer transformer implements one step of preconditioned gradient descent, on the least-squares lin- ear regression objective, with preconditioner Σ for some learning rate 17 > 0. Specifically,
given query token Un+1
In+1
the transformer outputs v) Ei= V: xx:) TTa::TTTnt Intt, where
F 17 X(X
x 2
1
17
"D-T
To prove this result, essentially perform change of variables to reduce this problem to the setting of the previous section - then, we directly apply Theorem The detailed proof is given in Appendix D
RESULTS FOR NONLINEAR TARGET FUNCTIONS
In this section, we extend to setting where the target function is nonlinear - our conditions on the target function are mild, and for instance allow the target function be a fully-connected neural network with arbitrary depth/width. However, we keep the model class the same (i.c. 1-layer trans- former with linear self-attention). We find that the transformer which minimizes the pre-training loss still implements one step of GD on the linear regression objective (Theorem 3). even though the tar- get function is nonlinear. This suggests that the distribution of ys |z, does not affect the algorithm learned by the transformer much as the distribution of 2;.
Data Distribution. In this section, we consider the same setup as Section but we change the distribution of the y,'s s. We now assume y/ =xx;;) + é,, where e; (0, oz) as before, but f is drawn from a family of nonlinear functions satisfying the following assumption:
Assumption 1. We assume that the target function I drawn from a family F, with a probability measure IP on F, such that the following conditions hold: for any fixed rotation matrix R E Raxd the distribution of functions ff is the same as the distribution of f R (where o denotes function composition). Moreover, the distribution of f is symmetric under negation. In other words, if E F is measurable under P. then P(E) P-EE,, where -E - (-fiff
For example, Assumption I is satisfied when f(x) is fully connected neural network, with arbitrary depth and width, where the first and last layers have i.i.d. /O0 1) entries see Appendix B for further discussion. Under this assumption, we prove the following result:
Theorem 3 (Global Minimum for 1-Layer 1-Head Linear Self-Attention with Nonlinear Target W²)
Function). Suppose Assumption I holds, and let (W k W Wi, h²) be global minimizer of the pre- pre-training loss. Then, the corresponding one-layer transformer implements one step of gradi- ent descent the least-squares linear regression objective, given (I1, Bl, -...... Tus yu). More
8

---

Under review as a conference paper at ICLR 2024
cretely, given a query token Un+1
xn-1
the transformer outputs 1) Ei=1 yeog In+ where
0
Rpln X 7 Bolg XX 3) " argmin, Benneenn [(u 2n+1 vn+1)² DJ. and D, D are in Section 2
The result is essentially the same as that of Theorem | note that the learning rate is potentially dif- ferent, as it may depend on the function family F. The proof is analogous to the proof of Theoreem First we prove the analogue of Lemma 1 defining L(w, M) as in Section 2
Lemma 5. There exists a constant C 2 0 such that L(w, M) C + E MI.M
"pllé
where "B argmin 11 [[ In+1 4n+1 D).
Next, in the proof of Lemma 2 we used the fact that odd-degree polynomials of the yi have expec- tation 0 - the corresponding lemma in our current setting is as follows:
Lemma 6. For even integers k and for replaced by an even-degree monomial of the Additionally. for odd integers k and for E [n]. Ely, X] This also holds with y replaced by an odd odd- degree monomial of the yi-
[m], Ely | X] O. This also holds with y
Proof of Lemma 0 This follows from Assumption [ This is because for each outcome (i.e. choice of f and Ei. EI....,En) En) which leads to (I1, 91..... Eni Bn), the corresponding outcome -f.-E1,... En which leads to (æ1, 181,-,1.... -............................... In. yn) is equally likely. The F which is ob- tained from the second outcome is the negative of the ū which is obtained from the first outcome. k is even, then y is the same under both outcomes since k is even, and the average of y " under these two outcomes is O. Additionally, if le is odd, then vf under the second outcome is the negative of y under the first outcome, and the average of y under these two outcomes is O. This completes the proof of the lemma.
Next, we show the analogue of Lemma
Lemma 7. E B- [XT
X] and E DNTT XT ju are scalar multiples of the identity. Thus, JJ
En XT X yu n B X Vy X). where 1j Boly X "f"
E B
The proof of Lemma is nearly identical to that of Lemma with Assumption I used where appropriate. We include the proof in Appendix We now state the analogue of Lemma 2
Lemma 8. There exists constant C1 2 O which is independent of w. M, such that Mw = 11X 5122 where 1)
Epluyxx Ep XX M7
C+
Theorem now follows from Lemma because M. G w is the weight vector for the effective linear predictor implemented by the transformer. All missing proofs are in Appendix
CONCLUSION
We thcoretically study one-layer transformers with linear self-attention trained on noisy linear gression tasks with randomly generated data. We confirm the empirical finding of von Oswald et al.) (022) by mathematically showing that the global minimum of the pre-training loss for one-layer transformers with linear self-attention corresponds to one step of GD on a least-squares linear re- gression objective, when the covariates are from an isotropic Gaussian distribution. We find that when the covariates are not from an isotropic Gaussian distribution, the global minimum of the pre-training loss instead corresponds to pre-conditioned GD. while the covariates are from an isotropic Gaussian distribution and the response is obtained from a nonlinear target function, then the global minimum of the pre-training loss will still correspond to one step of GD on least-squares linear regression objective. We study single-head linear self-attention layers - it is an interesting direction for future work to study the global minima the pre-training loss for a multi-head linear self-attention layer. Another interesting direction is to study the algorithms learned by multi-layer transformers when the response is obtained from a nonlinear target function. We note that Ahn et al.) 2023) have studied the case of multi- layer transformers when the target function is linear. They show that for certain restricted parameterizations of multi-layer linear transformers, the global min- ima critical points of the pre-training loss correspond interpretable sradient-based algorithms.

---

Under review as a conference paper at ICLR 2024
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 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 The Eleventh International Conference on Learning Representations, 2023,
net, forrmm?
id Nd-0gOX4RByN4I
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians:
Provable in-context learning with in-context algorithm selection, 2023.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, 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. CoRR, 0bs/2005.14116. 2020. URL https: http:///wwww.........//oo wwtpwwwwww.w///////....oo
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can trans- formers learn in-context? A case study of simple function classes. In NeurlPS, 2022. URL http://
html]
Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers, 2023.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representa- tions, 2023.
Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evalua- tion. Technical report, A121 Labs, 2021.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang Trans- formers learn shortcuts to automata. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https:://ppenrevier net peee/pdd?ic-eee4/jjie.
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, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019.
Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequence functions with infinite dimensional input, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2022.
Ben Wang and Aran Komatsuzaki. Gpt-j-6b: 6 billion parameter autoregressive language model, 2021.
10

---

Under review as a conference paper at ICLR 2024
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
https://openrews
net feet/fornmt forum?id RdJVE
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context, 2023.
A ADDITIONAL RELATED WORK
The work Takakura & Suzuki| (2023) theoretically shows that transformers can achieve low ap- proximation error when the target function is shift-equivariant on sequences of infinite length, sub- ject to certain regularity conditions. Bai al. (2023) propose various constructions through which transformers can solve problems such Bayesian linear regression and generalized linear models, achieving low test error they also show that transformers can represent form of model selec- tion. Guo et al. (2023) similarly provide constructions of transformers which represent certain algorithms. Compared to these works which give merooxmmation-therree constructions, we show that one step of gradient descent is the best predictor that can be implemented by layer of linear self-attention.
Another concurrent work by Tarzanagh al. (2023) analyzes the global minima and training dy- namics of transformer layers with general data. They show that the global minimum corresponds with type of max-mnrgin solution, and give conditions under which the training dynamics converge to this global minimum.
B DISCUSSION OF ASSUMPTION
Assumption 1 holds in the case where f is fully connected neural network, where the first and last layers have i.i.d. A 1) entries. For simplicity. we show this in the case of two-layer neural networks. We let our function family F be (x + W2o(W,x) W1, W2 have i.i.d. NN(0, 1) entries). Let fz) = "Var(Wii) be two-layer neural network drawn randomly from F. Let R be fixed rotation matrix. Then, (f Rx)x) W2o Since R is fixed rotation matrix, and W1 has i...d (0, 1) entries, W, R has the same distribution as W. Thus, f R has the same distribution as f-
C MISSING PROOFS FROM SECTION 3
Proof of Lemma 7
We can write
L(w, M) ED-T [(n+t1 w. G Munt2))
(27)
D. xn+1 Bynss [1 w. G Men+1) sm41, D7]
(28)
Let us simplify the inner expectation. For convenience, fix D and Un+1, and consider the function g RRd + RR given by
(2)) Ey... [wu In+1 3n+1)2 D, xn+1].
(29)
It is well-known fact that the minimizer of g(u) is given by wi where w TXT -
is the solution to ridge regression on X and y with regularization strength o², (See e.g. equation (19) Akkyrek et al. 2023).) Furthermore,
v.g(ūp) Ey... [2( Int1 Unt1)ann11 D,41l
(30)

---

Under review as a conference paper at ICLR 2024
and in particular, taking the dot product of both sides with - (for any vector IR") gives
Eyn+, [(0 In++ 1)+1)- (u In+1 - L In+1)] 0.
(31)
Thus, letting = w G M,,1aa we can simplify the inner expectation in Equation 28
E Ey [[nn+1 - w G M,, 11d Malttnn ) In+1, D]
(32)
E, [1 - un+1 1 + i D - w GB 1nt1, DJ
(33)
Eynti |(yn+1 2 [n+t1 )² In+1, D] (w D - w G.M,,ndun+1))
(34)
+ +2. yn [( 2 2 (0) o w G.M,Iatn+t) Tnt1,
(35)
Byn [V w ) Fat1, D) + (w (T+t w T G B M, 1dd
(By Equation 31
We can further write the final expression C C D + (m Dun+1 w T M. 1) where
C Cauti.D D is constant that depends 2n+1 and D but is independent of w and M. Thus, have
L(w, M) = E2 D. CC 20. + E D..._ 1 wI G G..MM... M,1
(36)
= =C + EC++Eplluy E - Gr
(B.c. Int1 ~ N(O, Iaxa))
where C is a constant which is independent of w and M.
Proof of Lemma For convenience let .(X) = E D.T yy X] - we use this notation to make the dependence on X clear. Then, the (i, j-th entry of (XX) is E((w + < 23 + +,))) X] and this is equal to [[(w (w: xz) X] for i + j. and [[[w x;)(w 2z) X] + o for i j. If we perform the change of variables m2 H Rx, for a fixed rotation matrix and all i E [n]. then
[w Rr;)(w (w- Rxy) X] [ x;)(w -2;) 22ÿ) XX] because w ~ NO, Taxa) and N(0, Taxa) is a rotationally invariant distribution. In other words, we have M(XRT) = M(X). Thus, for any rotation matrix R
Ex[X T M(X)X] = -xX(XX)) ((XRT)(X(
(By rotational invariance of dist. of x;))
=X M(XR )(XRT)XRR" )X]R
(37)
RREEx[X [XT M(X)X)R.
(Because M(XRT)= M(X))
This implies that [x[X ((X) EB-T(X Jy" X] is scalar multiple of the identity matrix.
Next, consider TENT(Ou X]. which write for convenience J(X). Similarly to the above proof, we use the observation that if make the change of variables 2; - Rrg, then the joint distribution of y1..... Nn. Un+1 is unchanged due to the rotational invariance of w. Addi- tionally, if we write úØ as W5 (21, V1....,, Ins yn) emphasize the fact that is a function of U1, ......... Tn, Un. then
w (RæI, HI, RoZ, y2,. Ren, yn) Roy (211 ; Vl, U2, y2. ... Iny Un)
(38)
because ú D is the minimizer of F(w) ||l Tuuxx X 7712 + 0 [|]l and if all the 2, rotated by R, then the minimizer of will also be rotated by R. Thus,
J(XRJ) E [ XR
15 R X = J(X)R
(39)
We can use this to show that E B rotation matrix, we obtain
XX gu IJ
is a scalar multiple of the identity. Letting R
E E [X ge I1
BB E[X Vi X
(40)
=BB-TI [X TTT X]]
(41)
=Ex(X T [(X)
(42)
=Ex((XRT) J(
(By rotational invariance of dist. of x;))
-EX(RXTJ(X))
(B.c. J(XR))= J
-REx ((X)RR
(43)
REE REE--TX T J.R...
(By Equation
42
This implies that E E (XT Tú
is scalar multiple of the identity.
12

---

Under review as a conference paper at ICLR 2024
Proof of Lemma 4 By Lemma we have E / [X' vy X] cil and E B [XT
c2l
for some scalars C2. Taking the traces of both matrices gives
cid = :od=:t(BBBpm)) i~
yy X]) = 1 [ur(X 8J X)] = 18 D~T T 1 xX 17
(44)
and similarly,
Cad tr( G [X
yr b.) = E [i yeo } E B 10 100 x
(45)
R D.t
Tn x 17
By the definition of 1 as
we have
E 10
XX 81
nB B-T [
5y
x]
E. BiTT
Xx1
ne1 I 022
(46)
E Bö 17
XX y1
E B [
"
1
(47)
d
d
B [10 I
ETx v1
E
E
[ 1 37
D
(48)
0,
(49)
completing the proof of the lemma.
Proof of Lemma 2 For convenience, write A == 1.dd
Then, wish to show that for all
A E (d+1) and w E Rd+1 we have
Ep |I u = C + E - T AG nx 772
(50)
for some constant C' independent of A and w. Define
J1 w) = 5 |G m. 5-
(51)
and
w) B To show that J1(A, w) and J2(A, w) are equal up constant, it suffices show that their gra- dients are identical, since J1 and J2 are defined everywhere on R++1 Idd+1)x(++)) and by the fundamental theorem of calculus.
(52)
Gradients With Respect to First, analyze the gradients with respect w. We have
V (A, w) 2E G
AT (AG B - vp
(53)
and
Vw.2(A,, w) = 2E 2E (AG n 17)
(54)
where we use the convention that the gradient with respect to w has the same shape as w. To show that 1 w) (2,) w). it suffices show that Bü.7G áū = n' G X g. Observe that we can write G as
GD GD= Σ
I
(55)
Given X, the expected value of any odd monomial of the y is equal to O, since and e, are independent of X and have mean O. Thus, we can ignore the blocks corresponding to xą and v/. since the entries of ún and X " contain odd monomials of the once X is fixed.
We now proceed in two steps. First, we deal with the terms corresponding to the lower left block of Gğ and show that their respective contributions to B G lū and nB G5 X y are equal. For B Din GD v. these terms contribute:
E B Σ yar (A 1id E JJ X XAT 1
(56)
= tr X XA_ (.1
(57)
tr (WT X X|A {1A,11a))
(58)
13

---

Under review as a conference paper at ICLR 2024
while for G A'' y, these terms contribute:
nR DNT E yer, A. A. == nB BNT o' X XA,,
(59)
-nE D tr( V0 X XA
(60)
i.w [nx]] "0 A)A
(61)
and since E "/ X] = In Vy
X] by Lemma these two contributions are equal.
Second, we deal with the terms corresponding to the upper right block of Gp- For En- Gl---ipi these terms contribute
G
E D~T Σ yi
(d+ = R D X JA (d.1))))
(62)
E = X ÿú J (d-11)
(B.c. dot product A A,1 is commutative)
while for nE D^ B GDA XT J. these terms contribute
17tE D~T
Σ yr ((t1). ブ= 1/ BNT X JA JA X X
(63)
iml
1 D W0 XA XA.,
(B.c. dot product A (,(111) X ÿ is commutative)
and again these contributions equal by Lemma In summary, have shown that for Eö-TGA A B and vBBB~GG X J. the contributions corresponding to the lower left and upper right blocks of G5 are equal, meaning that i Lū and nil GAA X y are equal. This shows that V (A, w) and (2,, w) are equal.
Gradients With Respect to A. We can compute the gradient of J1 with respect to A as follows:
AA w) Böur VallAG u 2
(64)
= =Enr V A úp) (AG -wp)
(655
BwT (AG w (
(66)
and we can similarly compute
V A2 EDTT V AlIAG D" n v
(67)
EDNT (AG 7x '7) (TT(A))) (AG n
(68)
(AG 11X (G
(69)
Thus, to show that the gradients with respect to A are equal, it suffices to show that
E D G D.T X "u T
(70)
for all
As before, we only need to consider the lower left and upper right blocks of Gp- since the entries of the other blocks have even powers of the yı- First let consider the lower left block. The contribution of the lower left block to u ED G is
E E W
Σ yau
EE u (a+1) T X w(a+ B BJ x.
E
(71)
while the contribution of the lower left block to nt D XT "ju" TWuTii
E n T1
Σ yat mu(d- T7 X,
(72)
14

---

Under review as a conference paper at ICLR 2024
and these two equal by Lemma 4 is
The contribution of the upper right block to B Gb
E B D Wid
yea,
iu w X x (ED-7 XWW1d
(73)
i=1
by the commutativity of the dot product x T J. while similarly, the contribution of the upper right block to nylE B
X jiu Gp is
E 7x yordd
Σ y: = B. yy X Wi:d 2=1
(74)
and these two quantities are equal by Lemma V w),
Thus, we have shown that (A, =
Summary. We have shown that I w) V w2, and VAJI(A,)) V A2,, w).
This completes the proof of the lemma.
Proof of Theorem 7 This theorem follows from Lemma and Lemma To see why, recall from Section that the output of the transformer on the last token Un+1 can be written as w G M where M = W&WQ and w = WVt. h. By Lemma 2 the pre-training loss is minimized when M. G Gpw nx y almost surely (over the randomness of X and y), and when this holds, the output of the transformer on the last token
w GGG w GpManadent (11xT Int1 Σ year Int1;
TT
(75)
as desired.
D MISSING PROOFS FOR SECTION
Proof of Theorem 2 As discussed in Section 4 we can write x; = E//2u; where u; ~ (O, Laxa). Furthermore, we can let U E RRaxd be the matrix whose 7th row is u." Note that X = UE since Σ is symmetric positive-definite matrix. Given D (x Vl..... Mn), we can define Wp = TU + 12 -U 7 - this would be the solution to ridge regression we were given the u, and yı. Using X Uy which implies U X we obtain
u パ.x = (E-1/2x² X 1/2 - -1
(76)
(77)
= Xx 1/2 - T
(78)
(79)
We now change variables in order to reduce Theorem to Theorem 1 using the fact that un and y/ together are from the same distribution as the data studied in Theorem We can write the loss
L(WK, Wo. Wv, = EEpTT ( AT 1
(80)
= =Ep [Bynes [Wa+1[ - hT 01 In+1, DI]
(81)
The solution to ridge regression given the u, and y/ 5 meaning Elya+1 | un+1,
D)
ai " "o,""n+ "n+1 and therefore Elyn+1 In+1, 55 = ub un+1, since un+1 Σ 1/2 10-1//2m44. En41, which an invertible function of In+1. Thus,
Ey... [[n++1 -h T 0nt1)² 11+1, D
(82)
Ey [(u++1 10 D.y un--n11 1) Ent1. D)
(83)
Eyn+: [1 D2 un+1 h bnt1) Int1, D)
(84)
+ 21E E1
[1
2
D.y
un+1) (10 I
D.
D unt1
h vn41) | 2n-1. D).
(85)
15

---

Under review as a conference paper at ICLR 2024
Now, note that Byntikyyo4 (b.24++)((0) D, un+ D
h un+1) | Intl, DD 0 since
T "n - hT bn41 fully determined by #n+1 and D. and [f - 1
D.S un+1 1n+1: DJ =
as mentioned above. Additionally, we can write Ey.. [1n+1 - uh h,, un+1)² In+1, DJ as C 1..
(denoting constant that depends only on 2n+1 and D) since independent of the parameters Wk. WQ. Wv, h. Thus, taking the expectation over 2n+1 and D, for some constant C which is independent of WK, Wq: Wv, we have
L(W,, Wo. Wv, h) = + +B,
(w DE D - h 12+1)
(86)
+ B
1 D. un+1 - W W W WoU++1)2
(87)
= + +BDD x0+1
1
D Un+1 - h W W W. In+1)
(88)
=C + CC+BB D,
DE
un+1
h W W Wo 1212 un+1)
(89)
To finish the proof, observe that we can write
GD
24
(90)
Vi
yi
1/2
0
X/2
0
(91)
34
y/:
1=1
HGP
'H
(92)
1//2
where we have defined H-
and
GD'
Ui
(93)
y/
Yi
Now, let us make the change of variables h = h,, Wv W WK W& and Wq W H-1 Then, we obtain
L(Wk. WQ. hv,
(94)
=C+ E E an+1
(vi D.y un+1 11 W W We. E (22un
(95)
=C + ED
(w D,y: Un+i - (W/) WV H HGGH
(96)
(Wk( Wa 1-1/2221/22 un+1)
(97)
=C+E (10 1,1 (W) WCG''(Wk) ( T
2n+1
W W 11+1)²
(98)
In other words, this is equal to the loss obtained by the transformer corresponding to hi, Wy. W W/ when the data is from the distribution studied in Section 3 The above proof also shows that the transformer corresponding to h, Wv, WK, Wq has the same output on Fa+1, as the transformer corresponding Ni, Wt, W K. W/ does th+1- By Theorem I the output of the latter transformer on u1+1 is
(nU (TT)"" Un-1 Tun- nh" Unn+1
(99)
Σ you! un-1
(100)
n
yi 1111/2 11)x1)TT² 1/2 T 1/2xn411
(101)
Σ
yo T Σ V-1an411 xn+1,
(102)
1=1
16

---

Under review as a conference paper at ICLR 2024
where the learning rate is
E EBnT
110
U
D.
" 18 DNTT
(103)
17 UU y
E Ont
17' X(X T + -
(By definition of (
E B D7 UU
E B
X X +
(B.c. U - 1/2)
B 17 XE-1125-1/XXX 57
E D
(X X + - -x
(104)
E D XE XX- XX y
This completes the proof.
L
E MISSING PROOFS FROM SECTION 5
Proof of Lemma S consider the function
We proceed similarly to the proof of Lemma
1 For convenience, fix D and
g(u) = E wa- [[u In+1 vn+1)² D].
(105)
Then, we have
Vug(u) = -Ezzi [2(u In+1 (n+1)aa411 D].
(106)
Therefore, if up is the minimizer of g(u) (note that üp depends Đ but not 2441 and fn+1),
then for any Rdd we have
Ex B2.11 llTB 1nt1 yo+ (w En+1 Up xn+1) [J (u -p)) Vug(up) 0,
(107)
meaning that for any E Rdd and some Cp which only depends on H and is independent of
E [(u Fa+1 2n+1) 55]
(108)
Bz (uu In+1 B5 In+1 UD Intl - Dn+1)² {)
(109)
E2,--1 1 [uu Int1
U5 xn+1) 1)2 D
(110)
Ba 11.6
[1
Inti yn+t)² D
(111)
+ 2E 2E
[(u In+1 - Up xn+1) ( (uD In+1 - Bo+1) [
(112)
E E2 1 [(u In+1
Up In+ |
(113)
Be. B...i....1 1.
[1 2n Dn+1)² D
(By Equation (
= ||u u + CD-
(114)
Here the last equality is because Jn+1 NO, Laxa), and because Ex,, Un- [ In+1 - yo+
)
D] is constant that depends D but not on u. We can apply this manipulation the loss function:
I(w, .)-
Ep-T
E [[ )n+1)² D)]
(115)
EBnT
Exoti
[1 w To G Mu )² [
(116)
EBiT
Ex [110411111111 w T M. 11 Di]
(117)
EBnT
[T M. M. G Cp]
(By Equation
(144)
+ Here, C' is constant independent of and M. This completes the proof of the lemma.
= EDwTlllz - M. n +C
(118)
17

---

Under review as a conference paper at ICLR 2024
Proof of Lemmal 7 We prove this using Assumption [ imitating the proof of Lemma 3 For con- venience, let (X) Ely [X. Then, the (i.j)-th entry of (X) is [f i = j and (zlf(z;)f we perform the change of variables Iú + Ray for fixed rotation matrix R and all i E [n], then by Assumption (Rr.)
+ o² if + j. since E[eg] and the e; are iid. and independent of X. If
(== meaning that for any rotation matrix R.
Ex[X M(X)X] = (xI(XRR) ) (XXR (XR
(119)
IEExIX M(XR' )XJRR
(120)
XX M(X)XJR
(121)
by the rotational invariance of the Tz- scalar multiple of the identity matrix.
This implies that Ex [X (MMM(X))- = E
(XT yy X] is
Next, we consider Ep-riX T.g. For convenience, let J(X) = Su X]. If we make the change of variables 2; -> Rri. then the joint distribution of y1.-... Nas In+1 does not change by Assumption 1. meaning that "h replaced by Rup- Thus, we can conclude that (X) is equivariant to rotations of all the by R:
J(XR_) E 5
XR E R X| J(X)R*
(122)
Thus,
E T "uu)
Ex(X J(X]
(123)
=Ex((XR)) J(XRT)
(By rotational invariance of NNO, Iaxd))
-REx [[X)] J(X)RR
(By Equation (122]
-RX yu R
(124)
Thus, D [XT Juu is scalar multiple of the identity matrix. The final statement of the lemma
follows by taking the trace of the left and right hand sides.
Proof of Lemma g This follows by the same argument as Lemma 2
here we use Lemma 6 in
order show that only need to consider the lower left and upper right blocks of Gp- and the rest of the proof follows from linear algebraic manipulations and applying Lemma 7 (in place of
Lemma which was used in the proof of Lemma 2).
Proof of Theorem 3 This follows directly from Lemma g since the effective linear predictor being nx J necessary and sufficient condition for minimizing the pre-training loss.
F LEARNING RATE SIMULATIONS
We follow setup similar to von Oswald et al. (2022). Our goal is to verify that trained linear self- attention layer will implement a single step of gradient descent (GD), with a learning rate which matches our theoretical prediction in Theorem [ As in yon Oswald et al. (2022). predictions of our trained linear self-attention layer with step of GD with learning rate y). One important difference between our experiment and that of von Oswald, al. is that we set 1 be our theoretical prediction from Theorem I rather than selecting using line search. As in von Oswald al. 22022), we evaluate the trained linear self-attention layer and the GD baseline on inputs which are scaled by a large amount compared to during training.
compare the
In our setting, the x; 's are 10 dimensional vectors chosen that each entry is uniformly random in [-1, 1] (we note that this is similar to the setting of von Oswald et al. (2022) but different from our theoretical setting). The weight vectors w for each sequence are random Gaussian vectors, chosen independently of the x, 's. (Thus, if we scale the az's by a certain at evaluation time, then this leaves the distribution of w unchanged, but the ys 's are also scaled by a factor of a at test time compared to during training.) In both the training and the evaluation data, d² 0.5.
Ep [ We estimate the learning rate for the GD baseline follows. We predict v) = Ep 18 XX T in Theorem [ We then sample 100,000 sequences in order to estimate the expectations in each of the numerator and the denominator.
18

---

Under review as a conference paper at ICLR 2024
LSA and 1-step GD Difference
1e17
3.5
LSA and GD Diff
LSA and Target Y Diff
2.5
2.0
MSE
1.5
1.0
0.5
0.0
2
16
64 128 200 300 400 500 600 700 800
Scaling X
Figure 1: As the 24 are scaled to larger, the difference between the linear self-attention layer and the ground truth increases much more than the difference between the linear self-attention layer and the GD baseline.
Since the x, 's are 10-dimensional vectors,
I; will be an 11-11mimmme 11-dimensional vector. Thus, to match
y1 theoretical setting, the key and query dimensions of our linear self-attention 11 well. We train the linear self-attention layer using Adam, with a learning rate of 1e-4. We use a batch size of 2048 and train for 12000 steps. We also perform gradient clipping with maximum norm of 10.
In all sequences (at both train and test time), there are 10 context vectors
2;
followed by one
y:
query vector
See Figure D for the results. Our experiment shows that, as the inputs are sealed to be larger, the trained linear self-attention layer becomes farther from the ground truth, while the difference be- tween the linear self-attention layer and the GD baseline does not grow as quickly, which suggests that the GD baseline with our predicted matches the algorithm that the trained linear self-attention layer learns. For instance, in the case where the scaling of the x; is 1, i.e. we do not rescale the I,, the mean-squared еттог between the linear self-attention layer and the ground truth is about 2.6, while the difference between the linear self-attention layer and the GD baseline is about 0.14. Additionally, Figure 2 shows that o² increases (shown in the x-axis) the theoretically predicted learning rate (shown in the y-axis) decreases.
19

---

Under review as a conference paper at ICLR 2024
Theoretical Prediction for Learning Rate using
Predicted LR
0.140
0.135
0.130
0.125
0.120
0.5
1.5
2
Output variance a
Figure 2: Change in learning rate as we vary o²
20


Reviews:
[{"review_id": "A85bbk0Vvm", "reviewer": "Reviewer_ksN8", "summary": "This paper theoretically analyzed the one-layer linear self-attention layer on the linear regression teacher model. The authors proved that after pertaining to this one-layer transformer under square loss, the minimizer we got is equivalent to a single-step gradient descent (GD) on the least-squares linear regression problem. This paper also considered covariate shifts for the data distribution, which correspond to preconditioned GD. Finally, the authors claimed for rotational invariant nonlinear teacher models, the global minimizer of the transformer is still equivalent to one step GD on least-squares linear regression.", "questions": "1. You may need to briefly explain the parameter $\\eta$ in Eq. (1).\n\n2. In the second paragraph on page 4, $v_n=\\\\begin{bmatrix}x_i\n\\\\\\\\ 0 \\\\end{bmatrix}$ should be $v_n=\\\\begin{bmatrix}x_{n+1}\n\\\\\\\\ 0 \\\\end{bmatrix}$.\n\n3. In Eq. (10), the number of training parameters is $d^2+d$ and we consider population squared loss for training. Does that mean this model is under-parameterized and has a unique global minimizer? And when could the constants in Lemma 1 and Eq. (13) in Lemma 2 be zero? More specifically, is the minimizer constructed in Theorem 1 unique and when will it attain zero training loss? \n\n4. How large the learning rate $\\eta$ is? Following the remark after Theorem 1, we know the global minimizer is equivalent to a step gradient descent on the empirical loss of the least squares problem with zero initialization and learning rate $\\eta$ defined in Theorem 1. How large $\\eta$ is, compared with the largest eigenvalue of the Hessian matrix $H$ of this least squares problem? Is it just close to or larger than the maximal learning rate $2/\\lambda_{\\max}(H)$?\n\n5. In Section 4, when we consider data covariance in $\\Sigma$, why do we renormalize back by $w\\sim\\mathcal{N}(0,\\Sigma^{-1})$? Can we consider $w$ has another different population covariance like [2]?\n\n6. In Theorem 3, when defining $\\eta$, what is $\\mathcal{D}$? No definition of this distribution.\n\n7. In the proof of Lemma 1, after Eq. (29), why is the minimizer of $g(u)$ given by $\\hat{w}_{\\tilde{D}}$? Here, for $g(u)$, you only have one data point. Can you explain more?\n\n8. It may be worthy to mention or compare with some of the references among [1] and [3-7].\n\n\n=================================================================================================\n\n[1] Bietti, et al. \"Birth of a Transformer: A Memory Viewpoint.\"  \n\n[2] Wu and Xu. \"On the Optimal Weighted $\\ell_2 $ Regularization in Overparameterized Linear Regression.\"  \n\n[3] Takakura and Suzuki. \"Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input.\" \n\n[4] Tarzanagh, et al. \"Margin Maximization in Attention Mechanism.\" \n\n[5] Tarzanagh, et al. \"Transformers as support vector machines.\" \n\n[6] Bai, et al. \"Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection.\"  \n\n[7] Guo, et al. \"How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations.\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is written very clearly in a way that highlights the main analysis techniques in the main body. It also provides enough summaries for some concurrent works and literature reviews. The main message of this paper is clear. The authors theoretically analyze the in-context learning capability for the self-attention layer in the linear regime and make the connection with the Bayes-optimal predictor of the linear regression model. This result provides a number of natural directions for future theoretical study in transformers.", "weaknesses": "1. The dataset assumption is simple and the authors only considered i.i.d. data sequence with linear teacher model. This setting helps the analysis but may not be able to fully capture the properties of the self-attention layer. Besides, the proofs rely on the rotational invariance of Gaussian distribution. It would be interesting to generalize the results in non-Gaussian datasets or consider more dependent structures in the data sequence, like the Bigram language model in [1].\n\n2. This paper focuses on the global minimizer of the population square loss of the self-attention layer which simplifies the analysis. It would be natural to consider the minimizer of the empirical loss during the pre-training process and how the minimizer of the GD or stochastic GD with finite step sizes generalizes in the test point. \n\n3. Further experiments and simulations should be presented for completeness. For instance, the training dynamic of nonlinear/multi-layer transformers with nonlinear target functions that are defined in Section 5. This will help us know the limitations of the current theory and potential interesting directions for future analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tuVq6420DI", "reviewer": "Reviewer_VQYN", "summary": "This paper considers a one layer self-attention model with linear attention and shows that one step of gradient descent is the optimal in-context learner in this case. Specifically, they consider a synthetic noisy linear regression task and show that when the covariates are drawn from a standard Gaussian, the model implements one step of GD, which is also the global minimizer of the pretraining loss. If the distribution of the covariates if changed to a non-isotropic Gaussian, it now implements pre-conditioned GD. On the other hand, when using a nonlinear model to generate the data, it still implements a single step of GD.", "questions": "Please see the weaknesses section. My main concern is that this paper does not offer new insights regarding in-context learning in transformers (that has not been discussed in one of the prior works), and also does not use any new proof techniques. It would be interesting to analyze multi-head attention or multilayer transformers, as the authors discuss in the conclusion.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper takes a step to improve the theoretical understanding of in-context learning in transformers, which is an important topic.", "weaknesses": "While this is an important topic, the paper does not seem to make a significant contribution. The main drawback is that it considers a one layer attention model, which has been studied extensively for the developing theoretical understanding of in-context learning.\n\nIn the first case, the only contribution seems to be that using an appropriate step size allows the resulting solution to be a global minimizer of the pretraining loss. This does not seem to add to the understanding of transformers, as it was already shown in [1] that transformers implement one step of GD. Similarly, the result in the third case is also not very informative. Given that this is a one layer model, it is not surprising that it implements one step of GD, even when the target function is nonlinear. \n\n[1] von Oswald et al. 'Transformers learn in-context by gradient descent', ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZR6ID5plzd", "reviewer": "Reviewer_1cqQ", "summary": "In this paper, the authors provide a theoretical analysis of transformers equipped with a single layer of linear self-attention, trained on synthetic noisy linear regression data. The primary focus of this paper lies in exploring in-context learning pretraining scenarios, where the training data consists of pairs (x_i, y_i) with associated ground truth, and the evaluation is based on the Mean Squared Error (MSE) metric for the test point (x', y').\n\nThe key findings presented in this paper can be summarized as follows: (1) Under the assumption of linear noisy ground truth, when x_i samples are drawn from an isotropic Gaussian distribution, the one-layer transformer model that minimizes the pretraining loss effectively corresponds to a single step of Gradient Descent (GD) applied to the least-squares linear regression problem. (2) When x_i samples are drawn from a non-isotropic Gaussian distribution, the optimization process becomes a preconditioned GD. The authors shed light on this aspect, showcasing the connection between the nature of the input distribution and the optimization approach. Furthermore, The paper goes beyond linear cases, demonstrating that the findings can be extended to non-linear scenarios under specific symmetric conditions.\n\nIn conclusion, I strongly recommend accepting this paper for the following reasons: (1) The paper demonstrates exceptional organization, making it highly accessible and comprehensible for the readers. (2) The topic addressed in this paper holds paramount significance within the Language Model (LLM) domain, contributing to our understanding of key theoretical aspects. (3) The paper introduces some innovative results, particularly in the sections related to preconditioning and non-linear extensions. These novel findings are likely to ignite further research and inspire intriguing follow-up studies.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Overall, this paper has the potential to inspire and stimulate further research in this area.\n1. The organization of the paper is well-structured, making it accessible to a broad readership.\n2. The paper addresses a crucial topic in the realm of Language Model (LLM) research, shedding theoretical insights on transformers under in-context learning scenarios.\n3. The results presented in the paper are noteworthy, particularly the connections made in Theorem 1, including the proof of global minimization and its equivalence to a single step of gradient descent. The exploration of non-isotropic Gaussian distributions leading to preconditioned GD is an interesting and novel aspect. Additionally, the extension to non-linear cases adds depth to the research.\n4. The paper is well-written and effectively communicates its findings and insights.", "weaknesses": "While the paper is commendable, there are a couple of minor questions and potential areas for further investigation:\n\n1. The usage of the statement \"(Wk, Wq, Wv, h) is a global minimizer\" in Theorem 1 raises questions about the specifics of this minimization process. Further clarification or details regarding this construction might be beneficial for readers.\n\n2. The reviewer suggests that, in in-context learning regimes, the downstream phase is crucial. Encouraging future research that delves into this aspect could be valuable for a more comprehensive understanding of the subject.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n3z3Y6l28Q", "reviewer": "Reviewer_nLRu", "summary": "This paper theoretically studies how one-layer transformers with linear attention implement one step of gradient descent on least-squares linear regression. The results include the cases when the covariates are from a standard Gaussian, a  non-standard Gaussian, and when the target function is nonlinear. The conclusion covers the global convergence of the network.\n\n\n---------------------------------------------------------------------\n\nAfter rebuttal, I tend to maintain the score of 5. The main concern is the significance to the community. \n\n**Practical Insight**: It is still a major concern. I am not clear on what I can learn from this paper. It makes the result less interesting and significant to me. For example, can this paper provide an explanation for any phenomenon in in-context learning in practice? How can this paper guide the training in practice?\n\n**Experiments**: Generally, I am satisfied with the efforts of the authors. Since I only specified one experiment, I will not treat the experiment part as a big weakness, although I expect to see more experiments. \n\n**Fully connected neural networks and Assumption 1**: Good, I am satisfied with this result.\n\n**Contribution compared to [Zhang, et al. 2023]**: OK. Although [Zhang et al., 2023] was posted online 3.5 months before ICLR submission deadline, I agree it can be treated as a concurrent work.", "questions": "1. Can you verify that the $\\eta$ in Theorem 1,2 are as predicted by experiments? Specifically, can you show how $\\sigma^2$ in Theorem 2 affects $\\eta$ by experiments?\n\n2. I don't know why fully-connected neural networks satisfy Assumption 1 (1). Can you provide a proof for this claim? \n\n3. Without section 5, the contribution compared with Zhang el at., 2023 will only be incremental. Why do you assign too much content to Section 3? I think it is better to enlarge the content of Section 5. \n\n[Zhang et al., 2023] \" Trained transformers learn linear models incontext.\"", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The significance is good since the studied problem is essential and interesting to the community. The paper is overall well-written with good clarity. This paper provides a comparison with existing works and concurrent works. The contributions include that it provides a global optimal analysis when constructing a linear-attention Transformer to implement gradient descent. Meanwhile, it shows analyses on non-standard Gaussian inputs and non-linear target functions.", "weaknesses": "1. This paper lacks empirical justification. \n2. I am not sure about the practical insight from the theoretical analysis of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)