Title: Transformers learn to implement preconditioned gradient descent for in-context learning
Venue: NeurIPS 2023
Accepted: yes
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

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

1 Introduction

In-context learning (ICL) is the striking capability of large language models: Given a prompt containing examples and a query, the transformer produces the correct output based on the context provided by the examples, without adapting its parameters (Brown et al., 2020; Lieber et al., 2021; Rae et 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 et al., 2022; von Oswald et al., 2023; Li and Malik, 2017; Min et 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 et al., 2021; Wei et al., 2022). Given an in-context prompt, Edelman et al. (2022); Olsson et al. (2022) argue that transformers are able to implement algorithms through the recurrence of multi-head attentions to extract coarse information from raw input prompts. Akyürek et al. (2022); von Oswald et al. (2023) assert that transformers can implement gradient descent on linear regression encoded in a given input prompt. It is thought provoking that transformers can 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.

More specifically, we investigate the learning of gradient-based methods. It is hard to mathematically formulate what it means to learn gradient descent for general functions with transformers. Yet, Garg et al. (2022) elegantly examine it in the specific setting of ICL 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 a single-layer linear transformer. In 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 us to formulate training transformers as a search over k -step adaptive gradient-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 data-dependent preconditioning, can be derived from a 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-based 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 is available at https://github.com/chengxiang/LinearTransformer .

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

A 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

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 on a 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 it the same as that in Theorem 1. Mahankali et al. (2023) also characterize the global minimizer of a single layer attention without softmax for a different data distribution. In addition to results for a 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 et 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 ( i ) ∈ R d be the covariates drawn i.i.d. from a distribution D X , and w ⋆ ∈ R d be drawn from D W . Let X ∈ R ( n +1) × d be the matrix of covariates whose row i contains tokens x ( i ) . Given x ( i ) 's and w ⋆ , the responses are defined as y = [ ⟨ x (1) , w ⋆ ⟩ , . . . , ⟨ x ( n ) , w ⋆ ⟩ ] ∈ R n . Define the input matrix Z 0 as

where zero in the above matrix is used to replace the unknown response variable corresponding to x ( n +1) . Then, our goal is to predict w ⊤ ⋆ x ( n +1) given Z 0 . In other words, the training data consists of pairs ( Z 0 , w ⊤ ⋆ x ( n +1) ) for x ( i ) ∼ D X and w ⋆ ∼ D W . We then consider training transformers over this data distribution.

Self-attention layer without softmax. Following (Schlag et al., 2021; von Oswald et al., 2023), we consider the linear self-attention layer. To motivate, we first briefly review the standard self-attention layer (Vaswani et al., 2017). Letting Z ∈ R ( d +1) × ( n +1) be the input matrix with n +1 tokens in R d +1 , a single-head self-attention layer denoted by Attn smax is a parametric map defined as

where W v , W k , W q ∈ R ( d +1) × ( 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 x ( n +1) is excluded from the input. To reflect this asymmetric structure, the mask matrix M is 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 v ∈ R ( d +1) × ( d +1) and Q := W k ⊤ W q ∈ R ( d +1) × ( d +1) and consider

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 to solve linear regression in-context.

Architecture for prediction. Wenowpresent the neural network architecture that will be used throughout this paper. For the number of layers L , we define an L -layer transformer as a stack of L linear self-attention blocks. Formally, denoting by Z ℓ the output of the ℓ th layer attention, we define

The scaling factor 1 / n is used only for ease of notation and does not influence the expressive power of the transformer. Given Z L , we define TF L ( Z 0 ; { P ℓ , Q ℓ } ℓ =0 , 1 ,...L -1 ) = -[ Z L ] ( d +1) , ( n +1) , i.e., the ( d +1 , n +1) -th entry of Z L . The reason for the minus sign is to be consistent with (von Oswald

et al., 2023), and we will clarify such a choice in Lemma 1. For training, the parameters are optimized to minimize in-context loss as

Goal: the landscape analysis of the training objective functions. We are interested in understanding how the optimization of f 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 i , Q i } even for a single layer transformer. (ii) 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 to multi-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 1 summarizes our results for various parameteric models.

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 1 ( Optimizing (5) vs. practical transformer optimization ) . Interestingly, a recent work by Ahn et al. (2023) reports that common optimization algorithms such as SGD/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 a good proxy for understanding practical transformer optimization.

3 The global optimum for a single-layer transformer

For the single layer case of L = 1 , the following result characterizes the optimal parameters P 0 and Q 0 for the in-context loss (5).

Theorem 1 ( Single-layer; non-isotropic data ) . Assume that vector x ( i ) is sampled from N (0 , Σ) , i.e., a Gaussian with covariance Σ = U Λ U ⊤ where Λ = diag( λ 1 , . . . , λ d ) . Moreover, assume that w ⋆ is sampled from N (0 , I d ) . Then, the following choice of parameters

is a global minimizer of f ( P, Q ) up to re-scaling, i.e., P 0 ← γP 0 and Q 0 ← γ -1 Q 0 for a scalar γ .

See Appendix A for the proof of Theorem 1. In the specific case when the Gaussian is isotropic, i.e., Σ = I d , the optimal Q 0 has the following simple form

Up to scaling, the above parameter configuration is equivalent to the parameters used by von Oswald et 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-isotropic, the transformer learns to implement one step of a preconditioned gradient descent as we shall detail in Lemma 1. Here the 'preconditioning matrix' given in (6) has interesting properties:

When the number of samples n is large, the first d × d submatrix of Q 0 approximates Σ -1 , the inverse of the data covariance matrix, which is also close to the Gram matrix formed from x (1) , . . . , x ( n ) . Hence the preconditioning can lead to considerably faster convergence rate when Σ is ill-conditioned.

Moreover, 1 n ∑ k λ k in (6) acts as a regularizer. It becomes more significant when n is small and variance of the x ( i ) 's is high. Such an adjustment resembles structural risk minimization (Vapnik, 1999) where the regularization strength is adapted to 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, we impose the following restrictions on the parameters:

The next lemma proves that a forward-pass of a L -layer transformer, with the parameter configuration (8) is the same as taking L steps of gradient descent, preconditioned by A ℓ .

Lemma 1 ( Forward pass as a preconditioned gradient descent ) . Consider the L -layer linear transfomer parameterized by A 0 , . . . , A L -1 as in (8) . Let y ( n +1) ℓ be the ( d + 1 , n + 1) -th entry of the ℓ -th layer output, i.e., y ( n +1) ℓ = [ Z ℓ ] ( d +1) , ( n +1) for ℓ = 1 , . . . , L . Then, it holds that y ( n +1) ℓ = -⟨ x ( n +1) , w gd ℓ ⟩ where { w gd ℓ } is defined as w gd 0 = 0 and as follows for ℓ = 1 , . . . , L -1 :

See Subsection C.1 for a proof. The iterative scheme (9) includes various optimization methods including gradient descent with A ℓ = γ ℓ I d , and (adaptive) preconditioned gradient descent, where the preconditioner A ℓ depends on the time step. In the upcoming sections, we characterize how the optimal { A ℓ } are linked to the input distribution.

4.1 Warm-up: optimal two-layer transformer with symmetric weights

For the rest of this section, we 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 a more general parameters. For a two-layer transformer, the next Theorem proves the optimal in-context loss obtains the simple gradient descent with adaptive coordinate-wise stepsizes.

Theorem 2 (Global optimality for the two-layer (symmetric) transformer) . Consider the optimization of in-context loss for a two-layer transformer with the parameter configuration in Eq. (8) , and additionally assume that A 1 , A 2 are symmetric matrices. More formally, consider

Assume x ( i ) i.i.d. ∼ N (0 , I d ) and w ⋆ ∼ N (0 , I d ) ; then, there are diagonal matrices A 1 and A 2 that are a global minimizer of f .

