Title: In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization
Venue: NeurIPS 2024
Accepted: yes
Keywords: In-Context Learning

Transformers

Approximation Theory

Optimization
Abstract: We study the \emph{in-context learning} (ICL) ability of a \emph{Linear Transformer Block} (LTB) that combines a linear attention component and a linear multi-layer perceptron (MLP) component. 
For ICL of linear regression with a Gaussian prior and a \emph{non-zero mean}, we show that LTB can achieve nearly Bayes optimal ICL risk. In contrast, using only linear attention must incur an irreducible additive  approximation error. 
Furthermore, we establish a correspondence between LTB and one-step gradient descent estimators with learnable initialization ($\mathsf{GD}-\beta$), in the sense that every $\mathsf{GD}-\beta$ estimator can be implemented by an LTB estimator and every optimal LTB estimator that minimizes the in-class ICL risk is effectively a $\mathsf{GD}-\beta$ estimator.
Finally, we show that $\mathsf{GD}-\beta$ estimators can be efficiently optimized with gradient flow, despite a non-convex training objective.
Our results reveal that LTB achieves ICL by implementing $\mathsf{GD}-\beta$, and they highlight the role of MLP layers in reducing approximation error.

In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

Abstract

1 Introduction

The recent dramatic progress in natural language processing can be attributed in large part to the development of large language models based on Transformers [34], such as BERT [11], LLaMA [32], PaLM [8], and the GPT series [28, 29, 6, 25]. In some pioneering pre-trained large language models, a new learning paradigm known as in-context learning (ICL) was observed (see, for example, [6, 7, 24, 20]). ICL refers to the capability of a pre-trained model to solve a new task based on a few in-context demonstrations without updating the model parameters.

A recent line of work quantifies ICL in tractable statistical learning setups such as linear regression with a Gaussian prior (see [14, 3] and references thereafter). Specifically, an ICL model takes a sequence ( x 1 , y 1 , ..., x M , y M , x query ) as input and outputs a prediction of y query , where ( x i , y i ) i M =1 and ( x query , y query ) are independent samples from an unknown task-specific distribution (where the task admits a prior distribution). The ICL model is pre-trained by fitting many empirical observations of such sequence-label pairs. Experiments show that Transformers can achieve an ICL risk close to that achieved by Bayes optimal estimators in many statistical learning tasks [14, 3].

For ICL of linear regression with a Gaussian prior, the data is generated as

where ˜ β is a task parameter that satisfies a Gaussian prior, ˜ β ∼ N ( 0 , ψ 2 I ) . In this setup, [14, 3] showed in experiments that Transformers can achieve nearly optimal ICL by matching the performance of the Bayes optimal estimator, that is, an optimally tuned ridge regression. Besides, [35] showed by construction that a single linear self-attention (LSA) can implement one-step gradient descent with zero initialization ( GD -0 ), offering an insight into the ICL mechanism of the Transformer.

Later, theoretical works showed that optimal LSA models effectively correspond to GD -0 models [1], trained LSA models converge to GD -0 models [38], and GD -0 models (hence LSA models) can provably achieve nearly Bayes optimal ICL in certain regimes [37].

̸

Our contributions. In this paper, we consider ICL in a more general setup, that is, linear regression with a Gaussian prior and a non-zero mean (that is, E [ ˜ β ] = β ∗ = 0 ). The non-zero mean in task prior captures a common scenario where tasks share a signal. In this setting, we show that the LSA models considered in prior papers [1, 38, 37] incur an irreducible additive approximation error. Furthermore, we show this approximation error is mitigated by considering a linear Transformer block (LTB) that combines a linear multi-layer perceptron (MLP) component and an LSA component. Our results highlight the important role of MLP layers in Transformers in reducing the approximation error, and they suggest that theories about LSA [1, 38, 37] do not fully explain the power of Transformers. Motivated by this understanding, we investigate LTB in depth and obtain the following additional results:

We show that LTB can implement one-step gradient descent with learnable initialization , referred to by us as GD -β . Additionally, we show that every optimal LTB estimator that minimizes the in-class ICL risk is effectively a GD -β estimator.

Moreover, we show that the optimal GD -β , hence also the optimal LTB, nearly matches the performance of the Bayes optimal estimator for linear regression with a Gaussian prior and a non-zero mean, provided that the signal-to-noise ratio is upper bounded. These two results together suggest that LTB performs nearly optimal ICL by implementing GD -β .

Finally, we show that the GD -β estimator can be efficiently optimized by gradient descent with an infinitesimal stepsize (that is, gradient flow) under the population ICL risk, despite the nonconvexity of the objective.

Paper organization. The remaining paper is organized as follows. We conclude this section by introducing a set of notations. We then discuss related papers in Section 2. In Section 3, we set up our ICL problems and define LTB and LSA models mathematically. In Section 4, we show a positive approximation error gap between LSA and LTB, which highlights the importance of the MLP component and motivates our subsequent efforts to study LTB. In Section 5, we connect LSA estimators to GD -β estimators that are more interpretable for ICL of linear regression. In Section 6, we study the in-context learning and training of GD -β estimators. We conclude our paper and discuss future directions in Section 7.

Notations. We use lowercase bold letters to denote vectors and uppercase bold letters to denote matrices and tensors. For a vector x and a positive semi-definite (PSD) matrix A , we write ∥ x ∥ A := √ x ⊤ A x . We write ⟨· , ·⟩ for the inner product, which is defined as ⟨ x , y ⟩ := x ⊤ y for vectors and ⟨ A , B ⟩ := tr( AB ⊤ ) for matrices. We write A [ i ] as the i -th row of the matrix A , A m,n as the ( m,n ) -th entry, and A -1 , -1 as the right-bottom entry. We write 0 n , 0 m × n , I n for the zero vector, zero matrix and identity matrix, respectively. We denote the Kronecker product as ⊗ . For two matrices A and B , A ⊗ B is a linear mapping which operates as ( A ⊗ B ) ◦ C = BCA ⊤ . For a positive semi-definite matrix A , we write A 1 2 as its principle square root, which is the unique positive semi-definite matrix B such that BB = BB ⊤ = A . We also write A -1 2 = ( A 1 2 ) + , where ( · ) + denotes the Moore Penrose pseudo-inverse. For two sets A,B we write A + B for the Minkowski sum, which is defined as { a + b : a ∈ A,b ∈ B } . We also write a + B = { a } + B for an element a and a set B. We write null ( · ) for the null set of a matrix or a tensor.

2 Related Works

Empirical results for ICL in controlled settings. The work by [14] first considered ICL in controlled statistical learning setups. For noiseless linear regression, [14] showed in experiments that Transformers match the performance of the optimal estimator (that is, Ordinary Least Square ). Subsequent works by [3, 18] extended their result to noisy linear regression with a Gaussian prior and showed that Transformers can match the performance of the Bayesian optimal estimator (that is, an optimally tuned ridge regression). Besides, [30] showed that the above holds even when Transformers are pretrained on a limited number of linear regression tasks. These papers only considered a Gaussian

prior with a zero mean. In contrast, we consider a Gaussian prior with a non-zero mean. Our setup better captures the common scenarios where the tasks share a signal.

The empirical investigation of ICL in tractable statistical learning setups goes beyond linear regression settings. For more examples, researchers empirically studied the ICL ability of Transformers for decision trees (see [14], they also considered two-layer networks), algorithm selection [4], linear mixture models [26], learning Fourier series [2], discrete boolean functions [5], representation learning [13, 16], and reinforcement learning [17, 19]. Among all these settings, Transformers can either compete with the Bayes optimal estimators or expert-designed strong benchmarks. These works are not directly comparable with our paper.

Transformer implements gradient descent. A line of work interpreted the ICL of Transformers by their abilities to implement gradient descent (GD) [35, 3, 10, 1, 38, 4, 37]. In experiments, [35, 10] showed that (multi-layer) Transformer outputs are close to (multi-step) GD outputs. When specialized to linear regression tasks, [35] constructed a single linear self-attention (LSA) that implements onestep gradient descent with zero initialization ( GD -0 ). Subsequently, [1] showed that optimal LSA models effectively correspond to GD -0 models, [38] proved that trained LSA models converge to GD -0 models, [37] showed that GD -0 models (hence LSA models) can provably achieve nearly Bayes optimal ICL in certain regimes and provided a sharp task complexity analysis of the pre-training. These papers focused on LSA models. Instead, we consider a linear Transformer block that also utilizes the MLP layer.

From an approximation theory perspective, [3, 4] showed that Transformers can implement multi-step GD under general losses. In comparison, we consider a limited setting of linear regression and show the LSA models can implement one-step GD with learnable initialization ( GD -β ), moreover, every optimal LTB model is effectively an GD -β model. Both of our constructions utilize MLP layers in Transformers, highlighting its importance in reducing approximation error. Different from their results, we also show a negative approximation result that reveals the limitation of LSA models.

3 Preliminaries

Model input. We use x ∈ R d and y ∈ R to denote a feature vector and its label, respectively. Throughout the paper, we assume a fixed number of context examples, denoted by M > 0 . We denote the context examples by ( X , y ) ∈ R M × d × R M , where each row represents a context example, denoted by ( x ⊤ i , y i ) , i = 1 , . . . , M . To formalize an ICL problem, the input of a model is a token matrix given by [1, 38]

The output of a model corresponds to a prediction of y .

A Transformer block. Modern large language models are often made by stacking basic Transformer blocks (see, for example, [34]). A basic Transformer block consists of a self-attention layer and a multi-layer perceptron (MLP) layer [34], E ↦→ MLP [ ATTN ( E )] . Here, the MLP layer is defined as MLP ( E ) := W ⊤ 2 ReLU ( W 1 E ) , where ReLU ( · ) refers to the entrywise rectified linear unit (ReLU) activation function, and W 1 , W 2 ∈ R d f × ( d +1) are two weight matrices. The self-attention layer (we focus on the single-head version in this paper) is defined as ATTN ( E ) := E + W ⊤ P W V EM · sfmx ( ( W K E ) ⊤ W Q E ) , where sfmx ( · ) refers to the row-wise softmax operator, W K , W Q ∈ R d k × ( d +1) are the key and query matrices, respectively, W P , W V ∈ R d v × ( d +1) are the projection and value matrices, respectively, and M is a fixed masking matrix given by

