Title: Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond
Venue: NeurIPS 2024
Accepted: yes
Keywords: In-context learning

linear attention

state-space model

optimization

RAG

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

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

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

a) Correlated features

b) Task-feature alignment:

Task and

feature vectors are

𝛼

-

correlated

Query

features

Database

Relevant

examples

Ex:

ATTN

(

𝑄𝐾

!

𝑉

)

Linear Attention

H3 (SSM)

×

SSM

×

Distributional Alignment

Architecture Choice

Low-rank Parameterization

Pretrained

Weights

rank

𝑟

+

𝑊

!

𝑊

"

a) Low-rank attention weights:

rank 𝑊

!

, 𝑊

"

≤ 𝑟

b) LoRA adaptation

0

10

20

30

40

50

# in-context samples

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Test risk

Linear Att

H3

Theorem 1

(a) Linear attention

=

H3

0

10

20

30

40

50

# in-context samples

0.2

0.4

0.6

0.8

Test risk

=0

=0.2

=0.4

=0.6

(b) RAG with

α

correlation

0

10

20

30

40

50

# in-context samples

0.1

0.3

0.5

0.7

0.9

Test risk

r

=1

r

=5

r

=10

r

=20

(c) LoRA

1 Introduction

Modern language models exhibit the 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 and 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].

ICL ability 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 , yi = x ⊤ i β + ξ i ) n i = 1 ⊂ R d × R and a test instance or query x n + 1 ∈ R d to the model, with d being the feature dimension, β ∈ R d being the task weight vector, and ( ξ i ) n i = 1 denoting the noise in individual labels. Given the in-context prompt, the model is tasked to predict ˆ yn + 1 - an estimate for yn + 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?

(Q2) In language modeling, ICL often works well with few-shot samples whereas standard linear estimation typically requires O ( d ) samples. How can we reconcile this discrepancy between classical learning and ICL?

(Q3) To our knowledge, existing works assume linear-attention 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 1) and the alignments in Fig. 1a verify that where the dotted curve represents the theoretical PGD result derived from Theorem 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 Appendix D.

(A2) Proposition 1 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 2 and Fig. 1b) and Task-Feature Alignment (c.f. Theorem 3), 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. 1c), and show that, these bounds accurately predict the empirical performance.

2 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] models when training with prompts containing in-context data following a linear model. We construct the input in-context prompt similar to Ahn et al. [2023], Mahankali et al. [2024], Zhang et al. [2024] as follows.

Linear data distribution. Let ( x , y ) ∈ R d × R be a (feature, label) pair generated by a d -dimensional linear model parameterized by β ∈ R d , i.e., y = x ⊤ β + ξ , where x and β are feature and task vectors, and ξ is the label noise. Given demonstrations ( x i , yi ) n + 1 i = 1 sampled from a single β , define the input in-context prompt

Here, we set z i = " x i yi # for i ≤ n and the last/query token z n + 1 = " x n + 1 0 # . Then, given Z , the goal of the model is to predict the correct label yn + 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 singlelayer 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

where W k , W q , W v ∈ R ( d + 1) × ( 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 = P i j = 1 f i -jaj . Note that our notation slightly differs from the original H3 model [Fu et al., 2023] in two ways:

SSMs provide efficient parameterization of f which would otherwise grow with sequence length. In essence, H3 utilizes a linear state-space model s i = As i -1 + B ui and yi = 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 = CA 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 1, 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.

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.

2.1 In-context Linear Estimation

We will next study the algorithms that can be implemented by the single-layer attention and statespace 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 optimallypreconditioned 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 , yi ) n i = 1 , define

Starting from β 0 = 0 d and letting η = 1 / 2 be the step size, a single-step GD preconditioned with weights W returns prediction

and a single-step sample-weighted GD given weights ω ∈ R n and W ∈ R d × d returns prediction

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.

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):

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.

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 cross-moments of order 2. Then, it motivates the following data assumptions.

Assumption 1 All cross moments of the entries of ( x i ) n + 1 i = 1 and β with odd orders are zero.

Assumption 2 The label noise ( ξ 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 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 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 ) , for all 1 ≤ i ≤ n + 1.

Retrieval augmented generation: β , x ∼ N (0 , I d ) , x i GLYPH<12> GLYPH<12> GLYPH<12> x ∼ N ( α x , (1 -α 2 ) I d ) , for all 1 ≤ i ≤ n .

Task-feature alignment: β ∼ N (0 , I d ) , x i GLYPH<12> GLYPH<12> GLYPH<12> β ∼ N ( α β , I d ) , for all 1 ≤ i ≤ n + 1 .

Next, we introduce the following result which establishes the equivalence among optimizing 1-layer linear attention (c.f. (6a)), 1-layer H3 (c.f. (6b)), and 1-step gradient descent (c.f. (5)).

Proposition 1 Suppose Assumptions 1 and 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,

Additionally, if the examples ( x i , yi ) n i = 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 A.1. Proposition 1 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 one-step PGD algorithm. Notably, under the independent, RAG and task-feature alignment data settings discussed above, examples ( x i , yi ) 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 1 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 ) and derive the subsequent lemmas.

