Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond

Authors: yingcong  li
ankit  singh rawat
samet  oymak
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:38:52

Abstract: Recent research has shown that Transformers with linear attention are capable of in-context learning (ICL) by implementing a linear estimator through gradient descent steps. However, the existing results on the optimization landscape apply under stylized settings where task and feature vectors are assumed to be IID and the attention weights are fully parameterized. In this work, we develop a stronger characterization of the optimization and generalization landscape of ICL through contributions on architectures, low-rank parameterization, and correlated designs: (1) We study the landscape of 1-layer linear attention and 1-layer H3, a state-space model. Under a suitable correlated design assumption, we prove that both implement 1-step preconditioned gradient descent. We show that thanks to its native convolution filters, H3 also has the advantage of implementing sample weighting and outperforming linear attention in suitable settings. (2) By studying correlated designs, we provide new risk bounds for retrieval augmented generation (RAG) and task-feature alignment which reveal how ICL sample complexity benefits from distributional alignment. (3) We derive the optimal risk for low-rank parameterized attention weights in terms of covariance spectrum. Through this, we also shed light on how LoRA can adapt to a new distribution by capturing the shift between task covariances. Experimental results corroborate our theoretical findings. Overall, this work explores the optimization and risk landscape of ICL in practically meaningful settings and contributes to a more thorough understanding of its mechanics.

Main paper: 
 introduction: 
 Modern language models exhibit remarkable ability to learn novel tasks or solve complex problems from the demonstrations provided within their context window (Brown et al., 2020; GeminiTeam et al., 2023; OpenAI, 2023; Touvron et al., 2023) . Such in-context learning (ICL) offers a novel and effective alternative to traditional fine-tuning techniques. It enables successful prediction across a wide range of tasks simply through a forward pass, eliminating the need for task-specific model weight updates. Since its introduction, ICL capability has become an important feature of LLM with its applications spanning retrieval-augmented generation (Lewis et al., 2020) , and reasoning via advanced prompting techniques, such as chain-of-thought (Wei et al., 2022) . While ICL already exhibits considerable benefits with a small number of demonstrations, i.e., few-shot data, there is a growing interest in extending its benefits to the manyshot settings, potentially realizing even more pronounced benefits (Agarwal et al., 2024) .ICL ability also presents an important research avenue to develop stronger theoretical and mechanistic understanding of large language models. To this aim, there has been significant recent interest in demystifying ICL through the lens of function approximation (Liu et al., 2023a) , Bayesian inference (Müller et al., 2021; Xie et al., 2022; Han et al., 2023) , and learning and optimization theory (Ahn et al., 2023; Mahankali et al., 2024; Zhang et al., 2024; Duraisamy, 2024) . The latter is concerned with understanding the optimization landscape of ICL, which is also crucial for understanding the generalization properties of the model. A notable result in this direction is the observation that linear attention models (Schlag et al., 2021; Von Oswald et al., 2023; Ahn et al., 2023) implement preconditioned gradient descent (PGD) during ICL (Ahn et al., 2023; Mahdavi et al., 2024) . While this line of works provide a fresh perspective to ICL, the existing studies do not address many questions arising from real-life applications nor provide guiding principles for various ICL setups motivated by practical considerations.To this aim, we revisit the theoretical exploration of ICL with linear data model where we feed an in-context prompt containing n examples (x i , the feature dimension, β ∈ R d is the task vector, and (ξ i ) n i=1 denote the noise in individual labels. Given the in-context prompt, the model is tasked to predict ŷn+1 -an estimate for y n+1 = x ⊤ n+1 β + ξ n+1 . We aim to provide answers to the following questions by exploring the loss landscape of ICL: (Q1) Is the ability to implement gradient-based ICL unique to (linear) attention? Can alternative sequence models implement richer algorithms beyond PGD?y i = x ⊤ i β + ξ i ) n i=1 ⊂ R d × R(Q2) In language modeling, ICL often works well with few-shot samples whereas standard linear estimation typically requires O (d) samples. How to reconcile this discrepancy between classical learning and ICL?(Q3) To our knowledge, existing works assume linearattention is fully parameterized, i.e., key and query projections W k , W q ∈ R d×d . What happens when they are low-rank? What happens when there is distribution shift between training and test in-context prompts and we use LoRA (Hu et al., 2022) for adaptation?In this work, we conduct a careful investigation of these questions. Specifically, we focus on ICL with 1-layer models and make the following contributions:(A1) We jointly investigate the landscape of linear attention and H3 (Fu et al., 2023) , a widely popular state-space model (SSM). We prove that under correlated design, both models implement 1-step PGD (c.f. Proposition 2.3) and the alignments in Fig. 1 (a) verify that where the dotted curve represents the theoretical result derived from Theorem 3.1. Our analysis reveals that the gating mechanism in H3 imitates attention. We also empirically show that H3 has the advantage of implementing sample-weighting which allows it to outperform linear attention in temporally-heterogeneous problem settings in Section 4 and Figure 4 .(A2) Proposition 2.3 allows for task and features to be correlated to each other as long as odd moments are zero. Through this, we can assess the impact of distributional alignment on the sample complexity of ICL. Specifically, we characterize the performance of Retrieval Augmented Generation (RAG) (c.f. Theorem 3.3 and Fig. 1(b) ) and Task-Feature Alignment (c.f. Theorem 3.4), where the in-context examples are α-correlated with either the query or the task vector.For both settings, we prove that alignment amplifies the effective sample size of ICL by a factor of α 2 d + 1, highlighting that aligned data are crucial for the success of ICL in few-shot settings.(A3) We show that, under low-rank parameterization, optimal attention-weights still implements PGD according to the truncated eigenspectrum of the fused task-feature covariance (see Section 3.2). We similarly derive risk upper bounds for LoRA adaptation (c.f. Eq. ( 14) and Fig. 1(c )), and show that, these bounds accurately predict the empirical performance. 