This mask matrix is included to reflect the asymmetric structure of a prompt since the label of query input x is not included in the token matrix [1, 21]. In the above formulation, d k , d v , d f are three hyperparameters controlling the key size, value size, and width of the MLP layer, respectively. In the single-head case, it is common to set d k = d v = d +1 and d f = 4( d +1) , where d +1 corresponds to the embedding size (see, for example, [34]). Our formulation of a basic transformer block ignores all bias parameters and some popular techniques (such as layer normalization and dropout) to focus solely on the benefits brought by the model structure.

A linear Transformer block. To facilitate theoretical analysis, we ignore the non-linearities in the Transformer block (specifically, sfmx and ReLU ) and work with a linear Transformer block (LTB) defined as

where E is the token matrix given by (3.1), [ · ] -1 , -1 refers to the bottom right entry of a matrix, M is the fixed masking matrix given by (3.2). Here, the trainable parameters are W P , W V ∈ R d v × ( d +1) , W K , W Q ∈ R d k × ( d +1) , and W 1 , W 2 ∈ R d f × ( d +1) . We use the bottom right entry of the transformed token matrix as the model output to form a prediction of the label y . The 1 /M factor is a normalization factor in linear attention and can be absorbed into trainable parameters. Finally, we denote the hypothesis class formed by LTB models as F LTB := { f LTB : W K , W Q , W V , W P , W 1 , W 2 , d k ≥ d, d v ≥ d +1 , d f ≥ 1 } , where f LTB is defined in (3.3).

A linear self-attention. We will also consider a linear self-attention (LSA) defined as

where W K , W Q , W P , W V are trainable parameters. An LSA model can be viewed as an LTB model without the MLP layer (setting d f = d +1 and W 1 = W 2 = I ). We remark that, unlike LTB, the residual connection in LSA plays no role because the bottom right entry of the prompt E is zero (see (3.1)). A variant of the LSA model has been studied by [1, 38, 37], where W ⊤ K W Q and W ⊤ P W V are respectively merged into one matrix parameter. Similarly, we denote the hypothesis class formed by LSA models as F LSA := { f LSA : W K , W Q , W V , W P , d k ≥ d, d v ≥ 1 } , where f LSA is defined in (3.4).

Linear regression tasks with a shared signal. Assume that data and context examples are generated as follows.

Assumption 3.1 (Distributional conditions) . Assume that ( X , y , x , y ) are generated by:

First, a task parameter is independently generated by ˜ β ∼ N ( β ∗ , Ψ ) .

The feature vectors are independently generated by x , x 1 , . . . x M i . i . d . ∼ N (0 , H ) .

Then, the labels are generated by y = ⟨ ˜ β , x ⟩ + ε, y i = ⟨ ˜ β , x i ⟩ + ε i , i = 1 , . . . , M, where ε and ε i 's are independently generated by ε, ε 1 , . . . , ε M i . i . d . ∼ N ( 0 , σ 2 ) .

Here, σ 2 ≥ 0 , H ⪰ 0 , Ψ ⪰ 0 , and β ∗ ∈ R d are fixed but unknown quantities that govern the data distribution. We denote ε = ( ε 1 , ..., ε M ) ⊤ .

We emphasize the importance of the mean of the task parameter β ∗ in Assumption 3.1. A non-zero β ∗ represents a shared signal across tasks, which is arguably common in practice. This assumption is implicitly used in [12] where they assumed task parameters are close to a meta parameter. In comparison, the prior works for ICL of linear regression [1, 38, 37] only considered a special case where β ∗ = 0 . In this special case, they showed that a single LSA layer can achieve nearly optimal ICL by approximating one-step gradient descent from zero initialization ( GD -0 ). In more general cases where β ∗ is non-zero, we will show in Section 4 that LSA is insufficient to learn the shared signal and must incur an irreducible approximation error compared to LTB models. This sets our results apart from the prior papers.

ICL risk. We measure the ICL risk of a model f by the mean squared error,

where E is defined in (3.1) and the expectation is over E (equivalent to over X , y , and x ) and y .

4 Benefits of the MLP Component

Our first main result separates the approximation abilities of the LTB and LSA models for ICL. Recall that an LSA model can be viewed as a special case of the LTB model with W 2 = W 1 = I . So the

best ICL risk achieved by LTB models is no larger than the best ICL risk achieved by LSA models. However, our next theorem shows a strictly positive gap between the best ICL risks achieved by those two model classes. This result highlights the benefits of the MLP layer for reducing approximation error in Transformer.

Theorem 4.1 (Approximation gap) . Consider the ICL risk defined by (3.5) and the two hypothesis classes F LTB and F LSA . Suppose that Assumption 3.1 holds. Then we have

inf f ∈F LTB R ( f ) is independent of β ∗ .

inf f ∈F LSA R ( f ) is a function of β ∗ . Moreover,

The proof of Theorem 4.1 is deferred to Appendix B. Theorem 4.1 reveals a gap in terms of the approximation abilities between the hypothesis set of LTB models and that of LSA models. Specifically, the best ICL performance achieved by LTB models is independent of β ∗ , while the best ICL performance achieved by LSA models is sensitive to the norm of β ∗ . In particular, when ∥ β ∗ ∥ 2 H = Ω( M ) , the best ICL risk of the former is smaller than that of the latter by at least a constant additive term. So when β ∗ is large, the hypothesis class formed by LSA models is restricted in its ability to perform effective ICL. We will show in Section 5 that the hypothesis class formed by LTB models can achieve nearly optimal ICL in this case.

We also remark that the Θ(1 /M ) factor in the lower bound in Theorem 4.1 is not improvable. This is because LSA models can implement a one-step GD algorithm that is consistent for linear regression tasks (that is, the risk converges to the Bayes risk as the number of context examples goes to infinity), with an excess risk bound of Θ(1 /M ) [38, 37]. So the approximation error gap is at most Θ(1 /M ) . Nonetheless, the Ω( ∥ β ∗ ∥ 2 H /M ) approximation gap between LSA models and LTB models shown by Theorem 4.1 suggests that β ∗ is not learnable by LSA models during pre-training. In contrast, we will show in Section 6 that LTB models can learn β ∗ during pre-training.

We emphasize that the ability of LTB to learn non-zero mean is a joint effect of an MLP component and a skip connection. Note that LSA also has a skip connection, but its skip connection is inactive. In comparison, the MLP component in LTB activates the skip connection. Therefore, we attribute the ability to learn non-zero mean to the MLP component, which is the only difference between LTB and LSA. Nonetheless, one can attribute the ability to learn non-zero mean to the skip connection: without a skip connection, LTB reduces to LSA with a potential rank constraint on the parameter, which cannot learn non-zero mean as we have proved. The above two explanations take different perspectives to interpret the same phenomenon. Finally, we remark that Theorem 4.1 holds even when H and Ψ in Assumption 3.1 are not full rank.

Does scratchpad help? We have demonstrated that employing a single-layer LSA introduces an additional approximation error in the in-context learning problem for the linear regression task defined in Assumption 3.1. Nonetheless, are there alternative structures, apart from MLPs, that could potentially reduce this approximation error? One plausible strategy is to include a 'scratchpad' in the input token. Specifically, we construct the token matrix as follows:

which we then input into the LSA layer. We discuss the limitation of this scheme in Appendix J. Notably, this method does not successfully recover the GD -β estimator defined in Section 5. We leave it as future work to see whether the token matrix with scratchpad could implement other types of estimators that more effectively address the linear regression tasks defined in Assumption 3.1, as well as whether additional structures could help alleviate this approximation error.

Experiments on GPT2. Theorem 4.1 shows the importance of the MLP component in LSA models for reducing approximation error. We also empirically validate this result by training a more complex GPT2 model [14] for the ICL tasks specified by Assumption 3.1. In the experiments, we use a GPT2-small model (with or without the MLP component) with 6 layers and 4 heads in each layer. The experiments follow the setting in [14], except that we train the model using a token matrix defined

Test Loss CurveAlongTrainingProcessforLTBlayerversusLSAlayer

5.5

5.0

Test Loss

4.5

4.0

3.5

0

50

100

150

200

250

300

Epoch

in (3.1). We considere two ICL settings, which instantiates Assumption 3.1 with β ∗ = (0 , 0 , ..., 0) ⊤ and β ∗ = (10 , 10 , ..., 10) ⊤ , respectively. We set d = 20 , M = 40 , σ = 0 , Ψ = H = I d in both settings. More experimental details are in Appendix I. In each setting, we train and test the model using the same data distribution. The experimental results are presented in Table 1. From Table 1, we observe that both models (with or without the MLP component) achieve a nearly zero loss when the task mean is zero. However, when the task mean is set away from zero, the GPT2 model with MLP component still performs relatively well while the GPT2 model without MLP component incurs a significantly larger loss. These empirical observations are consistent with our Theorem 4.1, indicating the benefits of MLP layers in reducing approximation error for ICL of linear regression with a shared signal.

Experiments on LSA and LTB. We trained both the LSA and LTB layers for β ∗ = (1 , 1 , ..., 1) ⊤ , and found that the trained LTB layer consistently achieved a significantly lower ICL risk than the LSA layer (see Figure 1). In these experiments, we adhered strictly to the previously defined LSA and LTB structures, setting the parameters as follows: d = 5 , M = 5 , σ = 0 , and Ψ = H = I d . At each training step, we sampled B = 128 new linear regression tasks.

In this part, we have shown that LSA models, the primary focus in previous ICL theory literature (see, e.g., [1, 38, 37] and references therein), are not sufficiently expressive for ICL of linear regression with a shared signal. In what follows, we will show LTB models are sufficient for this ICL problem.

5 LTB Implements One-Step GD with Learnable Initialization

