Transformers learn to implement preconditioned gradient descent for in-context learning
Authors: Kwangjun Ahn
Xiang Cheng
Hadi Daneshmand
Suvrit Sra
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformers
In-context learning
adaptive gradient methods

Abstract: Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: \emph{Can transformers learn to implement such algorithms by training over random problem instances?} To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances of linear regression. For a single attention layer, we prove the global minimum of the training objective implements a single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy.  For a transformer with $L$ attention layers, we prove certain critical points of the training objective implement $L$ iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.

Transformers learn to implement preconditioned gradient descent for in-context learning
Kwangjun Ahn" MIT EECS/LIDS kj jahhmmit.... edu
Xiang Cheng" MIT LIDS chengxOmit.edi
Hadi Daneshmand MIT LIDS/FODSI hdanesh@mit. edu
Suvrit Sra TU Munich / MIT suvrit@mit, edu
Abstract
Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: Can transformers learn to implement such algorithms by training over random problem instances? To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances linear regression. For a single attention layer, we prove the global minimum of the training objective implements single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy. For transformer with L attention layers, we prove certain critical points of the training objective implement L iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.
Introduction
In-context learning (ICL) is the striking capability large language models: Given a prompt containing examples and query, the transformer produces the correct output based the context provided by the examples, without adapting its parameters (Brown et al., 2020; Lieber et al., 2021; Rae al., 2021; Black et al., 2022). This property has become the focus of body of recent research that aims to shed light on the underlying mechanism of large language models (Garg et al., 2022; Akyürek al., 2022; Oswald al., 2023; Li and Malik, 2017; Min al., 2021; Xie et al., 2021; Elhage et al., 2021; Olsson et al., 2022).
A line of research studies ICL via the expressive power of transformers. Transformer architectures are powerful Turing machines, capable of implementing various algorithms (Pérez al., 2021; Wei et al., 2022). Given an in-context prompt, Edelman et al. (2022); Olsson et al. (2022) argue that transformers are able implement algorithms through the recurrence of multi-head attentions to extract coarse information from input prompts. Akyürek al. (2022); von Oswald et al. (2023) assert that transformers can implement gradient descent on linear regression encoded in a given input prompt. is thought provoking that transformers implement such algorithms.
Although transformers are universal machines to implement algorithms, they need specific parameter configurations for achieving these implementations. In practice, their parameters are adjusted via training using non-convex optimization over random problem instances. Hence, it remains unclear whether this non-convex optimization can be used to learn algorithms. The present paper investigates the possibility of learning algorithms via training over random problem instances.
Equul contribution, alphahetical order.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

More specifically, we investigate the learning of gradient-based methods. It is hard to mathematically formulate what means learn gradient descent for general functions with transformers. Yet. Garg et al. (2022) elegantly examine it in the specific setting of ICI. for learning functions. Empirical evidence suggests that transformers indeed learn to implement gradient descent, after training on random instances of linear regression (Garg et al., 2022; Akyürek et al., 2022; von Oswald et al., 2023). Motivated by these observations, we theoretically investigate the loss landscape of a simple transformer architecture based on attention without softmax (Schlag et al., 2021; von Oswald et al., 2023) (see Section 2 for details).
Summary of our main results. Our main contributions are the following:
We provide a complete characterization of the global optimum of single-layer linear transformer. particular, we observe that, with the optimal parameters, the transformer implements a single step of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the distribution of input data but also to the variance caused by data inadequacy. We present this result in Theorem 1 in Section 3.
Next, we focus on a subset of the transformer parameter space, defined by a special sparsity condition (8). Such a parameter configuration allows to formulate training transformers as a search over k-step adaptive gradient-bor based algorithms. Theorem 2 characterizes the global minimizers of the training objective of a two-layer linear transformer over isotropic regression instances, and shows that the optima correspond to gradient descent with adaptive stepsizes. For multilayer transformers, Theorem 3 demonstrates that gradient descent, with a depende preconditioning, can be derived from critical point of the training objective.
Finally, we study the loss landscape in the absence of the sparsity condition (8), which goes beyond searching over conventional gradient-based optimization methods. In this case, we prove and interpret the structure of a critical point of the training objective. We show that a certain critical point in parameter space leads to an intriguing gradient-basee algorithm that simultaneously takes gradient steps preconditioned by data covariance, and applies a linear transformation to further improve the conditioning. In the specific case when data covariance is isotropic, this algorithm corresponds to the GD++ algorithm of von Oswald et al. (2023) which is experimentally observed to be the outcome of training.
We empirically validate the critical points analyzed in Theorem 3 and Theorem 4. For a transformer with three layers, our experimental results confirm the structural of critical points. Furthermore, we observed the objective value associated with these critical points is close to 0, suggesting that the critical points might be global optima. These experiments substantiate our theoretical analysis and suggests that our theory indeed aligns with practice. Code for our experiments available at https://g/ thub. com/ ctt/gxiinnnnnn wwinearrransfee
1.1
Related works
The ability of neural network architectures to implement algorithms has been investigated in various context. The seminal work by Siegelmann and Sontag (1992) investigate the Turing completeness of recurrent neural networks. Despite this computational power, training recurrent networks remains a challenge. Graves et al. (2014) design an alternative neural architecture known as the neural Turing machine, building on attention layers introduced by Hochreiter and Schmidhuber (1997). Leveraging attention, Vaswani et al. (2017) propose transformers powerful neural architectures, capable of solving various tasks in natural language processing (Devlin et al., 2019). This capability inspired a line of research that examines the algorithmic power of transformers (Pérez et al., 2021; Wei et al., 2022: Giannou ct al., 2023; Akyürek et al., 2022: Olsson et al., 2022). What sets transformers apart from conventional neural networks is their impressive performance after training. In this work, we focus on understanding how transformers learn to implement algorithms by training over problem instances.
line of research investigates how deep neural networks process data across their layers. The seminal work by Jastrzebski et al. (2018) observes that hidden representations across the layers of deep neural networks approximately implement gradient descent. Recent observations provide novel insights into the working mechanism of ICL for large language models, showing they can implement optimization algorithms across their layers (Garg et al., 2022; Akyürek et al., 2022; von Oswald et al., 2023). Moreover, Zhao et al. (2023); Allen-Zhu and Li (2023) observe transformer perform
2

---