problem setup and preliminaries: 
 We begin with a short note on notation. Let bold lowercase and uppercase letters (e.g., x and X) represent vectors and matrices, respectively. The symbol ⊙ is defined as the element-wise (Hadamard) product, and * denotes the convolution operator. 1 d and 0 d denote the d-dimensional all-ones and all-zeros vectors, respectively; and I d denotes the identity matrix of dimension d × d. Additionally, let tr (W) denote the trace of the square matrix W.As mentioned earlier, we study the optimization landscapes of 1-layer linear attention (Katharopoulos et al., 2020; Schlag et al., 2021) and H3 (Fu et al., 2023) EQUATIONHere, we set z i = x i y i for i ≤ n and the last/query tokenz n+1 = x n+1 0 .Then, given Z, the goal of the model is to predict the correct label y n+1 corresponding to x n+1 . For cleaner notation, when it is clear from context, we drop the subscript n + 1 and set x = x n+1 , z = z n+1 . Different from the previous work (Ahn et al., 2023; Mahankali et al., 2024; Zhang et al., 2024; Mahdavi et al., 2024) where (x i ) n+1 i=1 and β are assumed to be independent, our analysis focuses on a more general linear setting that captures the dependency between (x i ) n+1 i=1 and β. Model architectures. To start with, we first review the architectures of both Transformer and state-space model (SSM). Similar to the previous work (Von Oswald et al., 2023; Ahn et al., 2023; Mahankali et al., 2024; Zhang et al., 2024) and to simplify the model structure, we focus on single-layer models and omit the nonlinearity, e.g., softmax operation and MLP activation, from the Transformer. Given the input prompt Z ∈ R (n+1)×(d+1) in (1), which can be treated as a sequence of (d + 1)-dimensional tokens, the single-layer linear attention ATT and H3-like single-layer SSM SSM are denoted by d+1) denote the key, query and value weight matrices, respectively. In (2b), the parameter f ∈ R n+1 is a 1-D convolutional filter that mixes tokens. The Hadamard product ⊙ is the gating mechanism (Dauphin et al., 2017) between key and query channels, which is crucial for attention-like feature creation. Thus, (2b) is more generally a gated-convolution layer. For f only, we use indexing f = [ f 0 . . . f n ] ⊤ ∈ R n+1 and given any vector a, denote convolution output (a * f ) i = i j=1 f i-j a j . Note that our notation slightly differs from the original H3 model (Fu et al., 2023) in two ways:ATT(Z) = (ZW q W ⊤ k Z ⊤ )ZW v (2a) SSM(Z) = (ZW q ) ⊙ ((ZW k ⊙ ZW v ) * f ) (2b) where W k , W q , W v ∈ R (d+1)×(1. SSMs provide efficient parameterization of f which would otherwise grow with sequence length. In essence, H3 utilizes a linear state-space models i = As i-1 + Bu i and y i = Cs i with parameters (A ∈ R d×d , B ∈ R d×1 , C ∈ R 1×d ) from which the filter f is obtained via the impulse response f i = C A i B for i ≥ 0.Here d is the state dimension and, in practice, A is chosen to be diagonal. Observe that, setting d = 1 and A = ρ, C = B = 1, SSM reduces to the exponential smoothing f i = ρ i for i ≥ 0. Thus, H3 also captures the all-ones filter as a special instance. As we show in Proposition 2.3, this simple filter is optimal under independent data model and exactly imitates linear attention. Note that, utilizing a filter f as in ( 2b) is strictly more expressive than the SSM as it captures all possible impulse responses.2. H3 also applies a shift SSM to the key embeddings to enable the retrieval of the local context around associative recall hits. We opted not to incorporate this shift operator in our model. This is because unless the features of the neighboring tokens are correlated (which is not the case for the typical independent data model), the entry-wise products between values and shifted keys will have zero mean and be redundant for the final prediction.We note that we conduct all empirical evaluations with the original H3 model, which displays exact agreement with our theory formalized for (6b), further validating our modeling choice.We will next study the algorithms that can be implemented by the single-layer attention and state-space models. Through this, we will show that training ATT and SSM with linear ICL data is equivalent to the prediction obtained from one step of optimally-preconditioned gradient descent (PGD) and sample-weighted preconditioned gradient descent (WPGD), respectively. We will further show that under mild assumption, the optimal sample weighting for SSM (e.g., f ) is an all-ones vector and therefore, establishing the equivalence among PGD, ATT, and SSM.Background: 1-step gradient descent. Consider minimizing squared loss and solving linear regression using one step of PGD and WPGD. Given n samples (x i , y i ) n i=1 , defineX = [x 1 • • • x n ] ⊤ ∈ R n×d and y = [y 1 • • • y n ] ⊤ ∈ R n .Starting from β 0 = 0 d and letting η = 1/2 be the step size, 1-step GD preconditioned with weights W returns prediction ŷEQUATIONand a single-step sample-weighted GD given weightsEQUATION)where Z is defined in (1) consisting of X, y and x. Our goal is to find the optimal W, as well as ω in (4) that minimize the population risks defined as follows.minW L PGD (W) where L PGD (W) = E (y -g PGD (Z)) 2 , (5a) min W,ω L WPGD (W) where L WPGD (W) = E (y -g WPGD (Z)) 2 . (5b)Here, the expectation is over the randomness in (x i , ξ i ) n+1 i=1 and β, and we use W to represent the set of corresponding trainable parameters. The search spaces for ω and W are R n and R d×d , respectively.As per (2), given input prompt Z ∈ R (n+1)×(d+1) , either of the underlying models outputs a (n + 1)-length sequence. Note that the label for the query x = x n+1 is excluded from the prompt Z. Similar to Ahn et al. (2023) ; Mahankali et al. (2024) , we consider a training objective with a causal mask to ensure inputs cannot attend to their own labels and training can be parallelized. Let Z 0 = [z 1 . . . z n 0] ⊤ be the features post-causal masking at time/index n + 1. Given weights W k , W q , W v and the filter f for SSM, predictions at the query token z = x 0 take the following forms following sequence-to-sequence mappings in (2):g ATT (Z) = (z ⊤ W q W ⊤ k Z ⊤ 0 )Z 0 W v v, g SSM (Z) = (z ⊤ W q ) ⊤ ⊙ ((Z 0 W k ⊙ Z 0 W v ) * f ) n+1 v,where v ∈ R d+1 is the linear prediction head and ((Z 0 W k ⊙ Z 0 W v ) * f ) n+1 returns the last row of the convolution output. Note that SSM can implement the mask by setting f 0 = 0. Now consider the meta learning setting and select loss function to be the squared loss, same as in (5). Thus, the objectives for both models take the following forms. min Wk,Wq,Wv,v L ATT (W) whereL ATT (W) = E (y -g ATT (Z)) 2 , (6a) min Wk,Wq,Wv,v, f L SSM (W) where L SSM (W) = E (y -g SSM (Z)) 2 . (6b)Here, similarly, the expectation subsumes the randomness of (x i , ξ i ) n+1 i=1 and β and W represents the set of trainable parameters. The search space for matrices W k , W q , W v is R (d+1)×(d+1) , for head v is R d+1 , and for f is R n+1 .Note that for all the optimization methods (c.f. ( 5), ( 6)), to simplify the analysis, we train the models without capturing additional bias terms. Therefore, in the following, we introduce the centralized data assumptions such that the models are trained to make unbiased predictions.To begin with, a cross moment of random variables is defined as the expectation of a monomial of these variables, with the order of the cross moment being the same as order of the monomial. For example, E[x ⊤ Wβ] is a sum of crossmoments of order 2. Then, it motivates the following data assumptions. Assumption 2.1. All cross moments of the entries of (x i ) n+1 i=1 and β with odd orders are zero. Assumption 2.2. (ξ i ) n+1 i=1 are independent of (x i ) n+1 i=1 and β, and their cross moments with odd orders are zero.Note that compared to (Ahn et al., 2023; Mahankali et al., 2024; Zhang et al., 2024) , Assumption 2.1 is more general which also subsumes the dependent distribution settings. In this work, we consider the following three linear models (omitting noise) satisfying Assumption 2.1. Let Σ β , Σ x ∈ R d×d represent the task and feature covariance matrices for independent data, and let 0 ≤ α ≤ 1 be the correlation level when considering data dependency. More specific discussions are deferred to Section 3.• Independent task and data:β ∼ N(0, Σ β ), x i ∼ N(0, Σ x ), i ≤ n + 1.• Retrieval augmented generation:β, x ∼ N(0, I d ), x i x ∼ N(αx, (1 -α 2 )I d ), i ≤ n.• Task-feature alignment:β ∼ N(0, I d ), x i β ∼ N(αβ, I d ), i ≤ n + 1.Next, we introduce the following result which establishes the equivalence among optimizing linear attention (c.f. (6a)), H3 (c.f. (6b)), and gradient descent (c.f. ( 5)). Proposition 2.3. Suppose Assumptions 2.1 and 2.2 hold. Consider the objectives as defined in (5) and (6), and let L ⋆ PGD , L ⋆ WPGD , L ⋆ ATT , and L ⋆ SSM be their optimal risks, respectively. Then,L ⋆ PGD = L ⋆ ATT and L ⋆ WPGD = L ⋆ SSM . Additionally, if the examples (x i , y i ) ni=1 follow the same distribution and are conditionally independent given x, β, then SSM/H3 can achieve the optimal loss using the all-ones filter and L ⋆ PGD = L ⋆ SSM .We defer the proof to Appendix B.1. Proposition 2.3 establishes that analyzing the optimization landscape of ICL for both single-layer linear attention and the H3 model can be effectively reduced to examining the behavior of a onestep PGD algorithm. Notably, under the independent, RAG and task-feature alignment data settings discussed above, examples (x i , y i ) n i=1 are independently sampled given x and β, and we therefore conclude that L ⋆ PGD = L ⋆ ATT = L ⋆ SSM . Leveraging this result, the subsequent section of the paper concentrate on addressing (5a), taking into account various linear data distributions.While Proposition 2.3 demonstrates the equivalence of optimal losses, we also study the uniqueness and equivalence of optimal prediction functions. To this end, we analyze the strong convexity of L PGD (W).Lemma 2.4. Suppose Assumption 2.2 holds and let ξ =[ξ 1 ξ 2 • • • ξ n ] ⊤ . Then the loss L PGD (W) in (5a) is strongly- convex iff. E[(x ⊤ W X ⊤ Xβ) 2 ] + E[(x ⊤ W X ⊤ ξ) 2 ] is strongly- convex.Additionally, let g ⋆ PGD , g ⋆ ATT be the optimal prediction functions of (5a), (6a). Then under the conditions of Assumptions 2.1 and 2.2, and the strong convexity, g ⋆ PGD = g ⋆ ATT . Lemma 2.5. Suppose that the label noise (ξ i ) n i=1 are i.i.d., zero-mean, variance σ 2 and independent of everything else, and that there is a decomposition x = x 1 + x 2 , X = X 1 + X 2 , and β = β 1 + β 2 such that either of the following holds • σ > 0, and (x 1 , X 1 ) have full rank covariance and are independent of each other and (x 2 , X 2 ).• (x 1 , β 1 , X 1 ) have full rank covariance and are independent of each other and (x 2 , β 2 , X 2 ).Then, the loss L PGD (W) in (5a) is strongly-convex.As mentioned above, we study three specific linear models: with general independent, RAG-related, and task-feature alignment data. Note that for all the three cases, according to Proposition 2.3, we haveL ⋆ PGD = L ⋆ ATT = L ⋆ SSM .Additionally, the second claim in Lemma 2.5 holds, and L PGD (W) is strongly convex. Therefore, following Lemma 2.4, we haveg ⋆ PGD = g ⋆ ATT .Thanks to the equivalence among PGD, ATT, and SSM, in the next section, we focus on the solution of objective (5a) under different scenarios, which will reflect the optimization landscapes of ATT and SSM models. 