Lemma 1 Suppose Assumption 2 holds and let ξ = [ ξ 1 ξ 2 · · · ξ n ] ⊤ . Then the loss L PGD ( W ) in (5a) is strongly-convex if and only if E [( x ⊤ WX ⊤ X β ) 2 ] + E [( x ⊤ WX ⊤ ξ ) 2 ] is strongly-convex. Additionally, let g ⋆ PGD , g ⋆ ATT be the optimal prediction functions of (5a) and (6a) . Then under the conditions of Assumptions 1 and 2, and the strong convexity, g ⋆ PGD = g ⋆ ATT .

Lemma 2 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, in this work, 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 1, we have L ⋆ PGD = L ⋆ ATT = L ⋆ SSM . Additionally, the second claim in Lemma 2 holds, and L PGD ( W ) is strongly convex. Therefore, following Lemma 1, we have g ⋆ 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.

3 Main Results

In light of Proposition 1, optimizing a single layer linear-attention 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. The latter refers to the scenario where the key and query matrices have rank restrictions, e.g., W k , W q ∈ R ( d + 1) × r , as well as LoRA-tuning when adapting the model under distribution shift.

3.1 Analysis of Linear Data Models

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

and the label is obtained via yi = x ⊤ i β + ξ i . Our following result characterizes the optimal solution of (5a). Note that the data generated from (7) satisfies the conditions in Proposition 1. Therefore, the same results can be applied to both linear-attention and H3 models.

Theorem 1 Consider independent linear data as defined 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 ) , and L ⋆ = L PGD ( W ⋆ ) . Additionally, let Σ = Σ 1 / 2 x Σ β Σ 1 / 2 x and M = tr ( Σ ) + σ 2 . Then W ⋆ and L ⋆ satisfy

where we define ¯ W ⋆ = GLYPH<16> ( n + 1) I d + M Σ -1 GLYPH<17> -1 .

Corollary 1 Consider noiseless i.i.d. linear data where Σ x = Σ β = I d and σ = 0 . Then, the objective in (5a) returns

See Appendix B.2 for proofs. Note that Theorem 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 1 and 2 continue to hold and Proposition 1 applies.

Retrieval augmented generation. 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 , yi ) n i = 1 as follows:

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 B.3. Specifically, when we carry out the analysis by assuming α = O GLYPH<16> 1 / √ d GLYPH<17> and d / n = O (1) where O ( · ) denotes proportionality, our derivation leads to the following result:

Theorem 2 Consider linear model as defined in (9) . Recap the objective from (5a) and let W ⋆ : = arg min W L PGD ( W ) , and L ⋆ = L PGD ( W ⋆ ) . Additionally, let κ = α 2 d + 1 and suppose α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1) and d is sufficiently large. Then W ⋆ and L ⋆ have approximate forms

Here, (10) is reminiscent of Corollary 1 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 1 whereas we need to set α = O GLYPH<16> 1 / √ d GLYPH<17> 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 4 in Appendix B.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 LLMcan 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 GLYPH<12> GLYPH<12> GLYPH<12> β 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 , yi ) n + 1 i = 1 as follows:

Above, κ -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 5 in Appendix B.4. Similar to the RAG setting, by assuming α = O GLYPH<16> 1 / √ d GLYPH<17> and d / n = O (1), we obtain the following results for the optimal parameter and risk.

Theorem 3 Consider linear model as defined in (11) . Recap the objective from (5a) and let W ⋆ : = arg min W L PGD ( W ) , and L ⋆ = L PGD ( W ⋆ ) . Additionally, given κ = α 2 d + 1 and suppose α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1) and d is sufficiently large. Then W ⋆ and L ⋆ have approximate forms

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

3.2 Low-rank Parameterization and LoRA

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 Consider independent linear data as defined in (7) . Recap the objective from (5a) and enforce rank ( W ) ≤ r and W ⊤ = W . Let Σ = Σ 1 / 2 x Σ β Σ 1 / 2 x and M = tr ( Σ ) + σ 2 . Denoting λ i to be the i'th largest eigenvalue of Σ , we have that

0

10

20

30

40

50

# in-context samples

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Test risk

Linear Att

H3

Theorem 1

(a)

Σ

x

=

Σ

β

=

I

d

and

σ

=

0

0

10

20

30

40

50

# in-context samples

0.4

0.6

0.8

1.0

1.2

Test risk

2

/

d

=0

2

/

d

=0.1

2

/

d

=0.2

2

/

d

=0.3

(b) Noisy label

0

10

20

30

40

50

# in-context samples

0.2

0.4

0.6

0.8

Test risk

=0

=0.3

=0.6

=0.9

(c) Non-isotropic task

Note that tr ( Σ ) = P d i = 1 λ i . Removing the rank constraint and considering noiseless data setting, this reduces to the following optimal risk L ⋆ = P d i = 1 λ i + M n + 1 + M /λ i . See Appendix C.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 population loss ˜ L ( W lora ) : = L ( W + W lora ) in (5a) with fixed W . Suppose tr ( Σ ) = tr ( Σ new ) = M , σ = 0 and W is jointly diagonalizable with Σ , Σ new , then LoRA's risk is upper-bounded by

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 .