To understand the expressive power of the LTB models, we build a connection between F LTB and its subset of models which we call one-step GD with learnable initialization ( GD -β ). We will first introduce GD -β models and then show that the best LTB models that minimize the ICL risk effectively belong to GD -β models.

The GD -β models. A GD -β model is defined as

where E is the token matrix given by (3.1), Γ ∈ R d × d and β ∈ R d are trainable parameters. Similarly, we define the function class formed by GD -β models as F GD -β := { f GD -β : β ∈ R d , Γ ∈ R d × d } .

A GD -β model computes a parameter that fits the context examples ( X , y ) and uses that parameter to make a linear prediction of label y on feature x . More specifically, the first step is by using one gradient descent step with a matrix stepsize Γ and an initialization β on the least square objective formed by context examples ( X , y ) .

LTB implements GD -β . A GD -β model (5.1) is a special case of an LTB model (3.3) by setting

where ∗ denotes the entries that do not affect the model output (hence can be set to anything). Note that Γ and β are free provided that d v ≥ d +1 , d k ≥ d and d f ≥ 1 (as required in F LTB ). In sum, we have proved the following lemma showing that the set of GD -β models belongs to the set of LTB models.

Optimal GD -β models for ICL. We now consider F GD -β and its optimal ICL risk. We have the following theorem that computes the globally minimal ICL risk over F GD -β and specifies the sufficient and necessary conditions for a global minimizer. The proof is deferred to Appendix C.

Theorem 5.2 (Optimal GD -β models) . Consider the ICL risk defined by (3.5) . Suppose that Assumption 3.1 holds. Then we have

The minimal ICL risk of F GD -β is

where Ω := [( M +1) H 1 2 ΨH 1 2 +(tr( HΨ ) + σ 2 ) I d ] /M.

The global optimal parameters for f GD -β that attain the minimum (5.2) take the following form:

and β ∗ is the mean of the task parameter in Assumption 3.1 Here, null ( H ) := { z ∈ R d : H z = 0 } is the null space of H , and null ( H ⊗ 2 ) = { Z ∈ R d × d : H Z H = 0 } is the null space of H ⊗ 2 . In particular, when H is positive definite, the global optimal parameter is unique, ( β , Γ ) = ( β ∗ , Γ ∗ ) .

Under Assumption 3.1, the global optimal f GD -β that attains the minimum (5.2) is unique as a function of E (given by (3.1) ) and takes the following form:

Theorem 5.2 characterizes the optimal GD -β models for ICL. In the above theorem, Γ ∗ → H -1 as the context length M goes to infinity (assuming that H is positive definite). In this case, the optimal GD -β function (5.2) implements one Newton step from initialization β ∗ . With a finite context length M , (5.2) implements one regularized Newton step from initialization β ∗ .

Optimal LTB models for ICL. Lemma 5.1 shows that the best ICL risk achieved by an GD -β model is no smaller than the best ICL risk achieved by an LTB model. Surprisingly, our next theorem shows that the best ICL risk achieved by a GD -β is equal to that achieved by an LTB model. Therefore, the hypothesis set F GD -β is diverse enough to match the approximation ability of the larger hypothesis set F LTB for ICL.

Theorem 5.3 (Optimal LTB models) . Consider the ICL risk defined by (3.5) . Suppose that Assumption 3.1 holds and that rank ( H 1 2 Ψ 1 2 ) ≥ 2 . Then we have

The minimal ICL risk of F LTB and of F GD -β are equal,

Rewrite an LTB model as f LTB in (3.3) with parameters ( ∗ denotes parameters that do not affect the output)

Then the sufficient and necessary conditions for f LTB ∈ arg min f ∈F LTB R ( f ) are

̸

where β ∗ is defined in Assumption 3.1 and Γ ∗ is defined in (5.3) . In particular, when H is positive definite, the globally optimal parameter represented by ( V 11 , v 12 , v 21 , v -1 , γ ) is unique up to a rescaling of v -1 .

Under Assumption 3.1, the globally optimal LTB model (that is, a function in arg min f ∈F LTB ) is unique as a function of E (given by (3.1) ) and takes the form of (5.4) almost surely.

Lemma 5.1 and Theorem 5.3 together show that F GD -β is a representative subset of F LTB that does not incur additional approximation error. In addition, every optimal LTB model is effectively an optimal GD -β model when restricted to all possible token matrices. Note that the optimal model parameters for LTB or GD -β are not unique because of redundant parameterization. But the optimal LTB and GD -β models are unique as a function of all possible token matrices. The above holds even when H and Ψ are potentially rank deficient.

Comparison with prior works. Aline of papers considers LSA models for ICL under Assumption 3.1 with β ∗ = 0 [35, 1, 38, 37]. They show that LSA models can (effectively) implement all possible GD -0 models that specialize GD -β models by fixing β = 0 . In addition, they show that every optimal LSA model is (effectively) a GD -0 model for ICL under Assumption 3.1 with β ∗ = 0 (see, for example, Theorem 1 in [1]). In comparison, we consider a harder ICL problem that allows a large shared signal in tasks (that is, a large β ∗ in Assumption 3.1). In this setting, our Theorem 4.1 shows that F LSA (hence its subset formed by GD -0 models), as a subet of F LTB , incurs an additional approximation error propotional to ∥ β ∗ ∥ 2 H compared with F LTB . In contrast, F GD -β , as a subset of F LTB , does not incur additional approximation error according to our Theorem 5.3. Thus the LSA and GD -0 models considered by [35, 1, 38, 37] are not capable of learning the shared signal β ∗ , while an LTB model can learn β ∗ through implementing GD -β and encoding β ∗ in the initialization parameter.

6 Training and In-Context Learning of GD-beta

We have shown that F GD -β is a representative subset of F LTB that effectively contains every optimal LTB model. We now examine the ICL and training of GD -β models.

Nearly optimal ICL with GD -β . We will compare the best ICL risk achieved by GD -β with the best ICL risk achieved by any estimator. The following lemma is an extension of Proposition 5.1 and Corollary 5.2 in [37] (which is based on [33]) that characterizes the Bayes optimal ICL risk among all estimators.

Lemma 6.1 (Bayes optimal ICL) . Given a task-specific dataset ( X , y , x , y ) sampled according to Assumption 3.1, let g ( X , y , x ) be an arbitrary estimator for y and measure the average linear regression risk by L ( g ; X ) := E [( g ( X , y , x ) -y ) 2 | X ] . It is clear that E L ( g ; X ) = R ( g ) . Then,

The optimal estimator that minimizes the average linear regression risk L ( · ; X ) is g ∗ ( X , y , x ) = x ⊤ β ∗ + x ⊤ Ψ 1 2 ( Ψ 1 2 X ⊤ XΨ 1 2 + σ 2 I d ) -1 Ψ 1 2 X ⊤ ( y -X β ∗ ) .

Assume the signal-to-noise ratio is upper bounded, that is, tr ( HΨ ) ≲ σ 2 , then with probability at least 1 -exp( -Ω( M )) over the randomness of X , it holds that L ( g ∗ ; X ) -σ 2 ≃ ∑ d i =1 min { ¯ ϕ, ϕ i } , where ¯ ϕ ≂ σ 2 M , and ( ϕ i ) i ≥ 1 are the eigenvalues of Ψ 1 2 HΨ 1 2 .

The proof is deferred to Appendix F. Lemma 6.1 shows that the Bayes optimal estimator is a ridge regression estimator centered at β ∗ . This is consistent with [37] where the Bayes optimal estimator

is a ridge regression estimator as they assumed β ∗ = 0 . The following corollary of Theorem 5.2 computes the rate of the ICL risk achieved by the optimal GD -β model.

Corollary 6.2. Under the setup of Theorem 5.2, additionaly assume the signal-to-noise ratio is upper bounded, that is, tr ( ΨH ) ≲ σ 2 . Then we have inf f ∈F GD -β R ( f ) -σ 2 ≃ ∑ d i =1 min { ¯ ϕ, ϕ i } , where ( ϕ i ) i ≥ 0 are the eigenvalues of Ψ 1 2 HΨ 1 2 and ¯ ϕ := [tr( Ψ 1 2 HΨ 1 2 ) + σ 2 ] /M ≃ σ 2 /M.

The optimal (expected) ICL risk achieved by F GD -β in Corollary 6.2 matches the (high probability) Bayes optimal ICL risk in Lemma 6.1 ignoring constant factors, provided that the signal-to-noise ratio is upper bounded. Therefore F GD -β achieves nearly Bayes optimal ICL risk. As a consequence, the larger hypothesis set F LTB also achieves nearly Bayes optimal ICL of linear regression under Assumption 3.1.

For the simplicity of discussion, we assume a fixed context length during pretraining and inference. Our discussions can be extended to allow a different context length during pretraining and inference using techniques in [37]. However, this is not the main focus of this work.

Optimization of GD -β with infinite tasks. We have shown that F GD -β is a representative subset of F LTB that covers the optimal LTB models and achieves nearly optimal ICL risk. We now consider the optimization in the parameter space specified by F GD -β . For simplicity, we follow [38] and consider gradient descent with an infinitesimal stepsize on the ICL objective with an infinite number of tasks. That is, we consider the optimization of gradient flow on the population ICL risk under the parameterization of GD -β ,

where R is defined by (3.5) and f GD -β is defined by (5.1).

The following theorem guarantees the global convergence of gradient flow. We introduce some notation to accommodate cases when H is rank deficient. Let P S be the orthogonal projection operator onto a subspace S . Let H = Im ( H ) be the image space of matrix H (viewing H as a linear map) and H ⊥ := null ( H ) be its orthogonal complement. Similarly, let Z := Im ( H ⊗ 2 ) = { H Z H , Z ∈ R d × d } be the image space of the operator H ⊗ 2 and Z ⊥ be its orthogonal complement. Then we have the following theorem.

Theorem 6.3. Consider the gradient flow defined by (6.1) with initialization β (0) , Γ (0) . We have,

as t →∞ . In particular, if H is positive definite, the gradient flow converges to the unique global minimizer of ICL risk over GD -β class, that is, β ( t ) → β ∗ and Γ ( t ) → Γ ∗ as t →∞ .