main results: 
 In light of Proposition 2.3, optimizing a single layer linearattention or H3 model is equivalent to solving the objective (5a). Therefore, in this section, we examine the properties of the one-step PGD in (5a). To this end, we consider multiple problem settings, including distinct data distributions and low-rank training.We first consider the standard independent data setting. We will then examine correlated designs. Independent data model. Let Σ x and Σ β be the covariance matrices of the input feature and task vectors, respectively, and σ ≥ 0 be the noise level. We assumeβ ∼ N(0, Σ β ), x i ∼ N(0, Σ x ), ξ i ∼ N(0, σ 2 ), i ≤ n + 1 (7)and the label is obtained via y i = x ⊤ i β + ξ i . Our following result characterizes the optimal solution of (5a). Note that the data generated from (7) satisfies the conditions in Proposition 2.3. Therefore, the same results can be applied to both linear-attention and H3 models.Theorem 3.1. Consider independent linear data in (7), and suppose the covariance matrices Σ x , Σ β are full rank. Recap the objective from (5a) and let W ⋆ := arg min W L PGD (W), andL ⋆ = L PGD (W ⋆ ). Additionally, let Σ = Σ 1/2 x Σ β Σ 1/2x and M = tr (Σ) + σ 2 . Then W ⋆ and L ⋆ satisfyW ⋆ = Σ -1/2 x W⋆ Σ -1/2 x and L ⋆ = M -ntr Σ W⋆ , (8)where we define W⋆ = (n + 1)I d + MΣ -1 -1 .Corollary 3.2. Consider noiseless i.i.d. linear data where Σ x = Σ β = I d and σ = 0. Then, the objective in (5a) returnsW ⋆ = 1 n + d + 1 I d and L ⋆ = d - nd n + d + 1 .See Appendix C.2 for proofs. Note that Theorem 3.1 is consistent with prior work (Ahn et al., 2023 , Theorem 1) when specialized to isotropic task covariance, i.e., Σ β = I d . However, their result is limited as the features and task are assumed to be independent. This prompts us to ask: What is the optimization landscape with correlated in-context samples? Toward this, we consider the following RAG-inspired and task-feature alignment models, where Assumptions 2.1 and 2.2 continue to hold and Proposition 2.3 applies.To provide a statistical model of the practical RAG approaches, given the query vector x n+1 = x, we propose to draw ICL demonstrations that are similar to x with the same shared task vector β.Modeling feature similarity through the cosine angle, RAG should sample the ICL examples x i , i ≤ n, from the original feature distribution conditioned on the event cos(x i , x) ≥ α where α is the similarity threshold. As an approximate proxy, under the Gaussian distribution model, we assume that β ∼ N(0, I d ), x ∼ N(0, I d ) and that RAG samples α-correlated demonstrations (x i , y i ) n i=1 as follows:x i x ∼ N(αx, (1 -α 2 )I d ), ξ i ∼ N(0, σ 2 ), 1 ≤ i ≤ n (9)and y i = x ⊤ i β + ξ i . Note that the above normalization ensures that the marginal feature distribution remains N(0, I d ).The full analysis of RAG is provides in Appendix C.3. Specifically, when we carry out the analysis by assumingα = O 1/ √ d and d/n = O (1)where O (•) denotes proportionality, our derivation leads to the following result: Theorem 3.3. Consider linear model in (9). Recap the objective from (5a) and let W ⋆ := arg min W L PGD (W), andL ⋆ = L PGD (W ⋆ ). Additionally, let κ = α 2 d + 1 and suppose α = O 1/ √ d , d/n = O (1)and d is sufficiently large. Then W ⋆ and L ⋆ have approximate formsW ⋆ ≈ 1 κn + d + σ 2 I d and L ⋆ ≈ d + σ 2 - κnd κn + d + σ 2 . (10)Here, ( 10) is reminiscent of Corollary 3.2 and has a surprisingly clean message. Observe that, α 2 d + 1 is the dominant multiplier ahead of n in both equations. Thus, we deduce that, RAG model follows the same error bound as the independent data model, however, its sample size is amplified by a factor of α 2 d + 1. α = 0 reduces to the result of Corollary 3.2 whereas we need to set α = O 1/ √ d for constant amplification. When α = 1, RAG achieves the approximate risk L ⋆ ≈ 2+σ 2 , where the constant bias is due to the higher order moments (e.g., the 4'th and 6'th moments) of the standard Gaussian distribution. As d increases, the normalized loss L ⋆ /d → 0. The full analysis of its optimal solution W ⋆ and loss L ⋆ are deferred Theorem C.1 in Appendix C.3.Task-feature alignment. We also consider another dependent data setting where task and feature vectors are assumed to be correlated. This dataset model has the following motivation: In general, an LLM can generate any token within the vocabulary. However, once we specify the task (e.g. domain of the prompt), the LLM output becomes more deterministic and there are much fewer token candidates. For instance, if the task is "Country", "France" is a viable output compared to "Helium" and vice versa when the task is "Chemistry". Formally speaking, this can be formalized as the input x having a diverse distribution whereas it becomes more predictable conditioned on β. Therefore, it can be captured through a linear model by making the conditional covariance of x β to be approximately low-rank. This formalism can be viewed as a spectral alignment between input and task, which is also well-established in deep learning both empirically and theoretically (Li et al., 2020; Arora et al., 2019; Canatar et al., 2021; Cao et al., 2019) . Here, we consider such a setting where the shared task vector is sampled as standard Gaussian distribution β ∼ N(0, I d ) and letting κ = α 2 d + 1, we sample the α-correlated ICL demonstrations (x i , y i ) n+1 i=1 as follows:EQUATION)andy i = κ -1/2 x ⊤ i β + ξ i .Here, κ -1/2 is a normalization factor to ensure that label variance remains invariant to α. To keep the exposition cleaner, we defer the full analysis of its optimal solution W ⋆ and loss L ⋆ to Theorem C.2 in Appendix C.4. Similar to the RAG setting, by assuming α = O 1/ √ d and d/n = O (1), we obtain the following results for the optimal parameter and risk.Theorem 3.4. Consider linear model in (11). Recap the objective from (5a) and let W ⋆ := arg min W L PGD (W), andL ⋆ = L PGD (W ⋆ ). Additionally, given κ = α 2 d + 1 and suppose α = O 1/ √ d , d/n = O (1)and d is sufficiently large. Then W ⋆ and L ⋆ have approximate formsW ⋆ ≈ 1 κn + (d + σ 2 )/κ I d and L ⋆ ≈ d + σ 2 - κnd κn + (d + σ 2 )/κ . (12)Similar to (10), ( 12) contains κ = α 2 + 1 multiplier ahead of n, which reduces the in-context sample complexity and setting α = 0 reduces to the results of Corollary 3.2.In this section, we investigate training low-rank models, which assume W k , W q ∈ R (d+1)×r where r is the rank restriction. Equivalently, we consider objective (5a) under condition rank (W) = r.Lemma 3.5. Consider independent linear data in (7). Recap the objective from (5a) and enforce rank (W) ≤ r andW ⊤ = W. Let Σ = Σ 1/2 x Σ β Σ 1/2x and M = tr (Σ) + σ 2 . Let λ i be the i'th largest eigenvalue of Σ, we have thatEQUATION)Note that tr (Σ) = d i=1 λ i . Removing the rank constraint and considering noiseless data setting, this reduces to the following optimal risk L ⋆ = d i=1 λ i +M n+1+M/λ i . See Appendix D.1 for more details.Impact of LoRA: Based on the above lemma, we consider the impact of LoRA for adapting the pretrained model to a new task distribution under jointly-diagonalizable old and new eigenvalues ofΣ, Σ new , (λ i ) d i=1 , (λ new i ) d i=1 .Consider adapting LoRA matrix to the combined key and value weights in attention, which reflects minimizing the popu-EQUATION)Note that, the right hand side is provided assuming the optimal LoRA-updated model W lora is also jointly diagonalizable with covariances Σ, Σ new , and W.  