Combining the above result with Lemma 1 concludes that the two iterations of gradient descent with coordinate-wise 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 its stepsize based on the individual problem instance, the algorithm learned adjusts its stepsize to the underlying data distribution.

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 a specific preconditioned gradient algorithm, where the preconditioning matrix is the inverse covariance of the input distribution. Before stating this result, let us first consider a motivating scenario in which the data-covariance matrix is non-identity:

Linear regression with distorted view of the data: Suppose that w ⋆ ∼ N (0 , I ) and the latent covariates are x (1) , . . . , x ( n +1) , drawn i.i.d from N (0 , I ) . We are given y (1) , . . . , y ( n ) , with y ( i ) = 〈 x ( i ) , w ⋆ 〉 . However, we do not observe the latent covariates x ( i ) . Instead, we observe the distorted covariates x ( i ) = Wx ( i ) , where W ∈ R d × d is a distortion matrix. Thus the prompt consists of ( x (1) , y (1) ) , . . . , ( x ( n ) , y ( n ) ) , as well as x ( n +1) . The goal is still to predict y ( n +1) . Note that this setting is quite common in practice, when covariates are often represented in an arbitrary basis.

Assume that Σ := WW ⊤ ≻ 0 . We verify from our definitions that for w ⋆ := Σ -1 / 2 w ⋆ , y ( i ) = 〈 x ( i ) , w ⋆ 〉 . Furthermore, x ( i ) ∼ N (0 , Σ) and w ⋆ ∼ N (0 , Σ -1 ) . From Lemma 1, the transformer with weight matrices { A 0 , . . . , A L -1 } implements preconditioned gradient descent with respect to R w ⋆ ( w ) = 1 2 n ( w -w ⋆ ) T XX ⊤ ( w -w ⋆ ) , with X = [ x (1) , . . . , x ( n ) ] . Under this loss, the Hessian matrix ∇ 2 R w ⋆ ( w ) = 1 2 n XX ⊤ (at least in the case of large n ). For any fixed prompt, Newton's method corresponds to A i ∝ ( XX ⊤ ) -1 , which makes the problem well-conditioned even if Σ is very degenerate. As we will see in Theorem 3 below, the choice of A i ∝ Σ -1 = E [ XX ⊤ ] -1 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 i } L -1 i =0 ∈ R L × d × d . We use f ( A ) to denote the in-context loss of f ( { P i , Q i } L -1 i =0 ) as defined in (5), when Q i depends on A i , and P i is a constant matrix, as described in (8).

Theorem 3. Assume that x ( i ) iid ∼ N (0 , Σ) and w ⋆ ∼ N (0 , Σ -1 ) , for i = 1 , . . . , n , and for some Σ ≻ 0 . Consider the optimization of in-context loss for a k -layer transformer with the the parameter configuration in Eq. (8) given by:

Let S ⊂ R L × d × d be defined as follows: A ∈ S if and only if for all i = 0 , . . . , L -1 , there exists scalars a i ∈ R such that A i = a i Σ -1 . Then

where ∇ A i f denotes derivative wrt the Frobenius norm ∥ A i ∥ F .

As discussed in the motivation above, under the setting of A i = a i Σ -1 , the linear transformer implements an algorithm that is reminiscent of Newton'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 is ill-conditioned. The proposed parameters A i in Theorem 3 are also similar to A i 's in Theorem 1 when n is large. However, in contrast to Theorem 1, there is no trade-off with statistical robustness; this is because w ⋆ has covariance matrix Σ -1 in the Theorem 3, while Theorem 1 has isotropic w ⋆ .

Unlike our prior results, Theorem 3 only guarantees that the set S of transformer prameters satisfying { A i ∝ Σ -1 } L -1 i =0 essentially 2 contains critical points of the in-context loss. However, in the next section, we show experimentally that this choice of A i 's does indeed seem to be recovered by training.

We defer the proof of Theorem 3 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

exploiting invariances in the underlying distribution of prompts, we construct a flow, contained entirely in S , whose objective value decreases as fast as gradient flow. Since f 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 an 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 x ( i ) ∼ N (0 , Σ) and w ⋆ ∼ N (0 , Σ -1 ) , where Σ = U T DU , where U is a uniformly random orthogonal matrix, and D is a 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 A 0 , A 1 , and A 2 are initialized by i.i.d. Gaussian matrices. Each gradient step is computed from a 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 to 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 A 0 , A 1 , A 2 and Σ -1 (up to scaling), we use the normalized Frobenius norm distance : Dist( M,I ) := min α ∥ M -α · I ∥ ∥ M ∥ F , (equivalent to choosing α := 1 d ∑ d i =1 M [ i, i ] ). This is essentially the projection distance of M / ∥ M ∥ F onto the space of scaled identity matrices.

We plot Dist ( A i , I ) , averaged over 5 runs, against iteration in Figures 1a,1b,1c. In each plot, the blue line represents Dist(Σ 1 / 2 A i Σ 1 / 2 , I ) , and we verify that the optimal parameters are converging to the critical point introduced in Theorem 3, which implements preconditioned gradient descent. The red line represents Dist( A i , I ) ; it remains constant indicating that the trained transformer is not implementing plain gradient descent. Figures 2a-2c visualize each Σ 1 / 2 A i Σ 1 / 2 matrix at the end of training to further validate that the learned parameter is as described in Theorem 3.

0

5000

10000

15000

20000

25000

30000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

0

1/2

A

0

1/2

(a)

Dist(Σ

1

/

2

A

0

Σ

1

/

2

, I

)

0

5000

10000

15000

20000

25000

30000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

1

1/2

A

1

1/2

(b)

Dist(Σ

1

/

2

A

1

Σ

1

/

2

, I

)

0

5000

10000

15000

20000

25000

30000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

2

1/2

A

2

1/2

(c)

Dist(Σ

1

/

2

A

2

Σ

1

/

2

, I

)

0

5000

10000

15000

20000

25000

30000

Iteration

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

log(Loss)

(d) log(Loss)

5 Multi-layer transformers beyond standard optimization methods

In this section, we study the more general setting of

Note that A i , B i are not constrained to be symmetric. Similar to Section 4, we introduce the following simplified notation: let A := { A i } L -1 i =0 ∈ R L × d × d and B := { B i } L -1 i =0 ∈ R L × d × d . We use f ( A,B ) to denote the in-context loss of f ( { P i , Q i } L -1 i =0 ) as defined in (5), when P i and Q i depend on B i and A i as described in (11).

0

1

2

3

4

0

1

2

3

4

-0.69

-0.00

-0.00

-0.00

-0.00

-0.00

-0.69

-0.00

-0.00

-0.00

-0.00

-0.00

-0.69

-0.00

-0.00

0.00

-0.00

-0.00

-0.69

-0.00

0.00

-0.00

-0.00

0.00

-0.69

0.6

0.5

0.4

0.3

0.2

0.1

0.0

0

1

2

3

4

0

1

2

3

4

-0.42

-0.01

0.00

-0.00

0.00

-0.00

-0.43

-0.00

0.00

-0.00

0.00

-0.00

-0.43

0.00

-0.00

-0.00

0.00

-0.00

-0.42

-0.00

-0.00

0.00

-0.00

-0.00

-0.43

0.40

0.35

0.30

0.25

0.20

0.15

0.10

0.05

0.00

0

1

2

3

4

0

1

2

3

4

-1.71

-0.01

0.00

-0.01

-0.00

-0.00

-1.72

-0.01

-0.00

-0.00

0.01

-0.00

-1.72

0.01

-0.00

0.00

-0.00

0.00

-1.69

0.00

-0.00

0.00

-0.01

0.01

-1.72

1.6

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

(a) Visualization of Σ 1 / 2 A 0 Σ 1 / 2

(c) Visualization of Σ 1 / 2 A 2 Σ 1 / 2