The proof, as well as the convergence rate, is deferred to Appendix G. We remark that (6.1) is a complex dynamical system on a non-convex potential function of β and Γ . We briefly discuss our proof techniques assuming that H is full rank. The rank-deficient cases can be handled in the same way by applying appropriate project operators. To conquer the non-convex optimization issue, we observe that for every fixed Γ , the potential as a function of β is smooth and strongly convex with a uniformly bounded condition number. This observation allows us to establish a uniform convergence for β . When β is sufficiently close to β ∗ , the potential as a function of Γ is approximately convex, allowing us to track the convergence of Γ .

Theorem 6.3 shows that optimization of GD -β can be done efficiently by gradient flow without suffering from non-convexity. However, as we have shown in previous sections, LTB utilizes a more complex parameterization than GD -β . So Theorem 6.3 does not imply optimization of LTB is easy. We leave it as future work to study the optimization and statistical complexity for directly learning LTB models.

7 Concluding Remarks

In this paper, we study the in-context learning of linear regression with a shared signal represented by a Gaussian prior with a non-zero mean. We show that although the linear self-attention layer

discussed in prior works is consistent for this more complex task, its risk has an inevitable gap compared to that of the linear Transformer block (LTB), which is a linear self-attention layer followed by a linear multi-layer perception (MLP) layer. Next, we show that the effectiveness of the LTB arises because it can implement the one-step gradient descent estimator with learnable initialization ( GD -β ). Moreover, all global minimizers in the LTB class are equivalent to the unique global minimizer in the GD -β class, which can achieve nearly Bayes optimal in-context learning risk. Finally, we consider training on in-context examples and prove global convergence over the GD -β class of gradient flow on the population loss. Several future directions are worth discussing.

Optimization and statistical complexity. This paper provides an approximation theory of LTB and an optimization theory of GD -β . However, the statistical complexity of learning LTB or GD -β is not considered. The work by [37] provided techniques for analyzing the statistical task complexity for pre-training GD -0 . An interesting direction is to extend their method to study the statistical complexity of learning GD -β . However, their method crucially relies on the convexity of the risk induced GD -0 , while we have shown that the risk induced by GD -β is non-convex. New ideas for dealing with non-convexity are needed here.

From LTB to Transformer block. We focus on LTB in this work, which simplifies a vanilla Transformer block by removing the non-linearities from the softmax self-attention and the ReLU activation in the MLP layers. This simplification allows us to obtain precise theoretical results for LTB (such as its connection to GD -β ). On the other hand, non-linearities are arguably necessary for Transformers to work well in practice. An important next step is to further consider the theoretical benefits of non-linearities based on our current results.

Roles of MLP layers. The work by [15] (and references thereafter) empirically found that MLP layers operate as key-value memories that store human-interpretable patterns in some pre-trained Transformers. Their work motivated a method for locating and editing information stored in language models by modifying their MLP layers (see, e.g., [9, 23]). Our work proves that the MLP component enables LTB to learn the shared signal in linear regression tasks, which cannot be done by a single LSA component. We leave it as future work to theoretically clarify the information stored in the MLP component.

Acknowledgments

We gratefully acknowledge the support of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and #814639, and of the NSF through grant DMS-2023505.

References

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Kabir Ahuja, Madhur Panwar, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations , 2024.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations , 2022.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and LLMs by learning to learn discrete functions. In The Twelfth International Conference on Learning Representations , 2024.

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. Advances in neural information processing systems , 33:1877-1901, 2020.

Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via language model in-context tuning. arXiv preprint arXiv:2110.07814 , 2021.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8493-8502, 2022.

Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023 , 2023.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

Jingwen Fu, Tao Yang, Yuwang Wang, Yan Lu, and Nanning Zheng. How does representation impact in-context learning: A exploration on a synthetic task. arXiv preprint arXiv:2309.06054 , 2023.

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.

Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , 2021.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. In The Twelfth International Conference on Learning Representations , 2024.

Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306.14892 , 2023.

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , pages 19565-19594. PMLR, 2023.

Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations , 2024.

Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys , 55(9):1-35, 2023.

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. In The Twelfth International Conference on Learning Representations , 2024.

AR Meenakshi and C Rajian. On a product of positive semidefinite matrices. Linear algebra and its applications , 295(1-3):3-6, 1999.

Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems , 35:17359-17372, 2022.

Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:2110.15943 , 2021.

OpenAI. Gpt-4 technical report, 2023.

Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. arXiv preprint arXiv:2311.08362 , 2023.

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

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arXiv:2306.15063 , 2023.

George AF Seber. A matrix handbook for statisticians . John Wiley & Sons, 2008.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Alexander Tsigler and Peter L Bartlett. Benign overfitting in ridge regression. J. Mach. Learn. Res. , 24:123-1, 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 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.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State-of-the-art natural language processing, 2020.

Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference on Learning Representations , 2024.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

A Notation and Variable Transformation

In order to simplify the proof, let's first describe a variable transformation scheme for our model and data coming from Assumption 3.1.

Definition A.1 (Variable Transformation) . We fix M as the length of the contexts. Recall that X , x are, respectively, the features in the context and the query input, while ˜ β and β ∗ are the true task parameter for the inference prompt and its expectation, respectively. Following the Assumption 3.1, we have ˜ β ∼ N ( β ∗ , Ψ ) . From the definition for multivariate Gaussian distribution, we know there exists a random vector ˜ θ such that where

Recall the noise vector is defined as ε = y -X β and is generated from ε ∼ N ( 0 , σ 2 · I M ) .

Rank deficient case Note that, even when Ψ is rank-deficient, the variable transformation in (A.1) and (A.2) still hold. This can be seen from the definition of the multivariate Gaussian distribution (Def 20.11 and the discussion below in [31]): A vector ˜ β ∈ R d with mean β ∗ and covariance matrix Ψ has a multivariate normal distribution, if it has the same distribution as A ˜ θ + β ∗ where A is any d × m matrix satisfying AA ⊤ = Ψ and ˜ θ ∼ N ( 0 m , I m ) . Here, we can recover the variable transformation above if we take m = d and A = Ψ 1 2 .

Notation Before we delve into the detailed proof, let's repeat the notation part with some additional notations. We use lowercase bold letters to note vectors and uppercase bold letters to denote matrices and tensors. For a vector x and a positive semi-definite (PSD) matrix A , we denote ∥ x ∥ A := √ x ⊤ A x . We denote ⟨· , ·⟩ as the inner product, which is defined as ⟨ x , y ⟩ := x ⊤ y for vectors and ⟨ A , B ⟩ := tr ( AB ⊤ ) for matrices. For a matrix A , we denote ∥ A ∥ op , ∥ A ∥ F as the operator norm and the Frobenius norm, respectively. We denote A [ i ] as the i -th row of the matrix A , A m,n as the ( m,n ) -th entry, and A -1 , -1 as the right-botom entry. We denote 0 n , 0 m × n , I n as the zero vector, zero matrix and identity matrix, respectively.

For a positive semi-definite matrix A , we denote A 1 2 for the principle square root of A , which is defined as the unique real matrix B such that B is positive semi-definite and BB = BB ⊤ = A . For positive definite A , its principle square root is also positive definite. We denote A + as the Moore-Penrose pseudo-inverse for any matrix A . We also denote A -1 2 = ( A 1 2 ) + . We denote ⊗ as the Kronecker product. For compatible matrices of proper size A , B , C , B ⊤ ⊗ A is a linear mapping which is defined by ( B ⊤ ⊗ A ) ◦ C = ACB . We denote Λ := Ψ 1 2 HΨ 1 2 and ϕ 1 ≥ ϕ 2 ≥ ... ≥ ϕ d ≥ 0 are its ordered eigenvalues. We also denote λ 1 ≥ λ 2 ≥ ... ≥ λ d ≥ 0 as the ordered eigenvalues of H , and λ -1 > 0 as its minimal positive eigenvalue. Finally, we denote another important matrix

Note that, under this definition and our assumption on H and Ψ , we have Ω is invertible.

B Proof of Theorem 4.1

Proof. The fact that inf f ∈F LTB R ( f ) does not depend on the vector β ∗ is subsumed in the Theorem 5.3, so we do not prove it here. We are going to prove the inequality in the theorem. First, from Theorem 5.3 we know that,

where Ω is defined in (A.3). Then, it suffices to lower bound inf f ∈F LSA R ( f ) . So let's take an arbitrary f ∈ F LSA , which is denoted as

where E is the token matrix defined in (3.1). Since the prediction is the right-bottom entry of the output matrix, we know only the last row of the product W ⊤ P W V attends the prediction. Similarly, only the last column of the E on the far right in the above equation atttends the prediction. Since the last column of E is ( x ⊤ 0 ) ⊤ , we know that only the first d rows of the product W ⊤ K W Q enter the calculation (since other parts are multipled by zero). Therefore, we denote

where U 11 ∈ R d × d , u 12 , u 21 ∈ R d × 1 , u -1 ∈ R , and ∗ denotes entries that do not enter the final prediction. The model prediction can be written as

Step 1: simplify the risk function. We use ˜ β to denote the task parameter. From the Assumption 3.1 and Definition A.1, we have

and

Then the model output can be written as

To simplify the presentation, we denote

Since x , X , ε , ˜ β are independent, we have

Note that z 1 does not contain ε , z 2 is a linear form of ε , and z 3 is a quadratic form of ε . Using ε ∼ N ( 0 , σ 2 I d ) , we have E [( z 1 -˜ β ) · z ⊤ 2 ] = 0 and E [ z 2 z ⊤ 3 ] = 0 . Therefore, we have

Step 2: compute S 1 . By Lemma H.4 and that X is independent of all other random variables, we have

We denote

then applying ˜ β = β ∗ + Ψ 1 2 ˜ θ , we get

Note that the first and third moments of ˜ θ are zero. So the above equation only involves the zeroth, second, and fourth moments of ˜ θ . Therefore we have