4 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 1 as well as Theorem 1, we run random linear regression instances where in-context samples are generated obeying (7). Fig. 2a is identical to Fig. 1a where we set Σ x = Σ β = I d and σ = 0. In Fig. 2b, set Σ x = Σ β = I and vary noise level σ 2 from 0 to 0 . 3 × d . In Fig. 2c, we consider noiseless labels, σ = 0, isotropic feature distribution Σ x = I d and set task covariance to be Σ β = γ 11 ⊤ + (1 -γ ) I d by choosing γ in { 0 , 0 . 3 , 0 . 6 , 0 . 9 } . Note that in Fig. 2c, we train a sufficient number of models (greater than 20) to ensure the optimal model is obtained. In all the figures, solid and dashed curves correspond to the

0

10

20

30

40

50

# in-context samples

0.2

0.4

0.6

0.8

Test risk

=0

=0.2

=0.4

=0.6

(a) RAG

0

10

20

30

40

50

# in-context samples

0.2

0.4

0.6

0.8

Test risk

=0

=0.2

=0.4

=0.6

(b) Task-feature alignment

0

10

20

30

40

50

# in-context samples

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Test risk

r

=1

r

=5

r

=10

r

=20

(c) Low-rank attention

0

10

20

30

40

50

# in-context samples

0.1

0.3

0.5

0.7

0.9

Test risk

r

=1

r

=5

r

=10

r

=20

(d) LoRA adaptation

ICL results from training 1-layer ATT and SSM models, respectively, and dotted curves are obtained from (8) in Theorem 1. The alignment of solid, dashed and dotted curves validates our Proposition 1 and Theorem 1.

