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

Authors: ruiqi  zhang
u c berkeley
jingfeng  wu
peter l bartlett
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:24:23

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.

Main paper: 
 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 ) M i=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 asx ∼ N (0, H), y | β, x ∼ N β ⊤ x, σ 2 ,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 ⊤ Ax. 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 = (A2 ) + , 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. 
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. 
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] E := X ⊤ x y ⊤ 0 ∈ R (d+1)×(M +1) . (3.1)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 asMLP (E) := W ⊤ 2 ReLU (W 1 E), where ReLU(•) refers to the entrywise rectified linear unit (ReLU) activation function, and d+1) are two weight matrices. The self-attention layer (we focus on the single-head version in this paper) is defined as d+1) are the key and query matrices, respectively, W P , W V ∈ R dv×(d+1) are the projection and value matrices, respectively, and M is a fixed masking matrix given byW 1 , W 2 ∈ R d f ×(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 ×(M := I M 0 0 0 ∈ R (M +1)×(M +1) . (3.2)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 setd 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 asEQUATION) 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 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 asW P , W V ∈ R dv×(d+1) , W K , W Q ∈ R d k ×(d+1) , and W 1 , W 2 ∈ R d f ×(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.A linear self-attention. We will also consider a linear self-attention (LSA) defined asEQUATIONwhere 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 (settingd 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 asF 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.• 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, R(fEQUATION) where E is defined in (3.1) and the expectation is over E (equivalent to over X, y, and x) and y. 
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, inf f ∈F LSA R (f ) -inf f ∈F LTB R (f ) ≥ max 2 3 (M + 1) , tr (HΨ) + σ 2 2 (M + 1) 2 tr ((HΨ) 2 ) ∥β * ∥ 2 H .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.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:E :=   X ⊤ x 1 ⊤ M 1 y ⊤ 0   ∈ R (d+2)×(M +1) (4.1)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 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. 
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 asEQUATIONwhere 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 asF 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 settingW ⊤ 2 W 1 = * * β ⊤ 1 , W ⊤ P W V = -I d 0 d 0 ⊤ d 1 , W ⊤ K W Q = Γ * 0 ⊤ d * ,where * denotes the entries that do not affect the model output (hence can be set to anything). Note that Γ and β are free provided thatd 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.Lemma 5.1. We haveF GD-β ⊆ F LTB . Therefore, inf f ∈F GD-β R(f ) ≥ inf f ∈F LTB R(f ).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 haveEQUATIONwhereΩ := [(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:β = β * + null (H) , Γ = Γ * + null H ⊗2 , where Γ * := ΨH 1 2 Ω -1 H -1 2 (5.3)and β * is the mean of the task parameter in Assumption 3.1 Here, null (H) := z ∈ R d : Hz = 0 is the null space of H, and null H ⊗2 = Z ∈ R d×d : HZH = 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:EQUATIONTheorem 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 H1 2 Ψ 1 2≥ 2. Then we have• The minimal ICL risk of F LTB and of F GD-β are equal,inf f ∈F LTB R(f ) = inf f ∈F GD-β R(f ) = RHS of (5.2)• Rewrite an LTB model as f LTB in (3.3) with parameters ( * denotes parameters that do not affect the output)W ⊤ 2 W 1 = * * γ ⊤ * , W ⊤ K W Q = V 11 * v ⊤ 12 * , W ⊤ 2 W 1 W ⊤ P W V = * * v ⊤ 21 v -1.Then the sufficient and necessary conditions forf LTB ∈ arg min f ∈F LTB R(f ) are v -1 ̸ = 0, v -1 v 12 ∈ null (H) , v 21 ∈ -v -1 β * + null (H) , γ ∈ β * + null (H) , v -1 V 11 ∈ Γ * ⊤ -v -1 β * v ⊤ 12 + null H ⊗2, 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. A line 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. 
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 EL (g; X) = R(g). Then,• The optimal estimator that minimizes the average linear regression riskL(•; 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) -EQUATION(ϕ 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 haveinf 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-β,EQUATIONwhere 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 = {HZH, 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,P H (β(t)) → P H (β * ) , P H ⊥ (β(t)) = P H ⊥ (β(0)) , P Z (Γ(t)) → P Z (Γ * ) , P Z ⊥ (Γ(t)) = P Z ⊥ (Γ(0))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. 
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.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β = β * + Ψ 1 2 θ, (A.1) where θ ∼ N (0, I d ) . (A.2)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. .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 ⊤ Ax. 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 byB ⊤ ⊗ A • C = ACB. We denote Λ := Ψ 1 2 HΨ 1 2and ϕ 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Ω := M + 1 M H 1 2 ΨH 1 2 + tr (HΨ) + σ 2 M • I d . (A.3)Note that, under this definition and our assumption on H and Ψ, we have Ω is invertible.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,inf f ∈F LTB R(f ) = σ 2 + tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 , (B.1)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 asf (E) = E + W ⊤ P W V EM E ⊤ W ⊤ K W Q E M -1,-1, 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 denoteW ⊤ P W V = * * u ⊤ 21 u -1 , W ⊤ K W Q = U 11 * u ⊤ 12 *, 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 asf (E) = u ⊤ 21 u -1 • EM M E ⊤ M • U 11 u ⊤ 12 • x = u ⊤ 21 • 1 M X ⊤ X • U 11 + u ⊤ 21 • 1 M X ⊤ y • u ⊤ 12 + u -1 • 1 M y ⊤ X • U 11 + u -1 • 1 M y ⊤ y • u ⊤ 12 • x.Step 1: simplify the risk function. We use β to denote the task parameter. From the Assumption 3.1 and Definition A.1, we havey = X β + ε, y = β, x + ε, β ∼ N (β * , Ψ) , β = β * + Ψ 1 2 θ;andX[i], x i.i.d. ∼ N (0, H) , ε[i], ε i.i.d. ∼ N 0, σ 2 , θ ∼ N (0, I d ) .Then the model output can be written asf (E) = u ⊤ 21 • 1 M X ⊤ X • U 11 + u ⊤ 21 • 1 M X ⊤ y • u ⊤ 12 + u -1 • 1 M y ⊤ X • U 11 + u -1 • 1 M y ⊤ y • u ⊤ 12 • x = u 21 + u -1 β ⊤ • 1 M X ⊤ X • U 11 + βu ⊤ 12 • x + u ⊤ 21 • 1 M X ⊤ ε • u ⊤ 12 + u -1 • 1 M ε ⊤ X • U 11 + u -1 • 2 M ε ⊤ X βu ⊤ 12 + 1 M ε ⊤ ε • u -1 u ⊤ 12 • x.To simplify the presentation, we denotez ⊤ 1 = u 21 + u -1 β ⊤ • 1 M X ⊤ X • U 11 + βu ⊤ 12 , z ⊤ 2 = u ⊤ 21 • 1 M X ⊤ ε • u ⊤ 12 + u -1 • 1 M ε ⊤ X • U 11 + u -1 • 2 M ε ⊤ X βu ⊤ 12 z ⊤ 3 = 1 M ε ⊤ ε • u -1 u ⊤ 12 .Since x, X, ε, β are independent, we haveEQUATION(ε is independent from other variables and zero-mean)= E z 1 + z 2 + z 3 -β, x 2 + σ 2 = H, E z 1 + z 2 + z 3 -β z 1 + z 2 + z 3 -β ⊤ + σ 2 .Note that z 1 does not contain ε, z 2 is a linear form of ε, andz 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 R (f ) -σ 2 = H, E z 1 -β z 1 -β ⊤ S1 + H, Ez 2 z ⊤ 2 S2 + H, Ez 3 z ⊤ 3 S3 + 2 H, E z 1 -β z ⊤ 3 S4 . (B.2)Step 2: compute S 1 . By Lemma H.4 and that X is independent of all other random variables, we haveH, Ez 1 z ⊤ 1 =Etr U 11 + βu ⊤ 12 ⊤ • 1 M X ⊤ X • u 21 + u -1 β u 21 + u -1 β ⊤ • 1 M X ⊤ X • U 11 + βu ⊤ 12 H = M + 1 M E β tr U 11 + βu ⊤ 12 ⊤ • H • u 21 + u -1 β u 21 + u -1 β ⊤ • H • U 11 + βu ⊤ 12 H + 1 M E β tr tr u 21 + u -1 β u 21 + u -1 β ⊤ H U 11 + βu ⊤ 12 ⊤ H U 11 + βu ⊤ 12 H . We denote b := u 21 + u -1 β * ∈ R d , A := U 11 + β * u ⊤ 12 ∈ R d×d , (B.3) then applying β = β * + Ψ 1 2 θ, we get H, Ez 1 z ⊤ 1 = M + 1 M E θ∼N (0,I d ) tr A + Ψ 1 2 θu ⊤ 12 ⊤ • H • b + u -1 Ψ 1 2 θ b + u -1 Ψ 1 2 θ ⊤ • H • A + Ψ 1 2 θu ⊤ 12 H + 1 M E θ∼N (0,I d ) tr b + u -1 Ψ 1 2 θ b + u -1 Ψ 1 2 θ ⊤ H tr A + Ψ 1 2 θu ⊤ 12 ⊤ H A + Ψ 1 2 θu ⊤ 12 H .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 haveH, Ez 1 z ⊤ 1 = M + 1 M tr A ⊤ Hbb ⊤ HAH + 1 M tr bb ⊤ H tr A ⊤ HAHThe leading term+T 2 + T 4 , (B.4)where T 2 and T 4 are the second order and the fourth order term, respectively. More concretely, the second order term isT 2 = M + 1 M Etr u -1 A ⊤ Hb θ ⊤ Ψ 1 2 HΨ 1 2 θu ⊤ 12 H + u -1 u 12 θ ⊤ Ψ 1 2 HΨ 1 2 θb ⊤ HAH + u -1 u 12 θ ⊤ Ψ 1 2 Hb θ ⊤ Ψ 1 2 HAH + u -1 A ⊤ HΨ 1 2 θb ⊤ HΨ 1 2 θu ⊤ 12 H + u 12 θ ⊤ Ψ 1 2 Hbb ⊤ HΨ 1 2 θu ⊤ 12 H + u 2 -1 A ⊤ HΨ 1 2 θ θ ⊤ Ψ 1 2 HAH + 1 M E u 2 -1 θ ⊤ Ψ 1 2 HΨ 1 2 θ • tr A ⊤ HAH + b ⊤ Hb • tr u 12 θ ⊤ Ψ 1 2 HΨ 1 2 θu ⊤ 12 H + 2u -1 b ⊤ HΨ 1 2 θ • tr u 12 θ ⊤ Ψ 1 2 HAH + 2u -1 b ⊤ HΨ 1 2 θ • tr A ⊤ HΨ 1 2 θu ⊤ 12 H = M + 1 M 2u -1 tr(Ψ 1 2 HΨ 1 2 ) • u ⊤ 12 HA ⊤ Hb + 2u -1 b ⊤ HΨHAHu 12 + b ⊤ HΨHb • u ⊤ 12 Hu 12 + u 2 -1 tr A ⊤ HΨHAH + 1 M u 2 -1 tr Ψ 1 2 HΨ 1 2 • tr A ⊤ HAH + tr Ψ 1 2 HΨ 1 2 • b ⊤ Hb • u ⊤ 12 Hu 12 + 4u -1 • b ⊤ HΨ 1 2 HAHu 12 = 2(M + 1) M u -1 b ⊤ [tr(HΨ)H + HΨH] AHu 12 + b ⊤ M + 1 M HΨH + 1 M tr(HΨ)H b • u ⊤ 12 Hu 12 + u 2 -1 tr A ⊤ M + 1 M HΨH + 1 M tr(HΨ)H AH + 4 M u -1 • b ⊤ HΨHAHu 12 . (B.5)The fourth order term isT 4 = M + 1 M Etr u 2 -1 u 12 θ ⊤ Ψ 1 2 HΨ 1 2 θ θ ⊤ Ψ 1 2 HΨ 1 2 θu ⊤ 12 H + 1 M Etr u 2 -1 θ ⊤ Ψ 1 2 HΨ 1 2 θ • tr u 12 θ ⊤ Ψ 1 2 HΨ 1 2 θu ⊤ 12 H = M + 2 M u 2 -1 u ⊤ 12 Hu 12 • E θ ⊤ Ψ 1 2 HΨ 1 2 θ 2 = M + 2 M u 2 -1 • 2tr (HΨHΨ) + tr (HΨ) 2 • u ⊤ 12 Hu 12 . (B.6)For the cross term in S 1 , we haveH, Ez 1 β ⊤ =E u 21 + u -1 β ⊤ • 1 M X ⊤ X • U 11 + βu ⊤ 12 H β (B.7) =E u 21 + u -1 β ⊤ H U 11 + βu ⊤ 12 H β (From the distribution of X) =E b + u -1 Ψ 1 2 θ ⊤ H A + Ψ 1 2 θu ⊤ 12 H β * + Ψ 1 2 θ (By (B.3)) =b ⊤ HAHβ * + E u -1 θ ⊤ Ψ 1 2 HΨ 1 2 θu ⊤ 12 Hβ * + u -1 θ ⊤ Ψ 1 2 HAHΨ 1 2 θ + b ⊤ HΨ 1 2 θu ⊤ 12 HΨ 1 2 θ =b ⊤ HAHβ * + u -1 tr (HΨ) • u ⊤ 12 Hβ * + u -1 tr (HAHΨ) + u ⊤ 12 HΨHb, (B.8)where the last row comes from the fact that θ ∼ N (0, I d ) . Moreover, we haveH, E β β ⊤ = β * ⊤ Hβ * + tr (HΨ) . (B.9)Combining (B.4), (B.5), (B.6), (B.8), (B.9), we haveS 1 = M + 1 M tr A ⊤ Hbb ⊤ HAH + 1 M tr bb ⊤ H tr A ⊤ HAH + M + 2 M u 2 -1 • 2tr (HΨHΨ) + tr (HΨ) 2 • u ⊤ 12 Hu 12 + 2(M + 1) M u -1 b ⊤ [tr(HΨ)H + HΨH] AHu 12 + u 2 -1 tr A ⊤ M + 1 M HΨH + 1 M tr(HΨ)H AH + 4 M u -1 • b ⊤ HΨHAHu 12 -2 b ⊤ HAHβ * + u -1 tr (HΨ) • u ⊤ 12 Hβ * + u -1 tr (HAHΨ) + u ⊤ 12 HΨHb + β * ⊤ Hβ * + tr (HΨ) + b ⊤ M + 1 M HΨH + 1 M tr(HΨ)H b • u ⊤ 12 Hu 12 .(B.10)Step 3: other terms. Let us compute S 2 , S 3 and S 4 . Using definitions, we rewrite z 2 asz ⊤ 2 = u ⊤ 21 • 1 M X ⊤ ε • u ⊤ 12 + u -1 • 1 M ε ⊤ X • U 11 + u -1 • 2 M ε ⊤ X βu ⊤ 12 = b ⊤ • 1 M X ⊤ ε • u ⊤ 12 + u -1 • 1 M ε ⊤ X • A + u -1 • 2 M ε ⊤ XΨ 1 2 θu ⊤ 12 . (B.11)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 haveS 2 = H, Ez 2 z ⊤ 2 = 1 M 2 E b ⊤ X ⊤ ε • u ⊤ 12 Hu 12 • ε ⊤ X • b + u 2 -1 ε ⊤ XAHA ⊤ X ⊤ ε + 2u -1 • ε ⊤ XAHu 12 • ε ⊤ Xb + 4u 2 -1 • ε ⊤ XΨ 1 2 θ • u ⊤ 12 Hu 12 • θ ⊤ Ψ 1 2 X ⊤ ε = σ 2 M b ⊤ Hb • u ⊤ 12 Hu 12 + u 2 -1 tr HAHA ⊤ + 2u -1 u ⊤ 12 HA ⊤ Hb + 4u 2 -1 tr (HΨ) u ⊤ 12 Hu 12 (B.12)For the other two terms, we haveEQUATION)andEQUATION)Step 4: combine all parts. Combining the four parts (B.10), (B.12), (B.13) and (B.14), we haveR (f ) -σ 2 = S 1 + S 2 + S 3 + S 4We remark 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 asEQUATION)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:Ω := M + 1 M H 1 2 ΨH 1 2 + tr (HΨ) + σ 2 M I d . (B.16)So we can write I asI = b ⊤ H 1 2 ΩH 1 2 b • u ⊤ 12 Hu 12 + u 2 -1 tr A ⊤ H 1 2 ΩH 1 2 AH + 2u -1 b ⊤ H 1 2 ΩH 1 2 AHu 12 -2u -1 tr (HAHΨ) -2u ⊤ 12 HΨHb = tr H 1 2 u 12 b ⊤ H 1 2 + u -1 H 1 2 A ⊤ H 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω • H 1 2 u 12 b ⊤ H 1 2 + u -1 H 1 2 A ⊤ H 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 ≥ -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 (B.17) = -tr H 1 2 ΨH 1 2 2 Ω -1 ,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 that4u -1 • b ⊤ HΨHAHu 12 ≥ -4 b ⊤ H 1 2 F • u -1 H 1 2 ΨH 1 2 F • H 1 2 AH 1 2 F • H 1 2 u 12 F ≥ -b ⊤ H 1 2 2 F H 1 2 AH 1 2 2 F -4 u -1 H 1 2 ΨH 1 2 2 F H 1 2 u 12 2 F ≥ -b ⊤ Hb • tr A ⊤ HAH -4u 2 -1 • tr (HΨHΨ) • u ⊤ 12 Hu 12 , (B.18)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 ⊤ Hb, which can be easily minimized. More concretely, we haveIII + M M + 1 β * -tr (HΨ) + σ 2 u -1 u 12 ⊤ H β * -tr (HΨ) + σ 2 u -1 u 12 = A ⊤ Hb - M M + 1 β * -tr (HΨ) + σ 2 u -1 u 12 ⊤ M + 1 M H • A ⊤ Hb - M M + 1 β * -tr (HΨ) + σ 2 u -1 u 12 ≥ 0. (B.19)Therefore, one hasIII ≥ - M M + 1 β * -tr (HΨ) + σ 2 u -1 u 12 ⊤ H β * -tr (HΨ) + σ 2 u -1 u 12 .Combining the three parts above, one hasR (f ) -σ 2 ≥ IV -tr H 1 2 ΨH 1 2 2 Ω -1 - M M + 1 β * -tr (HΨ) + σ 2 u -1 u 12 ⊤ H β * -tr (HΨ) + σ 2 u -1 u 12 = tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 inf f ∈F LTB R(f )-σ 2 + 1 M + 1 β * ⊤ Hβ * - 2 tr (HΨ) + σ 2 M + 1 β * ⊤ H (u -1 u 12 ) + 2tr (HΨHΨ) + 3M + 2 M (M + 1) tr (HΨ) + σ 2 2 (u -1 u 12 ) ⊤ H (u -1 u 12 ) .Therefore, one hasR (f ) -inf f ∈F LTB R(f ) ≥ 1 M + 1 β * ⊤ Hβ * - 2 tr (HΨ) + σ 2 M + 1 β * ⊤ H (u -1 u 12 ) + 2tr (HΨHΨ) + 3M + 2 M (M + 1) tr (HΨ) + σ 2 2 (u -1 u 12 ) ⊤ H (u -1 u 12 ) .The right hand side in the above inequality is a quadratic function of u -1 u 12 , so we can take its global minimizer:u -1 u 12 = (tr(HΨ)+σ 2 ) M +1 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2 β *to lower bound the risk gap asR (f ) -inf f ∈F LTB R(f ) ≥    1 M + 1 - (tr(HΨ)+σ 2 ) 2 (M +1) 2 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2    ∥β * ∥ 2 H .Finally, to simplify the results, we notice that on the one hand,(tr(HΨ)+σ 2 ) 2 (M +1) 2 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2 ≤tr (HΨ) + σ 2 2 2(M + 1) 2 tr (HΨHΨ) .On the other hand, one has(tr(HΨ)+σ 2 ) 2 (M +1) 2 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2 ≤ M (M + 1)(3M + 2) ≤ 1 3(M + 1).Therefore, we haveR (f ) -inf f ∈F LTB R(f ) ≥ max 2 3(M + 1) , 1 M + 1 - tr (HΨ) + σ 2 2 2(M + 1) 2 tr (HΨHΨ) • ∥β * ∥ 2 H . 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 R (f ) -inf f ∈F LTB R(f ) ≥ 1 M + 1 - tr (HΨ) + σ 2 2 2(M + 1) 2 tr (HΨHΨ) • ∥β * ∥ 2 H ≥ tr (HΨ) + σ 2 2 (M + 1) 2 tr (HΨHΨ) • ∥β * ∥ 2 H .Therefore, we finally haveR (f ) -inf f ∈F LTB R(f ) ≥ max 2 3(M + 1), tr (HΨ) + σ 2 2 (M + 1) 2 tr (HΨHΨ)• ∥β * ∥ 2 H .Note that this holds for an arbitrary f ∈ F LSA , so the proof finishes by taking infimum on the left hand side.Proof. Recall that from Assumption 3.1,x ∼ N (0, H), y = x ⊤ β + ε, X[i] ∼ N (0, H), y = X β + ε, where β ∼ N (β * , Ψ) , ε ∼ N 0, σ 2 , ε ∼ N 0, σ 2 • I M .Step 1: compute the risk function. From the independence of ε, ε with other random variables, we haveR M (β, Γ) = E x ⊤ β -β + ε + x ⊤ Γ • X ⊤ X M β -β -x ⊤ Γ • X ⊤ ε M 2 = E x ⊤ I d -Γ • X ⊤ X M • β -β 2 + 1 M 2 E x ⊤ ΓX ⊤ ε 2 + σ 2 = H, E I d -Γ • X ⊤ X M ⊗2 • E β -β ⊗2 + 1 M 2 E x ⊤ ΓX ⊤ ε 2 + σ 2 , (C.1)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β ∼ N (β * , Ψ) , we have E β -β ⊗2 = (β -β * ) ⊗2 + Ψ.Moreover, using the followingx ∼ N (0, H) , X[i] ∼ N (0, H) , ε ∼ N 0, σ 2 • I M ,we have thatE x ⊤ ΓX ⊤ ε 2 = σ 2 Etr XΓ ⊤ xx ⊤ ΓX ⊤ = M σ 2 • HΓH, Γ ⊤ .Bridging the two terms above into (C.1), we getR M (β, Γ) = H, E I d -Γ • X ⊤ X M ⊗2 • (β -β * ) ⊗2 + H, E I d -Γ • X ⊤ X M ⊗2 • Ψ + σ 2 M ΓHΓ ⊤ + σ 2 = E I d -Γ • X ⊤ X M ⊗2 • H, (β -β * ) ⊗2 V1 + H, E I d -Γ • X ⊤ X M ⊤ ⊗2 • Ψ + σ 2 M ΓHΓ ⊤ V2 +σ 2 (C.2)First, let's compute V 1 . From the definition above, we haveV 1 = (β -β * ) ⊤ H Γ (β -β * ) ,whereH Γ : = E I d -Γ • X ⊤ X M ⊤ ⊗2 • H (C.3) = E I d -Γ • X ⊤ X M ⊤ H I d -Γ • X ⊤ X M (the definition of the tensor product) = H -H Γ + Γ ⊤ H + tr HΓ ⊤ HΓ M H + M + 1 M HΓ ⊤ HΓH (X[i] ∼ N (0, H) and Lemma H.4) = (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH ⪰ 0. (C.4)Then, let's compute V 2 . We haveE I d -Γ • X ⊤ X M ⊗2 • Ψ + σ 2 M ΓHΓ ⊤ = E I d -Γ • X ⊤ X M Ψ I d -Γ • X ⊤ X M ⊤ + σ 2 M ΓHΓ ⊤ = Ψ -ΓHΨ -ΨHΓ ⊤ + Γ tr (HΨ) + σ 2 M • H + M + 1 M HΨH Γ ⊤ . (X[i] i.i.d.∼ N (0, H) and the Lemma H.4)From the definition of Ω in (A.3), we know that it is invertible. Moreover, it holds thatE I d -Γ • X ⊤ X M ⊗2 • Ψ + σ 2 M ΓHΓ ⊤ = Ψ -ΓHΨ -ΨHΓ ⊤ + ΓH 1 2 ΩH 1 2 Γ ⊤ = ΓH 1 2 -ΨH 1 2 Ω -1 Ω ΓH 1 2 -ΨH 1 2 Ω -1 ⊤ + Ψ -ΨH 1 2 Ω -1 H 1 2 Ψ.Therefore, we haveV 2 = tr H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ + tr (HΨ) -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2Step 2: solve the global minimum. Now we haveR (β, Γ) -σ 2 = (β -β * ) ⊤   (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH   (β -β * ) + tr H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ + tr (HΨ) -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 (C.5) ≥ tr (HΨ) -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 = tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 ,where the last line comes from the fact that Ω and H inf f ∈F GD-β R(f ) -σ 2 = inf β,Γ R (β, Γ) -σ 2 ≥ tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 .On the other hand, we takeβ = β * , Γ = Γ * := ΨH 1 2 Ω -1 H -1 2 , (C.6)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 haveR (β * , Γ * ) -σ 2 = tr H 1 2 Γ * H 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 Γ * H 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ + tr (HΨ) -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 . Since H 1 2 Γ * H 1 2 -H 1 2 -H 1 2 ΨH 1 2 Ω -1 = H 1 2 ΨH 1 2 Ω -1 H -1 2 H 1 2 -H 1 2 ΨH 1 2 Ω -1 = Ω -1 H 1 2 ΨH 1 2 -Ω -1 H 1 2 ΨH 1 2 H -1 2 H 1 2(Ω and H1 2 ΨH 1 2 commute) = Ω -1 H 1 2 ΨH 1 2 -Ω -1 H 1 2 ΨH 1 2 = 0 d×d .. (H 1 2 H -1 2 H 1 2 = H 1 2 )Therefore, we haveR (β * , Γ * )-σ 2 = tr (HΨ)-tr H 1 2 ΨH 1 2 2 Ω -1 ≥ inf f ∈F GD-β R(f )-σ 2 = inf β,Γ R (β, Γ)-σ 2 .Combining both directions, we conclude thatinf f ∈F GD-β R(f ) -σ 2 = inf β,Γ R (β, Γ) -σ 2 = tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 . (C.7)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 R (β, Γ) = inf f ∈F GD-β R(f ) ⇐⇒        V 1 = (β -β * ) ⊤ H Γ (β -β * ) = 0 tr H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ = 0Since Ω is invertible, the second equation is equivalent toH 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 = 0 d×d ,which is a linear system of Γ. Since Γ * is proved to be one solution, all solutions of this equation isΓ = Γ * + Z ∈ R d×d : H 1 2 ZH 1 2 = 0 d×d = Γ * + Im H ⊗2 ,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 thatH Γ = H -H Γ * + Γ * ⊤ M H + tr HΓ * ⊤ M HΓ * M H + M + 1 M HΓ * ⊤ M HΓ * H.Consider the following inequality:I d -H 1 2 Γ * + Γ * ⊤ H 1 2 + tr HΓ * ⊤ M HΓ * M I d + M + 1 M H 1 2 Γ * ⊤ M HΓ * H 1 2 ⪰ I d -H 1 2 Γ * + Γ * ⊤ H 1 2 + M + 1 M H 1 2 Γ * ⊤ M HΓ * H 1 2 = M M + 1 I d - M + 1 M H 1 2 Γ * H 1 2 • M M + 1 I d - M + 1 M H 1 2 Γ * H 1 2 ⊤ + M M + 1 I d ≻ 0 d×d .Therefore, we have(β -β * ) ⊤ H Γ (β -β * ) = 0 ⇐⇒ (β -β * ) ⊤ H 1 2 I d -H 1 2 Γ * + Γ * ⊤ H 1 2 + tr HΓ * ⊤ M HΓ * M I d + M + 1 M H 1 2 Γ * ⊤ M HΓ * H 1 2 H 1 2 (β -β * ) = 0 ⇐⇒ H 1 2 (β -β * ) = 0 d ⇐⇒ β = β * + null H 1 2 ⇐⇒ β = β * + null (H) ,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β = β * + null (H) , Γ = Γ * + Im H ⊗2 .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 havef β,Γ (E) = β - Γ M X ⊤ (Xβ -y) , x = β * - Γ * M X ⊤ (Xβ * -y) , x + h ⊤ x - 1 M x ⊤ ZX ⊤ Xβ * - 1 M x ⊤ ZX ⊤ Xh - 1 M x ⊤ Γ * X ⊤ Xh + x ⊤ ZX ⊤ y. Notice that X[i], x i.i.d.∼ N (0 d , H) , we haveE h ⊤ x 2 = h ⊤ Hh = 0 since h ∈ null (H). Therefore, we know h ⊤ x = 0 almost surely, and similarly, Xh = 0 d almost surely. Finally, we haveE 1 M x ⊤ ZX ⊤ • 1 M x ⊤ ZX ⊤ ⊤ = 1 M Etr [HZH] = 0,which imples x ⊤ ZX ⊤ = 0 d almost surely. Therefore, we concludef β,Γ (E) = f β * ,Γ * (E) almost surely.Proof. Recall the definition of Linear Transformer Block class:f LTB : R (d+1)×(M +1) → R E → W ⊤ 2 W 1 E + W ⊤ P W V EM E ⊤ W ⊤ K W Q E M -1,-1, where E is the token matrix defined by (3.1). Let's take an arbitrary function f in the LTB class with trainable matricesW K , W Q ∈ R d k ×(d+1) , W P , W V ∈ R dv×(d+1) , W 1 , W 2 ∈ R d f ×(d+1) .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 denoteW ⊤ 2 W 1 = * * γ ⊤ * , W ⊤ 2 W 1 W ⊤ P W V = * * v ⊤ 21 v -1 , W ⊤ K W Q = V 11 * v ⊤ 12 * , (D.1) 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 asf (E) = γ ⊤ x + v ⊤ 21 v -1 • EM M E ⊤ M • V 11 v ⊤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 haveV = tr H 1 2 v 12 b ⊤ H 1 2 + v -1 H 1 2 A ⊤ H 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 v 12 b ⊤ H 1 2 + v -1 H 1 2 A ⊤ H 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 ≥ -tr H 1 2 ΨH 1 2 Ω -1 H 1 2 ΨH 1 2 = -tr H 1 2 ΨH 1 2 2 Ω -1 ,where the last line comes from the fact that Ω and H1 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, since4v -1 • b ⊤ HΨHAHv 12 ≥ -4 b ⊤ H 1 2 F • v -1 H 1 2 ΨH 1 2 F • H 1 2 AH 1 2 F • H 1 2 v 12 F ≥ -b ⊤ H 1 2 2 F H 1 2 AH 1 2 2 F -4 v -1 H 1 2 ΨH 1 2 2 F H 1 2 v 12 2 F = -b ⊤ Hb • tr A ⊤ HAH -4v 2 -1 • tr (HΨHΨ) • v ⊤ 12 Hv 12 , (D.4)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 haveVII + M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 ⊤ H β * -γ -tr (HΨ) + σ 2 v -1 v 12 = A ⊤ Hb - M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 ⊤ M + 1 M H • A ⊤ Hb - M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 ≥ 0, (D.5) which implies VII ≥ - M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 ⊤ H β * -γ -tr (HΨ) + σ 2 v -1 v 12 .Combining the three parts above, one hasR (f ) -σ 2 ≥ VIII -tr H 1 2 ΨH 1 2 2 Ω -1 - M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 ⊤ H β * -γ -tr (HΨ) + σ 2 v -1 v 12 = tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 inf f ∈F GD-β R(f )-σ 2 + 1 M + 1 (β * -γ) ⊤ H (β * -γ) - 2 tr (HΨ) + σ 2 M + 1 (β * -γ) ⊤ H (v -1 v 12 ) + 2tr (HΨHΨ) + 3M + 2 M (M + 1) tr (HΨ) + σ 2 2 (v -1 v 12 ) ⊤ H (v -1 v 12 ) .Here, the global minimum of GD-β class is taken from Theorem 5.2, whose proof does not reply on the proof here. Therefore, one hasR (f ) -inf f ∈F GD-β R(f ) ≥ 1 M + 1 (β * -γ) ⊤ H (β * -γ) - 2 tr (HΨ) + σ 2 M + 1 (β * -γ) ⊤ H (v -1 v 12 ) + 2tr (HΨHΨ) + 3M + 2 M (M + 1) tr (HΨ) + σ 2 2 (v -1 v 12 ) ⊤ H (v -1 v 12 ) .The right hand side in the above inequality is a quadratic function of v -1 v 12 , so we can take its global minimizer:v -1 v 12 = (tr(HΨ)+σ 2 ) M +1 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2 (β * -γ)to lower bound the risk gap asR (f ) -inf f ∈F GD-β R(f ) ≥    1 M + 1 - (tr(HΨ)+σ 2 ) 2 (M +1) 2 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2    ∥β * -γ∥ 2 H ≥ 0.Note that, taking infimum on the left hand side, we getinf f ∈F LTB R (f ) -inf f ∈F GD-β R(f ) ≥ 0.On the other hand, in the main text we have showed thatF GD-β ⊂ F LTB , which implies inf f ∈F LTB R (f ) -inf f ∈F GD-β R(f ) ≤ 0.Therefore, we concludeinf f ∈F LTB R (f ) = inf f ∈F GD-β R(f ) = tr (HΨ) -tr H 1 2 ΨH 1 2 2 Ω -1 ,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 areH 1 2 v 12 b ⊤ H 1 2 + v -1 H 1 2 A ⊤ H 1 2 = H 1 2 ΨH 1 2 Ω -1 , (D.6) b ⊤ H 1 2 2 F H 1 2 AH 1 2 2 F = 4 v -1 H 1 2 ΨH 1 2 2 F H 1 2 v 12 2 F , (D.7) v -1 • b ⊤ HΨHAHv 12 = b ⊤ H 1 2 F • v -1 H 1 2 ΨH 1 2 F • H 1 2 AH 1 2 F • H 1 2 v 12 F , (D.8) H 1 2 A ⊤ Hb - M M + 1 β * -γ -tr (HΨ) + σ 2 v -1 v 12 = 0 d×d , (D.9) v -1 v 12 = (tr(HΨ)+σ 2 ) M +1 2tr (HΨHΨ) + 3M +2 M (M +1) (tr (HΨ) + σ 2 ) 2 (β * -γ) (D.10) ∥β * -γ∥ H = 0. (D.11)Let's first verify the necessary conditions of the system defined above. From (D.11) and Lemma H.5, we know γ ∈ β * + null H1 2 = β * + null (H) . Then, we have v -1 v 12 ∈ null (H) . Then, in (D.7), we know b ⊤ H 1 2 2 F H 1 2 AH 1 2 2 F = 4 v -1 H 1 2 ΨH 1 2 2 F H 1 2 v 12 2 F = 4 H 1 2 ΨH 1 2 2 F v -1 H 1 2 v 12 2 F = 0, which implies either H 1 2 b = 0 d or H 1 2 AH 1 2 = 0 d×d . If we assume H 1 2 AH 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(H1 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 H1 2 ΨH 1 2 Ω -1 ≥ 2, while H 1 2 v 12 b ⊤ H 1 2is a matrix of rank at most one, which contradicts with (D.6). Therefore, we haveH 1 2 b = 0, (D.12) v -1 H 1 2 A ⊤ H 1 2 = H 1 2 ΨH 1 2 Ω -1 . (D.13)Recall in Theorem 5.2, we have definedΓ * := ΨH 1 2 Ω -1 H -1 2 .Simple calculation showsH 1 2 Γ * H 1 2 = H 1 2 ΨH 1 2 Ω -1 H -1 2 H 1 2 = Ω -1 H 1 2 ΨH 1 2 H -1 2 H 1 2 = Ω -1 H 1 2 ΨH 1 2 = H 1 2 ΨH 1 2 Ω -1 ,where the second and the last equalities come from the fact that Ω and H1 2 ΨH 1 2commute, 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 isv -1 A ∈ Γ * ⊤ + Z : H 1 2 ZH 1 2 = 0 d×d = Γ * ⊤ + Im H ⊗2 .The solution to (D.12) isv 21 = -v -1 β * + null H 1 2 = -v -1 β * + null (H) .Therefore, we know the necessary conditions for (D.6) to (D.11) are               v -1 ̸ = 0, v -1 v 12 ∈ null (H) , v 21 = -v -1 β * + null (H) , v -1 V 11 ∈ Γ * ⊤ -v -1 β * v ⊤ 12 + Im H ⊗2 , γ = β * + null (H) (D.14)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 :v -1 ̸ = 0, v 12 = 0 d , v 21 = -v -1 β * , V 11 = 1 v -1 • Γ * ⊤ , γ = β * . (D.15)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 hasx ⊤ a = x ⊤ i a = 0 since x, x i i.i.d.∼ N (0 d , H) . We denotev 21 = -v -1 β * + a 1 , v -1 V 11 = Γ * ⊤ -v -1 β * v ⊤ 12 + Z, γ = β * + a 2 ,where a 1 , a 2 ∈ H, H1 2 ZH 1 2 = 0 d×d . Then, we have f (E) = W ⊤ 2 W 1 E + W ⊤ P W V EM E ⊤ W ⊤ K W Q E M -1,-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 inf f ∈F GD-β R(f ) -σ 2 = tr H 1 2 ΨH 1 2 -tr H 1 2 ΨH 1 2 2 Ω -1 ,whereΩ := M + 1 M H 1 2 ΨH 1 2 + tr (HΨ) + σ 2 M • I d .Therefore, we haveinf f ∈F GD-β R(f ) -σ 2 = tr H 1 2 ΨH 1 2 • Ω -H 1 2 ΨH 1 2 Ω -1 = 1 M tr Ω -1 H 1 2 ΨH 1 2 • H 1 2 ΨH 1 2 + tr H 1 2 ΨH 1 2 + σ 2 • I d . Since tr H 1 2 ΨH 1 2 + σ 2 •I d ⪯ H 1 2 ΨH 1 2 + tr H 1 2 ΨH 1 2 + σ 2 •I d ⪯ 2 tr H 1 2 ΨH 1 2 + σ 2 •I d , we have inf f ∈F GD-β R(f ) -σ 2 ≃ tr H 1 2 ΨH 1 2 + σ 2 M tr Ω -1 H 1 2 ΨH 1 2 = φ • d i=1 ϕ i M +1 M ϕ i + φ ≃ d i=1 min ϕ i , φ .Therefore, we conclude.F Proof of Theorem 6.1Proof. From the classical bias-variance decomposition, we knowL (g; X) := E (g(X, y, x) -y) 2 | X = E (g(X, y, x) -E [y | X, y, x]) 2 | X + E (E [y | X, y, x] -y) 2 | Xsince 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.,y Bayes = E [y | X, y, x] = E β | X, y, x , x .Since β ∼ N (β * , Ψ) , we have there exists a random vector θ ∼ N (0d , I d ) such that β = β * + Ψ 1 2θ almost surely. Therefore, one hasy Bayes = ⟨β * , x⟩ + E θ | X, y, x , Ψ 1 2 x .To compute E θ | X, y, x , it suffices to solve the posterior distribution of β given X, y. From θ ∼ N (0 d , I d ), we haveP θ | X, y, x ∝ P θ P y | X, θ ∝ exp   - y -X β * + Ψ 1 2 θ 2 2 2σ 2 - 1 2 θ ⊤ • θ    ∝ exp - 1 2σ 2 θ ⊤ Ψ 1 2 X ⊤ XΨ 1 2 + σ 2 I d θ -2 (y -Xβ * ) ⊤ XΨ 1 2 θ .Note that, the function above matches the probability density function of a multivariate Gaussian distrbution. Therefore, the posterior mean is given byE [θ | X, y, x] = Ψ 1 2 X ⊤ XΨ 1 2 + σ 2 I d -1 Ψ 1 2 X ⊤ (y -Xβ * ) .Therefore, we concludey Bayes = x ⊤ Ψ 1 2 Ψ 1 2 X ⊤ XΨ 1 2 + σ 2 I d -1 Ψ 1 2 X ⊤ (y -Xβ * ) + x ⊤ β * . (F.1)Proof. Let's first do a variable transformation. We denoteX = XΨ 1 2 , x = Ψ 1 2 x, y = y -Xβ * , Λ = Ψ 1 2 HΨ 1 2 . (F.2)Then, we knowX[i], x i.i.d.∼ N (0 d , Λ) . We can write the Bayesian optimal estimator in (F.1) asy Bayes = x ⊤ X ⊤ X + σ 2 I d -1 X y + x ⊤ β * .The true label isy = x ⊤ β + ε = x ⊤ β * + Ψ 1 2 θ .Therefore, we haveL ( y Bayes ; X) -σ 2 = E x ⊤ X ⊤ X + σ 2 I d -1 X y + x ⊤ β * -x ⊤ β * + Ψ 1 2 θ 2 = E x ⊤ X ⊤ X + σ 2 I d -1 X y -θ 2 = E θ -θ 2 Λ , where θ := X ⊤ X + σ 2 I d -1 X y.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θ -θ 2 H ≃ σ 2 + i>k * ϕ i M 2 θ 2 H -1 0:k * + θ 2 H k * :∞ + σ 2 k * M + M i>k * ϕ 2 i σ 2 + i>k * ϕ i ,whereϕ 1 ≥ ϕ 2 ≥ ... ≥ ϕ d ≥ 0 are ordered eigenvalues of Λ. k * := min k : ϕ k ≥ c • σ 2 + i>k * ϕ i Mand 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 ofH = Q • diag (ϕ 1 , ϕ 2 , ..., ϕ d ) • Q ⊤ ,where Q is an orthogonal matrix, then H 0:k * and H k * :∞ are given byH 0:k * = Q • diag (ϕ 1 , ϕ 2 , ..., ϕ k * , 0, 0, ..., 0) • Q ⊤ , H k * :∞ = H -H 0:k * .Taking expectation over θ ∼ N (0, I d ) , we haveL ( y Bayes ; X) -σ 2 = E θ∼N (0,I d ) θ Bayes -θ 2 H = σ 2 + i>k * ϕ i M 2 • i≤k * 1 ϕ i + i>k * ϕ i + σ 2 k * M + M i>k * ϕ 2 i σ 2 + i>k * ϕ iNow we simplify this expression. First, we defineφ := c • σ 2 + i>k * ϕ i M .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 φ ≃ σ 2 M . (F.3)Therefore, we haveL ( y Bayes ; X) -σ 2 ≃ φ2 • i≤k * 1 ϕ i + i>k * ϕ i + σ 2 M • k * + i>k * ϕ 2 i φ2 ≃ i min φ2 ϕ i , ϕ i + φ • i min 1, ϕ 2 i φ2 (from (F.3)) ≃ i min φ2 ϕ i , ϕ i + min φ, ϕ 2 i φ ≃ i min ϕ i , φ .This finishes the proof.G Proof of Theorem 6.3Now we consider doing gradient flow on the risk function, or equivalently, on the excess risk function which differs by only a constant:dβ dt = - 1 2 ∂ ∂β [R(β, Γ) -min R(•, •)] ; (G.1) dΓ dt = - 1 2 ∂ ∂Γ [R(β, Γ) -min R(•, •)] . (G.2)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 haveR(β, Γ) -min R(•, •) = (β -β * ) ⊤   (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH   (β -β * ) + tr H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ . (G.3)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) isdβ dt = -   (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH   (β -β * ) (G.4) dΓ dt = -HΓH 1 2 ΩH 1 2 -HΨH - M + 1 M HΓH (β -β * ) (β -β * ) ⊤ H - 1 M (β -β * ) ⊤ H (β -β * ) • HΓH + H (β -β * ) (β -β * ) ⊤ H. (G.5)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 have1 2 ∂ ∂Γ tr H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 ⊤ = H 1 2 H 1 2 ΓH 1 2 -H 1 2 ΨH 1 2 Ω -1 Ω 1 2 • Ω 1 2 H 1 2(the sixth equation in Lemma H.3)= HΓH 1 2 ΩH 1 2 -HΨH.∂ ∂Γ 1 2 (β -β * ) ⊤   (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH   (β -β * ) .Let's compute the partial derivativess separately. From Lemma H.3, we have∂ ∂Γ 1 2 (β -β * ) ⊤ -HΓ ⊤ H (β -β * ) = ∂ ∂Γ 1 2 (β -β * ) ⊤ (-HΓH) (β -β * ) = - 1 2 H (β -β * ) (β -β * ) ⊤ H; ∂ ∂Γ 1 2 M + 1 M (β -β * ) ⊤ HΓ ⊤ HΓH (β -β * ) = M + 1 M HΓ • H (β -β * ) (β -β * ) ⊤ H; ∂ ∂Γ 1 2   tr HΓ ⊤ HΓ M • (β -β * ) ⊤ H (β -β * )   = 1 M (β -β * ) ⊤ H (β -β * ) • HΓH.Summing over the three equations above and applying the definition of gradient flow, we conclude the dynamics of Γ in (G.5).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 defineH := Im (H)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 hasEQUATIONwhich impliesH 1 2 (β(t) -β * ) 2 2 ≤ λ 1 exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 , (G.7)and P H (β(t)) → P H (β * ) when t → ∞, from arbitrary initialization β(0) and Γ(0). Moreover, one has for any t > 0, it holds thatP H ⊥ (β(t)) = P H ⊥ (β(0)) (G.8)Proof. We first consider the orthogonal projection operator P. From Lemma H.5 and equation (G.4), we knowdP H (β(t) -β * ) dt = -HH + H Γ (β -β * ) ,whereH Γ := (I d -ΓH) ⊤ H (I d -ΓH) + tr HΓ ⊤ HΓ M H + 1 M HΓ ⊤ HΓH.From the property of pseudo-inverse, we knowdP H (β(t) -β * ) dt = -H Γ (β -β * ) = -H Γ H + H (β -β * ) = -H Γ P H (β -β * ) .Therefore, we haved dt 1 2 ∥P H (β(t) -β * )∥ 2 2 = -P H (β -β * ) ⊤ • H Γ • P H (β -β * ) .Notice thatH Γ ⪰ M M + 1 I - M + 1 M ΓH ⊤ H M M + 1 I - M + 1 M ΓH + 1 M + 1 H ⪰ 1 M + 1 H.whereA 1 = λ 2 -1 λ min (Ω) , A 2 = 2 + 2 M λ 2 1 ∥β(0) -β * ∥ 2 2 .(G. 14) 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 knowd dt 1 2 ∥P Z (Γ -Γ * )∥ 2 F = tr d dt P Z (Γ -Γ * ) • P Z (Γ -Γ * ) ⊤ = G 1 + G 2 + G 3 , (G.15)whereG 1 = -tr H • P Z (Γ -Γ * ) • H 1 2 ΩH 1 2 • P Z (Γ -Γ * ) ⊤ , G 2 = -tr M + 1 M H • P Z (Γ -Γ * ) • H (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ -tr 1 M (β -β * ) ⊤ H (β -β * ) • H • P Z (Γ -Γ * ) • H • P Z (Γ -Γ * ) ⊤ , G 3 = -tr 1 M (β -β * ) ⊤ H (β -β * ) • HΓ * H • P Z (Γ -Γ * ) ⊤ + 1 M tr HΓ * H (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ + tr 1 M H 1 2 Ω -1 H 1 2 ΨH 1 2 + tr H 1 2 ΨH 1 2 + σ 2 I d H 1 2 (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ .(G.16)From Lemma H.2, we know that G 2 ≤ 0. Then, we consider G 1 and G 3 . For G 1 , we haveG 1 = -tr H 1 2 P Z (Γ -Γ * ) ⊤ H 1 2• HEQUATION) in Lemma H.6 and (2) in Lemma H.2)≤ -λ 2 -1 λ min (Ω) ∥P Z (Γ -Γ * )∥ 2 F, where λ min (•) denote the minimal eigenvalue. Since Ω is positive definite, we know λ min (Ω) > 0.Finally, let's upper bound G 3 . First, we have-tr 1 M (β -β * ) ⊤ H (β -β * ) • HΓ * H • P Z (Γ -Γ * ) ⊤ ≤ √ d M (β -β * ) ⊤ H (β -β * ) • H 1 2 Γ * H 1 2 2 H 1 2 • P Z (Γ -Γ * ) ⊤ H 1 2 F (Lemma H.1) ≤ √ dλ 1 M exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 λ max H 1 2 Γ * H 1 2 • ∥H∥ F ∥P Z (Γ -Γ * )∥ F (from equation (G.7)) ≤ √ dλ 2 1 M exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 λ max H 1 2 Γ * H 1 2 • ∥P Z (Γ -Γ * )∥ F . For λ max H 1 2 Γ * H 1 2, we recallΓ * = ΨH 1 2 Ω -1 H -1 2 and obtain H 1 2 Γ * H 1 2 = H 1 2 ΨH 1 2 Ω -1 H -1 2 H 1 2 = Ω -1 H 1 2 ΨH 1 2 H -1 2 H 1 2(Ω and H1 2 ΨH 1 2 commute) = Ω -1 H 1 2 ΨH 1 2 . (H 1 2 H -1 2 H 1 2 = H 1 2 )Since Ω -1 and H -1 H 1 2 ΨH 1 2 are ϕ i M +1 M ϕ i + i ϕi+σ 2 M , i = 1, 2, ..., d.Note that, every eigenvalue is upper bounded by 1, so we simply getλ max H 1 2 Γ * H 1 2 ≤ 1.(G.17)Therefore, we have-tr 1 M (β -β * ) ⊤ H (β -β * ) • HΓ * H • P Z (Γ -Γ * ) ⊤ ≤ √ dλ 2 1 M exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 • ∥P Z (Γ -Γ * )∥ F . (G.18)Similarly, we have-tr 1 M HΓ * H (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ ≤ √ d M P Z (Γ -Γ * ) ⊤ F H 1 2 (β -β * ) (β -β * ) ⊤ H 1 2 F ∥H∥ F H 1 2 Γ * H 1 2 2 (Lemma H.1) ≤ √ dλ 2 1 M exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 • ∥P Z (Γ -Γ * )∥ F , (G.19)andtr 1 M H 1 2 Ω -1 H 1 2 ΨH 1 2 + tr H 1 2 ΨH 1 2 + σ 2 I d H 1 2 (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ ≤ √ d M Ω -1 H 1 2 ΨH 1 2 + tr H 1 2 ΨH 1 2 + σ 2 I d 2 H 1 2 (β -β * ) (β -β * ) ⊤ H • P Z (Γ -Γ * ) ⊤ H 1 2 F (Lemma H.1) ≤ √ d M • M • H 1 2 (β -β * ) (β -β * ) ⊤ H 1 2 F • ∥H∥ F • ∥P Z (Γ -Γ * )∥ F ≤ √ dλ 2 1 exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 • ∥P Z (Γ -Γ * )∥ F (G.20)Bridging (G.18), (G. 19 ) and (G.20) into (G.16), we getG 3 ≤ 2 + 2 M λ 2 1 exp -2λ -1 t M + 1 ∥β(0) -β * ∥ 2 2 • ∥P Z (Γ -Γ * )∥ F (G.21)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 hasP Z (Γ(t)) → P Z (Γ * ) , P Z ⊥ (Γ(t)) = P Z ⊥ (Γ(0))when t → ∞, from arbitrary initialization β(0) and Γ(0).Proof. We observe thatd dt 1 2 ∥P Z (Γ -Γ * )∥ 2 F = ∥P Z (Γ -Γ * )∥ F • d dt ∥P Z (Γ -Γ * )∥ F .Combining it with Lemma G.5, we haved dt ∥P Z (Γ -Γ * )∥ F ≤ -A 1 ∥P Z (Γ -Γ * )∥ F + A 2 exp -2λ -1 t M + 1 ,where A 1 > 0, A 2 > 0 are defined in (G.14). Simple calculation shows thatd dt [exp (A 1 t) ∥P Z (Γ -Γ * )∥ F ] ≤ A 2 exp A 1 - 2λ -1 M + 1 t . (G.22)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∥P Z (Γ(T ) -Γ * )∥ F ≤ ∥P Z (Γ(0) -Γ * )∥ F exp (A 1 T ) + A 2 A 1 -2λ-1 M +1 • exp A 1 -2λ-1 M +1 T -1 exp (A 1 T ) = ∥P Z (Γ(0) -Γ * )∥ F exp (A 1 T ) + A 2 A 1 -2λ-1 M +1 • exp - 2λ -1 T M + 1 -exp (-A 1 T ) → 0 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.Lemma H.1 (Von-Neumann's Trace Inequality). Let U, V ∈ R d×n with d ≤ n. We havetr U ⊤ V ≤ d i=1 σ i (U )σ i (V ) ≤ ∥U ∥ op × d i=1 σ i (V ) ≤ √ d • ∥U ∥ op ∥V ∥ F 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 ⊤ . • ∂x ⊤ Bx ∂x = B + B ⊤ x. • ∂a ⊤ Xb ∂X = ab ⊤ . • ∂a ⊤ X ⊤ b ∂X = ba ⊤ . • ∂b ⊤ X ⊤ DXc ∂X = D ⊤ Xbc ⊤ + DXcb ⊤ . • ∂ ∂X tr (AXB + C)(AXB + C) ⊤ = 2A ⊤ (AXB + C)B ⊤Lemma H.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. ThenE xx ⊤ Axx ⊤ = H A + A ⊤ H + tr(AH)H.If A is symmetric and the rows in X ∈ R M ×d are generated independently fromX[i] ∼ N (0, H), i = 1, . . . , M.Then, it holds thatE X ⊤ XAX ⊤ X = M • tr (HA) • H + M (M + 1) • HAH.Lemma H.5 (Linear Algebra). Suppose H is a (non-zero) positive semi-definite matrix in R d×d and H 12 is its principle square root. We denote λ -1 as the minimal non-zero eigenvector of H. Then, we have• 1. null (H) = null H 1 2 , Im (H) = Im H 1 2 . • 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 respectivelyP Im(H) (α) = HH + α = H + Hα, P null(H) (α) = I d -HH + α = I d -H + H α.• 4. For any vector α ∈ R d , we haveP Im(H) (α) ⊤ HP Im(H) (α) ≥ λ -1 P Im(H) (α)Proof. We consider the eigen-decomposition of matrix H, which isH = QDQ ⊤ , (H.1)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 knowH 1 2 = QD 1 2 Q ⊤ , (H.2)where D1 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 haveHα = d i=1 a i Hq i = r i=1 a i λ i q i ∈ H,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α = r i=1 a i λ i • λ i q i = H r i=1 a i λ i q i ∈ Im (H) ,which implies H ⊂ Im (H) . This showsIm (H) = H = span {q 1 , ...q r } ,and null (H) = H = span {q r+1 , ...q d } , 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 thatHH + = Q • diag(1, 1, ..., 1 r , 0, 0, .., 0) • Q ⊤ = H + H.To prove (3), it suffices to notice that actually HH + α ∈ Im (H) andHH + α, I d -HH + α = H + Hα, I d -HH + α = α ⊤ HH + -HH + HH + α = 0.Finally, to prove (4), we can write α = d i=1 a i q i for some a 1 , ..., a d . Then, by orthogonality we haveEQUATION. 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(A ⊗ B) • C = BCA ⊤ .Im (H ⊗ H) := (H ⊗ H) • Z : Z ∈ R d×d , null (H ⊗ H) := Z ∈ R d×d : (H ⊗ H) • Z = 0 , and define Im H 1 2 ⊗ H 1 2 and null H 1 2 ⊗ H 1 2similarly. Then, we have• 1. Im H 1 2 ⊗ H 1 2 = Im (H ⊗ H) , null H 1 2 ⊗ H 1 2 = null (H ⊗ H) . • 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 P Z (Z) = HH + ⊗2 • Z = HH + ZH + H = H 1 2 H -1 2 ZH -1 2 H 1 2 (H.3) P Z ⊥ (Z) = I ⊗2 d -HH + ⊗2 • Z = Z -HH + ZH + H = Z -H 1 2 H -1 2 ZH -1 2 H 1 2 . (H.4) • 3. For any matrix Z ∈ R d×d , it holds that H 1 2 ⊗ H 1 2 • P Z (Z) • H 1 2 ⊗ H 1 2 • P Z (Z) ⊤ ⪰ λ 2 -1 P Z (Z) • P Z (Z) ⊤ ,(H.5) 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 equation, it suffices to notice that for any matrix Z ∈ R d×d , the matrix HH + ZH + H is in Im H ⊗2 , together with the fact thatHH + ZH + H, Z -HH + ZH + H = tr HH + ZH + HZ ⊤ -tr HH + ZH + HHH + Z ⊤ H + H = tr HH + ZH + HZ ⊤ -tr HH + HH + ZH + HH + HZ ⊤ (HH + = H + H) = 0.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 haveP Im(H⊗H) (Z) = HH + ⊗2 • Z = HH + ZH + H = H 1 2 H -1 2 ZH -1 2 H 1 2 = P Im H 1 2 ⊗H 1 2 (Z) .Since this holds for any matrix Z ∈ R d×d , this suggests Im H :H = QDQ ⊤ , H 1 2 = QD 1 2 Q ⊤ ,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 H1 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 S := span {q i ⊗ q j : 1 ≤ i, j ≤ r} as the eigenspace spanned by all eigenvectors corresponding to positive eigenvalues. For any matrix Z, we can decompose it asZ = i,j a i,j • q i ⊗ q j , so (H ⊗ H) • Z = d i,j=1 a i,j (H ⊗ H) • (q i ⊗ q j ) = i,j (Hq i ) ⊗ (Hq j ) = d i,j=1a i,j λ i λ j q i ⊗ q j = r i,j=1 a i,j λ i λ j q i ⊗ q j ∈ S, 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 haveZ = r i,j=1 b i,j λ i λ j •(λ i q i )⊗(λ j q j ) = r i,j=1 b i,j λ i λ j •(Hq i )⊗(Hq j ) = (H ⊗ H)•   r i,j=1 b i,j λ i λ j (q i ⊗ q j )   ∈ Z,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 haveH 1 2 ⊗ H 1 2 • P Z (Z) = H 1 2 ⊗ H 1 2 • P S (Z) = H 1 2 ⊗ H 1 • r i,j=1 c i,j • q i ⊗ q j = r i,j=1 c i,j H 1 2 q i ⊗ H 1 2 q j = r i,j=1 c i,j λ i λ j • q i ⊗ q j .Therefore, we haveH 1 2 ⊗ H 1 2 • P Z (Z) H 1 2 ⊗ H 1 2 • P Z (Z) ⊤ =   r i,j=1 c i,j λ i λ j • q i ⊗ q j     r k,l=1 c k,l λ k λ l • q k ⊗ q l   ⊤ = r i,j=1 c 2 i,j λ i λ j (q i ⊗ q j ) (q i ⊗ q j ) ⊤ (By orthogonality) ⪰ λ 2 -1 r i,j=1 c 2 i,j (q i ⊗ q j ) (q i ⊗ q j ) ⊤ = λ 2 -1 P Z (Z) • P Z (Z) ⊤ .Therefore, we conclude.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.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 asE :=     X ⊤ x 1 ⊤ M 1 y ⊤ 0     ∈ R (d+2)×(M +1) . (J.1)where 1 M ∈ R M refers to the vector filled with ones.A LSA model f, is defined byf (E) := E + W ⊤ P W V EM E ⊤ W ⊤ K W Q E M -1,-1 = W ⊤ P W V EM E ⊤ W ⊤ K W Q E M -1,-1, 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. DenoteW ⊤ P W V = * * * w ⊤ a 1 a 2 , W ⊤ K W Q =     Q b * ⊤ 1 b 1 * q ⊤ 2 b 2 *     , where w ∈ R d , a 1 , a 2 ∈ R, Q ∈ R d×d , b, q 1 , q 2 ∈ R d , b 1 , b 2 ∈R, and * denotes entries that do not enter the final prediction. Then we havef (E) = w ⊤ a 1 a 2 EM M E ⊤ M     Q b * q ⊤ 1 b 1 * q ⊤ 2 b 2 *     x 1 0 = w ⊤ a 1 a 2 EM M E ⊤ M     Q b q ⊤ 1 b 1 q ⊤ 2 b 2     x 1 = w ⊤ a 1 a 2 1 M     X ⊤ X X ⊤ 1 M X ⊤ y 1 ⊤ M X M 1 ⊤ M y y ⊤ X y ⊤ 1 M y ⊤ y         Q b q ⊤ 1 b 1 q ⊤ 2 b 2     x 1 .Following the Assumption 3.1, we use β to refer to the task parameter, thenβ := β * + Ψ 1 2θ, where θ ∼ N (0 d , I d ) .We then decompose f (E) asf (E) = w ⊤ a 1 a 2 1 M     X ⊤ X X ⊤ 1 M X ⊤ y 1 ⊤ M X M 1 ⊤ M y y ⊤ X y ⊤ 1 M y ⊤ y         Q q ⊤ 1 q ⊤ 2     I •x + II,where terms I and II are independent of x. Therefore, we have R (f ) := E (f (E) -y) • 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. 
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?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). 
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). 
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). 
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.2 
 
 Appendix:  


Decision: This paper provides the benefit of an MLP layer on top of a linear transformer on an in-context learning task. This paper generalizes the problem setting to a non-zero mean task distribution while previous work mainly focused on mean 0 task distributions. The MLP layer is beneficial to deal with such a non-zero mean of the task distribution. Finally, it shows the convergence of gradient flow on the population risk .   

- The paper is well written. The message of the paper is very clear and thus the readers can easily grasp the main contribution.  
- Although there is not so much mathematical difficulty, this paper gives an interesting contribution to the line of research on this topic. Especially, it is nice that the paper provides not only statistical advantage but also optimization guarantee on the population risk.     

Indeed, the reviewers are overall positive on this paper. I also agree with their opinion and would like to recommend acceptance of this paper.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: I don’t see any major weaknesses in the paper as such, but I do have a few questions (see questions).
questions: questions: Have 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.

I 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?

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

It 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?
limitations: limitations: Yes, the authors have addressed the limitations.
paper_summary: paper_summary: In 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.
strength_weakness: strength_weakness: The 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.
questions: questions: 1)  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?

2) Does the optimisation dynamics of GD-$\beta$ offer any insights on the optimisation of the LTB ?   

3) Does the results qualitatively change if the input sequence is generated from uncentered Gaussian $\mathcal{N}( \mu, H )$ instead of zero mean Gaussian prior ?