where T 2 and T 4 are the second order and the fourth order term, respectively. More concretely, the second order term is

For the cross term in S 1 , we have

where the last row comes from the fact that θ ∼ N ( 0 , I ) . Moreover, we have

Step 3: other terms. Let us compute S 2 , S 3 and S 4 . Using definitions, we rewrite z 2 as

Since ε ∼ N ( 0 , σ 2 I M ) , ˜ θ ∼ N ( 0 , I d ) and they are independent, all terms in S 2 vanish except if the term contains even orders of ˜ θ or ε . So we have

For the other two terms, we have

and

Step 4: combine all parts. Combining the four parts (B.10), (B.12), (B.13) and (B.14), we have

Weremark that in the (B.10), (B.12), (B.13) and (B.14), the parameters are A , b , u -1 and u 12 , instead of the original parameter U 11 , u 12 , u 21 , u -1 . From the definitions of A and b , we know there is a bijective map between ( U 11 , u 12 , u 21 , u -1 ) and ( A , b , u 12 , u -1 ) , so these two parameterizations are equivalent when computing the minimal risk achieved by a model in a hypothesis class. From the expression of S 1 , S 2 , s 3 , S 4 , we can rewrite the risk as

where I, II, III, IV are defined as above.

Step 5: lower bound the risk function. Let's first define a new matrix, which by definition is invertible:

So we can write I as

I = b ⊤ H 1 2 ΩH 1 2 b · u ⊤ 12 Hu 12 + u 2 -1 tr ( A ⊤ H 1 2 ΩH 1 2 A H ) +2 u -1 b ⊤ H 1 2 ΩH 1 2 A Hu 12

where the last line comes from the fact that Ω and H 1 2 ΨH 1 2 commute.

Next, we claim II ≥ 0 . To see this, it suffices to notice that

where the last line comes from the fact that ∥ A ∥ 2 F = tr ( AA ⊤ ) for any matrix A .

Then, let's consider III. Notice that actually, III can be viewed as a quadratic function of A ⊤ H b , which can be easily minimized. More concretely, we have

Therefore, one has

Combining the three parts above, one has

Therefore, one has

The right hand side in the above inequality is a quadratic function of u -1 u 12 , so we can take its global minimizer:

to lower bound the risk gap as

Finally, to simplify the results, we notice that on the one hand,

On the other hand, one has

Therefore, we have

When 1 M +1 -( tr( HΨ )+ σ 2 ) 2 2( M +1) 2 tr( HΨHΨ ) ≥ 2 3( M +1) , we have 1 M +1 ≥ 3 ( tr( HΨ )+ σ 2 ) 2 2( M +1) 2 tr( HΨHΨ ) , which implies

Therefore, we finally have

Note that this holds for an arbitrary f ∈ F LSA , so the proof finishes by taking infimum on the left hand side.

C Proof of Theorem 5.2

Proof. Recall that from Assumption 3.1,

where

Step 1: compute the risk function. From the independence of ε, ε with other random variables, we have

where the last line comes from the independence between ˜ β and X , x , as well as the fact that x ∼ N ( 0 , H ) and the property of tensor product ◦ . Note that, β ad Γ are learnable parameters here, and we should set them apart from the task vector ˜ β or the prior mean vector ˜ β ∗ . Recall that

we have that

Bridging the two terms above into (C.1), we get

First, let's compute V 1 . From the definition above, we have

where

Then, let's compute V 2 . We have

From the definition of Ω in (A.3), we know that it is invertible. Moreover, it holds that

Therefore, we have

Step 2: solve the global minimum. Now we have

where the last line comes from the fact that Ω and H 1 2 ΨH 1 2 commute. Taking infimum over all β , Γ in the left hand side, we have

On the other hand, we take

where H 1 2 is the principle square root of H and H -1 2 is its Moore-Penrose pseudo inverse (see notation part in Section A). Then, we have

Since

Therefore, we have

Combining both directions, we conclude that

Step 3: identify all global minimizers Above we show that R ( β ∗ , Γ ∗ ) achieves the global minimum of the ICL risk over GD -β class. Now we will figure out the sufficient and necessary condition to achieve the global minimal risk. From the proof above, we know that for any β ∈ R d , Γ ∈ R d × d , it holds that

Since Ω is invertible, the second equation is equivalent to

which is a linear system of Γ . Since Γ ∗ is proved to be one solution, all solutions of this equation is

where the last equation comes from Lemma H.6. Here, + denotes Minkowski sum. This is defined as A + B = { a + b, a ∈ A,b ∈ B } for two sets A,B and a + B = { a } + B for an entry a and a set B. Under this condition, we know that

Consider the following inequality:

Therefore, we have

where null ( · ) denotes the null space of a matrix and + denotes the Minkowski sum. The last equivalence comes from Lemma H.6. Therefore, we conclude that the f β , Γ achieves the minimal ICL risk in GD -β class if and only if

Specially, if H is positive definite, there is unique global minimizer in GD -β class with parameters

Step 4: equivalence in the hypothesis class Finally, we show when H is rank-deficient, any global minimizer actually corresponds to one single function in F GD -β almost surely. For an arbitrary global minimizer, we assume β = β ∗ + h , Γ = Γ ∗ + Z , where h ∈ null ( H ) , Z ∈ Im ( H ⊗ 2 ) . Suppose we have a prompt X , y , x , y which follows the Assumption 3.1 and the token matrix is formed by (3.1), we have

Notice that X [ i ] , x i . i . d . ∼ N ( 0 d , H ) , we have

since h ∈ null ( H ) . Therefore, we know h ⊤ x = 0 almost surely, and similarly, Xh = 0 d almost surely. Finally, we have

which imples x ⊤ Z X ⊤ = 0 d almost surely. Therefore, we conclude

D Proof of Theorem 5.3

Proof. Recall the definition of Linear Transformer Block class:

where E is the token matrix defined by (3.1). Let's take an arbitrary function f in the LTB class with trainable matrices

Similar to the proof of the Theorem 4.1, we know that only the last row of W ⊤ 2 W 1 and the first d -columns of W ⊤ K W Q attend the prediction. Therefore, we denote

where γ , v 12 , v 21 , ∈ R d , v -1 ∈ R , V 11 ∈ R d × d , and ∗ denotes entries that do not enter the prediction. Then, the prediction of LTB function can be written as

Step 1: simplify the risk function. We use ˜ β to denote the task parameter. From the Assumption 3.1 and Definition A.1, we have

and

Then the model output can be written as

To simplify the presentation, we denote

Since x , X , ε , ˜ β are independent, we have

Note that z 1 does not contain ε , z 2 is a linear form of ε , and z 3 is a quadratic form of ε . Using ε ∼ N ( 0 , σ 2 I d ) , we have E [( z 1 + γ -˜ β ) · z ⊤ 2 ] = 0 and E [ z 2 z ⊤ 3 ] = 0 . Therefore, we have

Step 2: compute the risk function. Let's compute the risk function. Note that, the risk function is in the same form as the risk function of LSA layer, which is in the proof of Theorem 4.1 (see Section B and equation (B.2)). There are two differences: one is we replace U 11 , u 12 , u 21 , u -1 here with V 11 , v 12 , v 21 , v -1 . The second difference is that we replace ˜ β in (B.2) with ˜ β -γ . This is equivalent to replacing β ∗ with β ∗ -γ , since ˜ β -γ ∼ N ( β ∗ -γ , Ψ ) . Therefore, similar to (B.15), the risk function in (D.2) can be written as

where V, VI, VII, VIII are defined as above and Ω = M +1 M H 1 2 ΨH 1 2 + tr( HΨ )+ σ 2 M · I d , and b := v 21 + v -1 β ∗ ∈ R d , A := V 11 + β ∗ v ⊤ 12 ∈ R d × d ,

Step 3: solve the global minimum of the risk function. This is very similar to step 5 in Appendix B. The V term in (D.3) is actually equal to the I term in (B.15), so we have

where the last line comes from the fact that Ω and H 1 2 ΨH 1 2 commute. For the same reason, we know that the V I term above is equal to the II term in (B.15), which implies VI ≥ 0 , since

where the last line comes from the fact that ∥ A ∥ 2 F = tr ( AA ⊤ ) for any matrix A . The term VII above is equal to III term in (B.15), except that we replace β ∗ with β ∗ -γ . Therefore, we have

which implies

Combining the three parts above, one has

Here, the global minimum of GD -β class is taken from Theorem 5.2, whose proof does not reply on the proof here. Therefore, one has

The right hand side in the above inequality is a quadratic function of v -1 v 12 , so we can take its global minimizer:

to lower bound the risk gap as

Note that, taking infimum on the left hand side, we get

On the other hand, in the main text we have showed that F GD -β ⊂ F LTB , which implies

Therefore, we conclude

where the last equation is from Theorem 5.2.

Step 4: sufficient and necessary conditions for global minimizers. Let's now verify the conditions for the global minimizers. For a function in LTB class, the sufficient and necessary condition for it to be a global minimizer is that inequalities in the above step all hold, which are

Let's first verify the necessary conditions of the system defined above. From (D.11) and Lemma H.5, we know γ ∈ β ∗ + null ( H 1 2 ) = β ∗ + null ( H ) . Then, we have v -1 v 12 ∈ null ( H ) . Then, in (D.7), we know

which implies either H 1 2 b = 0 d or H 1 2 A H 1 2 = 0 d × d . If we assume H 1 2 A H 1 2 = 0 d × d , then (D.6) implies H 1 2 v 12 b ⊤ H 1 2 = H 1 2 ΨH 1 2 Ω -1 . From our assumption, we know rank ( H 1 2 Ψ 1 2 ) ≥ 2 , which implies rank ( H 1 2 ΨH 1 2 ) ≥ 2 , since for any matrix Z , it holds that rank ( Z ) = rank ( ZZ ⊤ ) . Since multiplication by an invertible matrix does not change the rank, we know rank ( H 1 2 ΨH 1 2 Ω -1 ) ≥ 2 , while H 1 2 v 12 b ⊤ H 1 2 is a matrix of rank at most one, which contradicts with (D.6). Therefore, we have