Distributional alignment experiments (Figs. 3a&3b). In Figs. 3a and 3b, we generate RAG and task-feature alignment data following (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 B.3 and (42) in Appendix B.4. The empirical alignments corroborate Theorems 4 and 5, further confirming that Proposition 1 is applicable to a broader range of real-world distributional alignment data.

Low-rank (Fig. 3c) and LoRA (Fig. 3d) 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. 3c, we set Σ x = I d , σ = 0 and task covariance to be diagonal with diagonal entries c [1 2 -1 · · · d -1 ] ⊤ for some normalization constant c = d / P 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. 3d, 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 ˆ W k , ˆ W q , ˆ W v ∈ R ( d + 1) × ( d + 1) . Phase 2: We generate new examples with task covariance Σ β being a diagonal matrix with diagonal entries c ′ [2 -1 2 -2 · · · 2 -d ] ⊤ for some normalization constant c ′ = d / P d i = 1 2 -i . Given the rank restriction r , we train additional LoRA parameters W up , W down ∈ R ( d + 1) × r where W lora : = W up W ⊤ down and (2a) becomes ATT ( Z ) = ( Z ( ˆ W q ˆ W ⊤ k + W up W ⊤ down ) Z ⊤ ) Z ˆ W v . Fig. 3d 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.

5 Related Work

There is growing interest in understanding the mechanisms behind ICL [Brown et al., 2020, Liu et al., 2023b, Rae et al., 2021] in LLMs due to its success in continuously enabling novel applications for LLMs [GeminiTeam et al., 2023, OpenAI, 2023, Touvron et al., 2023]. In the previous work, Garg et al. [2022] explored ICL ability of Transformers. In particular, they considered in-context prompts where each in-context example is labeled by a target function from a given function class, including linear models. A number of works have studied this and related settings to develop a theoretical understanding of ICL [von Oswald et al., 2023, Gatmiry et al., Collins et al., 2024, Lin and Lee, 2024, Li et al., 2024, Bai et al., 2024, Akyürek et al., 2023, Zhang et al., 2023, Du et al., 2023]. Akyürek et al. [2023] focus on linear regression and provide a construction of Transformer weights that can enable a single step of GD based on in-context examples. Along the similar line, Von Oswald et al. [2023] provide a construction of weights in linear attention-only Transformers that can emulate GD steps on in-context examples for a linear regression task. Similar to this line of work, Dai et al. [2023] argue that pre-trained language models act as meta-optimizer which utilize attention to apply meta-gradients to the original language model based on the in-context examples.

Building on these primarily empirical studies, Zhang et al. [2024], Mahankali et al. [2024], Ahn et al. [2023], Duraisamy [2024] focus on developing a theoretical understanding of Transformers trained to perform ICL. For single-layer linear attention model trained on independent in-context prompts for random linear regression tasks, Mahankali et al. [2024], Ahn et al. [2023] show that the resulting model implements a single step of PGD on in-context examples in a test prompt, thereby corroborating the findings of [Von Oswald et al., 2023]. Zhang et al. [2024] study the optimization dynamics of gradient flow while training a single-layer linear attention model on in-context prompts for random linear regression tasks. Similar to Mahankali et al. [2024], Ahn et al. [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.

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

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

References

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

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= 0g0X4H8yN4I .

Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332. PMLR, 2019.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems , 36, 2024.

Soumya Basu, Ankit Singh Rawat, and Manzil Zaheer. A statistical perspective on retrieval-based models. In International Conference on Machine Learning , pages 1852-1886. PMLR, 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communications , 12(1):2914, 2021.

Yuan Cao, Zhiying Fang, Yue Wu, Ding-Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint arXiv:1912.01198 , 2019.

Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts to function lipschitzness. arXiv preprint arXiv:2402.11639 , 2024.

Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizers. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Findings of the Association for Computational Linguistics: ACL 2023 , pages 4005-4019, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.findings-acl.247. URL https://aclanthology.org/2023.findings-acl.247 .

Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning , pages 933-941. PMLR, 2017.

Zhe Du, Haldun Balim, Samet Oymak, and Necmiye Ozay. Can transformers learn optimal filtering for unknown systems? IEEE Control Systems Letters , 7:3525-3530, 2023.

Karthik Duraisamy. Finite sample analysis and bounds of generalization error of gradient descent in in-context linear regression. arXiv preprint arXiv:2405.02462 , 2024.

Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/ forum?id=COZDy0WYGg .

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Forty-first International Conference on Machine Learning .

GeminiTeam, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arXiv preprint arXiv:2402.03170 , 2024.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 , 2023.

Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations , 2021a.

Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems , 34:572-585, 2021b.

Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305.12766 , 2023.

Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: A transformer that solves small tabular classification problems in a second. arXiv preprint arXiv:2207.01848 , 2022.

Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id= nZeVKeeFYf9 .

Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research , 24(251):1-43, 2023.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning , pages 5156-5165. PMLR, 2020.

Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Exploring the relationship between model architecture and in-context learning ability. arXiv preprint arXiv:2310.08049 , 2023.

Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems , 33: 9459-9474, 2020.

Mingchen Li, Mahdi Soltanolkotabi, and Samet Oymak. Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks. In International conference on artificial intelligence and statistics , pages 4313-4324. PMLR, 2020.

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

Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought: Compositionality through in-context filtering and learning. Advances in Neural Information Processing Systems , 36, 2024.

Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. arXiv preprint arXiv:2402.18819 , 2024.

Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations , 2023a. URL https://openreview.net/forum?id=De4FYqjFueZ .

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

Yue Lu, Mary I. Letey, Jacob A. Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention. arXiv preprint arXiv:2310.08391 , 2024.

Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id= 8p3fu56lKc .

Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Revisiting the equivalence of in-context learning and gradient descent: The impact of data distribution. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 7410-7414. IEEE, 2024.

Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. arXiv preprint arXiv:2112.10510 , 2021.

Samuel Müller, Matthias Feurer, Noah Hollmann, and Frank Hutter. Pfns4bo: In-context learning for bayesian optimization. In International Conference on Machine Learning , pages 25444-25470. PMLR, 2023.

Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, et al. Webgpt: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332 , 2021.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 , 2022.

OpenAI. Gpt-4 technical report. arXiv preprintarXiv:2303.08774 , 2023.

Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. International Conference on Machine Learning , 2024.

Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446 , 2021.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning , pages 9355-9366. PMLR, 2021.

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

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858 , 2023.

Shuohang Wang, Yichong Xu, Yuwei Fang, Yang Liu, Siqi Sun, Ruochen Xu, Chenguang Zhu, and Michael Zeng. Training data is more valuable than you think: A simple and effective method by retrieving from training data. arXiv preprint arXiv:2203.08773 , 2022.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems , 35:24824-24837, 2022.

Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391 , 2023.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=RdJVFCHjUMI .

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

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55, 2024.

Nicolas Zucchet, Seijin Kobayashi, Yassir Akram, Johannes Von Oswald, Maxime Larcher, Angelika Steger, and Joao Sacramento. Gated recurrent neural networks discover attention. arXiv preprint arXiv:2309.01775 , 2023.

Appendix

Table of Contents

A Equivalence among Gradient Descent, Attention, and State-Space Models

In this section, we present the proofs related to Section 2. Recap that given data

and corresponding prediction functions

Here, the expectation is over the randomness in ( x i , ξ i ) n i = 1 and β , and the search space for W is R d × d , for ω is R n , for W k , W q , W v is R ( d + 1) × ( d + 1) , for v is R d + 1 , and for f is R n + 1 .

we have objectives

A.1 Proof of Proposition 1

Consider the problem setting as discussed in Section 2, Proposition 1 can be proven by the following two lemmas.

Lemma 4 Suppose Assumptions 1 and 2 hold. Then, given the objectives (16a) and (16c) , we have

Proof. Recap the linear attention estimator from (15c) and denote

where ¯ W ∈ R d × d , w 1 , w 2 , v 1 ∈ R d , and w , v ∈ R . Then we have

We first show that for any given parameters W k , W q , W v , v ,

To this goal, we have

where we have decomposition

In the following, we consider the expectations of ( a ) , ( b ) , ( c ) , ( d ) sequentially, which return zeros under Assumptions 1 and 2. Note that since Assumption 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 therefore E [ x ⊤ β x ] = 0 d .

Combining the results with (19) returns that

which completes the proof of (18). Therefore, we obtain

We conclude the proof of this lemma by showing that for any W ∈ R d × d in g PGD , there exist W k , W q , W v , v such that g ATT ( Z ) = g PGD ( Z ). Let

Then we obtain

which completes the proof.

Lemma 5 Suppose Assumptions 1 and 2 hold. Then, given the objectives in (16) , we have

Additionally, if the examples ( x i , yi ) 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 and

Proof. Recap the SSM estimator from (15d) and let

where w qj , w k j , w v j ∈ R d + 1 for j ≤ d + 1, and let

Then we have

Next for all j ≤ d + 1, let

To start with, we obtain

where there is decomposition

In the following, similar to the proof of Lemma 4, we consider the expectations of ( a ) , ( b ) , ( c ) , ( d ) sequentially, which return zeros under Assumptions 1 and 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

where ¯ w qj , ¯ w k j , ¯ w v j ∈ R d and wqj , wkj , wvj ∈ R . Then we have

where

Then

where

Next we will show that for any W k , W q , W v , v ,

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.

Combining the results with (24) results that

Therefore we obtain,

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

Then we get and therefore

which completes the proof of (22).

Next, to show (23), for any W ∈ R d × d , let L ( ω ) = E h GLYPH<0> x ⊤ WX ⊤ ( y ⊙ ω ) -y GLYPH<1> 2 i . Then we have

Here since ( x i , yi ) n i = 1 follow the same distribution and are conditionally independent given x and β , for any i , j , j ′ , E h ( x ⊤ W yi x i ) 2 i = E h ( x ⊤ W yj x j ) 2 i and E h ( x ⊤ W yj x j )( x ⊤ W yi x i ) i = E h ( x ⊤ W yj ′ x j ′ )( x ⊤ W yi x i ) i . Then let

where ( c 1 , c 2 , c 3 ) : = ( c 1 ( W ) , c 2 ( W ) , c 3 ( W )). We get

If c 2 -c 1 = 0, then ∂ L ( ω ) ∂ω i ≡ 2 c 1 ω ⊤ 1 n -2 c 3 for all i ≤ n and any ω ∈ R n achieves the same performance.

If c 2 -c 1 , 0, setting ∂ L ( ω ) ∂ω i = 0 returns

Therefore the optimal loss is achieved via setting ω = C 1 n . Without loss of generality, we can update W → C W . Then ω = 1 n , and we obtain

which completes the proof of (23).

A.2 Proof of Lemma 1

Proof. Recap the loss L PGD ( W ) in (16a) and prediction g PGD ( Z ) in (15a), we have

where (25) follows Assumption 2. Since f 2 ( W ) is convex, L PGD ( W ) is strongly-convex if and only if f 1 ( W ) is strongly-convex, which completes the proof of strong convexity.

Next, (20) and (21) in the proof of Lemma 4 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 .

A.3 Proof of Lemma 2

Proof. According to Lemma 1, L PGD ( W ) is strongly-convex as long as either E [( x ⊤ WX ⊤ X β ) 2 ] or E [( x ⊤ WX ⊤ ξ ) 2 ] is strongly-convex. Therefore, in this lemma, the two claims correspond to the strong convexity of E [( x ⊤ WX ⊤ ξ ) 2 ] and E [( x ⊤ WX ⊤ X β ) 2 ] terms, respectively.

Suppose the decomposition claim holds. Without losing generality, we may assume ( x 1 , β 1 , X 1 ) are zero-mean because we can allocate the mean component to ( x 2 , β 2 , X 2) without changing the covariance.

Claim 1: Let ¯ Σ x = E [ x 1 x ⊤ 1 ], ¯ Σ β = E [ β 1 β ⊤ 1 ], and ¯ Σ X = E [ X ⊤ 1 X 1 ]. If the first claim holds, using independence, observe that we can write

where the last three terms of the right hand side are convex and the first term obeys

Since noise level σ > 0, using the full-rankness of covariance matrices ¯ Σ x and ¯ Σ X , we conclude with strong convexity of E [( x ⊤ WX ⊤ ξ ) 2 ].

Claim 2: Now recall that ¯ Σ X = E [ X ⊤ 1 X 1] and set A = X ⊤ 1 X 1 -¯ Σ X and B = X ⊤ 2 X 2 + ¯ Σ X . Observe that E [ A ] = 0. If the second claim holds, E [ X ⊤ X ] = E [ A + B ]. Note that ( A , β 1 , x 1) are independent of each other and ( B , β 2 , x 2). Using independence and E [ A ] = 0, similarly write

Now using E [ β 1 ] = E [ x 1 ] = 0 and their independence from rest, these terms obeys

In both equations, the last three terms of the right hand side are convex. To proceed, we focus on the first terms. Using independence and setting Σ X = E [ X ⊤ X ] ⪰ ¯ Σ X ≻ 0, we note that

where x 1 , β 1 , X are independent and full-rank covariance. To proceed, note that

Observing the convexity of the right hand side and focusing on the first term, we get

Using the fact that covariance matrices, ¯ Σ x , Σ X , ¯ Σ β , are full rank concludes the strong convexity proof of E [( x ⊤ WX ⊤ X β ) 2 ].

B Analysis of General Data Distribution

In this section, we provide the proofs in Section 3, which focuses on solving Objective (5a). For the sake of clean notation, let L ( W ) : = L PGD ( W ) and g : = g PGD in this section.

B.1 Supporting Results

We begin by deriving the even moments of random variables.

2 n 'th moment of a normally distributed variable: Let u ∼ N (0 , σ 2 ). Then we have

4 'th moment: Let u ∼ N (0 , I d ). Then for any W , W ′ ∈ R d × d , we have

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 have

6 'th moment: Let u ∼ N (0 , I d ). Then for any W , W ′ ∈ R d × d , we have

where (29) is obtained by following

8 'th moment: Let u ∼ N (0 , I d ). Then for any W , W ′ ∈ R d × d , we have

where (31) is obtained by following

B.2 Independent Data with General Covariance

Proof of Theorem 1. Consider a general independent linear model as defined in (7) where Σ x and Σ β are full-rank feature and task convariance matrices and

Let

To simplify and without loss of generality, let ¯ x = Σ -1 / 2 x x , ¯ X = X Σ -1 / 2 x , ¯ β = Σ 1 / 2 x β where we have

and

Then recap the loss from (5a), and we obtain

where the last equality comes from the independence of label noise ξ, ξ .

We first consider the following term

where Σ : = Σ 1 / 2 x Σ β Σ 1 / 2 x .

Let ¯ x i ∈ 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

Then we get

Therefore

Combining all together results in

where M : = tr ( Σ ) + σ 2 . Setting ∇ ¯ W L ( W ) = 0 returns

Then we have and

B.3 Retrieval Augmented Generation with α Correlation

In 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

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 sample

which results in (9) by setting γ 2 = 1 -α 2 .

Theorem 4 (Extended version of Theorem 2) 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 ⋆ satisfy

where

Suppose α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1) and d is sufficiently large. Let κ = α 2 d + 1 and γ 2 = 1 -α 2 . Then W ⋆ and L ⋆ have approximate forms

Proof. Here, for clean notation and without loss of generality, we define and rewrite (35) via

Then we obtain

To begin with, let

We first focus on the second term in (38)

Next, the first term in (38) can be decomposed into

In the following, we consider solving ( a )-( c ) sequentially.

Here, ( b ) utilizes the 4'th and 6'th moment results (27) and (30) and we define

Then combining all together results in

To find the optimal solution, set ∇L ( W ) = 0 and we obtain

Note that we have

Therefore, (39) returns

2 A 1 GLYPH<16> tr ( W ) I + W + W ⊤ GLYPH<17> + 2 A 2 W -2 A 3 + 2 n σ 2 ( α 2 ( tr ( W ) I + W + W ⊤ ) + γ 2 W ) I = 0 , (40) which implies that the optimal solution W ⋆ has the form of c I for some constant c . Then suppose W ⋆ = c I , we have tr ( W ) = cd and (40) returns

Then the optimal loss is obtained by setting W ⋆ = c I and

It completes the proof of (36). Now if assuming α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1) and sufficiently large dimension d , we have the approximate