Figure 2: Visualization of learned weights for the setting of Theorem 3. We visualize each Σ 1 / 2 A i Σ 1 / 2 matrix at the end of training. Note that the optimized weights match the stationary point discussed in Theorem 3.

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 Σ denote any PSD matrix. Assume that x ( i ) iid ∼ N (0 , Σ) and w ⋆ ∼ N (0 , Σ -1 ) , for i = 1 , . . . , n , and for some Σ ≻ 0 . Consider the optimization of in-context loss for a L -layer linear transformer with the the parameter configuration in Eq. (11) given by:

Let S ⊂ R 2 × L × d × d be defined as follows: ( A,B ) ∈ S if and only if for all i ∈ { 0 , . . . , k } , there exists scalars a i , b i ∈ R such that A i = a i Σ -1 and B i = b i I . Then

where ∇ A i f denotes derivative wrt the Frobenius norm ∥ A i ∥ F .

In words, parameter matrices in S implement the following algorithm: { A i = a i Σ -1 } L -1 i =0 plays the role of a distribution-dependent preconditioner for the gradient steps. At the same time, B i = b i I transforms the covariates themselves to make the Gram matrix have better condition number with each iteration. When the Σ = I , the algorithm implemented by A i ∝ I, b i ∝ I is exactly the GD++ algorithm proposed in (von Oswald et al., 2023) (up to stepsize).

The result in (12) says that the set S essentially 3 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 x ( i ) ∼ N (0 , Σ) and w ⋆ ∼ N (0 , Σ -1 ) , where Σ = U T DU , where U is a uniformly random orthogonal matrix, and D is a fixed diagonal matrix with entries (1 , 1 , 0 . 25 , 0 . 0625 , 1) . We train a three-layer linear transformer, under the constraints in (11) which is less restrictive than (8) in Subsection 4.3. We train the matrices A 0 , A 1 , A 2 , B 0 , B 1 4 using ADAM with the same setup as in Section Subsection 4.3. We repeat this experiment 5 times with different random seeds, each time we sample a different U (i.e. Σ ).

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.

0

20000

40000

60000

80000

100000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

B

0

(a)

Dist(

B

0

, I

)

0

20000

40000

60000

80000

100000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

B

1

(b)

Dist(

B

1

, I

)

0

20000

40000

60000

80000

100000

Iteration

6

4

2

0

2

log(Loss)

(c) log(Loss)

0

20000

40000

60000

80000

100000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

0

1/2

A

0

1/2

(d) Distances for

A

0

0

20000

40000

60000

80000

100000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

1

1/2

A

1

1/2

(e) Distances for

A

1

0

20000

40000

60000

80000

100000

Iteration

0.0

0.2

0.4

0.6

0.8

1.0

Distance to Id

A

2

1/2

A

2

1/2

(f) Distances for

A

2

We next verify that the parameters at convergence are consistent with Theorem 4. We will once again use Dist( M,I ) to measure the distance from M to the identity matrix, up to scaling (see Subsection 4.3 for definition of Dist ). Figures 3a and 3b show that B 0 and B 1 are close to identity, as Dist( B i , I ) appears to be decreasing to 0. Figures 3d, 3e and 3f plot Dist( A i , I ) (red line) and Dist(Σ 1 / 2 A i Σ 1 / 2 , I ) (blue line); the results here suggest that A i is converging to Σ -1 , up to scaling. In Figures 3a and 3b, we observe that B 0 and B 1 also converge to the identity matrix ( without left and right multiplication by Σ 1 / 2 ), consistent with Theorem 4.

We visualize each of B 0 , B 1 in Figure 4 and A 0 , A 1 , A 2 in Figure 5a-5c at the end of training. We highlight two noteworthy observations:

Let X k ∈ R d × n denote the first d rows of Z k , which are the output at layer k -1 defined in (4). Then the update to X k is X k +1 = X k + B k X k MX T k A k X k ≈ X k +1 = X k ( I -| a k b k | MX T k X k ) , where M is a mask defined in (2). As noted by von Oswald et al. (2023), this may be motivated by curvature correction.

As seen in Figures 5a-5c in the Appendix, ∥ A 0 ∥ ≤ ∥ A 1 ∥ ≤ ∥ A 2 ∥ that implies the transformer implements gradient descent with a small stepsize at the beginning and a large stepsize at the end. This makes intuitive sense as X 2 is better-conditioned compared to X 1 , due to the choice of B 0 , B 1 . This can be contrasted with the plots in Figures (2a)-(2c), where similar trends are not as pronounced because B i 's are constrained to be 0 .

6 Discussion

We take a 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 a single layer transformer, we prove that the global minimum corresponds to a 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 as a broad family of adaptive gradient-based algorithms.

0

1

2

3

4

0

1

2

3

4

0.90

-0.00

0.00

0.00

-0.00

0.00

0.91

0.00

-0.01

-0.00

0.01

0.01

0.90

-0.02

-0.01

-0.00

-0.00

0.00

0.90

-0.01

0.00

-0.00

-0.00

0.01

0.91

0.0

0.2

0.4

0.6

0.8

0

1

2

3

4

0

1

2

3

4

0.56

-0.00

0.00

0.00

-0.00

-0.00

0.57

-0.00

-0.00

0.00

-0.00

-0.00

0.57

0.00

0.00

0.00

0.00

-0.00

0.56

0.00

-0.00

-0.00

0.00

0.00

0.57

0.0

0.1

0.2

0.3

0.4

0.5

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 is to ask the effect of nonlinear activations for our main results. Empirically, von Oswald et 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 to match their performance with that of linear transformers.

As a first step analysis, we consider the nonlinear attention defined as

The following result is an analog of Theorem 1 for single-layer nonlinear attention. It characterizes a global minimizer for this setting with ReLU activation. Here, our choice of ReLU activation was motivated by Wortsman et al. (2023) who observed that ReLU attention matches the performance of softmax attention for vision transformers.

Theorem 5. Consider the single layer nonlinear attention setting with σ = ReLU . Assume that vector x ( i ) is sampled from N (0 , I d ) . Moreover, assume that w ⋆ is sampled from N (0 , I d ) . Consider the parameter configuration P 0 , Q 0 where we additionally assume that the last row of Q 0 is zero. Then, the following parameters form a global minimizer of the corresponding in-context loss:

The proof of Theorem 5 involves an 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 to the global minimum with linear attention, established in Theorem 1 (specifically the minimizer for the isotropic date given in (7)).

Refined landscape analysis for multilayer transformer. Theorem 4 proves that a 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 4 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)optimality.

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-20-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 acknowledges support from NSF CCF-2112665 (TILOS AI Research Institute). Hadi Daneshmand acknowledges support from NSF TRIPODS program (award DMS-2022448). Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute).

References

Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082 , 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.13673 , 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, et al. Gpt-neox-20b: 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 Linguistics: 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 L 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-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread , 2021. https://transformer-circuits.pub/2021/framework/index.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 S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 , 2023.

Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401 , 2014.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 1997.

Stanisław Jastrzebski, Devansh Arpit, Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual connections encourage iterative inference. In International Conference on Learning Representations , 2018. URL https://openreview.net/forum?id=SJa9iHgAZ .

Ke Li and Jitendra Malik. Learning to optimize. In International Conference on Learning Representations , 2017.

Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evaluation. White Paper. AI21 Labs , 2021.

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

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 Technologies , 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, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446 , 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 T 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 is 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: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 2022.

Mitchell Wortsman, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:2309.08586 , 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 L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117 , 2023.

Appendix

A Proofs for the single layer case

In this section, we prove our characterization of global minima for the single layer case (Theorem 1). We begin by simplifying the loss into a more concrete form. Throughout the proof, we will write P, Q instead of P 0 , Q 0 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:

From the definition of attention given in (3), one can further spell out the expression Z 0 + 1 n Attn P,Q ( Z 0 ) using the notation Z 0 = [ z (1) z (2) · · · z ( n +1) ] as follows:

Thus, the last column of the above matrix can be expressed as