Recall in Theorem 5.2, we have defined

Simple calculation shows

where the second and the last equalities come from the fact that Ω and H 1 2 ΨH 1 2 commute, and the third equality comes from the property of Moor Penrose pseudo-inverse. Therefore, one solution of (D.13) is v -1 A = Γ ∗⊤ . Since (D.13) is a linear equation to v -1 A , we know its full solution is

The solution to (D.12) is

Therefore, we know the necessary conditions for (D.6) to (D.11) are

̸

It is easy to verify these equations above are also sufficient conditions of (D.6) to (D.11) by directly replacing each variable with its value and validating equaions from (D.6) to (D.11).

Specially, if H is positive definite and hence, invertible, the global minimizer is unique up to a scaling to v -1 :

̸

Step 5: equivalence in the hypothesis class. Finally, we will prove that any global minimizer of F LTB is actually equivalent to one single function in F LTB almost surely. More concretely, let's take a function f ∈ F LTB with parameters W K , W Q , W P , W V , W 1 , W 2 and recall the parameter transformation in (D.1). We assume equations in (D.14) and Assumption 3.1 hold. Then, for any vector a ∈ null ( H ) , one has x ⊤ a = x ⊤ i a = 0 since x , x i i . i . d . ∼ N ( 0 d , H ) . We denote

where a 1 , a 2 ∈ H , H 1 2 Z H 1 2 = 0 d × d . Then, we have

where f β ∗ , Γ ( · ) is the GD -β function defined in (5.1).

E Proof of Corollary 6.2

Proof. We denote ϕ 1 ≥ ϕ 2 ≥ ... ≥ ϕ d ≥ 0 are ordered eigenvalues of Ψ 1 2 HΨ 1 2 . From Theorem 5.2, we have

where

Therefore, we have

Since

Therefore, we conclude.

F Proof of Theorem 6.1

F.1 Proof of first equation

Proof. From the classical bias-variance decomposition, we know

since the cross term vanishes. The second term does not depend on g and hence, the Bayesian optimal estimator is given by the posterior mean, i.e.,

Since ˜ β ∼ N ( β ∗ , Ψ ) , we have there exists a random vector ˜ θ ∼ N ( 0 d , I d ) such that ˜ β = β ∗ + Ψ 1 2 ˜ θ almost surely. Therefore, one has

To compute E [ ˜ θ | X , y , x ] , it suffices to solve the posterior distribution of β given X , y . From ˜ θ ∼ N ( 0 d , I d ) , we have

Note that, the function above matches the probability density function of a multivariate Gaussian distrbution. Therefore, the posterior mean is given by

Therefore, we conclude

F.2 Proof of second equation

Proof. Let's first do a variable transformation. We denote

Then, we know ˜ X [ i ] , ˜ x i . i . d . ∼ N ( 0 d , Λ ) . We can write the Bayesian optimal estimator in (F.1) as

The true label is

Therefore, we have

where

Note that, this is equivalent to the risk of estimating the ground true linear weight ˜ θ using ̂ θ under the Gaussian prior ˜ θ ∼ N ( 0 , I d ) . The estimator ̂ θ is a function of transformed input-output pairs in the context ( ˜ X , ˜ y ) and the transformed query input x , and takes the form of standard ridge estimator with regularization coefficient being σ 2 . We then revoke the standard results for the risk of a ridge estimator. Applying the Theorem 1 and 2 in [33], we know for a fixed weight vector ˜ θ , with probability at least 1 -exp( -Ω( M )) we have

where ϕ 1 ≥ ϕ 2 ≥ ... ≥ ϕ d ≥ 0 are ordered eigenvalues of Λ .

and c > 1 is an absolute constant. Here, H 0: k ∗ is SVD approximation with respect to the largest k ∗ singular values and H k ∗ : ∞ is the SVD approximation in the remaining singular values. Namely, if we have the eigen-decomposition of H = Q · diag ( ϕ 1 , ϕ 2 , ..., ϕ d ) · Q ⊤ , where Q is an orthogonal matrix, then H 0: k ∗ and H k ∗ : ∞ are given by

Taking expectation over ˜ θ ∼ N (0 , I d ) , we have

Now we simplify this expression. First, we define

From the definition, we see ¯ ϕ ≥ cσ 2 M . On the other hand, from the assumption that tr ( H ) = tr ( Ψ 1 2 HΨ 1 2 ) = ∑ d i =1 ϕ i ≲ σ 2 , we know that ¯ ϕ ≲ cσ 2 M . Combining two parts, we get

Therefore, we have

This finishes the proof.

G Proof of Theorem 6.3

G.1 Training dynamics

Now we consider doing gradient flow on the risk function, or equivalently, on the excess risk function which differs by only a constant:

First, we have the following corollary of Theorem 5.2, which computes the excess risk R ( β , Γ ) -min R ( · , · ) .

Corollary G.1 (Excess Risk) . We fix M as the context length. Consider the ICL risk in (3.5) , assume the data is generated following Assumption 3.1. Then, we have

Proof. This is obtained directly from equation (C.5) and (C.7).

Now, we write out the differential equations explicitly in the following lemma.

Lemma G.2 (Dynamical system) . The dynamical system of gradient flow described in (G.1) and (G.2) is

Proof. This can be obtained by directly calculating the derivatives over the excess risk (G.3). To write out the dynamics of β , it suffices to notice that β only attends the first term of the RHS of (G.3), which is a standard quadratic function. For the derivatives of Γ , we first have

Now it suffices to compute

Let's compute the partial derivativess separately. From Lemma H.3, we have

Summing over the three equations above and applying the definition of gradient flow, we conclude the dynamics of Γ in (G.5).

G.2 Proof of the global convergence

Let's now prove Theorem 6.3. Since the first part of Theorem 6.3 is directly implied by the second part, we only deal with the case with general H . In this section, we denote λ -1 > 0 as the minimal non-zero eigenvalue of H . As in the main text, we define

and H ⊥ as its orthogonal complement. First, let's prove the convergence of β .

Lemma G.3 (Convergence of β ) . Under the dynamical system (G.4) and (G.5) , one has

which implies and

when t →∞ , from arbitrary initialization β (0) and Γ (0) . Moreover, one has for any t > 0 , it holds that

Proof. We first consider the orthogonal projection operator P . From Lemma H.5 and equation (G.4), we know

where

From the property of pseudo-inverse, we know

Therefore, we have

Notice that

This implies

where the last line comes from Lemma H.5 and λ -1 is the minimal non-zero eigenvalue of H . Via standard integration method in ODE, we know this suggests

This indicates that P H ( β ( t )) →P H ( β ∗ ) when t →∞ . Finally, we have

Next, let's consider the convergence of Γ . As defined in the main text, we have

and Z ⊥ is its orthogonal complement. We have the following lemma.

Lemma G.4. Under the dynamical system (G.4) and (G.5) , one has

Proof. The first ODE is trivial since the RHS of (G.5) always lies in Z := Im ( H ⊗ H ) , so its projection on Z ⊥ always vanishes. To prove the second equation, we can rewrite (G.5) as

For term A, recalling Γ ∗ = ΨH 1 2 Ω -1 H -1 2 , we have

This suggests A = 0 . For B + C, we have

We can compute ( I d -HΓ ∗ ) H as

Therefore,

Bridging A = 0 and the result of B + C into the ODE, we have

Note that, all terms in the RHS above lie in Z = Im ( H ⊗ 2 ) , so this summation also equals d d t P Z ( Γ -Γ ∗ ) . Moreover, since

Bridging (G.12) into (G.11), we conclude.

Then, let's upper bound the dynamics of ∥P Z ( Γ -Γ ∗ ) ∥ 2 in the following lemma.

F

where

are two positive constant. Here, λ -1 > 0 is the minimal positive eigenvalue of H , λ 1 is the maximal eigenvalue of H , λ min ( Ω ) is the minimal eigenvalue of Ω (defined in (A.3) ) and is strictly positive.

Proof. From the dynamics of P Z Γ in (G.10), we know

where

From Lemma H.2, we know that G 2 ≤ 0 . Then, we consider G 1 and G 3 . For G 1 , we have

where λ min ( · ) denote the minimal eigenvalue. Since Ω is positive definite, we know λ min ( Ω ) > 0 . Finally, let's upper bound G 3 . First, we have

For λ max ( H 1 2 Γ ∗ H 1 2 ) , we recall Γ ∗ = ΨH 1 2 Ω -1 H -1 2 and obtain

Since Ω -1 and H 1 2 H -1 2 H 1 2 commute, they are simultaneously diagonalizable. The eigenvalues of Ω -1 H 1 2 ΨH 1 2 are

Note that, every eigenvalue is upper bounded by 1 , so we simply get

Therefore, we have

Similarly, we have

and

Bridging (G.18), (G.19) and (G.20) into (G.16), we get

Finally, we finish this section by proving the convergence of Γ .

Lemma G.6 (Convergence of Γ ) . Under the dynamical system (G.4) and (G.5) , one has

when t →∞ , from arbitrary initialization β (0) and Γ (0) .

Proof. We observe that

Combining it with Lemma G.5, we have

where A 1 > 0 , A 2 > 0 are defined in (G.14). Simple calculation shows that

̸

When A 1 = 2 λ -1 M +1 , we integrate both sides from t = 0 to t = T , then divide them by exp( A 1 T ) . This gives

when T → ∞ . Otherwise, if A 1 = 2 λ -1 M +1 , the right hand side of (G.22) reduces to a constant, which implies exp( A 1 t ) ∥P Z ( Γ ( t ) -Γ ∗ ) ∥ F grows at most at a linear rate, which implies ∥P Z ( Γ ( t ) -Γ ∗ ) ∥ F → 0 when t →∞ . Together, in all cases, we have ∥P Z ( Γ ( t ) -Γ ∗ ) ∥ F → 0 . From (G.9), we soon get P Z ( Γ ( t )) = P Z ( Γ (0)) foo all t > 0 . Therefore, we conclude.