and

B.4 Task-feature Alignment with α Correlation

In this section, we consider the task-feature alignment data model similar to (11), where we first draw task vector β via

Then we generate examples ( x i , yi ) n + 1 i = 1 according to the rule corr_coef( x i , β ) ≥ α ≥ 0 via

which results in (11) by setting γ 2 = 1 / ( α 2 d + 1).

Theorem 5 (Extended version of Theorem 3) Consider linear model as defined in (41) . Recap the objective from (5a) and let W ⋆ : = arg min W L PGD ( W ) , and L ⋆ = L PGD ( W ⋆ ) . Then W ⋆ and L ⋆ satisfy

where and

Suppose α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1) and d is sufficiently large. Let κ = α 2 d + 1 and γ 2 = 1 /κ . Then W ⋆ and L ⋆ have approximate forms

Proof. Here, for clean notation and without loss of generality, we define and rewrite (41) via

Similar to Appendix B.3, to begin with, let and additionally, given Λ = W ⊙ I

N 1 = tr ( W ) 2 + tr GLYPH<16> WW ⊤ GLYPH<17> + tr GLYPH<16> W 2 GLYPH<17> , N 2 = tr GLYPH<16> WW ⊤ GLYPH<17> , and N 3 = tr ( W ) , W , let

We first focus on the second term in (44)