dynamic programming to generate text. In this work, we theoretically study how transformer learns gradient-based algorithms for ICL.
We discuss here two related works (Zhang et al., 2023; Mahankali et al., 2023) that appeared shortly after publication of our original draft. Both of these studies focus single layer attention network (see Section 3). Zhang et al. (2023) prove the global convergence of gradient descent to the global optimum whose structure is analyzed independently from this study and the same as that in Theorem 1. Mahankali et al. (2023) also characterize the global minimizer of a single layer attention without softmax for different data distribution. In addition to results for single-layer attention, we analyze the landscape of two and multi-layer transformers.
2 Setting: training linear transformers over random linear regression
In order to understand the mechanism of ICL, we consider the setting of training transformers over the random instances of linear regression, following (Garg al., 2022; Akyürek et al., 2022; von Oswald et al., 2023). In particular, the random instances of linear regression are formalized as follows.
Data distribution: random linear regression instances. Let (x( Rd be the covariates drawn i.i.d.
from distribution Dx, and w, E R be drawn from Dw. Let X E [==-1)xd be the matrix of
covariates whose row i contains tokens g((), Given x:(0's and Ws, the responses are defined as y [(x(1) W.) w.)..... ...., (x(x), w.)] E R". Define the input matrix Zo as
[2)) (y())
x(2)
(()) x xln+1)] (n))
Zo (2 222) ((n) 2(+11)]
R R(d++)x(at)) (1)
y(2)
where zero in the above matrix is used to replace the unknown response variable corresponding to ((a+1), Then, our goal predict w.' 2(x+1) given Zy. In other words, the training data consists of pairs (Zo, w wTTa(2 ((*+1)) for (f) Dx and w, Dw. We then consider training transformers over this data distribution.
Self-attention layer without softmax. Following (Schlag et al., 2021; Oswald al., 2023), consider the linear self-attention layer. To motivate, we first briefly review the standard self-attention layer (Vaswani et al., 2017). Letting E p(d+1)x(x+1) be the input matrix with n + tokens in smax is parametric map defined as
Rd+1, a single-head self-attention layer denoted by Amtx
At.
smax
(Z)=
W.ZM smax Z WT W2Z), M
I 0
R (+++))x(++1)
where W.,, WR, W, E (d++1)) are the (value, key and query) weight matrices, and smax (-) is the softmax operator which applies softmax operation to each column of the input matrix. Note that the prompt is asymmetric since the label for ((*+1) is excluded from the input. To reflect this asymmetric structure, the mask matrix M included in the attention. In our setting, we consider the self-attention layer that omits the softmax operation in (2). In particular, we reparameterize weights as P = W. E x(d+1)x(d+)) and Q = W₁ W E R (dd+ (d+11) and consider
Attn Z) PZM (ZQZ).
(3)
At first glance, the omission of the softmax operation (3) might seem over-simplified. But, (von Oswald et al., 2023) proves such attention can implement gradient descent, and we will prove in Lemma 1 that it can also implement various algorithms solve linear regression in-context.
Architecture for prediction. We now present the neural network architecture that will be used through- out this paper. For the number of layers L, define L-layer transformer stack of L linear self-attention blocks. Formally, denoting by Z, the output of the eh layer attention, we define
1 n
Ze+1 Z5 Attmp, for l = 0, L,, 1,
The scaling factor 1/1 is used only for ease of notation and does not influence the expressive power of the transformer. Given Zz, define TF (Zo (Pt, Qc..-...... -[Zz]
(2+1). i.e. the (d + 1, n + 1) entry of Zz. The reason for the minus sign is to be consistent with (von Oswald
3

---

et al.. 2023), and we will clarify such choice in Lemma 1. For training. the parameters are optimized
minimize in-context loss
f ((Pt, Qc) (Zo, ( (P (1),O))) + w. (o]x((+t)))
(5)
Goal: the landscape analysis of the training objective functions. We are interested in understand- ing how the optimization leads to in-context learning. We investigate this question by analyzing its loss landscape. Such analysis is challenging due to two major reasons: (i) f is non-convex in parameters {P.. Qi3 even single layer transformer. (ti) The cross-product structures in attention makes f a highly nonlinear function in its parameters. Hence, we analyze a spectrum of settings from single-layer transformers multi nulti-layer transformers. For simpler settings such as single-layer transformers, we prove stronger results such as the full characterization of the global minimizers. For networks with more layers, we characterize the structure of critical points. Furthermore, we provide algorithmic interpretations of the critical points. Table summarizes our results for various parameteric models.
Results
xx(
Wx
Setting
Guarantees
Theorem
(0 Σ
(00,1) I)
single-layer
global minimizers
Theorem 2
(0,
N/
two-layer + symmetric (8)
global minimizers critical points critical points global minimizers
Theorem 3
N(O0 2
N(O, 2
multi-layer + (8)
Theorem 4
N(O,E)
N(0, E E-I)
multi-layer + (11)
Theorem I)
()O.
single-layer + ReLU activation
Table 1: Summary of our analyses for various models and input distributions. The additional
conditions (8) and (11) are about the sparsity structure of parameters. In addition, "symmetric (8)" means we additionally impose the weights to be symmetric.
Remark (Optimizing (5) vs. practical transformer optimization). Interestingly, a recent work by Ahn al. (2023) reports that common optimization algorithms such as SGI/ADAM behave remarkably similarly on the (linear Transformers linear regression) problem as they do on (practical transformers real language modeling tasks). In particular, they reproduce several distinctive features of transformer optimization under a simple shallow linear transformer. This work suggests that (linear transformer + linear regression) may serve as good proxy for understanding practical transformer optimization.
3 The global optimum for single-layer transformer
For the single layer case of L the following result characterizes the optimal parameters Po and Qo for the in-context loss (5).
Theorem (Single-layer; non-isotropic data). Assume that vector ((6) is sampled from (0, 2), i.e. a Gaussian with covariance Σ UAUT where A diag()1..... >a). Moreover, assume that W, is sampled from (0, Ia). Then, the following choice of parameters
Oaxd 0] O 1
Udiag
UT
Po =
Qo =
(En 2
(6)
a global minimizer of f(P. ()P.Q) up re-scaling, i.e., Po + Po and Qu x Qu for a sealar y
See Appendix A for the proof of Theorem 1. In the specific case when the Gaussian is isotropic, i.e..
Ia. the optimal Qu has the following simple form
1
Ta 0 0
Qo n + (
Up scaling, the above parameter configuration is equivalent to the parameters used by von Oswald al. (2023) to perform one step of gradient descent. Thus, in the single-layer setting, the in-context loss is indeed minimized by a transformer that implements the gradient descent algorithm.

---

More generally, when the in-context samples are non noniisotropi, the transformer learns to implement one step of a preconditioned gradient descent shall detail in Lemma Here the "preconditioning matrix" given in (6) has interesting properties:
When the number of samples is large, the first submatrix of Qu approximates Σ-1, the inverse of the data covariance matrix, which is also close to the Gram matrix formed from x(1) x((), Hence the preconditioning can lead considerably faster convergence rate when Σ is ill-conditioned.
Moreover, 7 Σx dς in (6) acts as regularizer. becomes more significant when is small and variance of the (3)) high. Such an adjustment resembles structural risk minimization (Vapnik, 1999) where the regularization strength is adapted the sample size.
4 Multi-layer transformers with sparse parameters
Theorem 1 proves a single layer of linear attention can implement a single step of preconditioned gradient descent. Inspired by this result, we investigate the algorithmic power of the linear transformer architecture. We show that the model can implement various optimization methods even under sparsity constraints. In particular, impose the following restrictions the parameters:
P. = Oaxd 0
Q. Q-
A
0 0 where A₁ E Rddd
(8)
The next lemma proves that forward pass of a L-layer transformer, with the parameter configuration (8) is the same as taking steps of gradient descent, preconditioned by At.
Lemma 1 (Forward pass a preconditioned gradient descent). Consider the I--laeer linear transfomer parameterized by Ag...., AL-1 as in (8). Let yi (n+1) be the (d 1, + 1)-th entry of the E-th layer output, i.e., Vi. (n+1) (Z)l(at1),(nt)) for l = 1....,A L. Then, it holds that
Vi (n+1) - (x((+1)), wli)) where (wtt) is defined as wh = 0 and as follows for l = l,..., 1:
W ui R...
(w)e)
where
Rw, (w)
(w x₁ W. T xx))2. 2n Σ
1
(9)
See Subsection C.I for proof. The iterative scheme (9) includes various optimization methods including gradient descent with A₁ = yela. and (adaptive) preconditioned gradient descent, where the preconditioner A depends on the time step. In the upcoming sections, we characterize how the optimal } are linked to the input distribution.
4.1 Warm-up: optimal two-layer transformer with symmetric weights
For the rest of this section, will study the optimal parameters for the in-context loss under the constraint of Eq. (8). Later in Section 5, we analyze the optimal model for more general parameters. For two-layer transformer, the next Theorem proves the optimal in-context loss obtains the simple gradient descent with adaptive coordinate-
stepsizes.
Theorem 2 (Global optimality for the two-layer (symmetric) transformer). Consider the optimization of in-context loss for two-layer transformer with the parameter configuration in Eq. (8), and A Ag are symmetric matrices. More formally, consider
additionally assume that AI,
{{ Pe = Oaxd
-:
0
-
min A1.Aq symmetric"
1
Assume (66 l.d. (O, La) and w. ((, Ia); then, there are diagonal matrices A and Ag that are a global minimizer of f.
Combining the above result with Lemma concludes that the two iterations of gradient descent with coordinate-we adaptive stepsizes achieve the minimal in-context loss for isotropic Gaussian inputs. Gradient descent with adaptive stepsizes such as Adagrad (Duchi et al., 2011) are widely used in machine learning. While Adagrad adjusts stepsize based on the individual problem instance, the algorithm learned adjusts its stepsize to the underlying data distribution.
5

---

4.2 Multi-layer transformers
We now turn to the setting of general L-layer transformers, for any positive integer L. The next theorem proves that certain critical points of the in-context loss effectively implement specific preconditioned gradient algorithm, where the preconditioning matrix is the inverse covariance of the input distribution. Before stating this result, let us first consider motivating scenario in which the data-covariance matrix is non-identity:
Linear regression with distorted view of the data: Suppose that w. ~ ~,,(,, and the latent covariates are æ(1),_... ,.,,,..,((tt)) I(n+1), drawn i.i.d from (0, We are given ,..,,..., (y)) with y (I(), W.) However, we do not observe the latent covariates ε Instead, we observe the distorted covariates 20()) Wx( where W E Rdxd is distortion matrix. Thus the prompt consists of (z(), y()),... ,..... (z("), () as well as x(*+1), The goal is still to predict y y(u+1), Note that this setting is quite common in practice, when covariates are often represented in an arbitrary basis.
Assume that Σ =- WWT O. We verify from our definitions that for w. :- E-1/2 y()) (I()), w.). Furthermore, 21)) ~ NO, Σ) and w, ~ ~V(0,, From Lemma 1, the transformer with weight matrices (Ao...., AL-I3 implements preconditioned gradient descent with respect to Rw, (w) 2 (w )² XX (w w.), with X = [xx(1) ................................ [.)]]. Under this loss, the Hessian matrix V (Z.,..(1)) (w)
XXT (at least in the case of large n)) For any fixed prompt, Newton's
method corresponds to A, x (XXT)=², which makes the problem well-conditioned even if Σ is
very degenerate. As we will see in Theorem 3 below, the choice of A; x Σ- E [XX]] appears to be a stationary point of the loss landscape, in expectation over prompts.
Before stating the theorem, we introduce the following simplified notation: let A := (A RLxxxd, We use [(A) to denote the in-context loss of / ((P,, 1 as defined (5), when Q; depends Ad, and P. is a constant matrix, described in (8).
Theorem 3. Assume that 16)) ~ (1,, E) and w. (0, for = 1...., and for some x O. Consider the optimization of in-context loss for a k-layer transformer with the the parameter configuration in Eq. (8) given by:
min S(A).
(A.
Let C. IR. Ixxxx be defined as follows: A E S if and only if for all O,... ......, L 1, there exists
scalars ag E R such that A; a.E Then
inf (A,B)ES
Σ AA,
B) 0,
t=0
where V f denotes derivative wrt the Frobenius norm ||A:l|
As discussed in the motivation above, under the setting of A; a:E-1, the linear transformer implements an algorithm that is reminiscent of Newton's 's method (as well as a number of other adaptive algorithms such as the full-matrix variant of Adagrad); these can converge significantly faster than vanilla gradient descent when the problem ill-conditioned. The proposed parameters A; in Theorem 3 are also similar to A,'s in Theorem when is large. However, in contrast to Theorem 1, there is no trade-off with statistical robustness; this is because w, has covariance matrix E- in the Theorem 3, while Theorem I has isotropic w..
Unlike our prior results, Theorem only guarantees that the set S of transformer prameters satisfying L-1 (A x Σ -ー1 i=o essentially? contains critical points of the in-context loss. However, in the next section, we show experimentally that this choice of A, 's does indeed seem to be recovered by training.
We defer the proof of Theorem to Subsection B.2. Due to the complexity of the transformer function, even verifying critical points can be challenging. We show that the in-context loss can be equivalently written as (roughly) a matrix polynomial involving the weights at each layer. By
AA subtle issue is that the infimum may not be attained, so it is possible that S contains points with arbitrarily small gradient, but does not contain point with exactly 0 gradient.

---

exploiting invariances in the underlying distribution of prompts, we construct flow, contained entirely in S, whose objective value decreases as fast as gradient flow. Since is lower bounded, we conclude that there must be points in S whos gradient is arbitrarily small.
4.3 Experimental validations for Theorem 3
We present here empirical verification of our results in Theorem 3. We consider the ICL loss for linear regression. The dimension is d = 5, and the number of training samples in the prompt is n = 20. Both g(()) ~ (O, 2) and w, ~ N0, )) where = U DU, where U is a uniformly random orthogonal matrix, and D is fixed diagonal matrix with entries (1, 1, 0.25, 0.0625, 1).
We optimizes f for a three-layer linear transformer using ADAM, where the matrices Ag. A1. and A2 are initialized by i.i.d. Gaussian matrices. Each gradient step is computed from minibatch of size 20000, and we resample the minibatch every 100 steps. We clip the gradient of each matrix to 0.01. All plots are averaged over 5 runs with different U (i.e. Σ) sampled each time.
Figure 1d plots the average loss. We observe that the training converges an almost 0 value, suggesting the convergence to global minimum. The parameters at convergence match the stationary point introduced in Theorem 3, and indeed appear to be globally optimal.
To quantify the similarity between Au, A2 and Σ-1 (up to scaling), we use the normalized Frobe- nius norm distance: Dist(M, I) mina
[|M aa-Ill (equivalent to choosing a := à Σi= M(i, i)).
[M] M This is essentially the projection distance M/II onto the space of scaled identity matrices.
We plot Dist (A, I). averaged over 5 runs, against iteration in Figures 1a, la,1b. In each plot, the blue line represents Dist(E A, I), and we verify that the optimal parameters are converging the critical point introduced in Theorem 3, which implements preconditioned gradient descent. The red line represents Dist (Ai, K /); it remains constant indicating that the trained transformer is not implementing plain gradient descent. Figures 2a-2c visualize each / A,E4/2 matrix at the end of training further validate that the learned parameter described in Theorem 3.
Ao
Detence
A S
tdl saico podne soleo vodaa
50c toles 20i0 soles vodoo
3id0e talec tales vodo
Beration
beration
leration
(a) 12 A (b) (2 A I) (c) D2 /
(d) log(Loss)
Figure 1: Plots for verifying convergence of general linear transformer, defined in Theorem 3. Figure (d) shows convergence of loss to O. Figures (a).(b)(c) illustrate convergence of A, 's to identity. More specifically, the blue line represents 1 1/21/2 T), which measures the convergence to the critical point introduced in Theorem 3 (corresponding to - ppecoondiiion gradient descent). The red line represents Dist(A,, I); it remains constant indicating that the trained transformer is not implementing plain gradient descent.
5 Multi-layer transformers beyond standard optimization methods
In this section, we study the more general setting of
P =
B 0
0
Q: =
A 0
where A,, B, E Bdddd
(11)
Note that A,, B, are not constrained to be symmetric. Similar to Section 4, we introduce the following simplified notation: let A := {A to denote the in-context loss of f ((P,, ).) as defined (5), when P. and Q2 depend on B, and A, described in (11).
L-1 e gxxxdd and B (B:) E [xxxxd We use f(A, B)
7

---

443
000
642
000
212
6.11
411
A00
200
6014
3.99
400
644
6:5
6:5
4.4
42
FM
2:
1:10
3:5
442
2:20
2:00
9.0
1.4
EM
2:
214
3.5
425
2.1
50
204
5.5
AM1
0:30
1:79
(a) Visualization of 5/4/2 A/ (h) Visualization of F,1/2 A, (c) Visualization of 5/ 5/321/2
Figure 2: Visualization of learned weights for the setting of Theorem 3. We visualize each Ed/2 A,E matrix at the end of training. Note that the optimized weights match the stationary point discussed in Theorem
With this relaxed parameter configuration, it turns out transformers can learn algorithms beyond the conventional preconditioned gradient descent. The next theorem asserts the possibility of learning a novel preconditioned gradient method. Let L be a fixed but arbitrary number of layers.
Theorem 4. Let E denote any PSD matrix. Assume that x() sid N(0 E) and w, ~ NO,EE for n,, n., and for some Σ> > Consider the optimization of in-context loss loss for L-layer linear transformer with the the parameter configuration in Eq. (11) given by:
i= L...,
min f (A,
(A..B.l:0 B.
Let S C. RR2xxxddd be defined follows: (A, B) E S if and only for all E (o..... .......,,, ..., there exists scalars a2 by E TR such that A₁ a.E- and B, = b, Then
EI + VV [A, B)lL
inf (A,B)ES
(12)
where V A denotes derivative wrt the Frobenius norm ||A:|]
In words, parameter matrices in S implement the following algorithm: -k==--111--- 1i=0 plays the role of a eessribution-dependdee transforms the covariates themselves to make the Gram matrix have better condition number with each iteration. When the Σ I, the algorithm implemented by A; x I, b, x / is exactly the GD++ algorithm proposed in (von Oswald et al., 2023) (up stepsize).
preconditioner for the gradient steps. At the same time, B; = b,
The result in (12) says that the set S essentially' contains critical points of the in-context loss f(A, B). In the next section, we provide empirical evidence that the trained transformer parameters do in fact converge to a point in S.
5.1 Experimental validations for Theorem 4
The experimental setup is similar to Subsection 4.3: we consider ICL for linear regression with n 10, d = 5, with (?) ~ (O, Σ) and w, (0, E-1), where Σ = U DU, where U is a uniformly random orthogonal matrix, and 1) is fixed diagonal matrix with entries (1, 1, 0.25, 0.0625, 1). We train a three-layer linear transformer, under the constraints (11) which is less restrictive than (8) in Subsection 4.3. We train the matrices Ag, A1, A2, Bo, B1 using ADAM with the same setup in Section Subsection 4.3. We repeat this experiment 5 times with different random seeds, each time we sample a different U (i.e. E)
Once again, similar to the case Theorem 3, the infimum may attained, so it possible that contains points with arbitrarily small gradient, hut does not contain point with exactly 0 gradient. "Ntte that the objective function does not depend on B2.

---

In Figure 3c, we plot the in-context loss through the iterations of ADAM; the loss appears to be converging to 0, suggesting that parameters are converging to the global minimum.
Bo
B1
0.5
Disiance to ld
Disiance to
6.3-
log0
20100
A3500 godoo
azdeo 1005co
20100
43500
53000
a3600
100520
20003
40000
53000
sodeo
200020
lteration
lteration
(a) (is 1)
(b) Dist B I)
(c) log(Loss)
Distance to id
0.2
0.2-
Distance to
0.5-
Ao 2 A3
Distance
0.5-
0.4
6.4
A
A
F
A,E
AA2
40deo
aodoo
03000
20300
42000
100bco
Iteration
Iteration
(d) Distances for Au
(c) Distances for A1
(f) Distances for A2
Figure 3: Plots for verifying convergence of general linear transformer, defined in Theorem 4. Figure (c) shows convergence of loss to 0. Figures (a). illustrate convergence of Bo, B1 to identity. Figures (f) illustrate convergence of A,'s to -ー1
We verify that the parameters convergence are consistent with Theorem 4. We will once again use Dist(M, f) to measure the distance from M to the identity matrix, to scaling (see Subsection 4.3 for definition Dist). Figures 3a and 3b show that By and B; are close identity, as Dist B,, 0,, 1) appears to be decreasing to O. Figures 3d, 3e and 3f plot Dist(A, I) (red line) and /,,,E///A/,/// II (blue line); the results here suggest that A, is converging to Σ Σ- up to scaling, In Figures 3a and 3b, we observe that Bo and B1 also converge to the identity matrix (without left and right multiplication by E1/22,, consistent with Theorem 4.
We visualize each of Bo, B1 Figure and Ag, A1, A2 in Figure 5a-5c at the end of training, We highlight two noteworthy observations:
1. Let X₁ E Rdx= denote the first d rows of Zz. which are the output at layer k - 1 de- fined in (4). Then the update to X is Xx+1 Xx BaXxM X Ag X X₃ ba MM Xx). where M is mask defined in (2). As noted by Oswald et al. (2023), this may be motivated by curvature correction.
2. As seen in Figures 5a-5c in the Appendix, ||Aoll < ||A,]| < ||A2ll that implies the trans- former implements gradient descent with a small stepsize at the beginning and a large stepsize at the end. This makes intuitive sense as X₂ is better-conditioned compared to X,, due the choice of Bu, B, This can be contrasted with the plots in Figures (2a)-(2c), where similar trends are not as pronounced because B; 's are constrained to be 0.
Discussion
We take first step toward proving that transformers can learn algorithms when trained over a set of random problem instances. Specifically, we investigate the possibility of learning gradient based methods when training on the in-context loss for linear regression. For single layer transformer, we prove that the global minimum corresponds to single iteration of preconditioned gradient descent. For multiple layers, we show that certain parameters that correspond to the critical points of the in-context loss can be interpreted broad family of adaptive gradient- based algorithms.

---

09
OBS
6:8
8.90
AJA
0.03
4.00
34
623
0.00
6:8
320
93
92
62
060
4.00
1.1
4.00
4006
8.09
043
301
Figure 4: Visualization of optimized weight matrices Bn (left) and B (right). One can see that the weight pattern matches the stationary point analyzed in Theorem Matrices Ao, A and A2 are similar to Figure 2, and are visualized in Figure 5 in Appendix D.
We discuss below two interesting future directions.
Beyond linear attention. The standard transformer architecture comes with nonlinear activations in attention. Hence, the natural question here to ask the effect of nonlinear activations for our main results. Empirically,
Oswald al. (2023) have observed that for linear regression task, softmax activations generally degrade the prediction performance, and in particular, softmax transformers typically need more attention heads match their performance with that of linear transformers.
As first step analysis, we consider the nonlinear attention defined as
A (Z) PZM (Z QZ) where σ R R applied entry-wise.
The following result is analog of Theorem for single-layer nonlinear attention. It characterizes global minimizer for this setting with ReLU activation. Here, our choice ReLU activation was motivated by Wortsman et al. (2023) who observed that ReLU attention matches the performance of softmax attention for vision transformers.
Theorem S. Consider the single layer nonlinear attention setting with a ReLU. Assume that from (0, Ig). Consider the parameter configuration Po. Qo where additionally assume that the last row of Qo is zero. Then, the following parameters form a global minimizer of the corresponding in-context loss:
vector (5) is sampled from (0, Ia). Moreover assume that w, is sampled from
Po = Oaxd
Ta 0
Qo n- + 2d+
0
1
The proof of Theorem 5 involves instructive argument and leverages tools from (Erdogdu et al., 2016); we defer it to Subsection A.4. Thus, for isotropic Gaussian data, the structure of global minimum under ReLU attention is similar the global minimum with linear attention, established in Theorem (specifically the minimizer for the isotropic date given in (7)).
Refined landscape analysis for multilayer transformer. Theorem 4 proves that stationary point of the in-context loss corresponds to implementing a preconditioned gradient method. However, we do not prove that all critical points of the non-convex objective lead to similar optimization methods. In fact, in Lemma in Appendix B, we prove that the in-context loss can have multiple critical points. It will be interesting to analyze the set of all critical points and try to understand their algorithmic interpretations, as well as quantify their (sub) sut))otimality.
Acknowledgments and Disclosure of Funding
We thank Ekin Akyürek, Johannes von Oswald, Alex Gu and Joshua Robinson for helpful discussions. Kwangjun Ahn was supported by the ONR grant (N00014-00-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. Kwangjun Ahn also acknowledges support from the Kwanjeong Educational Foundation. Xiang Cheng acknowledge wedgess support from NSF CCF-21 12665 (TILOS Al Research Institute). Hadi Daneshmand eeknowledgee support from NSF TRIPODS program (award DMS- 200--422448). Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute).
10

---

References
Kwangjun Ahn, Xiang Cheng, Minhak Song. Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Lin- car attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01002, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. International Conference on Learning Representations, 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305. 136773,, 2023.
Noga Alon and Joel H Spencer. The probabilistic method. John Wiley & Sons, 2016.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang. al. Gpt-neox-20bs An open-source autoregressive language model. Proceedings of BigScience - Workshop on Challenges & Perspectives in Creating Large Language Models, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Neural Information Processing Systems, 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. Volume 1, 2019.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Benjamin I. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning (ICML), 2022.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dddss Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https//transformer-icccisss cutprmerrrrrrrcuusssuu2222222222 wtt.riiiisssuuu//////a///ot/t/te h/ffr//wwor//idde...ie.. html.
Murat A Erdogdu, Lee H Dicker, and Mohsen Bayati. Scaled least squares estimator for glms in large-scale problems. Advances in Neural Information Processing Systems, 29, 2016.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30533-30598, 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Pa- pailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13199 2023.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.540) 2014.
arXiv preprint
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 1997.
Stanislaw Jastrzebski, Devansh Arpit, Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual connections encourage iterative inference. In International Conference Learning Representations, 2018 URL httpps///penneeieer: wwwpenrvvviweeeeeeeeeeeeeeeeeeee
Ke Li and Jitendra Malik, Learning to optimize. International Conference on Learning Represen- tations, 2017.
11

---

Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evaluation. White Paper. A121 Labs, 2021.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576. 2023.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl:
Learning to learn in context. Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Techmologies, 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 2021.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21 /2. 1446, 2021.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355-9366. PMLR, 2021.
Hava Siegelmann and Eduardo D Sontag. On the computational power of neural nets. In Proceedings of Workshop on Computational learning theory, 1992.
Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 1999.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 2017.
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 International Conference on Machine Learning, pages 35151-35174, PMLR, 2023.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:
2022.
Mitchell Wortsman, Jachoon Lee, Justin Gilmer, and Simon Korbllih.. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:2309.08588 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. International Conference on Learning Representations, 2021.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers lear linear models in-context. arXiv preprint 20X.:22306009927, 2023.
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.088 00 2023.
12

---

Appendix
Proofs for the single layer case
13
A.1 Rewriting the loss function
13
A.2 Warm-up: proof for the isotropic data
14
A.3 Proof for the non-isotropic case
16
A.4 Proof for non-linear attentions (Theorem 5)
18
B Proofs for the multi-layer case
21
B.1 Proof of Theorem 2
B.2 Proof of Theorem 3
B.3 Proof of Theorem 4
24
B.4 Equivalence under permutation
33
C Auxiliary Lemmas
33
C.I Proof of Lemma (Equivalence to Preconditioned Gradient Descent)
33
C.2 Reformulating the in-context loss
35
D Additional experimental results
36
Proofs for the single layer case
In this section, we prove our characterization of global minima for the single layer case (Theorem 1). begin by simplifying the loss into a more concrete form. Throughout the proof, we will write P, Q instead of Po, Qo for brevity.
A.1 Rewriting the loss function
Recall the in-context loss (5) for the single layer case f(P, Q) is defined as:
[ (30 Attn
f ( Q) Zo....
+ wT (x+1))
1),(+))
From the definition of attention given in (3), one can further spell out the expression Zo x Attnpq(Zo) using the notation Zn = (2) =(2) 2(x+1)] as follows:
(2) 2(+11]+ + P[
1
(1+1) (1=00) (0--11TQ(220 2(xt1)))
= =12() (n+1)) -p
(1)
(1)
Q(2(1) 2(+11)].
i=1
Thus, the last column of the above matrix can expressed
(x+1)
+ +
1
Σ
(0
(1)
Q
1=1
where note that the summation is for 1, 1, 2,...,, due to the mask matrix M. Therefore, letting
T be the last row of P, and A E I be the first d columns of Q (as we did in (14)), then
Zo + =tAttup: (x+1),(++1) can be written as
A (n+1)
bT
Σ
20 20000
200
(13)
13

---

in other words, f(P. Q) only depends on the parameter b and A. Henceforth, we will write f(P, Q) as fb, A). Let us summarize our conclusion so far since it's crucial for the analysis follow.
Conclusion so far: A careful inspection reveals that the in-context loss only depends on the last row of P and the first d columns of Q. Thus, consider the following parametrization
0 6 P
and Q
where gd+1 and A R [dd-1)xd
(14)
Now with this parametrization, the in-context loss can be written as f(b, A) = S((0 " 0)).
Now, us spell out l(b, A) based on (13) as follows:
[(b, A) Ezo,w,
b
Σ i
2 (0) 20)))()))
Aa(att) + w. wwT x(x+1)
1b (AA((a))))) + w x (nt1)]
(b GA + w )
B:ozo.w.
EZzo.,
(15)
where we used the notation G= 1 E, 2002 to simplify. We now analyze the global minima of this loss function. To illustrate the proof idea clearly, begin with the proof for the simpler case of isotropic data.
A.2 Warm-up: proof for the isotropic data
As warm-up, we first prove the result for the special case where x(1) is sampled from (O,
Decomposing the loss function into components. Writing (a] a1 ad). and use the fact for j + j' and ([z(*+ [ = 1, we get
that [j]
1 (b, Σ Ezo.w, Gay + w. [i]] Elg (1+
Σ
Ezo.., Ga, w.(l2
The key idea is to characterize the global minima of each component in the summation separately. Another key idea is reparametrize the cost function given the following identity:
E Ez,,w, [b Ga, + wi]² Ezo,u, at(Ga,bb)) w. will -Ez,,, [[G,, baj)) w. Lii]
where we use the notation (x, = T for two matrices X and Y here and below. Given the above identity, define each component in the summation follows.
fx) = Ezo.w. (G.X) + wbbl]
for X E ((d+1)x(d+1)
Characterizing global minima of each component. To characterize the global minima of each objective, prove the following result.
Lemma (Global minima each component). Suppose that (8 is sampled from N Ia) and w, is sampled from La). Consider the following objective (here, (X,Y) Tr(XY for two matrices X and Y
fx(X = Ezo... ((G, + bi]
Then a global minimum is given as
X, m-1 + (d + 2)
Fat1.j
where En is the matrix whose (i1, i2) entry is 1, and the other entries zero.
14

---

Proof of Lemma 2. Note first that f, is convex in X. Hence, in order to show that matrix X, is the global optimum of f, suffices to show that the gradient vanishes that point, other words, (X)())) = O. To verify this, let us compute the gradient of fz: for matrix X,
(X) E [(G, X)
G) + TE [., [i] G]
where we recall that G is defined as
2000000 200 y()²
x(0)
G
To verify that the gradient is equal to zero, let us first compute E [w [i] G]. For each = 1,..., n1, note that Elw., [j] (0)))))))) "J because E[w.] Moreover, E[w,(j] (1)0)
because ws
symmetric, i.e., w, d w,,, and y0()) 2(1). Lastly, for h: = 1, ...,, d. we have
E[w., [i] y (9) [ = E[w,1b [i] (w., (0) x [x]] = E [w-b] x [1] x [x]]]]] [a (16)
because E B/w,(e w. [1] 0 for i + j. Combining the above calculations, it follows that
E [ec bil GJ
Eatis + Ez,d11
(17)
In order to compute IE (G, X) G]. let us compute I [[G,, Expe) ) GJ for i, I,... d + 1. Without loss of generality, 2 z'. First of all We now compute compute [G, Ed+1J) GJ. Note first that
(G, Eat1a) = Σ (wo+, (x1) ())
Hence, it holds that
Σ E0maaO = [[Σ
(iwe, (w.,, 2(1) 20015] ) [ (00gge0)) O.
E
(G, Est1.)
because B(ws]] O. Next, we have
Σ (
{ス Σ (w)) ( 2(1111 [uroe) )
E
(G, Eatia)
because w, d W. Lastly, we compute
Σ E00
E
(G, Fatia)
To that end, note that for # j',
= [((()( , x0(0 x(' = ifi
I (we, 10)) x. [1] (w., (()) 1(05
( 1|2 2000 [i] -n =0 ifi="'
and
E []] (w., (w.,
[i] =
[E( 11) 11 E (w.,,(()) (c) (ز) = d ifi= (18)
ifi
where the last case follows from the fact that the fourth moment of Gaussian is and
E (. (z []) E || 112 (a (a)
1 1
Combining the above calculations together, we arrive at
1
[(G Ed+1,5) G] = (n(n - 1) + (d 2)m) (Eat1.5 + Es,041)
n² n - + (d + 2) (Eati.s + Ejst1).
(19)
Therefore, combining (17) and (19), the results follows.
15

---

3. Combining global minima of each component. From Lemma 2, it follows that
X3
Estla
(d 2)
is the unique global minimum of fj. Hence, b and A (aı ai adl achieve the global minimum of [b,, A) Σÿ= Si (baJ) if they satisfy
1
baj ( + (d+ 2)
Eg-1s for all 1,2, ....d
This can be achieve by the following choice:
T ed+1; aj
ej for = 1,2,.... d,
- + 2
where ej is the j-th coordinate vector. This choice precisely corresponds to
Ia
b = ed-l
n-1 + (d+ 2)
We next move on the non-isotropic case.
A.3
Proof for the non-isotropic case
Diagonal covariance case. We first consider the case where (0)) sampled from (0, A) where
= diag Ai,..., ,d) and w, is sampled from N (0, Ia). We prove the following generalization of Lemma 2.
Lemma 3. Suppose that (( is sampled from NO, where A = diag(A1. , and w, is sampled from (0, Ia). Consider the following objective
fx) = Ezo.w. X) + w. [i]2
Then global minimum is given
Xj n+ X3 + \x) Est
where Eigio the matrix whose (i,, i2) entry is 1, and the other entries are zero.
Proof of Lemma 3. Similarly to the proof of Lemma 2, it suffices to check that
2E [(G, Xo) G] + 2 I [w, [i] G] = 0,
where we recall that G is defined
20000 (1)
((g((1)) y(12
G
similar calculation as the proof Lemma 2 yields
E [w] G x, (Eat1.) + Ez,4+1).
(20)
Here the factor of X5 comes from the following generalization (16):
E Elw. [i] ( (0) [K]]
Elw, [i] (w., (0) ( [+]] = E w. (00) [i] [k]]]
X A
Next, we compute E (G, E G]. Again, follow similar calculation the proof of Lemma except that this time we use the following generalization of (18):
(w., (0) x L (w., (()) [0
x [ E (
i',
=,,
(0
=,
En d + +2
ifi=
16

---

where the last line follows since
E [w., 2012 20 [01)
20 E |1000011 x00
E xd 102 Σ
X3 Σ [41
d + +22
&
&
Therefore, we have
E KKGG, Ea+1.) G]=
1
ds nn(n 13123 +ndg E
+ + +nx23
(Eatis + Eza-1)
n
n2 + 23 (M3 Σ Ae)
(Eatiy + Ez,4t1).
(21)
n
Therefore, combining (20) and (21), the results follows.
Now we finish the proof. From Lemma 2, it follows that
X, , + (Ex As)
Edt1,5
is the unique global minimum of fz. Hence, and A = lai of f(b, A) Σラ= fy A,) if they satisfy
aal achieve the global minimum
buj] X3 - 2 (e) Batis for all 1,,,....
This can be achieve by the following choice:
bt ed+1:
a +
(Ex Ax)
for 1,2....,d.
where e, is the j-th coordinate vector. This choice precisely corresponds to
diag
ed+1, -
n+
Xy (E Av)
(Ex
0
Non-diagonal covariance case (the setting of Theorem 1). We finally prove the general result of Theorem 1, namely z(66 is sampled from Gaussian with covariance 2 UAU where A diag(A1,.... Xa) and w, is sampled from (O, I2). The proof works by reducing this case to the
previous case. For each i, define ()) == -Tx (6). Then (E()) = E(U (UAUT)UT A.
Now let us write the loss function (15) with this new coordinate system: since (x) = UZ0), we have
f(b, A) = EZ....
[OTGA + w. ) =
A Ezo.., ((bbGA GA + w UU) bil
j=1
Hence, let consider the vector (b
GA GA
w. YU. By definition of G, have
(1)
82
TbTGA GA + w. w.) JU = Σ T
AU + w. U
(2(() w.>
Σ bT
U (U)r(t) w.w
82
AU + I..U
Σ b
0
122
TT 0
AU w
0
1
(Ux((t)
w.
0
82
ET
TT
T (2(())
A A
w.>
17

---

U
U 0
where we define 8T = T
A
AU, and w. U UTw,.
w. By the rotational
0
1
symmetry, w. is also distributed as \o0, La). Hence, this reduces to the previous case, and global minimum is given as
diag
= ed+1,
Ā =
n x Ar
From the definition of b, A. it thus follows that a global minimum is given by
Udiag ({
} UT
b
ed+1.
A =
(2
0
as desired.
A.4 Proof for non-linear attentions (Theorem 5)
As mentioned in Theorem 5, we focus on the setting where the last row of Q is zero, i.e., let
A O
Q
for A Radd and Rd
0
We first rewrite the loss function and simplify it following Subsection A.1. Moreover, for simple notation will often write 2, x instead of 2(+11) ((++))
Rewriting loss function.
Following Subsection A.1, , let us write down the in-context loss (5). for the single-layer nonlinear attention denoted by [(P, Q):
- Aunp
( Q) EZoo.,
w
d+1,+11
Recalling the definition of the ReLU attention Attn (Z) = PZM (Z TZ), the data matrix Z (1) 2(n+1)], and the mask matrix M, the term Attn can be written as:
Inxn 0
0
Aton),o(Zo)
PZu ((d+1)x((+1)
QZo)
B(o (1)
Hence, follows that the (d+ 1, + 1)-th entry of Attn7, (Zo) equal to the product of the (d + 1)-th row PZo, the mask matrix M, and the (n + 1)-th column of a (Z) QZo). Hence, let us write them down explicitly:
Letting bT be the last row of the matrix P, it holds that the (d + 1)-th row of PZ, is equal to [(b, 200) (=)))|=_.....
The (n + 1)-th column of a (Zj QZo) is equal to [a ((=(0) ((=(0)TQz((=1)) i= Letting A the first d columns of Q, this vector is equal to [or ((z(1)) Ax ) because the last
row of Q is zero and the last row of 2(++1) is zero (since (z(=+1))T = [(x(n-1)) 1)T O).
Thus, the product of (bb, )= the mask matrix M, and Jo ((x(1)) Ax((+t))] :=1....
results in the following expression of the attention (writing z, instead of z (n+1) (n+1l))
[Attngg(Zo)) d-1,n+1
Σ (b,2(0)) o((x(0) Txx)
18

---

Since [Zold+1,nt1 = O, we therefore have
1 Zo + A (2))
(b, 2 ((x(1)) An)
d+1,n+1
Therefore, it follows that the in-context loss Q(P.Q) only depends on and A. Henceforth, let us write s(b, A) instead of f(P, Q) following Subsection A.I. In particular, writing bt = [bo bi] for bo E _ and bj E IR, the loss function can be expressed as
Σ (10), + b) o((x! Ax)) (w2))
f(b. =) =
(22)
2. Simplifying the loss function with symmetry.
Now, we use the fact that both (''' and are sampled from the isotropic Gaussian, (V La) in order to further simplify the loss function in (22). In particular, we use the following facts:
(a) For orthonormal matrices U.V V E Raxd, it holds that Ux(), Va and Uw, have the same distributions (0, Ta).
(b) Moreover, for diagonal matrix E diag(E) gdxd with the diagonal entries being random signs & {=1], it holds that (1) Er and Ew, have the same distributions as (O0, La).
Now let us fix matrix A Raxd and b 16 b1] for bu Rd and b1 R. Letting A UΣVT be the SVD of the matrix A, follows that
f(b, A)
E
(6 (1) + b,w, ( a((x!0)) Ax) w x
71
12
(a)
E
(b) U+( + w, UN(((())) a(xx(0))TT AVx) + + U Væ
(
(b) UE + b.w (( Txx) + " TE VEr
72
2
n
E
E5
(b0 VEx( + = d((xli)) ]200]] + w w.E.Tva V2
E
b T. (1) o (()(())) w. diag(U V)x
=: flower (b,, ED diag(U T V)).
where in the third line we use the fact that E ΣE = EIE = Σ; and the fourth line follows from the Jensen's inequality, Hence for the remainder of the proof, we will characterize the global minimizer of the lower bound, i.e., flower and then we will connect it back to the original objective.
Computation of the lower bound flower-
Let us now explicitly compute flower. Let us rewrite the definition of flower- In fact since, = ReLU homogenous, one can further simplify the lower bound by pushing the constant bj inside and write b, Σ as Σ. Hence, for two diagonal matrices Σ, D E IRdxdd flower is defined as:
2
flower (2, D) = E
(w., o(x 22000)) - w Dr 1=1
In particular, D is constrained to the diagonal part of an orthogonal matrix (since diag(U in the above derivation). Now focus characterizing the global minimizers of flower-
19

---

The main part of the argument is inspired by the elegant observation of Erdogdu et al. (2016). which says that the solution of least squares and generalized linear models are collinear for Gaussian inputs. We leverage the same proof technique (a la Stein's 's Lemma) to prove that the presence of ReLU only changes the scaling of global optimum.
First, since w. is isotropic Gaussian, we can take the expectation over w. to obtain
flower (2, D) =
71 o T
20) 211)) Dx
which after careful expansion becomes
E 2 Σ o( Ex (2 00 + 2 E o(x (a(a"Ex())aa))) Dr) + const.
Ex!0)
(23)
In order to compute (23), we will rely on the aforementioned argument of Erdogdu et al. (2016). In particular, from integration by parts, Stein's lemma (Erdogdu al., 2016) (since æ ~ N Ia)), we have
Eg (o T v)a} E, u)]o for a fixed v Rd.
We this to compute all the terms in (23) follows:
We first apply Stein's 's lemma the first term of (23) for i # j. This results
(1) (x x(0)) Eg(0 (xx Ex Jo En
Egui, B Jo' (x Ex! o1r T ExOi) x (1
Using the fact that (00) is a symmetric random variable, one can compute the expectation above as follows: one the one hand, we know Eg [a'(x T [x) = E lo'(=x (o'(==TZZ)))) On the other hand, we also know that for any scalar a, a' d-(a) + (a) = 1. Therefore, we conclude that Eg (x Ex()))))] 1/2. Thus, applying this technique twice,
obtain the following
1 E [x T = T+(=").
E20 in ox Ex (x (x (200) (0)) = 1
Similarly,
Stein's lemma to the second term of (23) to conclude
Eo(x' T Ex (xx()) Dx) = Eg x: TDD= Tr(ED).
2
2
Lastly, the computation of the first term of (23) for i j is straightforward. Using the fact that Va E RR (a) + (-aa = a², we get
2 (x T (1)")) ||
(0 ||]]
2 =11((x (0) (0
1
E ( |zx00) ||]
d 2 Tr(,)?).
2
Putting things all together (and ignoring the constant part in (23)), we have
fower (E, 2(d 2) + (n +
(24)
4n
4. Connecting back to the original loss function.
One can in fact write (24) solely in terms b, and A as follows:
2(d + 2) + (n 1) Tr(E?) + Tr(ED) = 2(d + 2) + (n 1)
[|b |b + Tr(b, Since the latter convex function in the matrix b, A, it follows that the minimizer corresponds to
4n
4n
2n
b 4 2)
Ta
(n 1)
In fact the choice bj 1, bn and A (n- achieves this, and more crucially, satisfies the property that flower = f for the corresponding parameters. Therefore, this shows that such choice is a global minimizer.
20

---

B Proofs for the multi-layer case
B.1 Proof of Theorem 2
The proof is based on probabilistic methods (Alon and Spencer, 2016). According to Lemma 5, the objective function can be written as (for more details check the derivations (25))
X A:XoM)Xo w. w. Xa II (1 MX (AX0)
A(A1, A2) = ETr
E
1=1
2=1
- - A X II MXT - A, 1)
= Tr
II
where we use the isotropy of w, and the linearity trace to get the last equation. Suppose that A and A denote the global minimizer of f over symmetric matrices. Since Ai is symmetric matrix, it admits the spectral decomposition A₁ UD where D1 is diagonal matrix and U is an orthogonal matrix. Remarkably, the distribution of Xo is invariant to a linear transformation by an orthogonal matrix, i.e, Xy has the same distribution as X,U This invariance yields
JUUD,UT = f(D,, U (T,,,UTTU)) AzU).
Thus, we can assume Ai diagonal without loss of generality. To prove AG also diagonal, we leverage a probabilistic proof technique. Consider the random diagonal matrix S whose diagonal elements are either 1 or -1 with probability J. Since the input distribution invariant orthogonal transformations, we have
2-
S(D1, J(S = (D)0:5)))_(((),)))) ((D1,
Note that we use SD,S D₃ in the last equation, which holds due to Dj and S are diagonal matrices and S has diagonal elements in (+1.-1}. Since f convex A2, straightforward application of Jensen's 's inequality yields
f(D1.A5) = E [f(,,, SA55) ((D,, = f(D1, diag(()))
Thus, there are diagonal D2 and diag(A}) for which (Di,,iag(AA)) < f[Ai, A5) holds for an optimal Aj and AG. This concludes the proof.
B.2 Proof of Theorem
Let us drop the factor of 1/n which was present in the original update (56). This is because the constant 1/n can be absorbed into A,'s Doing so does not change the theorem statement, but reduces notational clutter.
Let us consider the reformulation of the in-context loss f presented in Lemma 5. Specifically, let Zo be defined as
(x) (2)
(()) y(a) (m++1) E El(d-
[(++1]
Zo
(1) (22)
where ((++1) = (w.,, (w,,,r((+))) (n)))). Let 2; denote the output of the (i- 1) layer of the linear transformer (as defined in (56), initialized Zo). For the rest of this proof, we will drop the bar, and simply denote Ze by Z,5 Let X, dx++) denote the first d rows of Z; and let Y; E IR (x(x+)) denote the (d 1)th row of Zu. Under the sparsity pattern enforced in (8), we verify that, for any i E {O....,k,
X, = Xo.
Y = Y, + Y.X.MM AX, = Yo II (1+ MXJAAeX)
(25)
E=0
'This use of Z. differs the original definition in (1), But we will not refer to the original definition anywhere in this proof.
21

---

Inxn 0
where M
We adopt the shorthand A {Aili==
We adopt the shorthand = {4)]--- Let S C lx++)xxdxd and A E S and only if for all i E {O.....k) k), there exists scalars a, E IR such that A, = a,Σ and B, = b,I. We use f(A) to refer to the in-context loss of Theorem 3, that is,
f(A) = [A))= (f0:-
A 0 P
Oaxd P.
0 1 )
Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:
((A) (xo...) [Tr ((1 M) Y Yet1 (1 (M))]
(26)
The theorem statement is equivalent to the following:
inf
(A.f(A) 0,
(27)
AES
where VA f denotes derivative wrt the Frobenius norm |Aslle Towards this end, we establish the following intermediate result: if A E then for any R E I (x+1)xdxd there exists ft E S, such that, at L. = 0,
d f(A + tR) S
d
f(A tR). dt
(28)
In fact, we show that R2 := rıl, for r; 1 Ir (- ( This implies (27) via the following simple argument: Consider the "S "S-constrained gradient flow": let A(t) R+ + 1(++1)xxxx be defined as
d A((t) -n(()-- ri(1) = (((()EE
dt
for i = O..... k. By (28), we verify that
di ((A(t)) S - Σ VA. 1=0
(29)
We verify from its definition that fA) 2 if the infimum in (27) fails be zero, then inequality (29) will ensure unbounded descent as t + oc, contradicting the fact that f(A) is This concludes the proof.
Proof outline. The remainder of the proof will be devoted to showing (28), which we outline as follows:
In Step I, reduce the condition in (29) a more easily verified layer-wise condition. Specifically, we only need to verify (29) when RG are all zero except for R2 for some fixed j (see (30))
At the end of Step which roughly "a product of attention layer matrices". (31), we study the evolution of fAA(t) when A(t) moves in the direction of Rt as Xo is (roughly speaking) randomly transformed.
set up some additional notation, and introduce an important matrix G,
Step 2, we use the results of Step 2 to to study G (see (32)) and d G(A()) (see (33)) under random transformation of Xg. The idea in (33) is that "randomly transforming Xo" has the same effect as "randomly transforming S" (recall S is the perturbation to B).
de
In Step 3, we apply the result from Step 2 to the expression of 4 f(A(t) in (31). We verify that R in (28) is exactly the expected matrix after "randomly transforming S". This concludes our proof.
Reduction to layer-wise condition. To prove (28), it suffices to show the following simpler condi- tion: Let j E {O..... kJ. Let R, Rdxd be arbitrary matrices. For C' E ddd let A(tC, denote
22

---

the collection of matrices, where [A(tC,))l, A + and for # j. A(tC, = A. We show that for all j E (O,....,k] Rg Rdxd there exists R2 - - such that, at = 0,
d dt ((A(tR),5)) - dt ((A(tR))
(30)
We can verify that (28) is equivalent to (30) by noticing that for any R, at 0, + (R) Σj=o dh ()) We will now work towards proving (30) for some index j that is arbitrarily chosen but fixed throughout.
(A++
Let us define, for any C' E &Raxd GG =)====+)==>> and (26),
= I (I M (A(C,,)), By (25)
5(A(tR),5))
Tr ((I - -M) Yx+ Yx+1 - (I-M))
=E [Tr (II Aj - w, w. G(Xo +t,,) M))
=E [Tr (II-M) Aj - +R)) T G A₁ + (R2) (I -M)!
The second equality follows from plugging in (25). For the rest of this proof, let U denote uniformly randomly sampled orthogonal matrix. Let Us == E1/ Using the fact that Xn Xx, we can verify
di [(A(tR)))))
1=0
E Tr - (R)) + (R2) MM)]
di
1=
[Tr (1 dt Exo.U
A, + +tHz) - (1-)))
Ex Tr (U A)) E-) d di G(Ur:Xo, A₁ + +RR)
(I (1- ((-M))]
(31)
t=
G and d Gunder random transformation of Xo. We will now verify that G(UrXo,AA)) =
Ux)(X
G(U
-xxXo
II
(1+ M UU A.UEX)))
(=0
-X,,
(32)
where we use the fact that US A.Us: = UU (a:E~ = Ai. Next, we verify that
d
GlUn.o, A (R;)
dt
L [ + +T+MXX7A MXT A,Xo) MUSUS II (T+ MXAXAXX
1-1
UU>X0
1=/+1
-Us G(Xo, A, + +U
(33)
dt
where the first equality again uses the fact that U... U,UU As. Putting everything together. Let us continue from (31). Plugging (32) and (33) into (31).
(AA(R))))))
1=
23

---

Tr (1-MM)G(U2X.. A)) d
(I- (I-M) M) ]
=2Exo
G(Ug A, +tR,)
dt
1=0
( 22Exx U
Tr (Tx 5-
d
G(Xo, +
(1 - (4-M))
dt
1=0
-x
(( A,) Ey
-
G(Xo: A, UU RyU3:)
(I ( (T-
dt
2 Ex.
Tr (L --1 d
G(Xo Aj + tuu U1 R3Uss|
(1 (
dt
-2oxo
Tr
(1-M)G(XooAj)) E-1 dl
d
+ -
(1 (-MM)
1=
d.
=i [(Alt.r, (a..t.rT..11 (,.(,,,3)-''))) ())
where = 1 Tr T> (I1/2 R2.1/2). In the above, (i) uses 1. (32) and (33), as well as the fact that U- = -ー (it) uses the fact that d z G(X,, A₁ + (C) is affine in C. To see this, one can verify from the definition of G. e.g. using similar algebra as (33), that di G(x,, Ay + CC) affine in C. Thus Ey [G A₂ + tUş R,U))) = G(Xo, A, + t Ey UUS R,U+)].
B.3 Proof of Theorem 4
The proof of Theorem 4 is similar to that of Theorem 3, and with similar setup. However to keep the proof self-contained, we will restate the setup. Once again, we drop the factor of which was present in the original update (56). This is because the constant 1/n can be absorbed into A, 's. Doing so does not change the theorem statement, but reduces notational clutter.
Let us consider the reformulation of the in-context loss f presented in Lemma 5. Specifically, let Zo be defined as
(r( (2) x(n) ((m) g(++1) E =[(d-1)((+1)),
x(1+1)]
Zo
(1) y(2)
where y(n+1) (w., 1)). Let Z, denote the output of the (i (i- 1)th layer of the linear transformer (as defined in (56), initialized at 7o). For the rest of this proof, we will drop the bar, and simply denote Z, by Z,. Let X, Bdxn+1 denote the first rows of Z, and let Y, e Rlx++1 denote the (d + 1)th row of Zz. Under the sparsity pattern enforced in (11), we verify that, for any i E (0....,
X = X, + X A,X,
Y41 Y Y.XX A,, Y. II (T MXF
(34)
We adopt the shorthand A= (A.) and B = (B:) Let S C (x and (4, E S if and only if for all {O...., k), there exists scalars a; by E R such that A; agΣ and B, I. Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:
[A, = Ex....) [Tr (U--M)
M))]
(35)
(note that the only randomness in Zo comes from X, Yo is deterministic function of Xo). The theorem statement is equivalent to the following:
inf (A,S)ES
V VA,S(A + VV (A, B)
(36)
'This use of Z. differs the original definition in (1), But we will not refer to the original definition anywhere in this proof.
24

---

where A. denotes derivative wrt the Frobenius norm ||A
Our goal is to show that, if (A, S, then for any (R, E [xx((xx)xdxd there exists (R S, such that, t 0,
dt f(A + tR, B + t5) - d BRA
(37)
dt
In fact, we show that Ra rıl, for r. Ir (2 and S s,l, for s; Tr (2)) This implies (36) via the following simple argument: Consider the "S-
constrained gradient flow": let A(t) : R - R(++lixdxd and B(t) R - R(++1)xdxd be defined as
dt A,(t) = -r:(t)=- rt(t) =
B(1))
di (t) = -s((t)=- s,tt)
for 0, .., By (37), we verify that
de f((t), B()) S
Σ FVA)) B())II ((,,)
(38)
We verify from its definition that f(A, B) 0; if (36) does not hold then (38) will ensure unbounded descent as / + xo, contradicting the fact that [AA B) is lower-bounded. This concludes the proof. Proof outline. The remainder of the proof will be devoted to showing (37), which we outline as follows:
In Step 1, we reduce the condition in (37) to a more easily verified layer-wise condition. Specifically, we only need to verify (37) in one of the two cases: (I) when R, S; are zero except for R2 for some fixed (see (40)), or (II) when R,, S; are all zero except for SJ for some fixed (see (39)). We focus the proof of (II), as the proof of (I) is almost identical. At the end of Step 1, we set up some additional notation, and introduce an important matrix G, which is roughly "a product of attention layer matrices". In (41), study the evolution of [(A, B(C)) when B()) moves in the direction of S, as Xp is (roughly speaking) randomly transformed. This motivates the subsequent analysis in Steps 2 and 3 below.
In Step 2, we study how outputs of each layer (34) changes when Xo is randomly transformed. There are two main results here: First we provide the expression for X, (42). Second, we provide the expression for X2 (t)) in (43).
In Step 3. we use the results of Step 2 to to study G (see (47)) and d G(B(t)) (see (48)) under random transformation of Xg.
de
The idea in (48) is that "randomly transforming Xo" has the same effect as "randomly transforming S" (recall S is the perturbation to B).
In Step 4, we use the results from Steps 2 and 3 to the expression of [A, B(t)) in (41). We verify that Š in (37) is exactly the expected matrix after "randomly transforming S". This concludes our proof of (II).
In Step 5, we sketch the proof of (I), which is almost identical to Steps 2-4.
1. Reduction layer-wise condition. To prove (37), it suffices to show the following simpler con- dition: Let j E (O..... k}. Let Sy E Raxd be arbitrary matrices. For C E Raxd let A(tC, j) denote the collection of matrices, where A(tC,)), j); AN + tC, and for i # j. A(tC. f); = A, Define B(tC,j) analogously. We show that for all {0, (O...., k} and all R,, Sy dxdd there exists Rg = -- and 8, = S such that, at t = 0,
d ((A((R)
d
f( [A((IR), B) dt
(39)
dt
d
and
きt f(A < dt
f(A, B(tS,,
(40)
25

---

We can verify that (37) is equivalent to (39)+(40) by noticing that for any (R, E (2x(x++)xdx at t = 0, & J tR, B 1.) Eÿ () B) a [A, ()))
We will first focus on proving (40) (the proof of (39) is similar, and we present it in Step 5 at the end), for some index j that is arbitrarily chosen but fixed throughout. Notice that X; and Y; in (34) are in fact functions of A, B and Xo. For most of our subsequent discussion, A, (for all i) and B; (for all * f) can be treated as constant matrices. We will however make the dependence on Xo and B, explicit (as we consider the curve B, + (S),, i.e. we use (, C)) (resp Yı(X, C)) to denote the value of X; (resp Y, from (34), with Xo = X, and Bg = C.
By (35) and (34),
S(A, B()
=E ((T - M) (Xo, B2 +1) (X,, + (S}) (T - -M)] G(x, Bj + (s,) w, ux By + +S,) (T [M))]
=E [Tr ((T -
((T- M)
=E [Tr ((1 (L- B2 x Bj +S,) (1 - -M))]
where G(X,C) = x II (1 (X, C) (,X((X The second equality follows from plugging in (34).
For the rest of this proof, let U denote a uniformly randomly sampled orthogonal matrix. Let U2 = 1//2 Using the fact that X, d Us we can verify
d
d [A, B(t
dt
1=0
d Exo [Tr ((1- B, (S) (I-
di
1==0
d Exou [Tr (IT B, +15)) TE- B3 + (Sj) (1 - (1-M))]
dt
1=
-2Ex.,U Tr (I- G(UgX B,)
G(Us. By + tS₁) (I (1-M) M)
(41)
dt
2. Xx and 4X, under random transformation of Xo. In this step. prove that when Xo is trans- formed by Uz. X, for i 2 are likewise transformed in a simple manner. The first goal of this step is to show
(u, = UX, Bj).
(42)
We will prove this by induction. When this clearly holds by definition. Suppose that (42) holds for some i. Then
(Us, BB)
=X,(UxXo, B,) + B.X. -BKUCC...BAMMMXXXXX3XXXXXXX,,,,
-Xo,,(Xo, B,) + Us UX B₁) Axx,(XX. (X5x
where the second equality uses the inductive hypothesis, and the fact that A, a so that UT A,Us Ai. and the fact that B, b,1, from the definition of S and our assumption that (), S. This concludes the proof of (42).
We now present the second main result of this step. Let U := E1/2TT 1/2, so that it satisfies U- Us Ux I. For all i, d
d Xx By + +1 SUUz)
Us² dt X, By + (S;)
(43)
di
1=0
1=
To reduce notation, we will not write - explicitly in the subsequent proof. We first write down the dynamics for the hight-hand-side term of (43): From (34), for any E S j. and for any i 2 j 1, and
26

---

for any C' E Raxd
d
Xx (Xo, Bj [C) =0
dt
d
A,X, Xj+1 (Xo, B3 + +C²)= = CX, (Xo, B,)xX, (Xo, (Xo A,X,, (Xo.
dt
d
d
Xi+1 (Xo: B3 + +C)
Xi (Xo, B, +10)
di
di
+ B; X; (Xv, B,,,BBr + +10)) MX; B,)² A,X; (Xo,
d
dt
d
(Xo, By By+101) (Xo. X;
T
+ B,X, (Xo, B,)M
A,X; (Xo, (Xo, B,)
dt
d + B,X; (Xo, Bj) BB) Ai ( X: Bj + +10))
(44)
We ready prove (43) using induction. For the base case, verify that for l UZ de
d. X₁ Xo, B& + (S,) = 0 at Xe (Xo, B, + tU, +,,,,, (see first equation in (44)). For
index j + 1, we verify that
d Xit1 (Ux: B, + tS,)
dt
UF SJUn, ((0, (U B₁) (UX,( B,)
d X,+ (Xo: B, + S,Us)
(45)
di
where we use two facts: 1. X,(UsxXo, BB) (Xo, B,) from (42). 2. A, a so that USJ A:Us = Ai.
We verify by comparison to the second equation in (44) that UE di X, (og By + tS,) 0 = X (Xo, B, tU' "Ug"s,Us), These conclude the proof of the base case.
Now suppose that (43) holds for some i. We will now prove (43) holds for + 1. From (34),
d X (U2Xo, 1,+ tSp)
U
dt 1 d di
=
(x, (U>Xo. B, + +S,))
+ dt
(Us By + 1,,)) MX, (U=:o, Bj +S,)) A,X, (Ux: Bj + +1s,))
d
=U' dt
(X, (UnXo, B, + tS,))
d (Ug B, + (5,)) (U, BJ) (UxNo, B₁) d
+ B,
U B,X, (Ux B,M di By BB (3
A,X, (UsXo. Bj)
+ + (Us (UEXX. (Ux,,
① d
= U55 dt
X, (UrXo: Bj + t5,)
d X, (oxXo, By + +1,,)
+ B,
Us!
(xx, (xx
di
+ B,,,, (Xo, B,)M
US dt
x, (Us By (Us + Axx,
- B,,,, (Xo, (x,, (Xo, B,) A U5
X1 (UgXo B, + +(S))
dt
27

---

( d
X v (Xo, B, + W3 {5UUs)
dl
d
(Xo, B3 +W (X, B3)² A,X; (Xo,
di
x (Xo, B3 + t1E (UF'S,Us)) + B,X, (Xo, Bj)M
A,X, (Xo, (Xo, Bj)
B,X; (Xo, (.)MM,(X,, BB)" Ai (=xxxxxxxx (Xx:Xx.. Bj + (555)
(46)
- B; Xo Bj) (Xo, Bj) from (42), 3. A; a.E so that U A,Us A., Uzlús
In (i) above, we crucially the following facts: B3 be that Ug_
B 2.
US (ii) follows from our inductive hypothesis. The inductive proof is complete by verifying that (46) exactly matches the third equation of (44) when C= U5:" SUg.
G and dt G under random transformation of Xo- We now verify that G(UgXo. B,) G(o B,). This is a straightforward consequence of (42) as
G(UsXo, B2)
k
Xxxo
II
(1+ MX,(UXX (X, (.x,
t=0
(I+ MX, B,)² AXX((X.
i=0
Ur:G((Xo Bj):
(47)
where the second equality uses (42), well the fact that U. A,Us; A,, Next, will show that
d G(Xo B, + tU55 (1,55)
US G(UxXo. B3 +(,)
(48)
1= dt
1-0
To see this, expand
d
U
GoUs.Xo B₁ +t,))
dt
= -U=1 1 d
Us II (I MX, By +S,)² A,X, By ((,))
dt
t=
i-1
Xo Σ
II
[I (X+,,X ALX, Bb)
(=0
MM (Xi(UgXo, +t,)) AX, B,))
II (L+MX(UUXXo. AX, B:)
X
E
[ MX.(Xo, AXXXXXX Bi))
1=0
_e (15''
(15
d [x(UXXn,3,+sS))
A.X,(XX, B,) MX,(X.. B,) (US" d dt [x
&
II (1+ BJ) ALX, Bi))
1=+1
X
BX) Bi))
28

---

d M ((ax(xon))= T A.X, BB) + MX.(Xo.B:) B,) ( f (xxx By+ +1U5'
- S,Us)
II (T+ Axx Bil) A-i+1
(60) d dt
G(Xo; B3 + B
In (1) above, the following facts: 1. XUx B,) UX,, Bj) from (42), 2. A, =
so that US A,Ux, A,, 3. Us U3 Us: I. (fi) follows from (43). (iii) is by definition of G.
4. Putting everything together. Let us now continue from (41). We can now plug (47) and (48) into (41):
d dt"
fA, B(t
11=0
Tr (1-M))(USXX0,B) d
(1 - M) m1 ]
-2 Ex..U
di Co Bj + +5
1=0
① 2 Exo.U
Tr (I- B,) --1 G(Y, By + + SyUs:) SyUss)
(1 - M)]
dt
1=0
d
(1- (1-M) )
-2ox
Tr (I- E-) Ey
G(Xo, B, + :īss SyUx)
dt
1=0
(1)) 2 Ex.
Tr (I-M)G(X) --- d
G(Xo: Bj +Bu LUE SyUz])
(I (I- M) ]
dt
1==0
Tr (T- B d
-2 E2
G(X, Bj ts² k=
(I- M)
dt
d dt f(A, (AA,B(ts,I.))
1=0
where 8j = Tr (x- 12-1/2 In the above, (i) uses 1. (47) and (48), well as the fact that U₁ V -- (di) uses the fact that & G(Yo, B, + tt) is affine in C. To see this, one can verify from (44), using a simple induction argument, that æi (Xo, By + tC) is affine in C for all i. We can then verify from the definition of G, e.g. using similar algebra as the proof of (48), that d
di G(x., Bj + C') is affine in de (Xo, By tC). Thus Ey [G(Xu, By + tUss S,U²)] G(Xo, By + t Eu (UE)),
With this, we conclude our proof of (40).
5. Proof of (39), We will now prove (39) for fixed but arbitrary j. i.e. there is some ry such that d di ((A((Rg,)))BB)
di S(A(t ((A(t-rz--..)),
The proof is very similar the proof of (40) that we just saw, and will essentially repeat the same steps from Step 2-4 above.
Since now consider perturbations to A instead of to B, we will need to redefine some
notation: let Xx(X,C) (resp Yx(XC)) to denote the value of X, (resp Y,) from (34), with X, X, and A₁
C (previously it with By = C). Let G(X, Aj + C) II (I+ M (X,(X, A, + ()C, + C))), where recall that A(C.j) = A₁ C. and ACC,)) == A₁ for all i {0...k) {j}.
We first verify that
X,UUs A₁) Uso, Aj) GUU,,oo AJ) = UsG(Xo, A,).
(49)
29

---

The proof's are identical to the proofs of (42) and (47) so we omit them. Next, we show that for all i,
US dt Xo(Ux Ag + tR,)
dt X A, + +
(50)
We establish the dynamics for the right-hand- of (50):
X₂ (Xo: A, + +t0 =0
dt d
X5+1 (Xo; Aj +1C) = ByX, (Xu:
Ax),AA))c A,) cx, (Xo:
di
d dt
X1+1 (Xo: A, + +0) d X; (Xo, A,+ +10)
X; (Xu, A3 + +10) MX; (Xo, A,X, (Xo.AA)
+ B;
dt
d ( (Xo, Ay+ T
B,X, (Xo, A,)M
AXx: Aj)
dt
Ax) Ai (d XXo. + +10
B,X; (Xu, (.)MX;(Xo.. AT)
(51)
dt
Similar to (45), we show that for i S
Us dt
X; (Us: A5 + (R₂) =0 = Uz (Us A1 (Us
dt
and
U3 d
di X3+1 (UxXo A, +(2})
BUUs, (No, ()M AA) (Ux, Az)
dt X,,+ (Xo, A + UUS R2Ux).
Finally, for the inductive step, we follow identical steps leading up to (46) to show that
UZ X (UsNo, A3 + tR₂)
dt
d Xi (Xor A, + +UU R,U))
dt
Xi (Xo, Aj + UU (USR) (Xo A1) A:Xi (Xo,A.)
+ Bi
dl
( (Xo Aj + +1U Ax (1RRRRUU)) ArXi (Xo:
+ B,X; (Xo, A,)M
A ( (d (t (ax, (Xo. + +1U
+ B,X (Xo: (X.. AT)
(52)
The inductive proof is complete by verifying that (52) exactly matches the third equation of (51)
when C U3 SUs. This concludes the proof of (50).
Next, we study the time derivative of GUUxXo, A, + tRy) and show that
d US' " G(UgXo A, + (R,)
d G(Xo A, + 1U R,Us).
(53)
dt
This proof differs significantly from that (48) in a few places, so we provide the whole derivation below. By chain-rule, we can write
US' d dt
G(UxXo. A, +t2,) +
30

---

where
d = Uz dt
&
U2 II (I + A, + A, + +(R)))
i=0
and
II
A,
(UX RXX
k
II MX A,) A A.X.(UxXXo
We will separately simplify and V, and verify at the end that summing them recovers the right- hand-side of (53). We begin with and the steps almost identical to the proof of (48).
d
& Up II (I
= U3
AX + (X + +R2))
dt
1=0
=Xo
Σ
II AI P=
1=0
-M
(1) AX. AJ +(,)))
di
πI 1=++1
(T+ (X+ AI AX
X
(T+ M A,) AxXe(Xo, Ai))
10
P=
.
[ (1=5
d X.(Ux:X,,Ax+1)) A3 (
AiXi(Xo, AJ) + MX.(Xo A,)"
A, (x=3_x((,XX),)ttttt
di
di
II (I+ MX. A,) ArXi(Xo,
(=++1
(4)
X
Σ
II MX+ A,) Axxe(Xo:Ao)) P=
1=0
M
(
d Xi(Xo, A₁ Xx (++U3RRUss) dt
A:Xi(Xo, Aj) + MX,(Xo. A,) dt
d ( A3 + 18 A (D]RRUs))
πI (I U MX, A,) Ar (1+1
=X,
II (I+MX A, Axx:
1-0
P=
M
A A.X,(Xo, AJ + UJ R2Ux))
di
II (T+ A) A)
(54)
(=++1
31

---

In (i) above, the following facts: 1. (U, B,) UXo, Bj) from (4 A, =, so that U3 A,Ux A,, 3. Ux.U5²= U²-UU
Us I. (ii) follows from (50).
We will now simplify V
-Ux Xo
πI [1+MXXUUXX, AX,
MX)( AXX,
II ())
MX,(Xo, j-1 (T+ (T+ (TT-XX(76-44AXXX(70-) MMX, (X,
@ Xo
II (I+ (X,XX AxX(xx(X)
(55)
where (i) uses the fact that X, B3) UxX, B,) from (49) and the fact that A; a,E-!.
By expanding d G(Xo A, + tUg R, U2). we verify that
de
d
di G( A₁ + u R2U5) = (54) + (55) = + 0 = -= _h G(U A, + +tHg),
dt
this concludes the proof (53).
The remainder of the proof is similar to what was done in (41) in Step 4:
dt (A(tR),,)), B
11=0
Tr (I-M)) d AU Aj +tR₁)
(1- (M) ]
=2Ex.,U
dt
( Ex..U U
Tr (U-M))((X.,A) Aj) T-- d G(Xo, Ag A 115 R2Uss)
(I- (1-M))
dl
(1)) 22 Ex
d (I- A)) T G(Xo, A + (1, R(UU))
(1- [-M
dt
11=
d G(oo Aj +
(1 (1-M) M)]
-2Bxo
dl
1=1
2 (A(t-rr,--11 (,) ()),B) B)
where r; = 1 Tr (2//2 E+/2). the above, (i) uses 1. (49) and (53), as well as the fact that - = Σ Σ-1 (ii) uses the fact that के G A, + (C) is affine in C. To see this, one verify using simple induction argument, that d XX.o. A, + (C) is affine in C for all i. We can then verify from the definition of G, e.g. using similar algebra as the proof of (53), that á di G(oo, A, + C') is affine in d G(Xo, Aj + / Ev [U]] R.
de 4X,oXo Aj + tC') and C. Thus Ev [G A₁ + +UU R,Us)) =
This concludes the proof of (39), and hence of the whole theorem.
32

---

B.4 Equivalence under permutation
Lemma 4, Consider the same setup Theorem 3. Let A {Ail with A, =aa5- Let
f(A) : (fo-=
A
0 0 P, = 0axd
0
1 i=0
L,j E {O....,k) k) be any two arbitrary indices, and let A, = Ag, As = Ar, and let Ae = for all E ji Then f(A) = f(A)
Proof.
Following the same setup leading up to (26) in the proof of Theorem 3, we verify that the in-context loss is
f(A) E (r ((T M) A) (T M))]
where G(Xo. :) Xy II (I MxX) Consider any fixed index We will show that
(I+ MX2 AeXo) (I + MXO 44+1X) (I MX Aet1 Xo) (I MX2 AeXo)
The lemma can then be proven by repeatedly applying the above, so that indices of A, and Aj are swapped.
To prove the above equality,
(I+ + MX3 Ax) (I + MX. Ac-1
=I MX AeXo + MX A2+1 + MX. AeXaM A4
=I+ MX A,X0 + MX A MX. a X, a Σ- Xx
=I+ M XxXo + MXX A441Xo + MXO - 1xo
= (1 + MX. Xo) (I MX, A(Xo)
This concludes the proof. Notice that we crucially used the fact that Ae and Ap41 are the same matrix up scaling.
C Auxiliary Lemmas
C.1 Proof of Lemma (Equivalence to Preconditioned Gradient Descent)
Consider fixed samples (1)),... .....y (0),, and fixed w.. Let P = [P.].
o. Q = =Q::
i-o denote
fixed weights. Let Z, evolve described in (4), Let X, denote the first rows of Zy (under (8),
X; = Xu for all I) and let Y, denote the (d + 1)th row of Z,. Let g(z, k) IRd x IR x Z + R be
a function defined as follows: let =+1 = x and let yn n+1 [Yk]n+1
n+1 y. then g( y. k) == y 1+
Note that
We verify that, under (8), the formula for updating (u+1) is given by
Yk+1 = Y 1 -Y.M AxXo-
T
where M a mask given by
We can verify the following facts
1. g(z,v. k) g(x. k) + y. To see this, notice first that for all i E (1....,n).
(3)
1
x (1)
A1) (
33

---

In other words, Vi does not depend on Di
(n+1)
for any 1. Next, for y/ (n+1)
itself,
(x+1)
(x+1)
1
(n+1) Agx (f)
Mk+1
Uk n j=1
which depends on = only additively. We can verify under a simple induction that ( y, le y - g(, y, k)) y.
( is does not depend 2. ( 0,k) is linear in To see this, notice first that for j + + 1, yk (n+1) depends only linearly
on zt (x+1) for all t,j,k, k. Consequently, the update formula for yh++1
(n+1) 0 is linear in 2, the conclusion follows by on (m+1)) and yk (n+1), Finally, Yo induction.
With these two facts in mind, verify that for each k. there exists 02 Rd such that g(z, y. g(x, 0, k) (0k. x) y for all x, y. follows from definition that g y, 0) = so that (fb, = g(x y, 0) - = 0, so that 0g O.
now turn our attention to the third crucial fact: for all i, g(x yd k) = uk = (On, x111))) y(
To see this, suppose that we let (x+11) = (() for some E I...., Then
(6)
1
( A Apz(1y00 1/
Ve+1 = Uk
==1
(n+1)
(n+1)
(n+1)
Agx (0)
Metl
Vi
Mk+1 thus Vx+1 = (n+1) if yk (3) = 3/k
(8)
in+1) and the induction proof is completed by noting that (1)
(n+1) by definition. Let E πaxn be the matrix whose columns are x((), ".,,,..,,!,,,. x(m), leaving out yo
z(a+1), Let Ya Rxx= denote the vector of yb ..... 24 Then it follows that
(1)
(n)
Y Yo of X.
(n+1) can be written as
Using the above fact, the update formula for y
(n+1)
(n+1)
Axx T Yn, ((+tt)
yk+1
-P6
72
(04t1,0((tt)))
( -(0,, (at1))
(ALx (8T 0z + Po). ,x(ut)))
n
=(0e)
(AxX (X" (0x + w.)) 20((=1))
Since the choice of (1++)) is arbitrary, get the more general update formula
Oct1 0k X XT (Oz w.).
We can treat Aş as a preconditioner. Let f(0)= = 1 (0 + w.) XXX (O + w.), then
2n
Oc+1 02 A& (()).
Finally, let whth -kx. We verify that f(- w) Rw. (w), so that
gd #41
A₁ V R... (w
= w
n
We also verify that for any (n-)), the prediction of 1/4 g (m+1), ((++1), k)) ((++)) (0, x((tt))) ((+))) + wwee ,((+1)))
(x+1) is
This concludes the proof.
34

---

C.2 Reformulating the in-context loss
In this section, we will develop a re-formulation in-context loss, defined in (5), in more convenient form (see Lemma 5).
For the entirety of this section, we assume that the transformer parameters (Pi, Qi) are of the form defined in (11), which reproduce below for ease of reference:
B
Q = A !
P;
Recall the update dynamics in (4), which we reproduce below:
7441 7: +
PZ.MZ Q7+
(56)
n
where M is mask matrix given by M Lnxn 0
0 Let X Bdxn+1 denote the first rows
of Z& and let Y& IR1x++1 denote the (d + 1)th (last) row of Za. Then the dynamics in (56) is equivalent to
X411 = X; + - B.X,MXXAAN A.X,
1
Y41 Y, YYMXTA A,X,
(57)
n
We present below an equivalent form for the in-context loss from (5): (a+1) iid
Lemma 5. Let Pz and Pw denote distributions over Rd. Let x(1),
Pa and w. Pw-
Let Zo E Rdd 1xnt1 be as defined in (1);
(1) x(2)
((n) (n))
x(+41)]
Zo = (1))
E r ((d+))((++))
y(2)
Let Zu denote the output of the (k- (k-1) 1)th layer of the linear transformer (as defined in (56), initialized Za). Let f ((P:, Qi =-)) denote the in-context loss defined in (5), i.e.
f (1Ps 20) = -Bi ( (IZa((tt)))) + w. x (m+11) 2
(58)
Let Zo be defined as
[I) x(2) x(n) x(u+1)]
7, = (y(1) ((2)
((m) g(n+1) E ER (R(d-1)x(++1)
where g[n+1) = (w., x. (. Let Zx denote the output of the (I: 1)th layer of the linear
transformer (as defined in (56), initialized at Zu). Assume {P,, Qikk be of the form in (11). Then the loss in (5) has the equivalent form
f
= (IP,, = Ezzo,w.) Tr (I MMT-MYYAN k+1 (
where Y k++1 E Exn+1 is the (d + 1)th row of Zz-
Before proving Lemma 5, we first establish an intermediate result (Lemmaa 6 below). To facilitate discussion, let us define a function Fx ((A, B No, Yo and Fy (A, ; Xo. Yo to be the outputs, after & layers of linear transformers respectively. Le.
Fx (A, Bi) Xo. Yo) = =Xa+1
Fy (Ai, Bi) Xo, Yo) = Y
as defined in (57), given initialization X,, Yo.
We now prove a useful lemma showing that (V))211 = y(x+1) influences X3, Y, in very simple manner:
35

---

Lemma 6. Let Xs, follow the dynamics in (57). Then
1. [Xi] is are independent of [Yo],+r
2. For j + nn+1. 1, [Xilg is independent of [Yolatr
3. [Vil+1 depends additively on (Yol+tr-
In other words. for C == [0,0, S,.,0,...... 0,c| E R8 Rx((+11,
1: Fx (Ai, B.) Xo. Yo + C) = Bi) Vo) 2:3: Fy
Xo, Yo + +C) = (A, Xo, 1o) + +C
Proof of Lemma 6. The first and second items follows directly from observing that the dynamics for X; and Y; in (57) do not involve [Y:],41 due to the effect of M.
The third item again uses the fact that [Y++1 - Yil+11 does not depend on [Yil++1*
We are now ready prove Lemma 5
Proof of Lemma 5. Let Zo. Zr. Zo. Z& be as defined in the lemma statement. Let Xx and YR denote first d rows and last row of k. Then by Lemma Xx+1 Xx+1 and Y k+1 Yk+1 [0 0 (w., x(+t))] Therefore, (58) is equivalent to
(Z )
E (Zo.o.)
E (Zow..) )
([V+tl(nn1))
}"]
( M) 8 k+1 119
Tr (11- M) M) (T-I(T-
E (Zo,o.)
This concludes the proof.
D Additional experimental results
In this section, we present a few addition experimental results. We first present in Figure 5 a visualization of learned weights Ao; A1, for the setting of Theorem One can see that the weight pattern matches the stationary point analyzed in Theorem 4; hence, combining Figure 4 and Figure 5, we corroborate our results from Theorem 4. Interestingly, appears that the transformer implements a tiny gradient step using Xo (as Ao is small), and a large gradient step using X2 (as Ag is large). We believe that this due to X₂ being better-conditioned than X1, due the effects of Bg, B1.
4.4
0.20
0.42
0/1
SSX
0.00
0:11
0.21
4.00
+28
0.21
2030
415
0:33
(a) Visualization of // A05 (b) Visualization of /21/ 12/2 (c) Visualization of 1/1/2 A//21/2
Figure 5: Visualization of learned weights for the setting of Theorem 4. One can see that the weight pattern matches the stationary point analyzed in Theorem 4.
36

---

We next present some additional experiments that investigates the properties of the learned predictors of various algorithms. First, we plot the test losses against the number of examples provided in the prompt ("the number of ICI. examples"). We compare four different algorithms: (i) the predictor learned by a three-layered of linear transformer, (ii) three steps of GD, (iii) three steps of preconditioned GD, and (iv) the ordinary least-squared solution (OLS). For GD and preconditioned GD. the optimal stepsizes are found by gridsearch. For preconditioned GD, preconditioner is fixed to be -ー1 for comparison. In all cases, the dimension d = 5, and for each N, the linear Transformer is trained using Adam. The result presented in Figure 6.
4- 3-Step Preconditioned GD
Layer Linear Transformer
2-
1-
0-
2
6
10
14
18
22
Number Examples
Test loss comparison between (i) the predictor learned by a three-layered of linear trans- former, (ii) three steps of GD,, (iii) three steps of preconditioned GD,, and (iv) the ordinary least- squared solution (OLS).
Figure 6:
Lastly, in Figure 7, we plot the test losses against the number of layer L (or the number of steps in the case of gradient-based algorithms). For L = 1, 2, 3, 4, we compare between (i) the predictor learned by L-linear transformer and (i) L-steps of GD. (ii) L-steps of preconditioned GD. Again, the optimal stepsize is found by gridsearch, and for preconditioned GD, the preconditioner is fixed to be Σ-1 In all cases, the dimension d 5, and context length N 20. The linear transformer is trained with Adam.
-2-
los()
4-
GD
8
Proconditioned GD
Linear Transformer
Number of Layers/Steps
Figure 7: Test loss comparison between (i) the predictor learned by T.-layered linear transformer and (i) T-steps of GD, (ii) T-steps of preconditioned GD, for L. = 1, 1,2,33 2,
37


Reviews:
[{"review_id": "6qfcivlpvh", "reviewer": "Reviewer_m6sU", "summary": "The paper presents an in-depth analysis of in-context learning (ICL) in transformers, particularly focusing on linear transformers. The authors propose a novel interpretation of ICL as an implicit form of '**preconditioned**' gradient descent. They further extend this analysis to multi-layer transformers and provide insights into the optimal parameters for the in-context loss under certain constraints. The paper also includes empirical results on synthetic datasets to support their findings.", "questions": "1. In the current setup, I find there is no MLP block in the architecture, empirically that seems to be an important component in the transformer and ICL, could the authors provide more discussions/explanations on the roll of MLP (could also be listed as future directions)?\n2. Typo: missing '.' at the end of line189.", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper provides a novel perspective -- the preconditioning/adaptive perspective -- on in-context learning in transformers, which could potentially lead to new ways of understanding and improving these models. Given the growing interest in understanding ICL from a theoretical perspective, the `*preconditioning gradient descent*' can be served as an important building block. This preconditioning perspective is also new in the literature (as far as I know).\n\n2. The authors provide a rigorous mathematical analysis to support their claims, including single and multi-layer cases. The analysis tools developed in this paper could be helpful for future research on understanding ICL.\n\n3. The paper is well-structured and includes visualizations to help understand the concepts and results.", "weaknesses": "1. (minor) In the experiments, I did not find the ICL experiments --- x-axis represents the number of ICL samples, and the y-axis represents the test loss (similar to the experiments in Figure 2 of Garg et al. (2022)). It would be interesting to check those results besides the current ones on training dynamics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> I did not find the ICL experiments --- x-axis represents the number of ICL samples, and the y-axis represents the test loss (similar to the experiments in Figure 2 of Garg et al. (2022)). It would be interesting to check those results besides the current ones on training dynamics.\n\nThank you for the suggestion. We conducted additional experiments on ICL, plotting test loss against number of ICL samples, for Linear Transformer vs Gradient Descent vs Preconditioned Gradient Descent. You can find it in Figure 1 of the rebuttal PDF in the [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) above.\n\n\n\n\nAs noted by Oswald et al., the network analyzed in this paper is sufficiently expressive to implement gradient descent on the least-squares objective. Hence, one can zero-out all parameters of the MLP block and still implement gradient descent on ridge regression. Thus, in this work, we omit MLP block; notably, this approach was also taken in other recent works in the literature [1,2,3]. Note that ICL of kernel regression requires MLP blocks as shown in [1,4,5]. Extending our results to kernel regression is an interesting research topic for future work. We will add a detailed discussion of this in our final version as per your suggestion. \n\nThank you for catching the typos, we will correct them in the final version.\n\n\n\nThank you once again for taking the time to review our paper. We sincerely hope that the response to your concerns, as well as the overall response to other reviewers’ concerns, helps assuage your concerns, and view this paper in a more favorable light.\n\n**References.** \n\n[1] Von Oswald, Johannes, et al. “Transformers learn in-context by gradient descent.” International Conference on Machine Learning. PMLR, 2023.\n\n[2] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. \"Trained Transformers Learn Linear Models In-Context.\" arXiv preprint arXiv:2306.09927 (2023)\n\n[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. \"One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention.\" arXiv preprint arXiv:2307.03576 (2023)\n\n[4] Garg, Shivam, et al. \"What can transformers learn in-context? a case study of simple function classes.\" Advances in Neural Information Processing Systems 35 (2022)\n\n[5] Akyürek, Ekin, et al. \"What learning algorithm is in-context learning? investigations with linear models.\" arXiv preprint 2022."}, {"review_id": "USKVGtoPLO", "reviewer": "Reviewer_k8UL", "summary": "LLMs trained with transformers have shown remarkable abilities that have led to a lot of recent theoretical research towards understanding the origins of these abilities. In context learning is one such ability that refers to a model's ability to adapt in-context to new tasks without weight updates. In the work of Garg et al., the authors proposed a simple and powerful data generation process to understand the mechansims behind in-context learning. In Garg et al., it was empirically shown that transformers can learn to imitate ordinary least squares  algorithm. Since that work several works have taken steps to understand why and how transformers implement these algorithms. In recent work from Oswald et al., it was empirically shown that linear transformers implement one step of (modified) gradient descent per self-attention layer. In this work, the authors take an important step to explaining why do transformers implement gradient descent. The authors show that for shallow transformers pre-conditioned gradient descent is the global minima of the in-context training loss. For multi-layer transformers, the authors establish that pre-conditioned gradient descent form the critical points of the loss landscape.  ", "questions": "Please see the weakness section.", "limitations": "The authors discuss future directions, which stem from the limitations of the current study. The authors can elaborate further on limitations of the theory further and highlight some other potential gaps that can be addressed in future work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a very good paper. I was very happy to read it as it solves an important unresolved problem that has to come to fore in the recent investigations on in-context learning. \n\nOn the different axes, I would rate the paper as follows. \n\nOriginality: I really liked the approach that the authors have taken. It is quite a neat idea to consider sparsity driven constraints that permit tractability. \n\nQuality: This is a very good paper. The results resolve some very important questions towards understanding ICL. \n\nClarity: The paper is very well written and a joy to read. \n\nSignificance: . This paper serves as a crucial step towards the goal of understanding ICL is one of the central problems towards demistifying LLMs", "weaknesses": "I have some concerns that I would highly appreciate if authors can throw light on in the paper. \n\ni) The authors have studied linear attention mechanisms, which I very well understand are themselves non-trivial to understand. If we were to consider softmax attention and carry out the same numerical experiments with sparsity constraints of the form of equation 8, then do we continue to achieve global minima. If authors can add such an experiment, then this can serve as very solid motivation for someone to study ICL with softmax attention (and sparsity constraint) and develop theory for the same. \n\nii) The authors reparametrize the query and key matrix as $Q= W_k^{\\top} W_q$. This is not standard in transformers. What would happen if we were to consider the standard parametrization. Do we still obtain stationary points in the general case? Does the update dynamics still converge to the global minima (the verification step of Theorem 4 and 5)? \n\niii) The authors use LBFGS to learn the parameters. Do the results change under standard gradient descent or SGD or Adam?\n\niv) In the most general case of linear transformers, when $P$ and $Q$ are not constrained to be sparse at all, do the authors conjecture that similar results as the ones shown in the paper continue to hold.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n\nFor the linear attention vs softmax attention comment, we provide an extended response in the [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) above. Thank you for pointing it out! \n\n> What would happen if we were to consider the standard parametrization. Do we still obtain stationary points in the general case? Does the update dynamics still converge to the global minima (the verification step of Theorem 4 and 5)? \n\nThere are experimental and theoretical results confirming the optimization of reparameterized network will recovers the same solution characterized in this paper: $W_k W_q^\\top$ converges to $Q^*$ characterized in Theorems 1, 2 and 3 when $W_q$ and $W_k$ are square matrices. Experiments in [1] shows this convergence for single-layer and multi-layer linear attention. Theoretically, [2] recently proves the global convergence of GD to the optimal solution when  $W_q = W_k$ for a single attention layer. \n\n> The authors use LBFGS to learn the parameters. Do the results change under standard gradient descent or SGD or Adam?\n\nWe conjecture that various algorithms will converge to the same solution characterized in the Theorems. This is proven for SGD on single layer attention by (Zhang et al. 2023) [2]. Furthermore, [1] includes experimental results for AdamW on multilayer attention without sparsity constraint. We specifically chose LBFGS because it can get to _very small loss_ much more quickly than Adam, which is important as the predicted sparsity pattern only becomes apparent when the loss is very low (see e.g. Figure 3). Based on preliminary experiments, we believe that, SGD and Adam will both recover the same solutions given enough iterations (and small enough learning rate).\n\n> In the most general case of linear transformers, when $P$ and $Q$ are not constrained to be sparse at all, do the authors conjecture that similar results as the ones shown in the paper continue to hold.\n\n[1] conjectures the similar results hold without sparsity constraint based on experimental observations. In fact, in our Section 5, we did our best to relax such technical sparsity assumption required for theoretical proofs.  We will clarify this in the paper. \n\n\n\nThank you once again for taking the time to review our paper. We sincerely hope that the response to your concerns, as well as the overall response to other reviewers’ concerns, helps assuage your concerns, and view this paper in a more favorable light.\n\n\n\n**References**\n\n[1] Von Oswald, Johannes, et al. “Transformers learn in-context by gradient descent.”\n\n[2] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. “Trained Transformers Learn Linear Models In-Context.”\n \n\n\n"}, {"review_id": "HXB0NBKR1t", "reviewer": "Reviewer_iy8f", "summary": "This paper studies linear transformer and validates its ability in solving linear problems via applying conditioned gradient descent. The authors firstly state that there exists a unique global minimum for single attention layer. Then they explore multi-layer linear transformer and demonstrate that stationary points can be found that perform preconditioned gradient descent algorithm.", "questions": "1. Although $w^\\star$ is well-defined in the paper, it is not clear how $w$ relates to the model parameters or data. \n\nminor:\n\n2. Presuming you are studying least-squares loss, Eq.(5) may miss a square symbol.\n3. There seem to be notation typos: $k\\to i$ in Eq. (9) and $A_k\\to A_i$ in Lemma 2.", "limitations": "While the authors did not explicitly discuss limitations, they provide potential future directions in the discussion section that are related to their work and have not been explored.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper theoretically studies that a $k$-layer linear transformer can perform as a $k$-step preconditioned gradient descent algorithm when forwarding a linear ICL sequence. ", "weaknesses": "1. Related work is absent in this paper. \n2. This paper simplifies the problem a lot by treating all components as linear. However, for a standard transformer, the softmax and activation functions are two key components and introduce nonlinearity. A discussion on the relationship between linear and standard transformers is missing. \n3. Compared to the previous work (von Oswald et al., 2022), the novelty is insufficient except the introduction of non-isotropic data distribution. \n4. Although simulations are provided in this paper, it is not sufficient. It would be interesting to know: How performance varies with the number of layers? How a $k$-layer linear transformer performs compared with a direct gradient descent approach with different step sizes or a standard transformer (during training or after sufficiently trained)? ...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n Thank you for your constructive comments! We address your comments one by one as below.\n \n > Compared to the previous work (von Oswald et al., 2022), the novelty is insufficient except the introduction of non-isotropic data distribution.\n\nOur results in Theorems 1, 2, and 3 are novel even for isotropic inputs. Remarkably, [von Oswald et al. 2022] proves the recurrence of linear attentions is expressive enough to implement gradient descent. However, such expressivity result **do not imply anything about whether the trained transformers also attain such properties**. Here, we  theoretically address the conjecture of [von Oswald et al. 2022] which postulates that the optimization of parameters leads to the implementation of gradient descent. \n\nThe importance of going beyond expressivity can be highlighted by considering the setting of anisotropic data distribution: the network is still expressive to implement gradient descent but the optimal transformer (with a single-layer) is not one that implements gradient descent. This is proved in Theorem 1. Thus, it is important to analyze the landscape of training to characterize the outcome of training. \n\n\n> This paper simplifies the problem a lot by treating all components as linear. However, for a standard transformer, the softmax and activation functions are two key components and introduce nonlinearity. A discussion on the relationship between linear and standard transformers is missing.\n\n \nOur problem settings is taken from [1] since we are building on empirical observations in [1]. More recently, [2] and [3] have also considered similar settings\n - [MLP block with non-linear activations] We argue that MLP block are not necessary for learning linear models. Zeroing out the weights of MLP blocks allows skipping these blocks over residual connections. To implement gradient descent on ridge regression, one can omit MLP blocks and only rely on the attention. \n - [Linear attentions] We elaborate on this concern in the [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) and provide other references which also consider the same settings due to the expressively of linear attentions for linear models. Furthermore, we provide experimental and theoretical results for non-linear attentions. \n\n\n\n> Although simulations are provided in this paper, it is not sufficient. It would be interesting to know: How performance varies with the number of layers?\n\n\nTo address your concern, we actually conducted more experiments and share them in the [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) -- In Figure 2, we compare the linear transformer against Gradient Descent and Preconditioned Gradient Descent (with optimally chosen stepsize). We plot test loss against number of layers of linear transformer/steps of gradient descent.\n\nWe would like to also highlight that there have already been several **empirical works** in the literature [1,4,5] showing various properties of the model learned for in-context learning. Hence, our main focus was to provide \"**theoretical footing**\" for their interesting empirical observations. In particular, we refer the reviewer to these works for more empirical investigations.\n \n\n> Although $w^*$ is well-defined in the paper, it is not clear how $w$ relates to the model parameters or data.\n\n$w^*$ is a random vector, independent of the data distribution. Since the model is trained over random instances, the parameters of optimal model is independent of individual random $w^*$ and only depends on their distribution. We established the connection between distribution of $w^*$ and optimal model in Theorem 1. Although $w^*$ does not show up physically in the weights _the transformer's predicted label is close to the predicted label using $w^*$_ (and as number of layers $\\to \\infty$, the transformer's prediction converges to $y^{(n+1)} = -<w^*, x^{(n+1)}>$)\n\n  > Related work is absent in this paper. \n\nWe reviewed the related works in the introduction. To address your concern, we will add a broader scope of related works in the final version. Also, we will include any other related works that the reviewer recommend. \n > Presuming you are studying least-squares loss, Eq.(5) may miss a square symbol.\n> There seem to be notation typos: \n in Eq. (9) and \n in Lemma 2.\n \n Thank you very much for catching these typos and your comments.  \n \n \n \n\nThank you once again for taking the time to review our paper. We sincerely hope that the response to your concerns, as well as the overall response to other reviewers’ concerns, helps assuage your concerns, and view this paper in a more favorable light.\n\n **References**\n\n [1] Von Oswald, Johannes, et al. “Transformers learn in-context by gradient descent.” International Conference on Machine Learning. PMLR, 2023.\n\n\n[2] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. “Trained Transformers Learn Linear Models In-Context.” arXiv preprint arXiv:2306.09927 (2023)\n\n\n[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. “One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention.” arXiv preprint arXiv:2307.03576 (2023)\n\n\n[4] Garg, Shivam, et al. \"What can transformers learn in-context? a case study of simple function classes.\" Advances in Neural Information Processing Systems 35 (2022)\n\n\n[5] Akyürek, Ekin, et al. \"What learning algorithm is in-context learning? investigations with linear models.\" arXiv preprint 2022."}, {"review_id": "YTbPKtK1Jc", "reviewer": "Reviewer_vRHd", "summary": "The paper tries to tackle an important and relevant problem in LLM training: whether in-context learning properties can be explained by a form of gradient descent?\n\nFor the setting, the authors have chosen to focus specifically on the linear transformer with a single or multiple layers with an input that is given by the data generated using a linear model. Within this setting, the authors were able to show a series of results for various specially constructed transformers. \n\nSpecifically, they look at a (1) single-layer transformer, (2) multi-layer with sparse parameters and (3) multi-layer with more general parameters. For (1), they have provided a characterization of the unique global optimum of a single-layer transformer with linear attentions. They showed that foer the optimal parameters found, the transformer implements a single step of preconditioned gradient descent. For (2), they have shown that global minimizers of the training objective of a two-layer linear transformer over isotropic regression instances. For the most general  scenario (3), they found that certain critical points implement preconditioning using inverse-data-covariance of the input.\n\nExperimentally, the authors have confirmed using a toy-example that optimizing a three layer model indeed results in a minimizer that converges to an inverse of a covariance matrix.", "questions": "The equations for the gradient do not immediately pop out from the eq 7. I would encourage the authors to expose the equations for the gradient to the reader (maybe after eq 5) so that the connection between the gradient and the preconditioner would become more apparent. \n\nWhat are the properties of A in sections 3 and 4? Can a transformer find a positive definite matrix through training? \n\nFor the evaluation of the Theorem 4 and 5, it is a bit hard to interpret the numbers, since the used distance metric is close to zero, but not exactly zero. The obtained value of ~0.1-0.2 is quite hard to interpret, apart from the fact that it is smaller that the distance with respect to the identity. \n\nAlso, the metric that the authors are using involves the minima over the space of scalars $\\alpha$. It would be interesting if the authors can actually provide the scalars $a_i$ from Theorem 4 that they have found empirically. Do these scalars remain the same over multiple restarts of the algorithm?\n\nHow come the same $\\Sigma$ matrix is used for both data and the weights (theorem 4 and 5)? It doesn't sound very realistic that the data and the weights use the same covariance matrix. \n\nWould be great to add more information to Section 4.1 about adaptive coordinate-wise step sizes. If I understood correctly, Theorem 3 talks about the existence of the global minimizer under certain conditions. How exactly does adaptive coordinate-wise step sizes help find this solution?  \n\nI would encourage the authors to add a simple table at the end of the introduction clearly stating the assumptions that each section has. As written, It is quite hard to understand, especially when the notation change from section to section.\n\nSmall comments:\n- Typo l.5: “Can transformers can”\n- L.83: dimension of Z doesn’t match the dimension of W_k, and W_q in eq (1).\n- In Lemma 2 z^(1) etc are mentioned, but not properly defined nor used.\n- Gamma_k in l.155 plays a similar role to a_i from l.184. Maybe call both of these variables gamma?\n- In fig.1 (d) the log(loss) converges to -1.0, which is not a zero value, as the text suggests in l.213.\n- How come the authors choose L-BFGS as an optimizer? It is a fairly specialized algorithm that is not often used these days. What properties of L-BFGS are desirable in the settings that authors used it for? Can the same results be achieved with SGD or Adam?", "limitations": "It would be important to highlight limitations of the analysis. E.g. same covariance matrix, singe-head linear attention...", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper provides an interesting analysis that can be useful to understand the intriguing properties of in-context learning. It studies a series of transformer models, ranging from a simple single layer to a more general multi-layer ones and shows that for all the settings the in-context learning implements the preconditioned gradient descent. While by no means comprehensive, this paper can be useful as a stepping stone for any researcher that wants to further study intriguing properties of in-context learning.\n\nThe paper is easy to read and easy to understand. The assumptions and the math mainly seem to check out.", "weaknesses": "Generally, the paper restricts itself to a very narrow setting of a specific data that is centered at zero, linear model generating the labels, linear transformer, specific sparsity of the parameter matrix.\n\nThe transformers that they study are linear with a single head only. While I do appreciate the value of the analysis and I understand that this is just a step of many, it would be great to include a short paragraph outlining what the authors think about generalization of the proposed approach when the assumptions are removed.\n\nAlso, given the series of the assumptions, I would make sure that they are stated early on in the abstract and introduction.\n\nI also think that the experimental validation can be slightly improved by considering different choice of layers, dimension of x, conditioning factor of the input etc. It would also be great to have an empirical plot demonstrating how close A is to the Gram matrix as the number of points increase (the results from the end of section 3)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n\n**Regarding the setting** We agree with the reviewer that our setting is theoretical. However, we emphasize that despite our assumptions, the setting is expressive enough for ICL, and has been studied in a number of other papers.\n- [1,2,3] use linear attentions as linear attentions are expressive to implement gradient descent on linear regression.   \n- [1-5] use centered data and generated labels\n- [1,2,3,5] focus on linear models \n\nIn Section 5, we relax the sparsity constraint. As noted by reviewer `k8UL`: \"It is quite a neat idea to consider sparsity driven constraints that permit tractability.\" Without the minimal sparsity constraint in Eq. 10, the analysis becomes very difficult. \n\n\n**Generalization of our results**\n\n- **Multi-head attention.** We show here that learning Multi-head _linear_ attentions reduces to learning a single head attention: Recall notation $Attn_{P,Q}(Z)$ in Eq. 2. Due to the linearity of the attention in $P$ and $Q$, we have  \n $\\sum_{i} Attn_{P_i,Q_i}(Z) = Attn_{\\sum_i P_i,\\sum_{i}Q_i}(Z)$.  Thus the reparameterization $P' := \\sum_{i} P_i$ and $Q' = \\sum_{i} Q$ casts the problem to learning a single head attention. \n- **Linear attentions.** We elaborate on this concern in the [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) and provide other references which also consider the same settings due to the expressively of linear attentions for linear models. Furthermore, we provide experimental and theoretical results for non-linear attentions.  \n \n**Assumptions**\n\nIn the abstract, we narrow our scope to \"linear transformers trained over random instances of linear regression\". Our introduction further clarifies \"Akin to recent work, we too focus on the setting of linear regression encoded via ICL, for which we train a  transformer architecture with multiple layers of single-head self-attentions without softmax\". Linear regression encoded via ICL is introduced in [1,4,5] which encompasses the settings of generated labels, linear regression, centered input data and specific encoding of regression samples. We will gladly add more required details. \n\n**Additional experiments**\n\nThank you for your suggestions. We provide additional experiments in the PDF under [joint response](https://openreview.net/forum?id=LziniAXEI9&noteId=p362P0UrGI) above. In particular, in Figure 3, we plot \"Distance of A to Identity\" against \"number of ICL samples\".\n \n**Questions**\n \n- Regarding the scalars in Theorem 4, we have a theoretical result in Lemma 7 in the Appendix proving permutations of scalars obtains the same function value. Thus independent runs will not recover the same scalars. \n\n-  In experiments, we observe gradient descent converges to a positive definite matrix which depends on the covariance matrix of inputs (characterized in Theorems 4 and 5). We take steps towards proving this observation showing that the characterized parameters in Theorem 4 and 5 are stationary points of gradient descent. An interesting follow-up paper proves the convergence of gradients to the global optimum (in Theorem 1) for a single layer transformer [2]. (though the convergence for multi-layer transformer is still an open problem.)\n\n- We agree about the error ambiguity. In Figures 2 and 4, we visualize the outputs (after training) to illustrate the diagonal dominance of matrices which matches the result of Theorems 4 and 5. The error can be reduced by using a larger set of training set. \n\n-  In Theorem 1, we used different covariance matrices for data and $w^*$. However, the analysis for multi-layer attention is significantly more difficult when covariates are dependent to each other. We motivate this assumption in the practical scenario when \"observations of covariates are distorted by $\\Sigma^{1/2}$  (see line 190 of paper). The same covariance structure was also considered in a few other papers, e.g. [3].\n\n- Here, we do not analyze the convergence of adaptive coordinate-wise stepsizes to find the global optimum characterized in Theorem 3 (though it be an interesting future direction). Instead, we analyze the structure of a global minimizer of training objective and provide an algorithm interpretation for that. \n \n- We will improve readability: (i) we will add a table at the end of introduction and outline assumptions and main results of different sections (ii) add gradient formula. \n\n**Small comments:**\n\n- LBFGs enjoys considerably faster convergence rate compared to Adam in our settings. [1] has experiments with AdamW without sparsity constraints of Section 5. [2] proves the global convergence of gradient descent for a single-head attention. We will compare convergence rates of these algorithms to justify our choice for the optimizer. \n\n- In discussions, we included limitations of our analysis where we talked about linearity of attention. We will add more limitations such as covariance matrix structure, and single-head analysis. \n\nThank you once again for taking the time to review our paper. We sincerely hope that the response to your concerns helps assuage your concerns, and view this paper in a more favorable light.\n  \n **References**\n\n[1] Von Oswald, Johannes, et al. “Transformers learn in-context by gradient descent.” \n\n[2] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. “Trained Transformers Learn Linear Models In-Context.” \n\n[3] Mahankali, Arvind, et al. “One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention.” \n\n[4] Garg, Shivam, et al. \"What can transformers learn in-context? a case study of simple function classes.\" \n\n[5] Akyürek, Ekin, et al. \"What learning algorithm is in-context learning? investigations with linear models.\" \n\n[6] Imanol Schlag, et al. Linear transformers are secretly fast weight programmers. "}]

Author rebuttal:
We thank all the reviewers for very constructive comments. 
- We thank **Reviewer m6sU** recognizing the novelity of our work:  _"The paper provides a novel perspective -- the preconditioning/adaptive perspective -- on in-context learning in transformers... Given the growing interest in understanding ICL from a theoretical perspective, the "preconditioning gradient descent" can be served as an important building block. The analysis tools developed in this paper could be helpful for future research on understanding ICL."_
- We thank **Reviewer k8UL** for appreciating the value of our approach: _"It is quite a neat idea to consider sparsity driven constraints....The results resolve some very important questions towards understanding ICL....This paper serves as a crucial step towards the goal of understanding ICL is one of the central problems towards demistifying LLMs"_. 
- We thank **Reviewer vRHd** for appreciating the analysis in our paper:  _"The paper provides an interesting analysis that can be useful to understand the intriguing properties of in-context learning...this paper can be useful as a stepping stone for any researcher that wants to further study intriguing properties of in-context learning."_

We begin our response with the summary of our main contributions as summarized by the reviewers above.

**Primary technical achievements:**

- The inspiring previous works have shown that by setting the weights of transformers carefully, one can ICL the task of linear regression. We follow up on these works and show that in fact when you train your the transformer over the random instances of linear regression, **the global minima (and stationary points) correspond to interesting algorithms** (such as "preconditioned" gradient descent). 
- At a more technical level, **we propose an analytical techniques to analyze the training loss of transformer architecture**.
- We would also like to underscore the relevance of this work to this community by briefly mentioning concurrent works that were posted shortly after our submission. First, Zhang et al. [2] have shown that under a random initialization, gradient descent converges to the global minimum characterized in this work, contributing to our understanding training transformers. Moreover, Arvind et al. [3] establish a similar set of results to ours for the single layer case.
We thus believe that the **question we attempt to answer in this work is timely and relevant to the community,** and would hopefully spark further interests.

We next address a shared concern raised by the reviewers:

**Beyond Linear Attention:**

As noted by reviewers, we focus on theoretical analysis of transformers with _linear attention and no MLP layers_. Nevertheless, such networks are expressive enough for solving linear regression, and they have become the focus of theoretical studies in various recent works [1,2,3]. In particular, Lemma 2 proves these networks are expressive enough to implement a family of optimization methods. 

We observe empirically that softmax attention behaves similarly to the linear attention (consistent with observations in [Appendix A.9] in Oswald et al.) In fact, under the two-headed softmax attention, we empirically observe that the learned algorithm is similar to that of linear attention. Without going into details, the key intuition is the following linearization trick $\frac{1}{2} (e^x-e^{-x}) \approx x$. Indeed, in our experiments, we observe that the weights of the two attention heads have approximately opposite sign to each other. 

Since theoretical analysis of softmax is rather difficult, we instead analyze a variant where softmax is replaced by ReLU as suggested by [4]. Experimental observation, presented in the attached PDF showing that the sparsity pattern of the solution is the same as those in Thm. 1. Here, we provide a theoretical analysis for this observation. 

Following the theoretical settings of [1], we introduce a variant of non-linear attention where softmax is replaced by ReLU, i.e. $Attn_{Q}(Z) =  P ZM g(
Z^\top Q Z)$ 
where $g$ is the ReLU activation function, and $P = [0,\dots, 0, 1]^{\otimes 2} \in \mathbb{R}^{(d+1)\times (d+1)}$ and $Q = \begin{bmatrix} A & 0\\\\
0 & 0\end{bmatrix}$. 
The in-contex loss is defined as
$$f(A) = \mathbb{E} \left( \sum_{i} \langle x^{(i)}, w^* \rangle    g(x^\top A x^{(i)} )- \langle x, w^* \rangle \right)^2 $$
where we use the compact notation $x = x^{(n+1)}$. We omit the constant $1/N$ for the simplicity of notations.

**Lemma.** Suppose that the data and $w^*$ distribution is Gaussian and isotropic. Then there is a constant $c$ such that $c I = \arg\min_A f(A)$. 

This is a theoretical result derived in the limited rebuttal time, and hence has the following limitations (in addition to requiring further proof-checking):

- We fixed the parameter $P$ in the attention and only optimized $Q$. 
- We only analyzed ReLU instead of softmax
- We only analyzed isotropic inputs. 

The proof is based on an application of Stein's Lemma. [5] uses this technique to prove the solution of general linearized linear models for various non-linear functions and Gaussian inputs. Due to response limits, we could not include the proof. Upon request, we can add the proof in an official comment. 

---
**References.** 

[1] Von Oswald, Johannes, et al. "Transformers learn in-context by gradient descent."

[2] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. "Trained Transformers Learn Linear Models In-Context." 

[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. "One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention." 

[4] Zhao, H., Panigrahi, A., Ge, R., & Arora, S.. Do Transformers Parse while Predicting the Masked Word?

[5] Erdogdu, M. A., Dicker, L. H., & Bayati, M. (2016). Scaled least squares estimator for GLMs in large-scale problems.



Meta-review: The paper studies in-context learning, a phenomenon where Transformers can adjust their predictions based on additional data given in the input sequence itself. Prior work, e.g. von Oswald, showed that in-context learning in Transformers can approximate gradient-based learning within its forward pass.

The reviewers were in general rather positive about the paper, mentioning it provides an interesting perspective with a rigorous analysis. Some reviewers asked about the novelty compared to prior work, which in my view is adequately addressed in the rebuttal. For instance, the analysis of von Oswald is a constructive proof, stating a result about the potential expressivity of Transformers, but not what they learn. The paper therefore takes a significant step in the direction of showing that the optimization of parameters leads to the implementation of gradient descent.

Finally, the discussion between authors and reviewers has led several reviewers to increase their scores. Overall, the paper is clearly of interest to the community and I therefore recommend acceptance. I do encourage the authors to revise the paper according to the comments of the reviewers.
Decision: Accept (poster)