The Theorem 6.3 is proved by simplly combining Lemma G.3 and Lemma G.6.

H Technical lemmas

Lemma H.1 (Von-Neumann's Trace Inequality) . Let U, V ∈ R d × n with d ≤ n . We have

where σ 1 ( X ) ≥ σ 2 ( X ) ≥ · · · ≥ σ d ( X ) are the ordered singular values of X ∈ R d × n .

Lemma H.2 ([22]) . For any two positive semi-definite matrices A,B ∈ R d × d , we have

tr[ AB ] ≥ 0 .

AB ⪰ 0 if and only if A and B commute.

Lemma H.3 (Derivatives, [27]) . We denote A , B , C , D , X as matrices and a , b , x as vectors. Then, we have

∂ ∂ X tr [ X ⊤ BXC ] = BXC + B ⊤ XC ⊤ .

LemmaH.4 (Lemma D.2 in [38], Lemma 4.2 in [37]) . If x is Gaussian random vector of d dimension, mean zero and covariance matrix H , and A ∈ R d × d is a fixed matrix. Then

If A is symmetric and the rows in X ∈ R M × d are generated independently from

Then, it holds that

Lemma H.5 (Linear Algebra) . Suppose H is a (non-zero) positive semi-definite matrix in R d × d and H 1 2 is its principle square root. We denote λ -1 as the minimal non-zero eigenvector of H . Then, we have

2. HH + = H + H , where ( · ) + denotes the Moore-Penrose pseudo-inverse.

3. For any vector α ∈ R d , the orthogonal projection operator on Im ( H ) and null ( H ) are respectively

4. For any vector α ∈ R d , we have

Proof. We consider the eigen-decomposition of matrix H , which is

where D = diag ( λ 1 , λ 2 , ..., λ d ) is a diagonal matrix with diagonal entries being the eigenvalues of H , and Q is an orthogonal matrix. Then, from the definition of principle square root, we know

where D 1 2 = diag ( √ λ 1 , √ λ 2 , ..., √ λ d ) . We denote columns of Q as q 1 , ... q d . We know they are the eigenvectors of H . From the eigen-decomposition of H and H 1 2 , we know they share the same set of eigenvectors. Without loss of generality, we assume rank ( H ) = r and λ 1 ≥ λ 2 ≥ ... ≥ λ r > 0 , λ i = 0 , i = r +1 , ..., d. Then, we denote H as the linear vector space spanned by q 1 , ..., q r and H ⊥ as its orthogonal complement (which is the subspace spanned by q r +1 , ..., q d ). For any vector α ∈ R d , we can write its orthogonal decomposition as α = ∑ d i =1 a i q i , so we have

which implies Im ( H ) ⊂ H . On the other hand, we know for any α ∈ H , we can write it as α = ∑ r i =1 a i q i for some a 1 , ..., a r , then we have

which implies H ⊂ Im ( H ) . This shows

since null ( H ) is the orthogonal complement of Im ( H ) . Then, (1) is proved by noticing that H and H 1 2 share the same set of eigenvectors. To prove (2), it suffices to notice that

To prove (3), it suffices to notice that actually HH + α ∈ Im ( H ) and

Finally, to prove (4), we can write α = ∑ d i =1 a i q i for some a 1 , ..., a d . Then, by orthogonality we have

Therefore, we conclude.

Lemma H.6 (Tensor product) . Suppose H is a (non-zero) positive semi-definite matrix in R d × d and H 1 2 is its principle square root. We denote ⊗ as Kronecker product, which is defined as

and

We define

and define Im ( H 1 2 ⊗ H 1 2 ) and null ( H 1 2 ⊗ H 1 2 ) similarly. Then, we have

so

2. We denote Z = Im ( H 1 2 ⊗ H 1 2 ) = Im ( H ⊗ H ) and Z ⊥ = null ( H 1 2 ⊗ H 1 2 ) null ( H ⊗ H ) . For any matrix Z ∈ R d × d , we have

3. For any matrix Z ∈ R d × d , it holds that

where λ -1 is the minimal positive eigenvector of H .

Proof. Let's first prove the second part. From the definition of tensor product and the fact that HH + = H + H (see Lemma H.5), we know the second equation in H.3 holds. To prove the first equation, it suffices to notice that for any matrix Z ∈ R d × d , the matrix HH + Z H + H is in Im ( H ⊗ 2 ) , together with the fact that

The proof of (H.4) is similar. Then, from Lemma H.5, we know Im ( H ) = Im ( H 1 2 ) , which implies the projection operator onto those two subspace are identical, indicating HH + = H 1 2 H -1 2 . Therefore, we have

Since this holds for any matrix Z ∈ R d × d , this suggests Im ( H 1 2 ⊗ H 1 2 ) = Im ( H ⊗ H ) . Similarly, we also have null ( H 1 2 ⊗ H 1 2 ) = null ( H ⊗ H ) . Finally, to prove (3), we use the eigendecomposition of H and H 1 2 :

where Q = ( q 1 q 2 ... q d ) is an orthogonal matrix and D = diag ( λ 1 , ...., λ d ) , where λ 1 ≥ λ 2 ≥ ....λ d ≥ 0 are ordered eigenvalues. We assume rank ( H ) = r and λ r > 0 = λ r +1 . Then, from the property of Kronecker product (eg. see [27]), the eigenvalues of H 1 2 ⊗ H 1 2 are {√ λ i λ j : 1 ≤ i, j ≤ d } . The eigenvectors are { q i ⊗ q j : 1 ≤ i, j ≤ d } , which forms an orthogonal unit basis of R d × d . We define

as the eigenspace spanned by all eigenvectors corresponding to positive eigenvalues. For any matrix Z , we can decompose it as

which implies Z ⊂ S . On the other hand, for any Z ∈ S , we can write it as Z = ∑ r i,j =1 b i,j · q i ⊗ q j , so we have

which implies S ⊂ Z . Combining two directions, we have S = Z . Therefore, for any matrix Z ∈ R d × d , we can write it as Z = ∑ d i,j =1 c i,j · q i ⊗ q j for some c i,j . Then, we have

Therefore, we have

Therefore, we conclude.

I Experiment Details

Our experiments on GPT2 mostly follow the setting in [14], except that we use the token matrix defined in (3.1). In our experiments, we use d = 20 , M = 40 , Ψ = H = I d and σ = 0 . We train the GPT2 with and without MLP layers on two settings: β ∗ = (0 , 0 , ..., 0) ⊤ and β ∗ = (10 , 10 , ..., 10) ⊤ . For GPT2 with and without MLP layers, we initialize the W V and W 2 by normal distribution with standard deviation 0 . 02 / √ number of residual connections , where the number of residual connections equals the number of layers for GPT2 without MLP layers. For GPT2 with MLP layers, this is 2 times the number of layers. Initialization for other matrices follow the default setting in [36]. We use Adam with learning rate 0 . 0001 . We train the model for 200000 steps and we sample a batch of 256 new tasks for each step.

J Does scratchpad help?

In this section, we show the limitations of adding a scratchpad to the token matrix. We will show that by using a single LSA layer and the token matrix with scratchpad, one cannot recover the GD -β estimator defined in Section 5. We leave it as future work to see whether the token matrix with scratchpad could implement other types of estimators that more effectively address the linear regression tasks defined in Assumption 3.1, as well as whether additional structures could help alleviate this approximation error. We follow the notations in previous sections. The token matrix with scratchpad is defined as

where 1 M ∈ R M refers to the vector filled with ones.

A LSA model f, is defined by

where the second equality is because the bottom right entry of E is zero (see (3.1)). Note that the prediction is the bottom right entry of the output matrix. So only the last row of W ⊤ P W V and the last column of E attend the prediction. Denote

where

and ∗ denotes entries that do not enter the final prediction. Then we have

Following the Assumption 3.1, we use ˜ β to refer to the task parameter, then

We then decompose f ( E ) as

where terms I and II are independent of x . Therefore, we have

Note that the above equation holds if II = 0 , which can be easily achieved by setting b = 0 d , b 1 = b 2 = 0 . Therefore, without loss of generality, we can consider the LSA function which takes the following form:

Let's then try to determine whether such a function can effectively represent a GD -β function, which takes the form of

where β ∗ is the prior mean in Assumption 3.1 in our submission, and Γ ∗ = ΨH 1 2 Ω -1 H 1 2 and Ω = M +1 M H 1 2 ΨH 1 2 + σ 2 +tr( HΨ ) M I d is defined in (5.3) in our submission. In order to achieve this, one simple way is to let

However, this will incur some additive terms and will potentially enlarge the ICL risk. Inserting the parameters above, we have

This shows that the extended token with a single LSA cannot easily implement the GD -β function class and may incur an additive ICL risk depending on β ∗ (as shown in Theorem 4.1).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We believe our introduction and abstract are factually accurate in describing the contributions of the paper and of its result.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

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

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

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

2. Limitations

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

Answer: [Yes]

Justification: We discuss the limitation in Section 4 and the future work section.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

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

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

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: We provide the assumption in Assumption 3.1.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

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

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all detailed experiment setup in Appendix I.

Guidelines:

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

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

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

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

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

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

5. Open access to data and code

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

Answer: [Yes]

Justification: The data we use are simulated. We plan to release our code upon acceptance.

Guidelines:

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

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

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

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

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

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

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

6. Experimental Setting/Details

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

Answer: [Yes]

Justification: We include all details and hyperparameters in our experiments in Appendix I. Guidelines:

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

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

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

7. Experiment Statistical Significance

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

Answer: [No]

Justification: The cost for producing the experiments precludes us from reporting error bars. Moreover, the experiment section is not the main contribution of this paper.

Guidelines:

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

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

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

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

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

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

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

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

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

8. Experiments Compute Resources

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

Answer: [NA]

Justification: There is no strict requirement for the computer resources for our experiments.

Guidelines:

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

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: None to report.

Guidelines:

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

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

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This work focuses on the theory of in-context learning so it is not expected to have a direct societal impact.

Guidelines:

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

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

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

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

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The data used in this paper are fully simulated, so there is no data or model that have a high risk.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: We do not relase any asset.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not relase any asset.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

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

Answer: [NA]

Justification: We do not involve humans in our research.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

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

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the in-context learning (ICL) in the linear regression setting with a Gaussian prior with a non-zero mean. They first prove that linear Transformer block (LTB) will enjoy a smaller approximation error than linear self-attention (LSA) in this setting, where the non-zero mean of Gaussian emerges. Second, they show that the core benefit brought by linear MLP is that LTB can implement the one-step gradient descent with learnable initialization (GD-$\\beta$). Besides, all global minimizers in the LTB class are equivalent to the unique global minimizer in the GD-$\\beta$ class. Finally, they investigate the non-convex dynamics of the GD-$\\beta$ class and prove that it will converge to the unique global minimizer. Simulations on the GPT-2 support their theory.\n\nSTRENGTHS:\n1. Clarity: The paper is presented very clearly and is easy to follow.\n2. Significance: The paper characterizes the benefit brought by the linear MLP in the ICL for linear regression setting, which is undoubtedly an important topic in the ICL theory community.\n3. Quality: The paper investigates the additional effect of linear MLP from approximation and optimization perspectives in detail, and the derivation is solid.\n\nWEAKNESSES:\n1. In this paper, all non-linear parts in the transformer block (softmax and ReLU) are dropped. However, similar works on the non-linear transformer block have emerged in the ICL theory. [a] considers the ReLU MLP followed by a linear attention layer in the linear regression setting. [b] considers the standard transformer block in the in-context classification regime. Authors are suggested to discuss the potential impacts of these studies about non-linear transformers on the setting in this paper.\n2. Though the LTB class enjoys the same global minimizer as the GD-$\\beta$ class and the training dynamics of the GD-$\\beta$ class can converge to the global minima, whether the training dynamics of the LTB class can converge to the global minimizer is not discussed by the authors. The authors should clarify this point.\n\n[a] Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape, ICML, 2024\n\n[b] Training Nonlinear Transformers for Efficient In-Context Learning: A Theoretical Learning and Generalization Analysis, ICML, 2024\n\nQUESTIONS:\n1. In the LTB class, the paper does not merge $W_K^TW_Q$ and $W_P^TW_V$ into one matrix parameter, respectively. Are there any benefits? Does it will improve the approximation ability of the LTB?\n2. In line 272, I think the $\\Gamma$ is a precondition but not equivalent to the Newton step. Is it a typo?\n3. The pretraining goal of this paper is different from the autoregression training in practice. Are there any potential approaches to study the autoregression pretraining and few-shot regression inference?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nhis submission extends the earlier theoretical analyses of in-context learning with linear transformers by taking into account the data covariance and a non-zero mean on the task parameters. It is shown that a linear transformer with a linear MLP block (LTB) can implement the algorithm of one-step gradient descent with non-zero initialization. This additional flexibility of learned initialization allows LTB to match the Bayes optimal risk for linear regression, whereas the linear transformer studied in prior works cannot achieve such optimality at finite $M$. The authors also showed that the optimal one GD step solution can be obtained by training the parameters with gradient flow.\n\nSTRENGTHS:\nIn the past years, there have been many works that examined the ICL ability of linear transformers in the linear regression setting, where it is typically shown that transformers implement one preconditioned GD step on the squared error in-context. This submission goes beyond the prior results by showing that linear transformers can adapt to the non-zero mean of the task distribution with the aid of a linear attention block. This yields a separation in the ICL efficiency between transformers with or without an attention block, which is an interesting message in my opinion.\n\nWEAKNESSES:\nI have the following concerns and questions:\n\n1. Firstly, the obvious limitations of the analysis: the problem setting is still linear, and the authors only showed adaptivity of ICL to the task mean but not the full prior (see point below). Given the additional linear block, the approximation gap is not really surprising. The optimization analysis directly assumes the GD-$\\beta$ parameterization, and it is not clear if gradient descent on the standard LTB reaches the same solution. Also, convergence is only shown at a population level with no finite-sample guarantees.\n\n2. The Bayes optimal estimator in Lemma 6.1 is a generalized ridge regression estimator where the $\\ell_2$ regularization is anisotropic and adapts to the prior. On the other hand, when $\\beta^*=0$, the GD-$\\beta$ solution in Theorem 5.2 reduces to the one-step GD estimator studied in prior works, which does not incorporate the prior covariance.  \nIt is not clear to me why these two estimators always have the same statistical efficiency, as it is well known that anisotropic shrinkage provides a substantial advantage (Wu and Xu 2020) (Li et al. 2023). I guess the conclusion in Corollary 6.2 relies on $tr(H\\Psi)$ being bounded. Can the authors comment on the optimality of GD-$\\beta$ if this assumption is not satisfied?  \n(Wu and Xu 2020) *On the optimal weighted $\\ell_2$ regularization in overparameterized linear regression*.  \n(Li et al. 2023) *Transformers as algorithms: generalization and stability in in-context learning*.\n\n3. Some minor questions: \n* On page 5, why does the \"scratchpad\" take the specific form of (3.1)?\n* For the experiment reported in Table 1, what is the size of the trained model? Is the curriculum learning strategy in (Garg et al. 2022) also used?\n\n---\n**Post-rebuttal update:** the authors’ response addressed most of my concerns. I found that one of my criticisms is based on a misunderstanding of the optimal solution that the authors analyzed.  \nI have increased my score accordingly.\n\nQUESTIONS:\nSee Weaknesses above.\n\nLIMITATIONS:\nSee Weaknesses above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors investigate the in-context learning (ICL) of a linear transformer block (LTB), i.e., a single block comprised of a linear self-attention layer and a MLP layer. The task considered is ICL of linear regression with a Gaussian prior. Unlike earlier works that studied this problem, a Gaussian prior with a non-zero mean is used for generating the ICL tasks. For this setting, it is shown that the MLP layer is essential, as using only a self-attention layer results in a weaker approximation. Furthermore, the LTB at minimum is demonstrated to have a specific structure akin to one-step gradient descent with a learnable initialization, referred as GD-$\\beta$. The optimisation and statistical properties of the GD-$\\beta$ estimator is also widely studied.\n\nSTRENGTHS:\na) In-context learning of statistical tasks particularly the linear regression has received particularly interest. The authors have identified an interesting and relevant modification to this problem setting when the ICL tasks are generated from a more general setup Gaussian prior with non-zero mean. \n\nb) In this setting, an MLP layer and a skip connection are needed for approximation which is a novel finding. Even with these additional parameters it is shown that minimum of the task can be still be characterised and it corresponds to GD with learnable initialisation. \n\nc) The statistical and non-convex optimization aspects of the GD-$\\beta$ are interesting from a technical standpoint.\n\nWEAKNESSES:\nThe paper attributes that property of learning the MLP layer, however the skip connection is also equally important. Hence it is a bit of an misleading claim to attribute this to the MLP layer. Overall, the model restricted to only linear attention and also a single layer.\n\nQUESTIONS:\n1)  What happens if we include positional encoding (P) are also included?  It is possible to get an additional $P^{\\top} W_K^{\\top} W_QP$ term can come from the attention and can this overcome the need of positional encoding?\n\n2) Does the optimisation dynamics of GD-$\\beta$ offer any insights on the optimisation of the LTB ?   \n\n3) Does the results qualitatively change if the input sequence is generated from uncentered Gaussian $\\mathcal{N}( \\mu, H )$ instead of zero mean Gaussian prior ?\n\n\nMinor corrections: \n\n1) In l.251, it should be $ \\Tau^{\\top} $ in the expression of $W_K^{\\top} W_Q$.\n\n2) In l.263, an additional bracket in the expression of $\\Omega$.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the in-context learning of linear regression with a Gaussian prior that has a finite, non-zero mean across tasks. Specifically, the authors show that a linear transformer block (linear attention with an MLP layer) can achieve near-optimal Bayes risk for this task, whereas a linear attention-only block suffers from an approximation error. They further demonstrate that an LTB hypothesis class can be well-represented by a subset class of one-step gradient descent estimators, GD-$\\beta$, for the ICL task.\n\nSTRENGTHS:\nThis work delves into theoretically understanding the very relevant and hot topic of ICL. The problem motivation is clear and concise, and I really like how the authors essentially redo the previous results of LSA and GD-$0$ estimators with LTB and GD-$\\beta$, demonstrating the effective correspondence between them and showing that they achieve near Bayes optimal ICL risk. Overall, I think the authors do a great job of portraying the benefits of the MLP component in transformers, and it’s a good theoretical contribution, as most papers involve theory only from the LSA perspective.\n\nWEAKNESSES:\nI don’t see any major weaknesses in the paper as such, but I do have a few questions (see questions).\n\nQUESTIONS:\nHave the authors seen experimental evidence that trained LTB models converge to GD-$\\beta$ models, similar to how LSA converges to GD-$0$? I believe the correspondence established between the two is excellent, and theoretically studying the convergence is probably another work in itself, but seeing some experimental evidence would have been really cool.\n\nI believe the restrictions imposed on $d_k, d_v$ for the hypothesis classes are to ensure sufficient overparameterization, which is pretty common for theoretical works. However, I think in practice these dimensions are typically set to be less than $d$. Are these necessary conditions, or do the authors believe there are ways to override them for their analysis?\n\nThe constructions for Lemma $5.1$ and Theorem $5.3$ both operate on pairs of matrices composed together (if not more), which is also the common practice in the analysis that I’ve seen. Any reasons why the authors didn’t opt for analyzing with merged matrices, especially when the inner dimension of the product $W_p^TW_v​$, that is $d_v$​, is assumed to be larger than $d$? It doesn’t even enforce any additional rank constraints on the product matrix.\n\nIt was interesting to see that in your GPT-2 experiments, the full transformer model (with all the non-linearities) also relied heavily on the MLPs. How much of this do you think can be attributed to the presence of softmax only? Did you try experiments where you used LSA in GPT-2 training with MLPs?\n\nLIMITATIONS:\nYes, the authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-29"}