Next, the first term of (44) (omitting γ 2 ) returns the following decomposition:

Consider solving ( a )-( d ) sequentially as follows:

To begin with, we use the following decomposition for all ( a )-( d ):

Then, we have

where (45) and (47) utilize (30) and (32), and (46) is obtained via

which follows (27) and (28).

where (49) and (50) are obtained using (27), (30) and

Here we define

Then combining all together results in

and differentiating it results in

Similar to the proof in Appendix B.3, W ⋆ has the form of W ⋆ = c I and we have

Therefore, setting ∇L ( W ) = 0 returns

Then the optimal loss is obtained by setting W ⋆ = c I and

It completes the proof of (42). Now if assuming α = O GLYPH<16> 1 / √ d GLYPH<17> , d / n = O (1), γ 2 = 1 / ( α 2 d + 1) and sufficiently large dimension d , we have the approximate

and

C Analysis of Low-Rank Parameterization

C.1 Proof of Lemma 3

Proof. Recall the loss function from (34)

where ¯ W = Σ 1 / 2 x W Σ 1 / 2 x , Σ = Σ 1 / 2 x Σ β Σ 1 / 2 x 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

0

10

20

30

40

50

60

70

80

# in-context samples

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Test risk

=30

=50

=80

(a) Linear attention

0

10

20

30

40

50

60

70

80

# in-context samples

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Test risk

=30

=50

=80

(b) H3

30

40

50

60

70

80

0.45

0.50

0.55

0.60

Averaged test risk

Linear Att

H3

(c) Averaged risk

0

10

20

30

40

50

# in-context samples

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Test risk

Linear Att

H3

(d) Evolving

β

matrix. We will first treat U as fixed and optimize E . We will then optimize U . Fixing U , setting ¯ Σ = U ⊤ Σ U , we obtain