where note that the summation is for i = 1 , 2 , . . . , n due to the mask matrix M . Therefore, letting b ⊤ be the last row of P , and A ∈ R d +1 ,d be the first d columns of Q (as we did in (14)), then [ Z 0 + 1 n Attn P,Q ( Z 0 ) ] ( d +1) , ( n +1) can be written as

in other words, f ( P, Q ) only depends on the parameter b and A . Henceforth, we will write f ( P, Q ) as f ( b, A ) . Let us summarize our conclusion so far since it's crucial for the analysis to 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

Now with this parametrization, the in-context loss can be written as f ( b, A ) := f ([0 b ] ⊤ , [ A 0]) .

Now, let us spell out f ( b, A ) based on (13) as follows:

where we used the notation G := 1 n ∑ i z ( i ) z ( i ) ⊤ to simplify. We now analyze the global minima of this loss function. To illustrate the proof idea clearly, we begin with the proof for the simpler case of isotropic data.

A.2 Warm-up: proof for the isotropic data

As a warm-up, we first prove the result for the special case where x ( i ) is sampled from N (0 , I d ) .

̸

1. Decomposing the loss function into components. Writing A = [ a 1 a 1 · · · a d ] , and use the fact that E [ x ( n +1) [ j ] x ( n +1) [ j ′ ]] = 0 for j = j ′ and E [ x ( n +1) [ j ] 2 ] = 1 , we get

The key idea is to characterize the global minima of each component in the summation separately. Another key idea is to reparametrize the cost function given the following identity:

where we use the notation ⟨ X,Y ⟩ := Tr( XY ⊤ ) for two matrices X and Y here and below. Given the above identity, we define each component in the summation as follows.

Characterizing global minima of each component. To characterize the global minima of each objective, we prove the following result.

Lemma 2 ( Global minima of each component ) . Suppose that x ( i ) is sampled from N (0 , I d ) and w ⋆ is sampled from N (0 , I d ) . Consider the following objective (here, ⟨ X,Y ⟩ := Tr( XY ⊤ ) for two matrices X and Y )

Then a global minimum is given as

where E i 1 ,i 2 is the matrix whose ( i 1 , i 2 ) -th entry is 1 , and the other entries are zero.

Proof of Lemma 2. Note first that f j is convex in X . Hence, in order to show that matrix X j is the global optimum of f j , it suffices to show that the gradient vanishes at that point, in other words, ∇ f j ( X j ) = 0 . To verify this, let us compute the gradient of f j : for a matrix X ,

where we recall that G is defined as

To verify that the gradient is equal to zero, let us first compute E [ w ⋆ [ j ] G ] . For each i = 1 , . . . , n , note that E [ w ⋆ [ j ] x ( i ) x ( i ) ⊤ ] = O because E [ w ⋆ ] = 0 . Moreover, E [ w ⋆ [ j ] ( y ( i ) ) 2 ] = 0 because w ⋆ is symmetric, i.e., w ⋆ d = -w ⋆ , and y ( i ) = ⟨ w ⋆ , x ( i ) ⟩ . Lastly, for k = 1 , 2 , . . . , d , we have

̸

In order to compute E [ ⟨ G , X ⟩ G ] , let us compute E [ ⟨ G , E i,i ′ ⟩ G ] for i, i ′ = 1 , . . . , d +1 . Without loss of generality, i ≥ i ′ . First of all We now compute compute E [ ⟨ G , E d +1 ,j ⟩ G ] . Note first that

Hence, it holds that

because E [ w ⋆ ] = 0 . Next, we have

because w ⋆ d = -w ⋆ . Lastly, we compute

To that end, note that for j = j ′ ,

̸

and

̸

where the last case follows from the fact that the fourth moment of Gaussian is 3 and

Combining the above calculations together, we arrive at

Therefore, combining (17) and (19), the results follows.

̸

3. Combining global minima of each component. From Lemma 2, it follows that

is the unique global minimum of f j . Hence, b and A = [ a 1 a 1 · · · a d ] achieve the global minimum of f ( b, A ) = ∑ d j =1 f j ( ba ⊤ j ) if they satisfy

This can be achieve by the following choice:

where e j is the j -th coordinate vector. This choice precisely corresponds to

We next move on to the non-isotropic case.

A.3 Proof for the non-isotropic case

1. Diagonal covariance case. We first consider the case where x ( i ) is sampled from N (0 , Λ) where Λ = diag( λ 1 , . . . , λ d ) and w ⋆ is sampled from N (0 , I d ) . We prove the following generalization of Lemma 2.

Lemma 3. Suppose that x ( i ) is sampled from N (0 , Λ) where Λ = diag( λ 1 , . . . , λ d ) and w ⋆ is sampled from N (0 , I d ) . Consider the following objective

Then a global minimum is given as

where E i 1 ,i 2 is the matrix whose ( i 1 , i 2 ) -th 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

where we recall that G is defined as

A similar calculation as the proof of Lemma 2 yields

Here the factor of λ j comes from the following generalization of (16):

Next, we compute E [ ⟨ G , E d +1 ,j ⟩ G ] . Again, we follow a similar calculation to the proof of Lemma 2 except that this time we use the following generalization of (18):

̸

where the last line follows since

Therefore, we have

Therefore, combining (20) and (21), the results follows.

Now we finish the proof. From Lemma 2, it follows that

is the unique global minimum of f j . Hence, b and A = [ a 1 a 1 · · · a d ] achieve the global minimum of f ( b, A ) = ∑ d j =1 f j ( b, A j ) if they satisfy

This can be achieve by the following choice:

where e j is the j -th coordinate vector. This choice precisely corresponds to

2. Non-diagonal covariance case (the setting of Theorem 1). We finally prove the general result of Theorem 1, namely x ( i ) is sampled from a Gaussian with covariance Σ = U Λ U ⊤ where Λ = diag( λ 1 , . . . , λ d ) and w ⋆ is sampled from N (0 , I d ) . The proof works by reducing this case to the previous case. For each i , define ˜ x ( i ) := U T x ( i ) . Then E [ ˜ x ( i ) ( ˜ x ( i ) ) ⊤ ] = E [ U ⊤ ( U Λ U ⊤ ) U ] = Λ . Now let us write the loss function (15) with this new coordinate system: since x ( i ) = U ˜ x ( i ) , we have

Hence, let us consider the vector ( b ⊤ G A + w ⊤ ⋆ ) U . By definition of G , we have

where we define ˜ b ⊤ := b ⊤ [ U 0 0 1 ] , ˜ A := [ U ⊤ 0 0 1 ] AU , and ˜ w ⋆ := U ⊤ w ⋆ . By the rotational symmetry, ˜ w ⋆ is also distributed as N (0 , I d ) . Hence, this reduces to the previous case, and a global minimum is given as

From the definition of ˜ b, ˜ A , it thus follows that a global minimum is given by

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

We first rewrite the loss function and simplify it following Subsection A.1. Moreover, for simple notation we will often write z, x instead of z ( n +1) , x ( n +1) .

1. Rewriting loss function.

Following Subsection A.1, let us write down the in-context loss (5). for the single-layer nonlinear attention denoted by f ( P, Q ) :

Recalling the definition of the ReLU attention Attn σ P,Q ( Z ) := PZM σ ( Z ⊤ QZ ) , the data matrix Z := [ z (1) · · · z ( n +1) ] , and the mask matrix M , the term Attn σ P,Q ( Z 0 ) can be written as:

Hence, it follows that the ( d + 1 , n + 1) -th entry of Attn σ P,Q ( Z 0 ) is equal to the product of the ( d +1) -th row of PZ 0 , the mask matrix M , and the ( n +1) -th column of σ ( Z ⊤ 0 QZ 0 ) . Hence, let us write them down explicitly:

Letting b ⊤ be the last row of the matrix P , it holds that the ( d +1) -th row of PZ 0 is equal to [ ⟨ b, z ( i ) ⟩ ] i =1 ,...,n +1 .