Minor corrections: 

1) In l.251, it should be $ \Tau^{\top} $ in the expression of $W_K^{\top} W_Q$.

2) In l.263, an additional bracket in the expression of $\Omega$.
limitations: limitations: NA
paper_summary: paper_summary: his 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.
strength_weakness: strength_weakness: I have the following concerns and questions:

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

2. 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.  
It 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?  
(Wu and Xu 2020) *On the optimal weighted $\ell_2$ regularization in overparameterized linear regression*.  
(Li et al. 2023) *Transformers as algorithms: generalization and stability in in-context learning*.

3. Some minor questions: 
* On page 5, why does the "scratchpad" take the specific form of (3.1)?
* 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?

---
**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.  
I have increased my score accordingly.
questions: questions: See Weaknesses above.
limitations: limitations: See Weaknesses above.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.
2. 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.

[a] Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape, ICML, 2024

[b] Training Nonlinear Transformers for Efficient In-Context Learning: A Theoretical Learning and Generalization Analysis, ICML, 2024
questions: questions: 1. 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?
2. In line 272, I think the $\Gamma$ is a precondition but not equivalent to the Newton step. Is it a typo?
3. 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?
limitations: limitations: The authors have adequately addressed the limitations.

Comments: title: Reply to Author Rebuttal
Official Comment by Reviewer XdPR
comment: Thank you for answering my questions. Please add the figure regarding Q1, and the clarification regarding Q3 in the next version. I will keep my positive rating and recommend acceptance.
Thank you for the rebuttal and for addressing my questions. The rebuttal did not change my initial impression of the paper: it presents an interesting point but also has its share of limitations. Therefore, I maintain my rating.
Thank you for providing the rebuttal. The authors adequately addressed my concerns, so my confidence level regarding my review of this paper has increased.