Differentiating, we obtain

Setting ∇L ( E ) = 0 returns

Let ¯ λ i denote the i 'th largest eigenvalue of ¯ Σ . Plugging in this value, we obtain the optimal risk as a function of U is given by

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 P 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 where

The equality holds by setting U to be the topr eigenvectors of Σ and E = E ⋆ ( U ) to be the diagonal matrix according to (51).

D Additional Experiments

In this section, we present additional experiments demonstrating that the H3 model can outperform the linear attention model under different training or data settings. The implementation details are consistent with those outlined in Section 4.

· 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. 4a and 4b, instead of training separate models to fit the different context lengths, we train a single model with 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) with Σ x = Σ β = I d and σ = 0, and train 1-layer linear attention (Fig. 4a) and H3 (Fig. 4b) models with different max-lengths n max = 30 , 50 , 80. Comparison between Fig. 4a and 4b shows that 1-layer attention and H3 implement different algorithms in solving the averaged linear regression problem and H3 is more consistent in generalizing to longer context lengths. In Fig. 4c, we plot the averaged risks for each model and H3 outperforms linear attention. Furthermore, in Fig. 4d, we focus on the setting where in-context examples are generated using evolving task vector β . Specifically, consider that each sequence corresponds to two individual task parameters β (1) ∼ N (0 , I d ) and β (2) ∼ N (0 , I d ). Then the i 'th sample is generated via x i ∼ N (0 , I d ) and yi = β ⊤ i x i where β i = λ i β (1) + (1 -λ i ) β (2) and λ i = i / n . The results are reported in Fig. 4d 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 1.

E Extended Related Work

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]. Towards this, Xie et al. [2022] explain ICL by language model's ability to perform implicit Bayesian inference where, under specific assumptions on the pre-training data distribution, the model infers a shared latent concept among the in-context examples and leverages the concept to make a prediction. Müller et al. [2021], Hollmann et al. [2022], Müller et al. [2023] introduce prior-data fitted network (PFN) to approximate Bayesian inference on synthetic datasets and use it to perform downstream tasks such as tabular dataset classification. On the other hand, Olsson et al. [2022] posit induction heads as the key mechanism enabling ICL in Transformers. Park et al. [2024] study how various distributional properties of training data aid in the emergence of ICL in Transformers.

In the previous work, Garg et al. [2022] explored ICL ability of Transformers. In particular, they considered in-context prompts where each in-context example is labeled by a target function from a given function class, including linear models. A number of works have studied this and related settings to develop a theoretical understanding of ICL [von Oswald et al., 2023, Gatmiry et al., Collins et al., 2024, Lin and Lee, 2024, Li et al., 2024, Bai et al., 2024, Akyürek et al., 2023, Zhang et al., 2023, Du et al., 2023]. Akyürek et al. [2023] focus on linear regression and provide a construction of Transformer weights that can enable a single step of GD based on in-context examples. They further show that Transformers trained on in-context prompts exhibit behaviors similar to the models recovered via explicit learning algorithm on the in-context examples in a prompt. Along the similar line, V on Oswald et al. [2023] provide a construction of weights in linear attention-only Transformers that can emulate GD steps on in-context examples for a linear regression task. Interestingly, they find similarity between their constructed networks and the networks resulting from training on in-context prompts corresponding to linear regression tasks. Similar to this line of work, Dai et al. [2023] argue that pre-trained language models act as meta-optimizer which utilize attention to apply meta-gradients to the original language model based on the in-context examples. Focusing on various NLP tasks, they further connect it to a specific form of explicit fine-tuning that performs gradient updates to the attention-related parameters. Inspired by the connection between linear attention and GD, they developed a novel attention mechanism that mirrors the behavior of GD with momentum. Beyond Transformers, existing work [Lee et al., 2023, Zucchet et al., 2023, Grazzi et al., 2024] demonstrate that other model architectures, such as SSM and RNNs, are also capable of in-context learning (ICL).

Building on these primarily empirical studies, Zhang et al. [2024], Mahankali et al. [2024], Ahn et al. [2023], Duraisamy [2024] focus on developing a theoretical understanding of Transformers trained to perform ICL. For single-layer linear attention model trained on in-context prompts for random linear regression tasks with isotropic Gaussian features and isotropic Gaussian weight vectors, Mahankali et al. [2024], Ahn et al. [2023] show that the resulting model implements a single step of GD on in-context examples in a test prompt, thereby corroborating the findings of [Von Oswald et al., 2023]. They also show that the learned model implements a PGD step, when faced with anisotropic Gaussian features, with Mahankali et al. [2024] also considering anisotropic Gaussian weight vectors. Ahn et al. [2023] further study multi-layer model and show that the trained model can implement a generalization of GD++ algorithm, supporting an empirical observation in Von Oswald et al. [2023]. On the other hand, Mahankali et al. [2024] extend their single-layer setup to consider