experiments: 
 We now conduct synthetic experiments to support our theoretical findings and further explore the behavior of different models of interest under different conditions. The experiments are designed to investigate various scenarios, including independent data, retrieval-augmented generation (RAG), task-feature alignment, low-rank parameterization, and LoRA adaption.Experimental setting. We train 1-layer attention and H3 models for solving the linear regression ICL. As described in Section 2, we consider meta-learning setting where task parameter β is randomly generated for each training sequence. In all experiments, we set the dimension d = 20. Depending on the in-context length (n), different models are trained to make in-context predictions. We train each model for 10000 iterations with batch size 128 and Adam optimizer with learning rate 10 -3 . Since our study focuses on the optimization landscape, and experiments are implemented via gradient descent, we repeat 20 model trainings from different initialization and results are presented as the minimal test risk among those 20 trails. In all the plots, theoretical predictions are obtained via the corresponding formulae presented in Section 3 and the test risks are normalized by the dimension d.• Equivalence among L ⋆ PGD , L ⋆ ATT and L ⋆ SSM (Figure 2 ). To verify Proposition 2.3 as well as Theorem 3.1, we run random linear regression instances where in-context samples are generated obeying (7). Fig. 2 9) and ( 11), respectively, by setting σ = 0 and varying α from 0 to 0.6. Attention training results are displayed in solid curves, and we generate theory curve (dotted) via the L ⋆ formula as described in (36) in Appendix C.3 and (42) in Appendix C.4. The empirical alignments corroborate Theorems C.1 and C.2, further confirming that Proposition 2.3 is applicable to a broader range of real-world distributional alignment data.• Low-rank (Fig. 3(c )) and LoRA (Fig. 3(d) ) experiments. We also run simulations to verify our theoretical findings in Section 3.2. Consider the independent data setting as described in (7). In Fig. 3 (c), we set Σ x = I d , σ = 0 and task covariance to be diagonal with diagonal entries c[12 -1 • • • d -1 ] ⊤ for some normalization constant c = d/ d i=1 i -1, and parameterize the attention model using matrices W k , W q ∈ R (d+1)×r and vary r across the set {1, 5, 10, 20}. Results show that empirical (solid) and theoretical (dotted, c.f. ( 13)) curves overlap. In Fig. 3 (d), we implement two phases of training. Phase 1: Setting Σ x = Σ β = I d and σ = 0, we pretrain the model with full rank parameters and obtain weights Ŵk , Ŵq , Ŵv ∈ R (d+1)×(d+1) . Phase 2: We generate new examples with task covariance Σ β being a diagonal matrix with diagonal entries c Fig. 3(d) presents the results after two phases of training where dotted curves are drawn from the right hand side of ( 14) directly. Here, note that since Σ, Σ new are diagonal, the right hand side of ( 14) returns the exact optimal risk of LoRA and the alignments verify it.′ [2 -1 2 -2 • • • 2 -d ] ⊤ for some normalization constant c ′ = d/ d i=1 2 -i . Given the rank restriction r, we train additional LoRA parame- ters W up , W down ∈ R (d+1)×r where W lora := W up W ⊤ down and (2a) becomes ATT(Z) = (Z( Ŵq Ŵ⊤ k + W up W ⊤ down )Z ⊤ )Z Ŵv .• H3 outperforms linear attention (Figure 4 ). Until now, our analysis has established the equivalence between linear attention and H3 models in solving linear ICL problem. Furthermore, we also investigate settings where H3 could outperform linear attention due to its sample weighting ability. In Figs. 4(a fixed max-length n max and loss is evaluated as the average loss given samples from 1 to n max . Such setting has been wildly studied in the previous ICL work (Garg et al., 2022; Akyürek et al., 2023; Li et al., 2023) . We generate data according to ( 7 Then the i'th sample is generated viax i ∼ N(0, I d ) and y i = β ⊤ i x i where β i = λ i β 1 + (1 -λ i )β 2 and λ i = i/n.The results are reported in Fig. 4(d) which again shows that H3 achieves better performance compared to linear attention, as H3 may benefit from the additional convolutional filter (c.f. f in (2b)). Here, dotted curve represent the theoretical results under i.i.d. and noiseless setting, derived from Corollary 3.2. 
discussion: 
 In this work, we revisited the loss landscape of in-context learning with 1-layer sequence models. We have established a general connection between ICL and gradient methods that accounts for correlated data, non-attention architectures (specifically SSMs), and the impact of low-rank parameter-ization including LoRA adaptation. Our results elucidate two central findings: (1) The functions learned by different sequence model architectures exhibit a strong degree of universality and (ii) Dataset and prompt design, such as RAG, can substantially benefit ICL performance.Future directions and limitations. The results of this work fall short of being a comprehensive theory for ICL in LLMs and can be augmented in multiple directions. First, while the exact equivalence between H3 and linear attention is remarkable, we should examine whether it extends to other SSMs. Secondly, while empirically predictive, our RAG and LoRA analyses are not precise and fully formal. Thirdly, it is desirable to develop a deeper understanding of multilayer architectures and connect to iterative GD methods as in (Ahn et al., 2023; Von Oswald et al., 2023) . Finally, we have studied the population risk of ICL training whereas one can also explore the sample complexity of pretraining (Wu et al., 2023; Lu et al., 2024) . Moving beyond the theoretically tractable setup of this work, our simplified models are trained on in-context prompts from random initialization. Therefore, this theoretical study doesn't address more challenging in-context learning tasks, such as question answering, where both in-context demonstration and general knowledge from pretraining are required. Future work in this area could also shed light on how certain contexts might elicit undesirable behaviors acquired by an LLM during pretraining, an aspect not covered in our current analysis. This work also studies a theoretical model for retrieval augmentation-based ICL. In a real-life retrieval augmentation-based ICL, one needs to account for the quality of the collection of the retrievable demonstrations and its (negative) impacts on the final predictions. There is growing interest in understanding the mechanisms behind ICL (Brown et al., 2020; Liu et al., 2023b; Rae et al., 2021) in large language models (LLMs) due to its success in continuously enabling novel applications for LLMs (GeminiTeam et al., 2023; OpenAI, 2023; Touvron et al., 2023) 2023), they show that the trained model implements a single step of GD and PGD for isotropic and anisotropic Gaussian features, respectively. In addition, they also characterize the test-time prediction error for the trained model while highlighting its dependence on train and test prompt lengths. Interestingly, Zhang et al. (2024) further explore the effect of various distributional shifts, including the shift in task weight vector distributions between train and test time as well as the covariate shifts among train and test in-context prompts. Interestingly, they find that while linear-attention models are robust to most shifts, they exhibit brittleness to the covariate shifts.While our work shares similarities with this line of works, as discussed in our contributions in the introduction, we expand the theoretical understanding of ICL along multiple novel dimensions, which includes the first study of LoRA adaptation for ICL in the presence of a distributional shift. Furthermore, we strive to capture the effect of retrieval augmentation (Lewis et al., 2020; Nakano et al., 2021) on ICL through our analysis. Retrieval augmentation allows for selecting most relevant demonstration out of a large collection for a test instance, e.g., via a dense retrieval model (Izacard et al., 2023) , which can significantly outperform the typical ICL setup where fixed task-specific demonstrations are provided as in-context examples (Wang et al., 2022; Basu et al., 2023) . Through a careful modeling of retrieval augmentation via correlated design, we show that it indeed has a desirable amplification effect where the effective number in-context examples becomes larger with higher correlation which corresponds to preforming a successful retrieval of query-relevant demonstrations in a practical retrieval augmented setup.Recently, state space models (SSMs) (Gu et al., 2021b; a; Fu et al., 2023; Gu & Dao, 2023) have appeared as potential alternatives to Transformer architecture, with more efficient scaling to input sequence length. Recent studies demonstrate that such SSMs can also perform ICL for simple non-language tasks (Park et al., 2024; Grazzi et al., 2024) as well as complex NLP tasks (Grazzi et al., 2024) . That said, a rigorous theoretical understanding of ICL for SSMs akin to (Zhang et al., 2024; Mahankali et al., 2024; Ahn et al., 2023) is missing from the literature. In this work, we provide the first such theoretical treatment for ICL with SSMs. Focusing on H3 architecture (Fu et al., 2023) , we highlight its advantages over linear attention in specific ICL settings.In this section, we present the proofs related to Section 2. Recap that given dataX = [x 1 • • • x n ] ⊤ ∈ R n×d , ξ = [ξ 1 • • • ξ n ] ⊤ ∈ R n , y = [y 1 • • • y n ] ⊤ = Xβ + ξ ∈ R n , Z 0 = [z 1 . . . z n 0 d+1 ] ⊤ = x 1 . . . x n 0 d y 1 . . . y n 0 ⊤ ∈ R (n+1)×(d+1) ,and corresponding prediction functionsEQUATIONEQUATIONEQUATIONEQUATIONwe have objectives minW L PGD (W) where L PGD (W) = E (y -g PGD (Z)) 2 , (16a) min W,ω L WPGD (W) where L WPGD (W) = E (y -g WPGD (Z)) 2 , (16b) min W k ,W q ,W v ,v L ATT (W) where L ATT (W) = E (y -g ATT (Z)) 2 , (16c) min W k ,W q ,W v ,v, f L SSM (W) where L SSM (W) = E (y -g SSM (Z)) 2 . (16d)Here, the expectation is over the randomness in (x i , ξ i ) n i=1 and β, and the search space for d+1) , for v is R d+1 , and for f is R n+1 .W is R d×d , for ω is R n , for W k , W q , W v is R (d+1)×(Consider the problem setting as discussed in Section 2, Proposition 2.3 can be proven by the following two lemmas.Lemma B.1. Suppose Assumptions 2.1 and 2.2 hold. Then, given the objectives (16a) and (16c), we havemin W q ,W k ,W v ,v L ATT (W) = min W L PGD (W).Proof. Recap the linear attention estimator from (15c) and denoteW q W ⊤ k = W w 1 w ⊤ 2 w and W v v = v 1 v ,where W ∈ R d×d , w 1 , w 2 , v 1 ∈ R d , and w, v ∈ R. Then we haveEQUATIONwhere W := v W + w 1 v ⊤ 1 . We first show that for any given parametersEQUATION)To this goal, we haveEQUATION)where we have decomposition(g ATT (Z) -y)ε = (x ⊤ W X ⊤ y -y)x ⊤ W X ⊤ Xv 1 + v ∥y∥ 2 ℓ 2 w 1 = y ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 + v ∥y∥ 2 ℓ 2 w 1 -yx ⊤ W X ⊤ Xv 1 + v ∥y∥ 2 ℓ 2 w 1 = y ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 (a) + v ∥y∥ 2 ℓ 2 y ⊤ X W⊤ xx ⊤ w 1 (b) -yx ⊤ W X ⊤ Xv 1 (c) -vy ∥y∥ 2 ℓ 2 x ⊤ w 1 (d).In the following, we consider the expectations of (a), (b), (c), (d) sequentially, which return zeros under Assumptions 2.1 and 2.2. Note that since Assumption 2.1 holds, expectation of any odd order of monomial of the entries of X, x, β returns zero, i.e., order of x ⊤ βx is 3 and thereforeE[x ⊤ βx] = 0 d . (a) : E y ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 = E (Xβ + ξ) ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 = E β ⊤ X ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 + E ξ ⊤ X W⊤ xx ⊤ W X ⊤ Xv 1 = 0. (b) : E v ∥y∥ 2 ℓ 2 y ⊤ X W⊤ xx ⊤ w 1 = E v(Xβ + ξ) ⊤ (Xβ + ξ)(Xβ + ξ) ⊤ X W⊤ xx ⊤ w 1 = E v ∥ξ∥ 2 ℓ 2 ξ ⊤ X W⊤ xx ⊤ w 1 = 0. (c) : E yx ⊤ W X ⊤ Xv 1 = E (x ⊤ β + ξ)x ⊤ W X ⊤ Xv 1 = E β ⊤ xx ⊤ W X ⊤ Xv 1 + E ξx ⊤ W X ⊤ Xv 1 = 0. (d) : E vy ∥y∥ 2 ℓ 2 x ⊤ w 1 = v E (β ⊤ x + ξ)(Xβ + ξ) ⊤ (Xβ + ξ)x ⊤ w 1 = v E ξ ∥ξ∥ 2 ℓ 2 x ⊤ w 1 = 0.Combining the results with (19) returns thatE (g ATT (Z) -y) 2 -E (g ATT (Z) -y) 2 = E[ε 2 ] ≥ 0 (20)which completes the proof of ( 18). Therefore, we obtain minW q ,W k ,W v ,v E (g ATT (Z) -y) 2 ≥ min W E (g ATT (Z) -y) 2 = min W E (g PGD (Z) -y) 2 .We conclude the proof of this lemma by showing that for any W ∈ R d×d in g PGD , there existW k , W q , W v , v such that g ATT (Z) = g PGD (Z). Let W k = W v = I d+1 , W q = W 0 d 0 ⊤ d 0 , and v = 0 d 1 .Then we obtainEQUATION)which completes the proof. □ Lemma B.2. Suppose Assumptions 2.1 and 2.2 hold. Then, given the objectives in (16), we haveEQUATION)Additionally, if the examples (x i , y i ) n i=1 follow the same distribution and are conditionally independent given x and β, then SSM/H3 can achieve the optimal loss using the all-ones filter andEQUATION)Proof. Recap the SSM estimator from (15d) and letW q = w q1 w q2 • • • w q,d+1 , W k = w k1 w k2 • • • w k,d+1 , W v = w v1 w v2 • • • w v,d+1 ,where w q j , w k j , w v j ∈ R d+1 for j ≤ d + 1, and letv =               v 1 v 2 • • • v d+1               , and f =               f 0 f 1 • • • f n               . Then we have g SSM (Z) = (z ⊤ W q ) ⊤ ⊙ ((Z 0 W k ⊙ Z 0 W v ) * f ) n+1 v = n i=1 f n+1-i • v ⊤                       w ⊤ q1 z • • • w ⊤ q,d+1 z            ⊙           w ⊤ k1 z i w ⊤ v1 z i • • • w ⊤ k,d+1 z i w ⊤ v,d+1 z i                      = n i=1 f n+1-i • v ⊤            w ⊤ q1 zw ⊤ k1 z i w ⊤ v1 z i • • • w ⊤ q,d+1 zw ⊤ k,d+1 z i w ⊤ v,d+1 z i           .Next for all j ≤ d + 1, let w q j = wq j w q j , w k j = wk j w k j , w v j = wv j w v j where wq j , wk j , wv j ∈ R d and w q j , w k j , w v j ∈ R. Then we have w ⊤ q j zw ⊤ k j z i w ⊤ v j z i = w⊤ q j x w⊤ k j x i + w k j y i w⊤ v j x i + w v j y i = x ⊤ wq j w v j w⊤ k j + w k j w⊤ v j x i y i + w⊤ q j x w⊤ k j x i w⊤ v j x i + w k j w v j w⊤ q j xy 2i = x ⊤ W ′ j x i y i + δ j (x, x i , x i ) + w ′ j ⊤ xy 2 i where W ′ j := wq j w v j w⊤ k j + w k j w⊤ v j ∈ R d×d , w ′ j := w k j w v j wq j ∈ R d , δ j (x, x i , x i ) := w⊤ q j x w⊤ k j x i w⊤ v j x i ∈ R.Theng SSM (Z) = n i=1 f n+1-i • d+1 j=1 v j x ⊤ W ′ j x i y i + δ j (x, x i , x i ) + w ′ j ⊤ xy 2 i = x ⊤         d+1 j=1 v j W ′ j         X(y ⊙ f ) + n i=1 f n+1-i • d+1 j=1 v j • δ j (x, x i , x i ) +         d+1 j=1 v j w ′ j ⊤         xy ⊤ (y ⊙ f ) = x ⊤ W X ỹ gSSM (Z) + δ(x, X, X) ε 1 + w⊤ xy ⊤ ỹ ε 2.wheref := [ f n • • • f 1 ] ⊤ ∈ R n , ỹ := y ⊙ f ∈ R n , W := d+1 j=1 v j W ′ j ∈ R d×d , w := d+1 j=1 v j w ′ j ∈ R d , δ(x, X, X) := n i=1 f n+1-i • d+1 j=1 v j • δ j (x, x i , x i ) ∈ R.Next we will show that for anyW k , W q , W v , v, E (g SSM (Z) -y) 2 ≥ E (g SSM (Z) -y) 2 .To start with, we obtainEQUATION)where there is decomposition(g SSM (Z) -y)(ε 1 + ε 2 ) = δ(x, X, X) • x ⊤ W X ỹ (a) -δ(x, X, X)y (b) + w⊤ xy ⊤ ỹ • x ⊤ W X ỹ (c) -y • w⊤ xy ⊤ ỹ (d).In the following, similar to the proof of Lemma B.1, we consider the expectations of (a), (b), (c), (d) sequentially, which return zeros under Assumptions 2.1 and 2.2. Note that δ j (x, x i , x i )'s and δ(x, X, X) are summation of monomials of entries of (x, X, β) with order 3, and entries of y and y are summation of monomials of entries of (x, X, β) with even orders: e.g., y = x ⊤ β + ξ where ξ is of oder 0 and x ⊤ β is of order 2.(a) : E δ(x, X, X) • x ⊤ W X ỹ = E δ(x, X, X) • x ⊤ W X(Xβ ⊙ f ) + E δ(x, X, X) • x ⊤ W X(ξ ⊙ f ) = E δ(x, X, X) • x ⊤ W X E ξ ⊙ f = 0. (b) : E δ(x, X, X)y = E δ(x, X, X)(x ⊤ β + ξ) = E δ(x, X, X)x ⊤ β + E δ(x, X, X)ξ = 0. (c) : E w⊤ xy ⊤ ỹ • x ⊤ W X ỹ = E w⊤ x(Xβ + ξ) ⊤ (Xβ ⊙ f + ξ ⊙ f ) • x ⊤ W X(Xβ ⊙ f + ξ ⊙ f ) = 0. (d) : E y • w⊤ xy ⊤ ỹ = E (x ⊤ β + ξ) • w⊤ x(Xβ + ξ) ⊤ (Xβ ⊙ f + ξ ⊙ f ) = 0.Combining the results with (24) results thatE (g SSM (Z) -y) 2 -E (g SSM (Z) -y) 2 = E (ε 1 + ε 2 ) 2 ≥ 0.Therefore we obtain, minW q ,W k ,W v ,v, f E (g SSM (Z) -y) 2 ≥ min W, f E (g SSM (Z) -y) 2 = min W,ω E (g WPGD (Z) -y) 2 .Next we show that for any choices of W and ω in g WPGD , there are W q,k,v , v, f such that g SSM ≡ g WPGD . To this end, givenω = [ω 1 . . . ω n ] ⊤ , let W q = I d+1 , W k = W ⊤ 0 d 0 ⊤ d 0 , W v = 0 d×d 0 d 1 ⊤ d 0 , v = 1 d 0 and f =               0 ω n • • • ω 1              .Then we get((Z 0 W k ⊙ Z 0 W v ) * f ) n+1 = XW ⊤ 0 n 0 d 0 ⊙ y1 ⊤ d 0 n 0 d 0 * f n+1 = n i=1 ω i • y i W x i 0 = W X ⊤ (y ⊙ ω) 0 ,and thereforeg SSM (Z) = x ⊤ W X ⊤ (y ⊙ ω) = g WPGD (Z),which completes the proof of ( 22).Next, to show (23), for anyW ∈ R d×d , let L(ω) = E x ⊤ W X ⊤ (y ⊙ ω) -y 2 . Then we have ∂L(ω) ∂ω i = E         2         x ⊤ W n j=1 ω j y j x j -y         x ⊤ Wy i x i         = 2 n j=1 ω j E (x ⊤ Wy j x j )(x ⊤ Wy i x i ) -2 E yx ⊤ Wy i x i .Here since (x i , y i ) n i=1 follow the same distribution and are conditionally independent given x and β, for any i j j ′ ,E (x ⊤ Wy i x i ) 2 = E (x ⊤ Wy j x j ) 2 and E (x ⊤ Wy j x j )(x ⊤ Wy i x i ) = E (x ⊤ Wy j ′ x j ′ )(x ⊤ Wy i x i ) . Then let E (x ⊤ Wy j x j )(x ⊤ Wy i x i ) =        c 1 , i j c 2 , i = j and E yx ⊤ Wy i x i = c 3 , where (c 1 , c 2 , c 3 ) := (c 1 (W), c 2 (W), c 3 (W)). We get ∂L(ω) ∂ω i = 2c 1 ω ⊤ 1 n + 2(c 2 -c 1 )ω i -2c 3 . If c 2 -c 1 = 0, then ∂L(ω) ∂ω i ≡ 2c 1 ω ⊤ 1 n -2c 3 for all i ≤ n and any ω ∈ R n achieves the same performance. If c 2 -c 1 0, setting ∂L(ω) ∂ω i = 0 returns ω i = c 3 -c 1 n j=1 ω j c 2 -c 1 := C for all i ≤ n.Therefore the optimal loss is achieved via setting ω = C1 n . Without loss of generality, we can update W → CW. Then ω = 1 n , and we obtain minW,ω E x ⊤ W X ⊤ (y ⊙ ω) -y 2 = min W E (x ⊤ W X ⊤ y -y) 2which completes the proof of (23). □Proof. Recap the loss L PGD (W) in (16a) and prediction g PGD (Z) in (15a), we haveL PGD (W) = E[(y -g PGD (Z)) 2 ] = E x ⊤ β + ξ -x ⊤ W X ⊤ (Xβ + ξ) 2 = E (x ⊤ β -x ⊤ W X ⊤ Xβ) 2 + 2(x ⊤ β -x ⊤ W X ⊤ Xβ)(ξ -x ⊤ W X ⊤ ξ) + (ξ -x ⊤ W X ⊤ ξ) 2 = E (x ⊤ β -x ⊤ W X ⊤ Xβ) 2 + (ξ -x ⊤ W X ⊤ ξ) 2 + 2 E[(x ⊤ β -x ⊤ W X ⊤ Xβ)(ξ -x ⊤ W X ⊤ ξ)] = E (x ⊤ β -x ⊤ W X ⊤ Xβ) 2 + (ξ -x ⊤ W X ⊤ ξ) 2 (25) = E (x ⊤ W X ⊤ Xβ) 2 + (x ⊤ W X ⊤ ξ) 2 f 1 (W) -2 E[β ⊤ xx ⊤ W X ⊤ Xβ + ξx ⊤ W X ⊤ ξ] f 2 (W) + E[(x ⊤ β) 2 + ξ 2 ]constant where (25) follows Assumption 2.2. Since f 2 (W) is convex, L PGD (W) is strongly-convex if and only if f 1 (W) is stronglyconvex, which completes the proof of strong convexity.Next, ( 20) and ( 21) in the proof of Lemma B.1 demonstrate that the optimal loss is achievable and is achieved at ε = 0. Subsequently, (17) indicates that g ⋆ ATT has the same form as g ⋆ PGD . Under the strong convexity assumption, g ⋆ PGD is unique, which leads to the conclusion that g ⋆ PGD = g ⋆ ATT . □We begin by deriving the even moments of random variables.• 2n'th moment of a normally distributed variable: Let u ∼ N(0, σ 2 ). Then we haveEQUATION• 4'th moment: Let u ∼ N(0, I d ). Then for any W, W ′ ∈ R d×d , we haveEQUATION• 4'th cross-moment: Let u, v ∼ N(0, I d ) and for any W ∈ R d×d , let Λ W = W ⊙ I d . Then we haveEQUATION)• 6'th moment: Let u ∼ N(0, I d ). Then for any W, W ′ ∈ R d×d , we haveEQUATIONwhere ( 29) is obtained by followingE         u 4 i         d i ′ =1 u 2 i ′                 = E[u 6 ] + (d -1) E[u 4 ] E[u 2 ] = 3(d + 4), E         u 2 i u 2 j         d i ′ =1 u 2 i ′                 = 2 E[u 4 ] E[u 2 ] + (d -2) E[u 2 ] E[u 2 ] E[u 2 ] = d + 4.• 8'th moment: Let u ∼ N(0, I d ). Then for any W, W ′ ∈ R d×d , we haveEQUATION)where ( 31) is obtained by followingE         u 4 i         d i ′ =1 u 4 i ′ + i ′ j ′ u 2 i ′ u 2 j ′                 = E[u 8 ] + (d -1) E[u 4 ] E[u 4 ] + 2(d -1) E[u 6 ] E[u 2 ] + (d -1)(d -2) E[u 4 ] E[u 2 ] E[u 2 ] = 105 + 9(d -1) + 30(d -1) + 3(d -1)(d -2) = 3(d + 4)(d + 6), E         u 2 i u 2 j         d i ′ =1 u 4 i ′ + i ′ j ′ u 2 i ′ u 2 j ′                 = 2 E[u 6 ] E[u 2 ] + (d -2) E[u 4 ](E[u 2 ]) 2 + 2 E[u 4 ] E[u 4 ] + 4(d -2) E[u 4 ](E[u 2 ]) 2 + (d -2)(d -3)(E[u 2 ]) 4 = 30 + 3(d -2) + 18 + 12(d -2) + (d -2)(d -3) = (d + 4)(d + 6).Proof of Theorem 3.1. Consider a general independent linear model as defined in (7) where Σ x and Σ β are full-rank feature and task convariance matrices andx ∼ N(0, Σ x ), β ∼ N(0, Σ β ), ξ ∼ N(0, σ 2 ), and y = x ⊤ β + ξ.Let X = [x 1 • • • x n ] ⊤ , ξ = [ξ 1 • • • ξ n ] ⊤ , and y = [y 1 • • • y n ] ⊤ = Xβ + ξ.To simplify and without loss of generality, letx = Σ -1/2 x x, X = XΣ -1/2 x , β = Σ 1/2x β where we havex ∼ N(0, I), β ∼ N(0, Σ 1/2 x Σ β Σ 1/2 x ) and y = x⊤ β + ξ, y = X β + ξ.Then recap the loss from (5a), and we obtainEQUATION)where the last equality comes from the independence of label noise ξ, ξ.We first consider the following termE (x ⊤ W X ⊤ ξ) 2 = E ( x⊤ (Σ 1/2 x WΣ 1/2 x ) X⊤ ξ) 2 = nσ 2 • tr W W⊤where we define W = Σ 1/2 x WΣ 1/2x . Next, focus on the followingE (x ⊤ β -x ⊤ W X ⊤ Xβ) 2 = E ( x⊤ β -x⊤ W X⊤ X β) 2 = E x⊤ I -W X⊤ X β 2 = tr E I -W X⊤ X Σ I -W X⊤ X ⊤ = tr (Σ) -tr Σ( W + W⊤ ) E[ X⊤ X] + tr W⊤ W E[ X⊤ XΣ X⊤ X] = tr (Σ) -2n • tr Σ W + tr W⊤ W E[ X⊤ XΣ X⊤ X] ,whereΣ := Σ 1/2 x Σ β Σ 1/2x . Let xi ∈ R n be the i'th column of X and Σ i j be the (i, j)'th entry of Σ. Then the (i, j) entry of matrix X⊤ XΣ X⊤ X is( X⊤ XΣ X⊤ X) i j = d k=1 d p=1 Σ kp x⊤ i xk x⊤ p x j .Then we geti j : E X⊤ XΣ X⊤ X i j = Σ i j E[ x⊤ i xi x⊤ j x j ] + Σ ji E[ x⊤ i x j x⊤ i x j ] = n 2 Σ i j + nΣ ji i = j : E X⊤ XΣ X⊤ X ii = Σ ii E x⊤ i xi x⊤ i xi + j i Σ j j E x⊤ i x j x⊤ j xi = Σ ii E (x 2 i1 + • • • + x 2 in ) 2 + n j i Σ j j = Σ ii (3n + n(n -1)) + n j i Σ j j = n         Σ ii (n + 1) + d j=1 Σ j j         = n (Σ ii (n + 1) + tr (Σ)) . Therefore E[ X⊤ XΣ X⊤ X] = n(n + 1)Σ + n • tr (Σ) I.Combining all together results inEQUATIONwhere M := tr (Σ) + σ 2 . Setting ∇ W L(W) = 0 returns-2n • Σ + 2n(n + 1) • Σ W + 2nM W = 0 =⇒ W⋆ = (n + 1)I + MΣ -1 -1 .Then we haveW ⋆ = Σ -1/2 x (n + 1)I + MΣ -1 -1 Σ -1/2 x and L ⋆ = L(W ⋆ ) = M -ntr ((n + 1)Σ -1 + MΣ -2 ) -1 . □ C.3In this section, we consider the retrieval augmented generation (RAG) linear model similar to (9), where we first draw the query vector x and task vector β via x ∼ N(0, I) and β ∼ N(0, I).We then draw data (x i ) n i=1 to be used in-context according to the rule corr_coef(x, x i ) ≥ α ≥ 0. Hence, for i ≤ n we sampleEQUATIONwhich results in (9) by setting γ 2 = 1 -α 2 .Theorem C.1 (Extended version of Theorem 3.3). Consider linear model as defined in (35). Recap the objective from (5a) and let W ⋆ := arg min W L PGD (W), and L ⋆ = L PGD (W ⋆ ). Then W ⋆ and L ⋆ satisfyW ⋆ = cI and L ⋆ = d + σ 2 -cnd(α 2 (d + 2) + γ 2 ) (36) where c = α 2 (d + 2) + γ 2 α 4 n(d + 2)(d + 4) + α 2 γ 2 (d + 2)(d + 2n + 3) + γ 4 (d + n + 1) + σ 2 (α 2 (d + 2) + γ 2 ).Suppose α = O 1/ √ d , d/n = O (1)and d is sufficiently large. Let κ = α 2 d + 1 and γ 2 = 1 -α 2 . Then W ⋆ and L ⋆ have approximate formsEQUATION)Proof. Here, for clean notation and without loss of generality, we define and rewrite (35) viag i ∼ N(0, I), ξ i ∼ N(0, σ 2 ) and x i = αx + γ g i , y i = (αx + γ g i ) ⊤ β + ξ i .Then we obtainEQUATION)To begin with, letN 1 = tr (W) 2 + tr WW ⊤ + tr W 2 , N 2 = tr WW ⊤ , and N 3 = tr (W) .We first focus on the second term in ( 38)E (x ⊤ W X ⊤ ξ) 2 = E                n i=1 ξ i x ⊤ W(αx + γ g i )        2         = nσ 2 E x ⊤ W(αx + γ g)(αx + γ g) ⊤ W ⊤ x = nσ 2 α 2 E[x ⊤ W xx ⊤ W ⊤ x] + γ 2 E[x ⊤ W g g ⊤ W ⊤ x] = nσ 2 α 2 N 1 + γ 2 N 2 .(It follows ( 27) and independence of x, g.)Next, the first term in (38) can be decomposed intoE (x ⊤ β -x ⊤ W X ⊤ Xβ) 2 = E (x ⊤ β) 2 (a) + E (x ⊤ W X ⊤ Xβ) 2 (b) -2E x ⊤ βx ⊤ W X ⊤ Xβ (c).In the following, we consider solving (a)-(c) sequentially.(a) :E (x ⊤ β) 2 = d. (b) : E (x ⊤ W X ⊤ Xβ) 2 = E                x ⊤ W n i=1 (αx + γ g i )(αx + γ g i ) ⊤ β        2         = E                n i=1 x ⊤ W(α 2 xx ⊤ + γ 2 g i g ⊤ i + αγx g ⊤ i + αγ g i x ⊤ )β        2         = α 4 n 2 E (x ⊤ W xx ⊤ β) 2 + γ 4 E                n i=1 x ⊤ W g i g ⊤ i β        2         + α 2 γ 2 E                n i=1 x ⊤ W xg ⊤ i β        2         + α 2 γ 2 E                n i=1 x ⊤ W g i x ⊤ β        2         +2α 2 γ 2 n 2 E x ⊤ W xx ⊤ ββ ⊤ gg ⊤ W ⊤ x + 2α 2 γ 2 n E x ⊤ W x g ⊤ βx W gx ⊤ β = α 4 n 2 (d + 4)N 1 + γ 4 n(d + n + 1)N 2 + α 2 γ 2 ndN 1 + α 2 γ 2 n(d + 2)N 2 + 2α 2 γ 2 n 2 N 1 + 2α 2 γ 2 nN 1 = α 4 n 2 (d + 4) + α 2 γ 2 n(2n + d + 2) N 1 + α 2 γ 2 n(d + 2) + γ 4 n(d + n + 1) N 2 = A 1 N 1 + A 2 N 2 . (c) : E x ⊤ βx ⊤ W X ⊤ Xβ = E        n i=1 x ⊤ βx ⊤ W(αx + γ g i )(αx + γ g i ) ⊤ β        = E        n i=1 x ⊤ βx ⊤ W(α 2 xx ⊤ + γ 2 g i g ⊤ i + αγx g ⊤ i + αγ g i x ⊤ )β        = α 2 n E x ⊤ βx ⊤ W xx ⊤ β + γ 2 n E x ⊤ βx ⊤ W g g ⊤ β = α 2 n(d + 2)tr (W) + γ 2 ntr (W) = α 2 n(d + 2) + γ 2 n N 3 = A 3 N 3 .Here, (b) utilizes the 4'th and 6'th moment results ( 27) and (30) and we defineA 1 = α 4 n 2 (d + 4) + α 2 γ 2 n(2n + d + 2) A 2 = α 2 γ 2 n(d + 2) + γ 4 n(d + n + 1) A 3 = α 2 n(d + 2) + γ 2 n.Then combining all together results inL(W) = A 1 N 1 + A 2 N 2 -2A 3 N 3 + nσ 2 (α 2 N 1 + γ 2 N 2 ) + d + σ 2 .To find the optimal solution, set ∇L(W) = 0 and we obtainEQUATION)Note that we have∇N 1 = ∇ tr (W) 2 + tr WW ⊤ + tr W 2 = 2tr (W) I + 2W + 2W ⊤ ∇N 2 = ∇tr WW ⊤ = 2W ∇N 3 = ∇tr (W) = I.Therefore, (39) returnsEQUATIONwhich implies that the optimal solution W ⋆ has the form of cI for some constant c. Then suppose W ⋆ = cI, we have tr (W) = cd and (40) returns2A 1 (d + 2)cI + 2A 2 cI -2A 3 I + 2nσ 2 (α 2 (d + 2)cI + γ 2 cI) = 0Proof. Here, for clean notation and without loss of generality, we define and rewrite (41) viag i ∼ N(0, I), ξ i ∼ N(0, σ 2 ) and x i = αβ + g i , y i = γx ⊤ i β + ξ i = γ • (αβ + g i ) ⊤ β + ξ i .Recap the loss function from (5a), we obtainEQUATION)Similar to Appendix C.3, to begin with, letN 1 = tr (W) 2 + tr WW ⊤ + tr W 2 , N 2 = tr WW ⊤ , and N 3 = tr (W) ,and additionally, givenΛ W = W ⊙ I, let N 4 = 3tr Λ 2 W + (d + 4)tr WW ⊤ + tr W 2 .We first focus on the second term in ( 44)E (x ⊤ W X ⊤ ξ) 2 = E                (αβ + g) ⊤ W n i=1 ξ i (αβ + g i )        2         = nσ 2 E (αβ + g) ⊤ W(αβ + g ′ ) 2 = nσ 2 α 4 E (β ⊤ Wβ) 2 + 2α 2 E (β ⊤ W g ′ ) 2 + E (g ⊤ W g ′ ) 2 = nσ 2 α 4 tr (W) 2 + tr W 2 + tr WW ⊤ + (2α 2 + 1)tr WW ⊤ = nσ 2 α 4 N 1 + (2α 2 + 1)N 2 .(It follows ( 27) and independence of β, g, g ′ .)Next, the first term of (44) (omitting γ 2 ) returns the following decomposition:EQUATION.Consider solving (a)-(d) sequentially as follows:To begin with, we use the following decomposition for all (a)-(d):X ⊤ Xβ = n i=1 x i x ⊤ i β = n i=1 (αβ + g i )(αβ + g i ) ⊤ β = n i=1 α 2 ββ ⊤ β + αβ g ⊤ i β + αg i β ⊤ β + g i g ⊤ i β.Then, we have (a) : E[(β ⊤ W X ⊤ Xβ) 2 ] = E                n i=1 α 2 β ⊤ Wββ ⊤ β + αβ ⊤ Wβ g ⊤ i β + αβ ⊤ W g i β ⊤ β + β ⊤ W g i g ⊤ i β        2         = α 4 n 2 E β ⊤ Wββ ⊤ β 2 + α 2 E                n i=1 β ⊤ Wβ g ⊤ i β        2         + α 2 E                n i=1 β ⊤ W g i β ⊤ β        2         + E                n i=1 β ⊤ W g i g ⊤ i β        2         +2α 2 n E        n i=1 β ⊤ Wββ ⊤ ββ ⊤ W g i g ⊤ i β        + 2α 2 E        n i=1 β ⊤ Wβ g ⊤ i ββ ⊤ W g i β ⊤ β        = α 4 n 2 E β ⊤ Wββ ⊤ β 2 + α 2 n E β ⊤ Wβ g ′⊤ β 2 + α 2 n E β ⊤ W g ′ β ⊤ β 2 + E                n i=1 β ⊤ W g i g ⊤ i β        2         +2α 2 n 2 E β ⊤ Wββ ⊤ ββ ⊤ W g ′ g ′⊤ β + 2α 2 n E β ⊤ Wβ g ⊤ i ββ ⊤ W g i β ⊤ β = α= B 1 N 1 + B 2 N 2 + nN 4 ,where ( 45) and ( 47) utilize ( 30) and (32), and ( 46) is obtained viaE                n i=1 β ⊤ W g i g ⊤ i β        2         = n E β ⊤ W g ′ g ′⊤ β 2 + n(n -1) E β ⊤ W g ′ g ′⊤ ββ ⊤ W g ′′ g ′′⊤ β = nN 4 + n(n -1)N 1 ,which follows ( 27) and ( 28).(b) : where (49) and (50) are obtained using ( 27), ( 30) and (d) : = γ 2 α 2 B 1 N 1 + (α 2 B 2 + B 3 )N 2 -2(α 2 B 4 + B 5 )N 3 + α 2 nN 4 + nσ 2 (α 4 N 1 + (2α 2 + 1)N 2 ) + γ 2 d α 2 (d + 2) + 1 + σ 2 and differentiating it results in ∇L(W) = γ 2 α 2 B 1 ∇N 1 + (α 2 B 2 + B 3 )∇N 2 -2(α 2 B 4 + B 5 )∇N 3 + α 2 n∇N 4 + nσ 2 (α 4 ∇N 1 + (2α 2 + 1)∇N 2 ). x WΣ 1/2 x , Σ = Σ 1/2 x Σ β Σ 1/2EQUATIONE                n i=1 g ⊤ W g i g ⊤ i β        2         = n E g ⊤ W g ′ g ′⊤ β 2 + n(n -1) E g ⊤ W g ′ g ′⊤ β g ⊤ W g ′′ g ′′⊤ β = n(d + 2)N 2 + n(n -1)N 2 = n(n + d + 1)N 2 . (c) : E β ⊤ ββ ⊤ W X ⊤ Xβ = n E β ⊤ ββ ⊤ W(αβ + g ′ )(αβ + g ′ ) ⊤ β = α 2 n E β ⊤ ββ ⊤ Wββ ⊤ β + n E β ⊤ ββ ⊤ WEQUATIONx and M = tr (Σ) + σ 2 . For any W, let us parameterize W = UEU ⊤ where U ∈ R d×r denotes the eigenvectors of W and E ∈ R r×r is a symmetric square matrix. We will first treat U as fixed and optimize E. We will then optimize U. Fixing U, setting Σ = U ⊤ ΣU, we obtain L(E) = M -2ntr ΣE + n(n + 1)tr ΣE 2 + nMtr E 2 .Differentiating, we obtain 0.5n -1 ∇L(E) = -Σ + (n + 1) ΣE + ME.Setting ∇L(E) = 0 returnsEQUATIONLet λi denote the i'th largest eigenvalue of Σ. Plugging in this value, we obtain the optimal risk as a function of U is given byEQUATION)Now observe that, the right hand side is strictly decreasing function of the eigenvalues λi of Σ = U ⊤ ΣU. Thus, to minimize L ⋆ (U), we need to maximize r i=1 λi n+1+M λ-1 i . It follows from Cauchy interlacing theorem that λ j ≤ λ i where λ i is the i'th largest eigenvalue of Σ since Σ is an orthogonal projection of Σ on U. Consequently, we find the desired bound whereL ⋆ = M -n r i=1 λ i n + 1 + Mλ -1 i .The equality holds by setting U to be the top-r eigenvectors of Σ and E = E ⋆ (U) to be the diagonal matrix according to (51). □ 
 
 Appendix:  


Decision: This paper shows that linear transformers can perform ICL for a linear regression setting. In particular, this paper considers SSM as well as the usual linear attention model. It is shown that SSM can also perform ICL even when the distribution of input and task has non-diagonal covariance.  

I think the PGD part of this paper is not new. That is already provided in the previous papers. The novelty of this paper is to show the ICL ability of SSM. This is not trivial.    
The writing is clear and precise. The authors provided detailed results in some concrete situations such as RAG and LoRA. This provides better understanding on the authors' contribution.  

Overall, this paper provides an interesting contribution on the literature especially about understanding learning capability of SSM. The reviewers are all positive on this paper. Hence, I recommend acceptance of this paper.

Reviews: review: paper_summary: paper_summary: This paper studies how transformers can use ICL to solve linear regression problems. It is shown that state space models, transformers are both capable of performing linear regression as well as gradient descent (which implements the least squares solution). There are results about LORA and RAG.
strength_weakness: strength_weakness: The analysis of vanilla regression is not novel, I think, and has been done for instance in Theorem 1 of Ahn et al (transformers learn to do preconditioned gradient descent) where it is even shown that a pretrained transformer (with GD) learns to do this in a similar setting (which is stronger than showing that it can learn it). Maybe for SSMs this result is new but I am not sure how strong of a contribution this is considering that it only establishes the existence of such a solution.
questions: questions: What is the significance of the RAG result? It makes sense that giving the learner access to $\beta$ through $X$ directly improves the error. A comparison to some reasonable baseline might be a nice story if it came about that transformers are able to leverage this kind of side information particularly well. What i mean is to compare this with a loss like $\Vert y-X\beta\Vert^2 -\Vert X\beta\Vert^2 + \Vert \beta\Vert^2$ which is aware of this kind of side information.

What is the significance of the LoRA result? A decomposition of the optimal error in terms of the singular values of a covariance is shown and then it is observed that if the covariance changes in a specific way (I think the singular vectors have to remain fixed), then LoRA can help. What role does Lemma 1 play? This is not about LoRA, but rather for an entirely low rank model.

Why does linear attention perform better for shorter contexts than it was trained on (Figure 3a)? How does Figure 3b show that H3 is better at generalizing to a longer context length if the plot is discontinued at the context it was trained on? I noticed that the performance on varying $\beta$ in Figure 3d better than the iid noiseless case. Can you elaborate on that? Is this for a fixed (that is, rather than varying) $\beta$?
limitations: limitations: Yes
paper_summary: paper_summary: The authors examine the capabilities of Transformers with linear attention in performing in-context learning (ICL) by implementing a linear estimator through gradient descent. The existing studies mostly consider IID task and feature vectors and fully parameterized attention weights. This work expands on these studies by analyzing:

1. The landscape of 1-layer linear attention and 1-layer H3 (a state-space model), proving both can implement 1-step preconditioned gradient descent.

2. New risk bounds for retrieval augmented generation (RAG), showing benefits from distributional alignment.

3. The optimal risk for low-rank parameterized attention weights, illustrating how LoRA adapts to new distributions.
strength_weakness: strength_weakness: 1. The analysis is limited to single-layer linear attention and linear distribution data, which might not fully capture the complexities of multi-layer architectures. 2. The RAG and LoRA analyses are not precise and fully formal.
questions: questions: 1. How does the distributional alignment in RAG quantitatively affect the sample complexity of ICL? Are there specific metrics or case studies that illustrate these benefits in practical applications?

2. What are the limitations of LoRA in your cases? I believe LoRA significantly underperforms compared to fully parameterized models. If there are no limitations compared with the fully parameterized models because of the linear distribution data, can you talk about the relationship between the data distribution and LoRA.

3. The paper focuses on single-layer (linear) models. How do the findings extend to multi-layer (linear) architectures?

4. The paper discusses LoRA adaptation for distribution shifts. Are there practical examples or case studies where LoRA has been successfully applied to handle real-world distribution shifts?

5. The study primarily considers linear regression tasks. How would the insights gained from this study apply to more complex tasks, such as natural language processing or computer vision tasks that involve non-linear relationships?
limitations: limitations: As mentioned in Weakness: 1. The analysis is limited to single-layer linear attention and linear distribution data, which might not fully capture the complexities of multi-layer architectures. 2. The RAG and LoRA analyses are not precise and fully formal.
paper_summary: paper_summary: The paper studies the in-context learning capabilities of linear attention (ATT) and linear state space layers (SSM) on a linear regression task. It shows that for both ATT and SSM there exists a parametrization such that they perform as well as one step of preconditioned gradient descent (PGD) with optimal preconditioning in expectation, and that they cannot perform better.

The authors then propose a model for retrieval-augmented generation where in-context examples and the final query are correlated. For PGD, they provide an approximate expression for the optimal preconditioning weights and the related loss and relate it to the optimal weights in the noiseless, i.i.d. case, showing how the correlations between in-context examples and query translate into an increased effective sample size.

Next, the authors give an upper bound on the optimal expected loss (for 1-step PGD <=> ATT <=> SSM) in a low-rank adapter setting.

Experiments validate the theoretical results. 

A final experiment shows some initial results hinting that the SSM model studied (H3) performs slightly better than linear attention in an online setting and with a non-stationary task.
strength_weakness: strength_weakness: * The main result (12) is not stated as a theorem, and the derivation in the appendix is not completely rigorous. In particular, it is not clear under which conditions the approximation in l.558 is justified. Unless the conditions for the approximations to be valid are stated more clearly, it is difficult to extrapolate from the experimental results to the full claim (A2).

* The result on LoRA is not interpreted, and it is not immediately clear what the implications are from the statement of the result. I am not sure if the result implies claim (A3).

* The additional experimental results (Fig. 3) seem somewhat preliminary and detached from the rest of the paper. Based on these results, it is not clear that the part of claim (A1) claiming an advantage of H3 is fully supported by the evidence.
questions: questions: * Under which assumptions are the approximation in (12) valid? What is the nature of the approximations made in the proof of the statement (high probability, asymptotic, ...)?

* Can you elaborate on how claim (A3) follow from (14)?

* Define $\alpha$-correlated
* l.342: wildly studied -> widely studied (?)
limitations: limitations: The authors address the limitations of their work and acknowledge that their analyses are "not precise and fully formal".

Comments: title: Revised ranking to Accept
comment: Many thanks for your encouraging assessment and positive recommendations! Your suggestions have been very valuable, and we will revise and further improve our manuscript based on them.
We thank the reviewer for reassessing our work and for raising their rating. We would like to provide further clarifications on their comments:
1. For low-rank parameterization, our main technical result is Lemma 1. This result does not rely on the assumption of diagonalization and establishes a natural generalization of fully-parameterized results (Theorem 1) to the low-rank setting. In LoRA analysis, we assume diagonalization mostly to be able to provide closed form interpretable bounds on its impact in terms of eigenspace. We anticipate one can establish similar upper bounds on the impact of LoRA in terms of the angle between the eigenspaces of old and new covariance matrices. Finally, it worths noting that the analysis of general covariance is challenging even for vanilla ICL (fully parameterized) as both Ahn et al. and Mahankali et al. make assumptions about the covariance of $\beta$ and/or $x$ while we allow for arbitrary $x$, $\beta$ covariances in our main results.
2. Our RAG model is inspired from real RAG systems [[Karpukhin et al.](https://arxiv.org/pdf/2004.04906), [Lewis et al.](https://arxiv.org/pdf/2005.11401)], which utilize cosine similarity of the document (feature) embeddings. However, we agree with the reviewer that exploring more complex RAG models, either with general covariance models or allowing for nonlinear feature/document representations learned via a separate retriever model, would be valuable. This is certainly an exciting direction for future work and we appreciate their suggestion.
I had a misunderstanding about the RAG section. Please disregard what I said. I thought $x$ was sampled as $\mathcal(\alpha \beta, (1-\alpha)I)$.

Because of the clarifications on the experiments, and the contributions relating to SSMs (equivalences to PGD), I change my score from 3 to 5.

However, I still feel the following about the Transformer results:
1. The LoRA discussion holds only when the new and old $\Sigma$ are jointly diagonalizable and this makes the message a little weak, but it is possible that this is truly the only benefit from LoRA on a single layer. 
2. Using the present model to study RAG still seems somewhat simplistic to me.

These are difficult to fix, since this is a critique of the model itself. The results themselves seem reasonable to me.
Thank you for your response. My concerns have been addressed and I updated my ranking to Accept.