The ( n +1) -th column of σ ( Z ⊤ 0 QZ 0 ) is equal to [ σ ( ( z ( i ) ) ⊤ Qz ( n +1) )] i =1 ,...,n +1 . Letting A the first d columns of Q , this vector is equal to [ σ ( ( x ( i ) ) ⊤ Ax ( n +1) )] i =1 ,...,n +1 because the last row of Q is zero and the last row of z ( n +1) is zero (since ( z ( n +1) ) ⊤ = [( x ( n +1) ) ⊤ 0] ).

Thus, the product of [ ⟨ b, z ( i ) ⟩ ] i =1 ,...,n +1 , the mask matrix M , and [ σ ( ( x ( i ) ) ⊤ Ax ( n +1) )] i =1 ,...,n +1 results in the following expression of the attention (writing z, x instead of z ( n +1) , x ( n +1) ):

Since [ Z 0 ] d +1 ,n +1 = 0 , we therefore have

Therefore, it follows that the in-context loss f ( P, Q ) only depends on b and A . Henceforth, let us write f ( b, A ) instead of f ( P, Q ) following Subsection A.1. In particular, writing b ⊤ = [ b ⊤ 0 , b 1 ] for b 0 ∈ R d and b 1 ∈ R , the loss function can be expressed as

2. Simplifying the loss function with symmetry.

Now, we use the fact that both x ( i ) 's and w ⋆ are sampled from the isotropic Gaussian, i.e., N (0 , I d ) in order to further simplify the loss function in (22). In particular, we use the following facts:

( a ) For orthonormal matrices U, V ∈ R d × d , it holds that Ux ( i ) , V x and Uw ⋆ have the same distributions as N (0 , I d ) .

( b ) Moreover, for a diagonal matrix Ξ = diag( ξ i ) ∈ R d × d with the diagonal entries being random signs ξ i ∼ {± 1 } , it holds that Ξ x ( i ) , Ξ x and Ξ w ⋆ have the same distributions as N (0 , I d ) .

Now let us fix a matrix A ∈ R d × d and b ⊤ = [ b ⊤ 0 , b 1 ] for b 0 ∈ R d and b 1 ∈ R . Letting A = U Σ V ⊤ be the SVD of the matrix A , it follows that

where in the third line we use the fact that Ξ ⊤ ΣΞ = ΞΣΞ = Σ ; 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., f lower and then we will connect it back to the original objective.

3. Computation of the lower bound f lower .

Let us now explicitly compute f lower . Let us rewrite the definition of f lower . In fact since, σ = ReLU is homogenous, one can further simplify the lower bound by pushing the constant b 1 inside and write b 1 Σ as Σ . Hence, for two diagonal matrices Σ , D ∈ R d × d , f lower is defined as:

In particular, D is constrained to be the diagonal part of an orthogonal matrix (since D = diag( U ⊤ V ) in the above derivation). Now we focus on characterizing the global minimizers of f lower .

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 ( à la Stein'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

which after a careful expansion becomes

In order to compute (23), we will rely on the aforementioned argument of Erdogdu et al. (2016). In particular, from integration by parts, or Stein's lemma (Erdogdu et al., 2016) (since x ∼ N (0 , I d ) ), we have

We use this to compute all the terms in (23) as follows:

̸

We first apply Stein's lemma to the first term of (23) for i = j . This results in

Using the fact that x ( i ) is a symmetric random variable, one can compute the expectation above as follows: one the one hand, we know E x ( i ) [ σ ′ ( x ⊤ Σ x ( i ) )] = E x ( i ) [ σ ′ ( -x ⊤ Σ x ( i ) )] . On the other hand, we also know that for any scalar α , σ ′ ( -α ) + σ ′ ( α ) = 1 . Therefore, we conclude that E x ( i ) [ σ ′ ( x ⊤ Σ x ( i ) )] = 1 / 2 . Thus, applying this technique twice, we obtain the following

Similarly, we can use Stein's lemma to the second term of (23) to conclude

Lastly, the computation of the first term of (23) for i = j is straightforward. Using the fact that ∀ α ∈ R , σ 2 ( α ) + σ 2 ( -α ) = a 2 , we get

Putting things all together (and ignoring the constant part in (23)), we have

4. Connecting back to the original loss function.

One can in fact write (24) solely in terms of b 1 and A as follows:

Since the latter is a convex function in the matrix b 1 A , it follows that the minimizer corresponds to

In fact the choice b 1 = 1 , b 0 = 0 , and A = -2 n 2( d +2)+( n -1) · I d achieves this, and more crucially, satisfies the property that f lower = f for the corresponding parameters. Therefore, this shows that such choice is a global minimizer.

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 in (25))

where we use the isotropy of w ⋆ and the linearity of trace to get the last equation. Suppose that A ∗ 0 and A ∗ 1 denote the global minimizer of f over symmetric matrices. Since A ∗ 1 is a symmetric matrix, it admits the spectral decomposition A 1 = UD 1 U ⊤ where D 1 is a diagonal matrix and U is an orthogonal matrix. Remarkably, the distribution of X 0 is invariant to a linear transformation by an orthogonal matrix, i.e, X 0 has the same distribution as X 0 U ⊤ . This invariance yields

Thus, we can assume A ∗ 1 is diagonal without loss of generality. To prove A ∗ 2 is also diagonal, we leverage a probabilistic proof technique. Consider the random diagonal matrix S whose diagonal elements are either 1 or -1 with probability 1 2 . Since the input distribution is invariant to orthogonal transformations, we have

Note that we use SD 1 S = D 1 in the last equation, which holds due to D 1 and S are diagonal matrices and S has diagonal elements in { +1 , -1 } . Since f is convex in A 2 , a straightforward application of Jensen's inequality yields

Thus, there are diagonal D 1 and diag( A ∗ 2 ) for which f ( D 1 , diag( A ∗ 2 )) ≤ f ( A ∗ 1 , A ∗ 2 ) holds for an optimal A ∗ 1 and A ∗ 2 . This concludes the proof.

B.2 Proof of Theorem 3

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 i '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 Z 0 be defined as

where y ( n +1) = 〈 w ⋆ , x ( n +1) 〉 . Let Z i denote the output of the ( i -1) th layer of the linear transformer (as defined in (56), initialized at Z 0 ). For the rest of this proof, we will drop the bar, and simply denote Z i by Z i . 5 Let X i ∈ R d × ( n +1) denote the first d rows of Z i and let Y i ∈ R 1 × ( n +1) denote the ( d +1) th row of Z k . Under the sparsity pattern enforced in (8), we verify that, for any i ∈ { 0 , . . . , k } ,

We adopt the shorthand A = { A i } k i =0 . Let S ⊂ R ( k +1) × d × d , and A ∈ S if and only if for all i ∈ { 0 , . . . , k } , there exists scalars a i ∈ R such that A i = a i Σ -1 and B i = b i I . We use f ( A ) to refer to the in-context loss of Theorem 3, that is,

Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:

The theorem statement is equivalent to the following:

where ∇ A i f denotes derivative wrt the Frobenius norm ∥ A i ∥ F . Towards this end, we establish the following intermediate result: if A ∈ S , then for any R ∈ R ( k +1) × d × d , there exists ˜ R ∈ S , such that, at t = 0 ,

In fact, we show that ˜ R i := r i I , for r i = 1 d Tr ( Σ 1 / 2 R i Σ 1 / 2 ) . This implies (27) via the following simple argument: Consider the " S -constrained gradient flow": let A ( t ) : R + → R ( k +1) × d × d be defined as

for i = 0 , . . . , k . By (28), we verify that

We verify from its definition that f ( A ) ≥ 0 ; if the infimum in (27) fails to be zero, then inequality (29) will ensure unbounded descent as t →∞ , contradicting the fact that f ( A ) is lower-bounded. This concludes the proof.

Proof outline. The remainder of the proof will be devoted to showing (28), which we outline as follows:

In Step 1, we reduce the condition in (29) to a more easily verified layer-wise condition. Specifically, we only need to verify (29) when R i are all zero except for R j for some fixed j (see (30)) 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 (31), we study the evolution of f ( A ( t )) when A ( t ) moves in the direction of R , as X 0 is (roughly speaking) randomly transformed.

In Step 2, we use the results of Step 2 to to study G (see (32)) and d dt G ( A ( t )) (see (33)) under random transformation of X 0 . The idea in (33) is that "randomly transforming X 0 " has the same effect as "randomly transforming S " (recall S is the perturbation to B ).

In Step 3, we apply the result from Step 2 to the expression of d dt f ( A ( t )) in (31). We verify that ˜ R in (28) is exactly the expected matrix after "randomly transforming S ". This concludes our proof.

1. Reduction to layer-wise condition. To prove (28), it suffices to show the following simpler condition: Let j ∈ { 0 , . . . , k } . Let R j ∈ R d × d be arbitrary matrices. For C ∈ R d × d , let A ( tC, j ) denote

̸

the collection of matrices, where [ A ( tC, j )] j = A j + tC , and for i = j , A ( tC, j ) i = A i . We show that for all j ∈ { 0 , . . . , k } , R j ∈ R d × d , there exists ˜ R j = r j Σ -1 , such that, at t = 0 ,

We can verify that (28) is equivalent to (30) by noticing that for any R , at t = 0 , d dt f ( A + tR ) = ∑ k j =0 d dt f ( A ( tR j , j )) . We will now work towards proving (30) for some index j that is arbitrarily chosen but fixed throughout.

Let us define, for any C ∈ R d × d , G ( X,A j + C ) := X ∏ k i =0 ( I -MX ⊤ [ A ( C, j )] i X ) . By (25) and (26),

The second equality follows from plugging in (25). For the rest of this proof, let U denote a uniformly randomly sampled orthogonal matrix. Let U Σ := Σ 1 / 2 U Σ -1 / 2 . Using the fact that X 0 d = U Σ X 0 , we can verify

2. G and d dt G under random transformation of X 0 . We will now verify that G ( U Σ X 0 , A j ) = U Σ G ( X 0 , A j ) :

where we use the fact that U ⊤ Σ A i U Σ = U ⊤ Σ ( a i Σ -1 ) U Σ = A i . Next, we verify that

where the first equality again uses the fact that U ⊤ Σ A i U Σ = A i .

Putting everything together. Let us continue from (31). Plugging (32) and (33) into (31),

where r j := 1 d Tr ( Σ 1 / 2 R j Σ 1 / 2 ) . In the above, ( i ) uses 1. (32) and (33), as well as the fact that U ⊤ Σ Σ -1 U Σ = Σ -1 . ( ii ) uses the fact that d dt G ( X 0 , A j + tC ) ∣ ∣ t =0 is affine in C . To see this, one can verify from the definition of G , e.g. using similar algebra as (33), that d dt G ( X 0 , A j + C ) is affine in C . Thus E U [ G ( X 0 , A j + tU ⊤ Σ R j U Σ ) ] = G ( X 0 , A j + t E U [ U ⊤ Σ R j U Σ ) ] .

B.3 Proof of Theorem 4

The proof of Theorem 4 is similar to that of Theorem 3, and with a similar setup. However to keep the proof self-contained, we will restate the setup. Once again, we 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 i '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 Z 0 be defined as

where y ( n +1) = 〈 w ⋆ , x ( n +1) 〉 . Let Z i denote the output of the ( i -1) th layer of the linear transformer (as defined in (56), initialized at Z 0 ). For the rest of this proof, we will drop the bar, and simply denote Z i by Z i . 6 Let X i ∈ R d × n +1 denote the first d rows of Z i and let Y i ∈ R 1 × n +1 denote the ( d +1) th row of Z k . Under the sparsity pattern enforced in (11), we verify that, for any i ∈ { 0 , . . . , k } ,

We adopt the shorthand A = { A i } k i =0 and B = { B i } k i =0 . Let S ⊂ R 2 × ( k +1) × d × d , and ( A,B ) ∈ S if and only if for all i ∈ { 0 , . . . , k } , there exists scalars a i , b i ∈ R such that A i = a i Σ -1 and B i = b i I . Throughout this proof, we will work with the following formulation of the in-context loss from Lemma 5:

(note that the only randomness in Z 0 comes from X 0 as Y 0 is a deterministic function of X 0 ). The theorem statement is equivalent to the following:

where ∇ A i f denotes derivative wrt the Frobenius norm ∥ A i ∥ F .

Our goal is to show that, if ( A,B ) ∈ S , then for any ( R,S ) ∈ R 2 × ( k +1) × d × d , there exists ( ˜ R, ˜ S ) ∈ S , such that, at t = 0 ,

In fact, we show that ˜ R i := r i I , for r i = 1 d Tr ( Σ 1 / 2 R i Σ 1 / 2 ) and ˜ S i = s i I , for s i = 1 d Tr ( Σ -1 / 2 S i Σ 1 / 2 ) . This implies (36) via the following simple argument: Consider the " S -constrained gradient flow": let A ( t ) : R + → R ( k +1) × d × d and B ( t ) : R + → R ( k +1) × d × d be defined as

for i = 0 , . . . , k . By (37), we verify that

We verify from its definition that f ( A,B ) ≥ 0 ; if (36) does not hold then (38) will ensure unbounded descent as t →∞ , contradicting the fact that f ( A,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 i , S i are all zero except for R j for some fixed j (see (40)), or (II) when R i , S i are all zero except for S j for some fixed j (see (39)). We focus on 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), we study the evolution of f ( A,B ( t )) when B ( t ) moves in the direction of S , as X 0 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 X 0 is randomly transformed. There are two main results here: First we provide the expression for X i in (42). Second, we provide the expression for d dt X i ( B ( t )) in (43).

In Step 3, we use the results of Step 2 to to study G (see (47)) and d dt G ( B ( t )) (see (48)) under random transformation of X .

0 The idea in (48) is that "randomly transforming X 0 " 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 d dt f ( A,B ( t )) in (41). We verify that ˜ S 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 to layer-wise condition. To prove (37), it suffices to show the following simpler condition: Let j ∈ { 0 , . . . , k } . Let R j , S j ∈ R d × d be arbitrary matrices. For C ∈ R d × d , let A ( tC, j ) denote the collection of matrices, where A ( tC, j ) j = A j + tC , and for i = j , A ( tC, j ) i = A i . Define B ( tC, j ) analogously. We show that for all j ∈ { 0 , . . . , k } and all R j , S j ∈ R d × d , there exists ˜ R j = r j Σ -1 and ˜ S j = s j Σ -1 , such that, at t = 0 ,

We can verify that (37) is equivalent to (39)+(40) by noticing that for any ( R,S ) ∈ R 2 × ( k +1) × d × d , at t = 0 , d dt f ( A + tR, B + tS ) = ∑ k j =0 ( d dt f ( A ( tR j , j ) , B ) + d dt f ( A,B ( tS j , j )) ) .

̸

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 i and Y i in (34) are in fact functions of A,B and X 0 . For most of our subsequent discussion, A i (for all i ) and B i (for all i = j ) can be treated as constant matrices. We will however make the dependence on X 0 and B j explicit (as we consider the curve B j + tS ), i.e. we use X i ( X,C ) (resp Y i ( X,C ) ) to denote the value of X i (resp Y i ) from (34), with X 0 = X , and B j = C .

By (35) and (34),

where G ( X,C ) := X ∏ k i =0 ( I -MX i ( X,C ) T A i X i ( X,C ) ) . The second equality follows from plugging in (34).

For the rest of this proof, let U denote a uniformly randomly sampled orthogonal matrix. Let U Σ := Σ 1 / 2 U Σ -1 / 2 . Using the fact that X 0 d = U Σ X 0 , we can verify

X i and d dt X i under random transformation of X 0 . In this step, we prove that when X 0 is transformed by U Σ , X i for i ≥ 1 are likewise transformed in a simple manner. The first goal of this step is to show