suitable non-linear target functions, showing that learned Transformer again implements a single step of GD on lineare regression objective. For a single-layer linear attention model, Zhang et al. [2024] study the optimization dynamics of gradient flow while training such a model on in-context prompts for random linear regression tasks. Despite the non-convexity of the underlying problem, they show the convergence to the global minimum of the population objective. Similar to Mahankali et al. [2024], Ahn et al. [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 and 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.

NeurIPS Paper Checklist

1. Claims

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

Answer: [Yes]

Justification: All the theoretical contributions claimed in the abstract and introduction along with the underlying data model are presented in Section 2 and Section 3.

Guidelines:

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

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

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

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

2. Limitations

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

Answer: [Yes]

Justification: This is a theoretical study which (similar to prior studies in the field) relies on a precise but simplified data model to draw quantitatively precise conclusions. All the assumptions on the data model are clearly stated in Section 2 and Section 3. We have also added a paragraph after the conclusion to specifically highlight various limitations of our work.

Guidelines:

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

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

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

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

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

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

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

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

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: As discussed above, for all of our theoretical results and proofs we state the precise setup and assumptions in Section 2 and Section 3. Due to page limit, proofs are deferred to the supplemental material.

Guidelines:

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

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

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

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

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

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

4. Experimental Result Reproducibility

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

Answer: [Yes]

Justification: This is primarily a theoretical work where detailed synthetic experiments (on the same data model studied in our theoretical analysis) have been conducted to corroborate our theoretical findings. We provide sufficient details in Section 4 for reproducing these experiments.

Guidelines:

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

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

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

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

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

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

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

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

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

5. Open access to data and code

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

Answer: [No]

Justification: As discussed above, this paper conducts small scale synthetic experiments to corroborate our theoretical findings. We have provided sufficient details to reproduce these experiments in Section 4.

Guidelines:

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

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

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

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

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

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

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

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

6. Experimental Setting/Details

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

Answer: [Yes]

Justification: All the relevant details for our small scale experiments are provided in Section 4.

Guidelines:

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

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

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

7. Experiment Statistical Significance

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

Answer: [No]

Justification: This work is a theoretical work studying the optimization landscape of linear attention/H3 under population risk. Then our goal of simulations is to find the optimal solution corresponding to the minimal risks. Therefore, we do not report the error bars and we have included the discussion in the experiment section.

Guidelines:

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

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

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

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

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

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

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

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

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

8. Experiments Compute Resources

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

Answer: [No]

Justification: Our work only focuses on 1-layer attention/H3 model training with hidden dimension 21 and maximal context length < 100 , which can be implemented easily.

Guidelines:

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

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

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

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: Yes, the authors confirm that the research conducted in the paper conform wiht the NeurIPS Code of Ethics.

Guidelines:

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

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

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

10. Broader Impacts

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

Answer: [Yes]

Justification: In its current form, we don't see any specific negative impacts of our theoretical study. However, we have discussed potential broader impacts of the future extensions of this work, e.g., the ones tied to eliciting undesirable behavior of LLMs with in-context learning, while discussing the limitations of the work after the conclusion section.

Guidelines:

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

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

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

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

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

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

11. Safeguards

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

Answer: [NA]

Justification: The synthetic setup studied in the paper does not pose such risks.

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: The paper does not rely on existing assets.

Guidelines:

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

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

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

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

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

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

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

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

13. New Assets

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

Answer:[NA]

Justification: The paper does not release new assets such as code, data, or models.

Guidelines:

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

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

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

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

14. Crowdsourcing and Research with Human Subjects

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

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

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

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

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

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

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

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

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

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

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

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nThe 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.\n\nNext, the authors give an upper bound on the optimal expected loss (for 1-step PGD <=> ATT <=> SSM) in a low-rank adapter setting.\n\nExperiments validate the theoretical results. \n\nA 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.\n\nSTRENGTHS:\nIn terms of results, the main original contribution is the analysis of the model with correlated training examples and query points. However, the analysis of the SSM model in the appendix seems non-trivial and should therefore also be considered as a significant element of originality.\n\nThe section on related work is well-written and covers relevant work.\n\nThe problem setup is very clean, making it very easy to read and understand the results. All results are clearly stated with the necessary assumptions and accompanied by clear explanations.\n\nThe questions studied are highly relevant to contemporary machine learning research and applications of AI.\n\nWEAKNESSES:\n* 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).\n\n* 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).\n\n* 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.\n\nQUESTIONS:\n* 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, ...)?\n\n* Can you elaborate on how claim (A3) follow from (14)?\n\n* Define $\\alpha$-correlated\n* l.342: wildly studied -> widely studied (?)\n\nLIMITATIONS:\nThe authors address the limitations of their work and acknowledge that their analyses are \"not precise and fully formal\".", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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:\n\n1. The landscape of 1-layer linear attention and 1-layer H3 (a state-space model), proving both can implement 1-step preconditioned gradient descent.\n\n2. New risk bounds for retrieval augmented generation (RAG), showing benefits from distributional alignment.\n\n3. The optimal risk for low-rank parameterized attention weights, illustrating how LoRA adapts to new distributions.\n\nSTRENGTHS:\nThe authors provide a comprehensive theoretical analysis of the landscape for ICL, extending existing results to more practical settings. It offers new insights into the benefits of distributional alignment and the capabilities of low-rank parameterization in attention mechanisms. The theoretical findings are supported by experimental results, enhancing the credibility of the conclusions.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\n2. 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.\n\n3. The paper focuses on single-layer (linear) models. How do the findings extend to multi-layer (linear) architectures?\n\n4. 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?\n\n5. 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?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe presentation is clear.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nWhat 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.\n\nWhat 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.\n\nWhy 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$?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}