We will prove this by induction. When i = 0 , this clearly holds by definition. Suppose that (42) holds for some i . Then

where the second equality uses the inductive hypothesis, and the fact that A i = a i Σ -1 , so that U T Σ A i U Σ = A i , and the fact that B i = b i I , from the definition of S and our assumption that ( A,B ) ∈ S . This concludes the proof of (42).

We now present the second main result of this step. Let U -1 Σ := Σ 1 / 2 U T Σ -1 / 2 , so that it satisfies U Σ U -1 Σ = U -1 Σ U Σ = I . For all i ,

To reduce notation, we will not write ·| t =0 explicitly in the subsequent proof. We first write down the dynamics for the right-hand-side term of (43): From (34), for any ℓ ≤ j , and for any i ≥ j +1 , and

d × d

We are now ready to prove (43) using induction. For the base case, we verify that for ℓ ≤ j , U -1 Σ d dt X ℓ ( U Σ X 0 , B k + tS j ) = 0 = d dt X ℓ ( X 0 , B j + tU -1 Σ S j U Σ ) (see first equation in (44)). For index j +1 , we verify that

where we use two facts: 1. X i ( U Σ X 0 , B j ) = U Σ X i ( X 0 , B j ) from (42), 2. A i = a i Σ -1 , so that U ⊤ Σ A i U Σ = A i . We verify by comparison to the second equation in (44) that U -1 Σ d dt X j ( U Σ X 0 , B j + tS j ) = 0 = d dt X j ( X 0 , B j + tU -1 Σ S j U Σ ) . These conclude the proof of the base case.

Now suppose that (43) holds for some i . We will now prove (43) holds for i +1 . From (34),

In ( i ) above, we crucially use the following facts: 1. B i = b i I so that U -1 Σ B i = B i U -1 Σ , 2. X i ( U Σ X 0 , B j ) = U Σ X i ( X 0 , B j ) from (42), 3. A i = a i Σ -1 , so that U ⊤ Σ A i U Σ = A i , 4. U Σ U -1 Σ = U -1 Σ U Σ = I . ( ii ) follows from our inductive hypothesis. The inductive proof is complete by verifying that (46) exactly matches the third equation of (44) when C = U -1 Σ SU Σ .

d

where the second equality uses (42), as well as the fact that U ⊤ Σ A i U Σ = A i . Next, we will show that

To see this, we can expand

In ( i ) above, we the following facts: 1. X i ( U Σ X 0 , B j ) = U Σ X i ( X 0 , B j ) from (42), 2. A i = a i Σ -1 , so that U ⊤ Σ A i U Σ = A i , 3. U Σ U -1 Σ = U -1 Σ U Σ = I . ( ii ) 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):

where s j := 1 d Tr ( Σ -1 / 2 S j Σ 1 / 2 ) . In the above, ( i ) uses 1. (47) and (48), as well as the fact that U ⊤ Σ Σ -1 U Σ = Σ -1 . ( ii ) uses the fact that d dt G ( X 0 , B j + tC ) ∣ ∣ t =0 is affine in C . To see this, one can verify from (44), using a simple induction argument, that d dt X i ( X 0 , B j + 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 dt G ( X 0 , B j + C ) is affine in d dt X i ( X 0 , B j + tC ) . Thus E U [ G ( X 0 , B j + tU -1 Σ S j U Σ ) ] = G ( X 0 , B j + t E U [ U -1 Σ S j U Σ ) ] .

With this, we conclude our proof of (40).

Proof of (39) . We will now prove (39) for fixed but arbitrary j , i.e. there is some r j such that

The proof is very similar to the proof of (40) that we just saw, and we will essentially repeat the same steps from Step 2-4 above.

Since we now consider perturbations to A instead of to B , we will need to redefine some notation: let X i ( X,C ) (resp Y i ( X,C ) ) to denote the value of X i (resp Y i ) from (34), with X 0 = X , and A j = C (previously it was with B j = C ). Let G ( X,A j + C ) := X ∏ i i =0 ( I + M ( X i ( X,A j + C ) T A ( C, j ) i X i ( X,A j + C ) )) , where recall that A ( C, j ) := A j + C , and A ( C, j ) ℓ := A ℓ for all ℓ ∈ { 0 ...k } \ { j } .

We first verify that

The proofs are identical to the proofs of (42) and (47) so we omit them. Next, we show that for all i ,

We establish the dynamics for the right-hand-side of (50):

Similar to (45), we show that for i ≤ j ,

and

Finally, for the inductive step, we follow identical steps leading up to (46) to show that

The inductive proof is complete by verifying that (52) exactly matches the third equation of (51) when C = U -1 Σ SU Σ . This concludes the proof of (50).

Next, we study the time derivative of G ( U Σ X 0 , A j + tR j ) and show that

This proof differs significantly from that of (48) in a few places, so we provide the whole derivation below. By chain-rule, we can write

where and

We will separately simplify ♠ and ♡ , and verify at the end that summing them recovers the righthand-side of (53). We begin with ♠ , and the steps are almost identical to the proof of (48).

In ( i ) above, we the following facts: 1. X i ( U Σ X 0 , B j ) = U Σ X i ( X 0 , B j ) from (49), 2. A i = a i Σ -1 , so that U ⊤ Σ A i U Σ = A i , 3. U Σ U -1 Σ = U -1 Σ U Σ = I . ( ii ) follows from (50). We will now simplify ♡ .

where ( i ) uses the fact that X i ( U Σ X 0 , B j ) = U Σ X i ( X 0 , B j ) from (49) and the fact that A i = a i Σ -1 .

By expanding d dt G ( X 0 , A j + tU ⊤ Σ R j U Σ ) , we verify that

this concludes the proof of (53).

The remainder of the proof is similar to what was done in (41) in Step 4:

where r j := 1 d Tr ( Σ 1 / 2 R j Σ 1 / 2 ) . In the above, ( i ) uses 1. (49) and (53), as well as the fact that U ⊤ Σ Σ -1 U Σ = Σ -1 . ( ii ) uses the fact that d dt G ( X 0 , A j + tC ) ∣ ∣ t =0 is affine in C . To see this, one can verify using a simple induction argument, that d dt X i ( X 0 , A j + 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 (53), that d dt G ( X 0 , A j + C ) is affine in d dt X i ( X 0 , A j + tC ) and C . Thus E U [ G ( X 0 , A j + tU ⊤ Σ R j U Σ ) ] = G ( X 0 , A j + t E U [ U ⊤ Σ R j U Σ ] ) .

This concludes the proof of (39), and hence of the whole theorem.

B.4 Equivalence under permutation

Lemma 4. Consider the same setup as Theorem 3. Let A = { A i } k i =0 , with A i = a i Σ -1 . Let

Let i, j ∈ { 0 , . . . , k } be any two arbitrary indices, and let ˜ A i = A j , ˜ A j = A i , and let ˜ A ℓ = A ℓ for all ℓ ∈ { 0 , . . . , k } \ { i, j } . 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

where G ( X 0 , A ) := X 0 ∏ k ℓ =0 ( I + MX T 0 A ℓ X 0 ) .

Consider any fixed index ℓ . We will show that

The lemma can then be proven by repeatedly applying the above, so that indices of A i and A j are swapped.

To prove the above equality,

This concludes the proof. Notice that we crucially used the fact that A ℓ and A ℓ +1 are the same matrix up to scaling.

C Auxiliary Lemmas

C.1 Proof of Lemma 1 (Equivalence to Preconditioned Gradient Descent)

Consider fixed samples x (1) , . . . , x ( n ) , and fixed w ⋆ . Let P = { P i } k i =0 , Q = { Q i } k i =0 denote fixed weights. Let Z i evolve as described in (4). Let X i denote the first d rows of Z k (under (8), X i = X 0 for all I ) and let Y i denote the ( d +1) th row of Z i . Let g ( x, y, k ) : R d × R × Z → R be a function defined as follows: let x n +1 = x and let y n +1 0 = y , then g ( x, y, k ) := y n +1 k . Note that y n +1 k = [ Y k ] n +1 .

We verify that, under (8), the formula for updating y ( n +1) k is given by

where M is a mask given by [ I 0 0 0 ] . We can verify the following facts

1. g ( x, y, k ) = g ( x, 0 , k ) + y . To see this, notice first that for all i ∈ { 1 , . . . , n } ,

In other words, y ( i ) k does not depend on y ( n +1) t for any t . Next, for y ( n +1) k itself,

which depends on y n +1 k only additively. We can verify under a simple induction that g ( x, y, k +1) -y = g ( x, y, k ) -y .

̸

g ( x, 0 , k ) is linear in x . To see this, notice first that for j = n +1 , y ( j ) k is does not depend on x ( n +1) t for all t, j, k . Consequently, the update formula for y ( n +1) k +1 depends only linearly on x ( n +1) and y ( n +1) k . Finally, y ( n +1) 0 = 0 is linear in x , so the conclusion follows by induction.

With these two facts in mind, we verify that for each k , there exists a θ k ∈ R d , such that

for all x, y . It follows from definition that g ( x, y, 0) = y , so that ⟨ θ 0 , x ⟩ = g ( x, y, 0) -y = 0 , so that θ 0 = 0 .

We now turn our attention to the third crucial fact: for all i ,

To see this, suppose that we let x ( n +1) := x ( i ) for some i ∈ 1 , . . . , n . Then

thus y ( i ) k +1 = y ( n +1) k +1 if y ( i ) k = y ( n +1) k , and the induction proof is completed by noting that y ( i ) 0 = y ( n +1) 0 by definition. Let ¯ X ∈ R d × n be the matrix whose columns are x (1) , . . . , x ( n ) , leaving out x ( n +1) . Let ¯ Y k ∈ R 1 × n denote the vector of y (1) k , . . . , y ( n ) k . Then it follows that

Using the above fact, the update formula for y ( n +1) k can be written as

Since the choice of x ( n +1) is arbitrary, we get the more general update formula

We can treat A k as a preconditioner. Let f ( θ ) :== 1 2 n ( θ + w ⋆ ) T ¯ X ¯ X T ( θ + w ⋆ ) , then

Finally, let w gd k := -θ k . We verify that f ( -w ) = R w ⋆ ( w ) , so that

We also verify that for any x ( n +1) , the prediction of y ( n +1) k is

This concludes the proof.

C.2 Reformulating the in-context loss

In this section, we will develop a re-formulation in-context loss, defined in (5), in a more convenient form (see Lemma 5).

For the entirety of this section, we assume that the transformer parameters { P i , Q i } k i =0 are of the form defined in (11), which we reproduce below for ease of reference:

Recall the update dynamics in (4), which we reproduce below:

where M is a mask matrix given by M := [ I n × n 0 0 0 ] . Let X k ∈ R d × n +1 denote the first d rows of Z k and let Y k ∈ R 1 × n +1 denote the ( d + 1) th (last) row of Z k . Then the dynamics in (56) is equivalent to

We present below an equivalent form for the in-context loss from (5):

Lemma 5. Let p x and p w denote distributions over R d . Let x (1) , . . . , x ( n +1) iid ∼ p x and w ⋆ ∼ p w . Let Z 0 ∈ R d +1 × n +1 be as defined in (1) :

Let Z k denote the output of the ( k -1) th layer of the linear transformer (as defined in (56) , initialized at Z 0 ). Let f ( { P i , Q i } k i =0 ) denote the in-context loss defined in (5) , i.e.

Let Z 0 be defined as

where y ( n +1) = 〈 w ⋆ , x ( n +1) 〉 . Let Z k denote the output of the ( k -1) th layer of the linear transformer (as defined in (56) , initialized at Z 0 ). Assume { P i , Q i } k i =0 be of the form in (11) . Then the loss in (5) has the equivalent form

Before proving Lemma 5, we first establish an intermediate result (Lemma 6 below). To facilitate discussion, let us define a function F X ( { A i , B i } k i =0 , X 0 , Y 0 ) and F Y ( { A i , B i } k i =0 , X 0 , Y 0 ) to be the outputs, after k layers of linear transformers respectively. I.e.

as defined in (57), given initialization X 0 , Y 0 .

We now prove a useful lemma showing that [ Y 0 ] n +1 = y ( n +1) influences X i , Y i in a very simple manner:

Lemma 6. Let X i , Y i follow the dynamics in (57) . Then

[ X i ] is are independent of [ Y 0 ] n +1 .

̸

For j = n +1 , [ Y i ] j is independent of [ Y 0 ] n +1 .

[ Y i ] n +1 depends additively on [ Y 0 ] n +1 .

In other words, for C := [0 , 0 , 0 , . . . , , 0 , c ] ∈ R 1 × ( n +1)

Proof of Lemma 6. The first and second items follows directly from observing that the dynamics for X i and Y i in (57) do not involve [ Y i ] n +1 , due to the effect of M .

The third item again uses the fact that [ Y i +1 -Y i ] n +1 does not depend on [ Y i ] n +1 .

We are now ready to prove Lemma 5

Proof of Lemma 5. Let Z 0 , Z k , Z 0 , Z k be as defined in the lemma statement. Let X k and Y k denote first d rows and last row of Z k . Then by Lemma 6, X k +1 = X k +1 and Y k +1 = Y k +1 + [ 0 0 · · · 0 〈 w ⋆ , x ( n +1) 〉] . Therefore, (58) is equivalent to

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 A 0 , A 1 , A 2 for the setting of Theorem 4. 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, it appears that the transformer implements a tiny gradient step using X 0 (as A 0 is small), and a large gradient step using X 2 (as A 2 is large). We believe that this is due to X 2 being better-conditioned than X 1 , due to the effects of B 0 , B 1 .

0

1

2

3

4

0

1

2

3

4

-0.38

0.00

0.01

-0.01

-0.00

-0.00

-0.38

0.01

-0.00

0.00

-0.01

-0.01

-0.38

-0.01

-0.00

0.01

0.00

0.01

-0.38

0.01

0.00

-0.00

0.00

-0.01

-0.38

0.35

0.30

0.25

0.20

0.15

0.10

0.05

0.00

0

1

2

3

4

0

1

2

3

4

-2.40

-0.00

0.00

-0.00

-0.00

0.00

-2.40

0.01

0.01

0.00

0.00

0.00

-2.39

0.01

0.00

-0.00

0.01

0.01

-2.38

0.01

0.00

0.00

0.00

0.01

-2.42

2.0

1.5

1.0

0.5

0.0

0

1

2

3

4

0

1

2

3

4

-9.59

-0.01

-0.01

-0.03

0.02

-0.01

-9.54

0.01

0.05

0.03

-0.00

0.01

-9.55

0.05

0.01

-0.03

0.06

0.04

-9.46

0.02

0.01

0.03

0.02

0.02

-9.65

8

6

4

2

0

(a) Visualization of Σ 1 / 2 A 0 Σ 1 / 2 (b) Visualization of Σ 1 / 2 A 1 Σ 1 / 2 (c) Visualization of Σ 1 / 2 A 2 Σ 1 / 2

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 ICL 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 is presented in Figure 6.

2

6

10

14

18

22

Number of ICL Examples

0

1

2

3

4

Loss

3-Step GD

3-Step Preconditioned GD

3-Layer Linear Transformer

OLS

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.

1

2

3

4

Number of Layers/Steps

8

6

4

2

0

log(Loss)

Preconditioned GD

Linear Transformer

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nGenerally, 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)\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nIt would be important to highlight limitations of the analysis. E.g. same covariance matrix, singe-head linear attention...", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n1. 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)? ...\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLLMs 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.\n\nSTRENGTHS:\nThis 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\n\nWEAKNESSES:\nI 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\nQUESTIONS:\nPlease see the weakness section.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. (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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}