Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond
Authors: Yingcong Li
Ankit Singh Rawat
Samet Oymak
Venue: NeurIPS 2024 poster (2024)
Keywords: In-context learning
linear attention
state-space model
optimization
RAG
LoRA
Primary area: optimization_for_deep_networks
TL;DR: We study the loss landscape of in-context learning for single-layer linear attention and state-space models under general linear tasks model while delineating the effect of distributional alignments (e.g., RAG), low-rank constraints, and 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
Yingcong Li University of Michigan yingcong@umi ch. edu
Ankit Singh Rawat Google Research NYC anki mikitsraatt@poggee.com
Samet Oymak University of Michigan oymakθumi ch. edu
Abstract
Recent research has shown that Transformers with linear attention are capable in-context learning (ICL) by implementing linear estimator through gradient descent steps. However, the existing results 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, develop stronger characterization of the optimization and generalization landscape ICL through contributions architectures, low-rank parameterization, and correlated designs: (1) We study the landscape of I-layer linear attention and 1 1-layer H3, a state- space model. Under suitable correlated design assumption, 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 suitable settings. (2) By studying correlated designs, we provide new risk bounds for retrieval augmented generation (RAG) and task-fcature 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 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 more thorough understanding of its mechanics.
prove that both
Architecture Choice
Distributional Alignment
Low-rank Parameterization
Correlated features
Low-rank attention weights: rank(W, W,)
ATTIN
Retrieval
(OKCTV)
SSM
b) LORA adaptation
Query features
Relevant examples
W
Pretrained Weights
Database
rank
b) festure a alignment: Task and feature vectors e-comellaeed
Wr
Uneer Attention
HB (SSM)
Uncer Att
=
0.2
0.7
r18k
Test rick
0.6
0.6
Test risk
Fer 10
0.5
0.3
0.1
samples
samples
samples
Linear attention=H3
(b) RAG with correlation
(c) LoRA
Figure 1: We investigate the optimization landscape of in-context learning from the lens of archi-
tecture choice, the role of distributional alignment, and low-rank parameterization. The empirical performance (solid curves) are aligned with our theoretical results (dotted curves) from Section 3. More experimental details and discussion are deferred to Section 4.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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, DeminiTeam et al., 2023, OpenAI, 2023, Touvron et al., 2023]. Such in-context learning (ICL) offers novel and effective alternative to traditional fine-tuning techniques and has become an important feature of LLM generation [Lewis et al., 2020], and reasoning via advanced prompting techniques, such as chain-of-thoughe [Wei et al., 2022].
with its applications spanning etrreval--uggeed
ICL ability presents an important research avenue to develop stronger theoretical and mechanistic understanding of large language models. To this aim, there has been sigmificant 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 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, revisit the theoretical exploration of ICL with linear data model where we feed an in-context prompt containing n examples (x ; = + C R x R and a test instance or query Ka+l E Rd to the model, with d being the feature dimension, B E Idd being the task weight vector, and (E))" denoting the noise in individual labels. Given the in-context prompt, the model is tasked to predict 9n+1 - an estimate for Ya+1 = x B++ Eu+1. We aim to provide answers to the following questions by exploring the loss landscape of ICL:
(Q1) Is the ability 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 (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 Wx, W E Raxd, 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 I-layer models and make the following contributions:
(A1) We jointly investigate the landscape of linear attention and H3 [Fu al., 2023], widely popu- lar 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. la 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 camporaaly-heterrggeeeeo problem settings in Appendix D.
(A2) Proposition I 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 and Fig. 1b) and Task-Feature Alignment (c.f. Theorem 3), where the in-context examples are a-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 a²d + 1, highlighting that aligned data are crucial for the success of ICI. 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. Ic), 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., and X) represent vectors and matrices, respectively. The symbol @ is defined as the element- wise (Hadamard) product, and denotes the convolution operator. 1a and 0g denote the d-dimensional all-ones and all-zeros vectors, respectively; and la denotes the identity matrix of dimension d x 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 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 [2023], Mahankali et al. [2024], Zhang et al. [2024] as follows.
Linear data distribution. Let (x. y) E Rdd x R be a (feature, label) pair generated by a d-dimensional linear model parameterized by B e R i.e. y = x" B where x and ß are feature and task vectors, and & is the label noise. Given demonstrations (1) sampled from a single B, define the input in-context prompt
Z = [z]
Zn Zn+1 1i=
x1 y1
Xa yn
Nnt1 0
E R (n+ (x(dd)) l)
x for i n and the last/query token Zn-1 yi
Kn+1
Here, we set z,
Then, given Z, the goal of
0
the model is to predict the correct label Yx+1 corresponding to kn+1 For cleaner notation, when it is clear from context, we drop the subscript n 1 and set x = Xn+l, Z zæ+1- Different from the previous work [Ahn et al., 2023, Mahankali et al., 2024, Zhang et al., 2024, Mahdavi et al., 2024] where (x:) and B are assumed to be independent, our analysis focuses on a more general linear setting that captures the dependency between (x)) and B.
Model architectures. To start with, we first review the architectures of both Transformer and state-space model (SSM). Similar the previous work [Von Oswald et al., 2023, Ahn al., 2023, Mahankali et al., 2024, Zhang et al., 2024] and to simplify the model structure, we focus on single- layer models and omit the nonlinearity, e.g., softmax operation and MLP activation, from the Transformer. Given the imput prompt Z E (==1))(d+ in (1), which can be treated as sequence of (d + D-dimensional tokens, the single-layer linear attention ATT and H3-like single-layer SSM SSM are denoted by
ATT(2) (ZW, 4 W z JZW, SSM(Z) (ZZW,) o (ZWW, @ ZW,) f)
(2a)
(2b)
where W.. Wg W, R(d+1xd++1) denote the key, query and value weight matrices, respectively.
In (2b), the parameter f E R=*1 is 1-D convolutional filter that mixes tokens. The Hadamard product O 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 I only, we use indexing f Ifo fa] E R and given any vector a, denote convolution output (a f); = Σ/ ay Note that our notation slightly differs from the original H3 model [Fu et al., 2023] in two ways:
1. SSMs provide efficient parameterization of f which would otherwise grow with sequence length. In essence, H3 utilizes linear state-space model s, = AA Bu, and y = Cs) with parameters (A E R B E Rdxl C E R. !xd) from which the filter f is obtained via the impulse response
f== CA B for i 2 Here d is the state dimension and, in practice, A is chosen to be diagonal. Observe that, setting d 1 and A p. C. C= B = 1, SSM reduces to the exponential smoothing f = for 2 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 filter f as in (2b) strictly more expressive than the SSM as it captures all possible impulse responses.
2. H3 also applies a shift SSM to the key embeddings to enable the retrieval of the local context around associative recall hits. We opted not to incorporate this shift operator in our model. This is because unless the features of the neighboring tokens correlated (which is not the case for
3

---

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 state- space models. Through this, will show that training ATT and SSM with linear ICL data is equivalent to the prediction obtained from one step of pniimallyprecoondiioneee gradient descent (PGD) and cample-weighter 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 step of PGD and WPGD. Given samples (x1, yi) define
X = [x] x,,]² Rand and y [y] ya] e R..
Starting from Bo O₂ and letting y 1/2 be the step size, single-step GD preconditioned with weights W returns prediction
y = x WX == (Z).
(3)
and single-step sample-weightee GD given weights w E R and W E Raxd returns prediction
y WX"
(X (w (wPwo
where Z is defined in (1) consisting of X.y and x. Our goal to find the optimal W, as well w in (4) that minimize the population risks defined as follows.
min (Pco((W))) where LPwo(")) = B B((0- (PcD(()))
(5a)
W
min (W) where Lurc((") = E (v gureo (Z)) We
(5b)
Here, the expectation is the randomness in (x,,5)) and B, and use W to represent the set of corresponding trainable parameters. The search spaces for and W are R= and Raxd respectively.
As per (2), given input prompt Z E R(++Ixxd-1) either the underlying models outputs (n+ I)-length sequence. Note that the label for the query x == Fx++ is excluded from the prompt Z. Similar to Ahn al. [2023], Mahankali al. [2024], we consider training objective with a causal mask to ensure inputs cannot attend to their own labels and training
be parallelized. Let Zo [z] ... 2n 0J be the features post-causal masking at time/index n 1. Given weights W.. W., W, and the filter f for
x
SSM, predictions at the query token
take the following forms following sequence-to-sequene 0
mappings in (2):
BATT (2) (z W. W 2 W.r..
BSSH Z) W,) O ((ZoW, o ZoW,) fm+1)
where Rd++ is the linear prediction head and ((Z, ZoW,) ) fa+1 returns the last of the convolution output. Note that SSM can implement the mask by setting fo = 0. Now consider the meta learning setting and select loss function to be the squared loss, same in (5). Thus, the objectives for both models take the following forms.
min W.W. W.W.W..... W..
LAT where LATT(W) = B (0 8 (Z))
(6a)
min W.W. W.F
(s) where Lssu(W) B B 8ssx Z 2))
(6b)
Here, similarly, the expectation subsumes the randomness of (x)+ and B and W represents the set of trainable parameters. The search space for matrices Wt. W. W, is R(d+Ixid+1) for head is dd+1 and for f is 1*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, these variables, with the order of the cross moment being the same as order of the monomial. For example, B[x' WB] is sum of cross-moments of order Then, it motivates the following data assumptions.
moment of random variables is defined as the expectation of monomial of
Assumption All cross moments of the entries of (x:) and B with odd orders are zero.
Assumption? The label noise (5i)m) are independent of (x))=) and B. and their cross moments with odd orders are zero.
Note that compared to Ahn al. [2023], Mahankali et al. [2024), Zhang 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 Σg. My E Raxd represent the task and feature covariance matrices for independent data, and let 0 s a s be the correlation level when considering data dependency. More specific discussions are deferred Section 3.
Independent task and data: B ~ N(O, Eg), X; N(O, Ex), for all 1 isisnti
Retrieval augmented generation: B. N(O. Id) x; N(ax, (1 o2)12), for all isis
Task-feature alignment: B N(O, la, B N(aB, Ia). for all s i +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 I-step gradient descent (c.f. (5)).
Proposition 1 Suppose Assumptions and hold. Consider the objectives defined in (5) and (6), and let Lpar Lwvcrr LAND and L be their optimal risks, respectively. Then,
= = and LNOD L
Additionally, the examples Vili follow the same distribution and are conditionally independent given B, then SSM/H3 can achieve the optimal loss using the all-ones filter and L' L
We defer the proof to Appendix A.1. Proposition establishes that analyzing the optimization landscape ICL for both single-layer linear attention and the H3 model can 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,, 1)), are independently sampled given and B, and we therefore conclude that LpOo LAT L 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, also study the uniqueness and equivalence of optimal prediction functions. To this end, we analyze the strong convexity of ('W) and derive the subsequent lemmas.
Lemma 1 Suppose Assumption holds and let f [81 &2 En Then the loss (P)(W) in (5a) {)
strongly-conver if and only Bl WXT x " is strongly-conver Additionally, let Byer Shrr be the optimal prediction functions of (5a) and (6a). Then under the conditions of Assumptions and and the strong convexity, 8 = 8AT
Lemma 2 Suppose that the label noise (E:),. are i.i.d., zero-mean, variance o² and independent of everything else, and that there a decomposition x X1 X2, X X, X2, and B = B1 + BB such that either of the following holds
or > O, and (x1, X1) have full rank covariance and are independent of each other and (x2, X2).
(x1.Br, X1) have full rank covariance and are independent of each other and (x2. X2).
Then, the loss (P) in (5a) is strongly-convex.
5

---

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 Proposi- tion 1, we have Ppx LAT LSs Additionally, the second claim in Lemma 2 holds, and Epco("W) strongly convex. Therefore, following Lemma 1, we have &pop ВÂTT: 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., We Wg E Rdd+1)xr 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 Σy and Σg be the covariance matrices of the input feature and task be the noise level. We assume
vectors, respectively, and
B NO Ep). X, E.). & - ~(0,00))) 1sisn+1
(7)
and the label is obtained via y; = B &. Our following result characterizes the optimal solution of (5a). Note that the data generated from (7) satisfies the conditions Proposition Therefore, the same results can be applied to both linear-attention and H3 models.
Theorem Consider independent linear data defined in (7), and suppose the covariance matrices Σx, Eg are full rank. Recap the objective from (5a) and let W. := arg minw LV), and L. 11/2 and M = tr + 0². Then W, and L. satisfy L. - M-ntr(wW.)
L Additionally, let Σ =-x E3=12
W. -- E-1/ 1-1/ and
(8)
where define W. (11 1)La ME
Corollary Consider noiseless i..dd linear data where Σ, = Li la and or 0. Then, the objective (5a) returns
nd
W. = n1 + d+1 la and
L. = n 1+
See Appendix B.2 for proofs. Note that Theorem is consistent with prior work [Ahn al., 2023, Theorem I] when specialized to isotropic task covariance, i.e., Lp = la- However, their result is limited as the features and task are assumed to be independent. This prompts us ask: What the optimization landscape with correlated in-context samples? Toward this, we consider the following RAG-inspired and task-fcature alignment models, where Assumptions and continue to hold and Proposition I applies.
Retrieval augmented generation. To provide statistical model of the practical RAG approaches, given the query vector x,,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; S n, from the original feature distribution conditioned on the event cos(x,, a where is the similarity threshold. As an approximate proxy, under the Gaussian distribution model, we assume that B N(O, I.). x N(0, Ia) and that RAG samples a-correlated demonstrations (x,, )) follows:
x; N(xx, (1 d2) & &~N(0,03) and y; = x Isisn.
(9)
Note that the above normalization ensures that the marginal feature distribution remains N(O, Iz). The full analysis of RAG is provides in Appendix B.3. Specifically, when we carry out the analysis by assuming a = O (1/ vœ) and d/n (1) where (c) 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 minw L W), and L. (w.) Additionally, let K = a²d + 1 and suppose 0 (1/ Vd). d/n (1) and d is sufficiently large. Then W. and La have approximate forms
1 kn + d + oz la
knd
W.
and
L. d
(10)
Kn d o²
Here, (10) is reminiscent of Corollary and has a surprisingly clean message. Observe that, a²d + 1 the dominant multiplier ahead of 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 d + 1. = O reduces to the result of Corollary 1 whereas we need to set a = O (1/ vd) for
constant amplification. When a RAG achieves the approximate risk L. = 2 2 o², where the constant bias is due to the higher order moments (e.g., the 4 ''th and 6 6'th moments) of the standard Gaussian distribution. As d increases, the normalized loss L./d -> The full analysis of its optimal solution Wa 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 LLM can generate any token within the vocabulary. However, once we specify the task (e.g. domain of the prompt), the LLM output becomes more deterministic and there are much fewer token candidates. For instance, if the task is "Country". "France" is a viable output compared "Helium" and vice versa when the task is "Chemistry". Formally speaking, this can be formalized as the input x having diverse distribution whereas becomes more predictable conditioned on B. Therefore, it can be captured through a linear model by making the conditional covariance of B to be approximately low-rank. This formalism can be viewed as a spectral alignment between imput 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 B N(0, Ia) and letting x a²d + 1, we sample the a-correlated ICL demonstrations (x; Ya) as follows:
x ~ (P-N(aB.La) Ia). & &~N(0, and y; x " 1sisn+1.
(11)
Above, x- 1/2 is normalization factor to ensure that label variance remains invariant to a. 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 the RAG setting, by assuming a (1/ vd) and d/n (1), we obtain the following results for the optimal parameter and risk.
Theorem 3 Consider linear model defined in (11), Recap the objective from (5a) and let W. == arg minw Lea((W) W), and L. = (W.).W.). Additionally, given = a²d 1 and suppose 0 |/ vd). d/n (1) and d is sufficiently large. Then W. and L. have approximate forms
knd
W. kn + (d+o2)/K la and L.x
(12)
Kn + xd
Similar to (10), (12) contains a² + multiplier ahead of 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 We, W, E Rid+)err 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) S and W W. Let Σ 2x i ''th largest eigenvalue of E, we have that
Eg= == and M tr + 02. Denoting d, to be the
n
min WW
L(W) = M
(13)
(n + 1)4, + M
7

---

0.9
Lineat
1.2
HB
-'1-0.1
0.8
0.3
0.7
Theorem
Test risk
ofdd - 0.2
Test risk
0.6
0.6
0.5
0.6
0.4
0.2
10 20
10
20
30
40
50
# in-context samples
in-context samples
# in-context samples
(a) I Ip and
(b) Noisy label
(c) Non-isotropic task
Figure 2: Empirical evidence validates Theorem and Proposition 1. We train 1-layer linear attention
and H3 models with prompts containing independent demonstrations following linear model, and dotted curves are the theory curves following Eq. (8). (a): We consider noiseless i.i.d. setting where Mg = Σ la and o O, with results presented in red (attention) and blue (H3) solid curves. (b): We conduct noisy label experiments by choosing o # O. (c): Consider non-isotropic task by setting Ep = y11 (1 y) Solid and dashed curves in (b) and (c) represent attention and H3 results, respectively. The alignments in (a), (b) and (c) show the equivalence between attention and H3, validating Theorem and Proposition 1. More
experimepeeett experimental details are discussed in Section
Note that tr (2) E 4. Removing the rank constraint and considering noiseless data setting, this reduces to the following optimal risk L. Ei
A,+ See Appendix C.1 for more details.
e+1+M;A.
Impact of LoRA: Based on the above lemma, we consider the impact of LoRA for adapting the pretrained model /i Consider adapting LoRA matrix to the combined key and value weights in attention, which reflects minimizing the population loss L(W := L(W Wiora) in (5a) with 0 and jointly diagonalizable with Σ, ynew then
task distribution under jointly-diagonalizable old and eigenvalues of
I, nee _n) (A
fixed Suppose tr (2)= = tr (Σ")")= = M, o
LoRA's risk is upper-bounded by
min L(Wiona)
min Miscc
Σ iE
2; + M
Σ IET
New M
(14)
rank
n+ M/A:
n+
M/
Note that, the right hand side is provided assuming the optimal LoRA-updated model W is also jointly diagonalizable with covariances Σ, yeew, and
Experiments
We now conduct synthetic experiments to support 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-aute ettrivall-ugmmenee 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 ICI., As described 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 make in-context predictions. 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, repeat 20 model trainings from different initialization and results 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 and the test risks normalized by the dimension d.
train
Equivalence among Li LAT and LSx (Figure 2). verify Proposition as well as Theorem 1, we run random linear regression instances where in-context samples are generated obeying (7). Fig. 2a identical to Fig. la where we set Σ, = Σp = La and o O. Fig. 2b, set Σ, = Eg = and vary noise level 2 from to 0.3 x d. Fig. 2c, we consider noiseless labels, or = O, isotropic feature distribution Σ, = J, and set task covariance to be Σn y11 + (1 La by choosing y 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.9
0=0
0-0
0.02
0.8
0=02
ra5
0.04
0m
Text rick
0.04
0.
20
Tost PGd
02
022
20 20
20
30
10
20
30
40
50
10
20 n-context samples
20
50
10 20 50 context samples
59
context samples
in-context samples
(a) RAG
(b) Task- feature alignment
(c) Low-rank attention
(d) LoRA adaptation
Figure 3: Distributional alignment and low-rank parameterization experiments. (a) and (b) show the ICL results using data generated via (9) and (11), respectively, by changing a from to 0.6. In (c), we train low-rank linear attention models by setting Wz, Wg y R(d+1) and in (d), we apply the low-rank LoRA adaptor, Wiora- == Wap WT down where W W down E Rd-1)xr to pretrained linear attention models and adjust the LoRA parameters under different task distribution. Solid and dotted curves correspond to the linear attention and theoretical results (c.f. Section 3), respectively, and the alignments validate our theorems in Section 3. More experimental details are discussed Section 4.
ICL results from training I-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- mask-feaurr alignment data following (9) and (11), respectively, by setting o = 0 and varying from O to 0.6. Attention training results are displayed in solid curves, and 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 is applicable to 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). Fig. 3c, set Σ, La. 0 and task covariance to be diagonal with diagonal entries c[1 2-1 d for some normalization constant
d/ Σ r|, and parameterize the attention model using matrices W1. W E R(d+ 1xxr 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 I: Setting Σ, Σp Ia and o O, we pretrain the model with full rank parameters and obtain weights W.. Wq W. I(d+1)x Phase 2: We generate new examples with task covariance Σg being diagonal matrix with diagonal entries c'[2-1 2 ... 2-²JT for some normalization constant 2- Given the rank restriction r, we train additional LoRA parameters W W Woown R(d+1Jer where Wiora == Wup WT down and (2a) becomes ATT(Z)= (Z(W, W. Wup WT down )Z W.W, 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 Σ, nnew diagonal, the right hand side of (14) returns the exact optimal risk of LoRA and the alignments verify it.
é' = E
Related Work
There is growing interest in understanding the mechanisms behind ICL [Brown al., 2020, Liu et al., 2023b, Rae et al., 2021] ] in LLMS due to its success in continuously enabling novel applications for LLMs [GeminiTTemm al., 2023, OpenAI, 2023, Touvron al., 2023]. In the previous work, Garg et al. [2022] explored ICI. ability of Transformers. In particular, they considered in-context prompts where each in-context example is labeled by a target function from 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 al., Collins 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 al. [2023] focus on linear regression and provide a construction Transformer weights that can enable single step of GD based on in-context examples. Along the similar line, Von Oswald al. [2023] provide a construction of weights in linear attention-only Transformers that can emulate GD steps on in-context examples for 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 gattiggradents 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 al. [2023]. Duraisamy [2024] focus on developing theoretical understanding of Transformers trained to perform ICI. 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 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 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 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 test instance, e.g., via 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 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 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, 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 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: different sequence model architectures exhibit strong degree of universality and (ii) Dataset and prompt design, such as RAG, can substantially benefit ICL performance.
(1) The functions Iearned by
Future directions and limitations. The results of this work fall short of being 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 '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-baat baaed ICI. In real-life retrieval augmentation- based ICI,, one needs to account for the quality of the collection of the retrievable demonstrations and its (negative) impacts on the final predictions.
10

---

Acknowledgements
This work was supported in part by the National Science Foundation grants CCF-2046816, CCF- 2403075, the Office of Naval Research award N000142412289, an Adobe Data Science Research award, and a gift by Google Research.
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 in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https:: //pps/////p/eeeeee///e net/ formfformm?i id= N40X
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 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, 13:1877-1901, 2020.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model align- ment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1):2914, 2021.
Yuan Cao, Zhiying Fang, Yue Wu, Ding- spectral bias of deep learning. arXiv preprint arXiv: 2019.
Ding-Xuan Xuan Zhou, and Quanquan Gu. Towards understanding the
Liam Collins, Advait Parulckar, 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 meta-optimizers. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Findings of the Associa- tion for Computational Linguistics: ACI. 2023, pages 4005-4019, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. (0.18653/11/0233 htt.://///22....nnnnnaaaa/////// findings-acl.227. URL https://aclantholle org/2023. findingsn-nnna acl 247.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. 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.02-422 2024.
Daniel 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. wwt.hh::////pereee/// net/ forum?i forumPid=-OOZDy0WYGG.

---

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 Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn 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: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arXiv preprint 20XIv:2402.03110, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:21/2.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 second. arXiv preprint arXiv:2207. 01 048, 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: pestps://opentt //pp/nevview.neettteteeeeeeeeeee mnne/fooum?ii id= nZeVKeeFYf99
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-44, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference machine learning, pages 5156-5165. PMLR, 2020.
Ivan Lee, Nan Jiang, and Taylor Kinn-KKrkkatrikk. Exploring the relationship between model architecture and in-context learning ability. arXiv preprint arXiv:23 2023.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmented genera- tion for wwwllede-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.
12

---

Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. arXiv preprint arXiv:2402.18889, 2024.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers
learn shorteuts to automata. In The Eleventh International Conference on Learning Representations, 2023a. URL ht htttts//oppereiiiieee www..neeoooommmi--eFFFoooo..
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.0833) 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. The Twelfth International Conference on Learning Representations, 2024. URL htttss:/openrevieeet net forum?id= 8p3fu561Kc.
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 arXi::2112.005100 2021.
Samuel Müller, Matthias Feurer, Noah Hollmann, and Frank Hutter. Pfns4bo: In-context learning for bayesian optimization. In International Conference 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 preprini
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.
OpenAl. Gpt-4 technical report. arXiv prepiintaXXv:2333.06744,, 2023.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how 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 Machine Learning, pages 9355-9366. PMLR, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Marie-Anm 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.13977 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 S.XXi:2330905850 2023.
13

---

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 urXiv:220308877. 2022.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. tf-thuuhtt prompting elicits reasoning in large language models. Advances in 2022.
neural information processing systems,
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 0rXIv: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 httph::t////. /tttts:///opnrevieeee ew. wttp:/ewwweeeetmmm/ieiieaae hett/forummdddRddJvvCCCj..n
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.0999 09927, 2023.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers leam 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 preprini arXiv:2309.01777 775, 2023.
14

---

Appendix
Table of Contents
A Equivalence among Gradient Descent, Attention, and State-Space Models
15
A.1 Proof of Proposition
16
A.2 Proof of Lemma
20
A.3 Proof of Lemma
21
Analysis of General Data Distribution
21
Supporting Results
22
B.2 Independent Data with General Covariance
24
B.3 Retrieval Augmented Generation with Correlation
25
B.4 Task feature Alignment with Correlation
28
C Analysis of Low-Rank P Parameterization
31
C.I Proof of Lemma 3
31
D Additional Experiments
32
Extended Related Work
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
X = [x] ... x,] E RRdd
5 = K
En E
IR""
y = [y
yal X
Zo lz Zn 1
xi
X, 0g
R
ine Dxid+l)
y1
yo
and corresponding prediction functions
SPco (Z) x WXTx y.
(15a)
Sweao (2) = x W (w),
(15b)
8ATT (Z) (z W. W 2 W.,
(15c)
8ssm 7.) (w) (lZoW, o ZoW,) Dmt)
(15d)
we have objectives
min (P) where Leco (W) B (1: geGO 22)²
(16a)
min (w) where ( (y Swpco (2))
(16b)
Was
min W.W.W.W...
Larr(W) where
Lan(WW E ( BATr(Z))²?
(16c)
min
(ss where LssM(W) ( - gssm (2)
(16d)
W.W.W.W... W.y
Here, the expectation is over the randomness in (x,, Ei and B, and the search space for W is Rddd for is R"" for W.. W., W, is (xd+1x) for v is Rd+1 and for f is B
15

---

A.1 Proof of Proposition 1
Consider the problem setting as discussed in Section 2, Proposition can be proven by the following two lemmas.
Lemma 4 Suppose Assumptions and hold. Then, given the objectives (16a) and (16c), we have
min W.WWi.W...
Lar((W) min LP))
Proof.
Recap the linear attention estimator from (15c) and denote
W W
W
Wi
and
W.r
V1
w
v
where We Radd, WT, W2, V1 E Rd, and W, V E R.
Then we have
SATT = Wg W 2
W.
= [xx 0]
w
X
0g
y
F1
W2
y
0
0
0
= (x WW + x why (Xv1 + yy)
xx(WWWxx y "wyy" Xvi x (WX" Xvi rllvlie, wi)
= = wir
(WX Xv1 v IIvell w:)
=x' (WX Xv + +lblllf wi).
(17)
(Z)
where W := vW + W
We first show that for any given parameters W. W 40 W,,,,
E (BATT(Z) - y) 2 (Z) - -
(18)
To this goal, we have
B [((ATT(()) - -)) B (BATT(Z) -> = B (B (2) + = y) -11 (8)m(()) y>²
= 2 2B (SATT(Z) y)a]
(19)
where we have decomposition
(BATr(Z) - y)c = (x" WXT -)xx (WX" Xv + v |1vll wi)
= = xx
[Wx XP1 + +bbll wi) yx (Wx") Xv + |1rll wi)
y XW xx WXT
Xv |||l y' XW xx w yx W Xv vy |i
wi
(e)
(c)
In the following, we consider the expectations of (a), (b), (c), (d) sequentially, which return zeros under Assumptions and 2. Note that since Assumption holds, expectation of any odd order of monomial of the entries of X, B returns zero, i.e., order of x fx is 3 and therefore E[x" Bxx = 0a-
(a): E B XWT xx WXT. Xri
B (XB - XWT xx' WXT Xvi]
E B x XW xx WX Xv, + + XWTxx xx (VXTXr]
0.
(b): E B|Pllylle y XW xxxxxTT w1]
= B|I(XB= + E) (XB + EXXB+ +) XW xx wr]
= B(vl1ETR, HT xW xx xxTT] W1
O.
16

---

(c): B B W Xvi]
= B B + E)x W Xvi
= B B xx WX Xri] + + fxx W Xvi)
O.
(d) B [ x
=v (B x (XX +) (XB =xx"'" Twi]
= -v | x wi
0.
Combining the results with (19) returns that
E (EATT((Z) - y> =[E²] 0
(20)
which completes the proof of (18). Therefore, obtain
min B (2AT(()) y) y> W.WWW.WW.. W.
2 min (ZATT(Z) ->> min (8)c - y>² W
We conclude the proof of this lemma by showing that for any W e Rada in gEGD. there exist W.. W. W such that gATT (Z) grcD( Let
W
0g
0g 1
W
W, Ionts
Wg
and
0
0
Then obtain
(Z) = WWX = 8PCD
(21)
which completes the proof.
Lemma Suppose Assumptions and hold. Then, given the objectives in (16), we have
min W...W.W.....
Lssx((W min ("P
(22)
W.o
Additionally, if the examples (x,, YoY; follow the same distribution and are conditionally independent given x and then SSM/I13 can achieve the optimal loss using the all-ones filter and
min (x)a Wo
min LPa((W).
(23)
W
Proof. Recap the SSM estimator from (15d) and let
W
Wel
Wyy
Wyde1
W;
WA1
Wx2
Wk.+11
W,
Wyl
Wo2
Wydeel
where Woth Wky,
W.j Wky.
Rd+1 for s 1, and let
Vi
fo f
V2
and
fi
Va+1
Then have
Essn (Z) ((z WT T @ ((Z W. @ ZoW,)
D
w "
W zan Ze
fool-i-p T
w 1
w Wate Zi vidti
W zw
Ni Zi
foll-i T
WT
a.t0 2W R.dul Ziw vidal
17

---

Next for all j و d 1, let
Wal Wass
Wx2
Wyj
wxj
Wej =
where Waji Wej. Woj E Rd and Way Wejo Woj E B. Then we have
w 22 2, Miz (w | + WaSNi) W Wijaa mossi)
x Was (hr (Wo) + wxy xay (w (MTTer) (1 (w.Txx)) + (w
W =x"W
x + xx, + w xy?
where
wy= == Was (Wo, weyn E Raxd
w/. weyWirWWWr R.d
6 6,(x,x,,x)) = xxx,,,,x (,) (1 (w) ER.
Then
d+1
Bssn ()) = Σ
fa+1-i
Σ
v/ (x W''WYYy -0,(x, w,,xx,
F11
F=l
d+1
d+1
d+1
T
W,W X Σ SMLI- Σ
(x, x,, +
Σ
r xy xyT
WXy
8(x X. X) w xy'y
(Z)
where
: Un ... fil E ER"
5 yo f ER..
d+1
W
vy W E Rand
w = Σ Vjnyy E
d+1
8(x, X)
fual-i Σ
v.. (,), R.
iml
Next we will show that for any Wt. W W,.v,
E (gssn (2) - y>² 2 B ( 1) -
To start with, we obtain
B (Essx(Z) y> E (gssn (Z) + + -
E (Zssn(Z) - +B +B( 82) (Zssxx)))()) y) + {2)
(24)
where there is decomposition
(Essn ()) ()) 82) =
WXy (w)
5(x, X, X)
xy j WXy w xy y.
(b)
(4)
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 and 2. Note that ,,(x, x; xx)'s and ,(x X, X) summation of monomials of entries of (x, X,B) with order 3, and entries of and y summation
18

---

of monomials of entries of (x, X, with even orders: e.g. x' x B+ where & is of oder 0 and x x²B of order 2.
(a): E 0(x, X, X) x Wxj
= =B( X) x W ( O A3] +B B x T N)
= B X, x W 5 i
O.
(b): B flxx X, X)y
= B Blx X, X)(xxxxx B +)
= B f(xx X, X)x B) + ooxx x, X)
= 0.
(o: B B xy x WXy
= - x(XB E) x *Wxxxof++o
= O.
(d): B Byy- To xyy }|
= 3 [ B + E) x"xxXx++)xXX0f=++
= O.
Combining the results with (24) results that
B -2) (y)) + + 20. 0.
Therefore we obtain,
B (gsse (2) -y>²
2 min B (ā)):(() - y>²
min E (EWPGD (Z) (Z) y)
min
W.
W.o
Next we show that for any choices of and w gueGD, there are Warne f such that gssa gWPGD- To this end, given as [w)] wn]" let
0
W,
Lasis
W.
WT
0g
W.
0axd
0g]
1d
and
Wn
07
1
0
Then we get
W2 @ Z.W.) fn+
XWT 0,
y1
0,
%
*
0
n+1
Σ we Wx: O
WXT (y w
and therefore
82s = WXT (y 2 w) = BWPGD (Z),
which completes the proof of (22).
19

---

Next, to show (23), for any W E Raxd let L(w) = E (xx WX ( ()0
Then we have
TL(w)
2
B 2 x W Σ w - (x'Wsyxz))
dwi
2 Σ Wj B Wyjjyy)) Wyxx)] - 2 "ax"Wyyxk
Here since (x).) follow the same distribution and are conditionally independent given and B. for any j J'. (y;)) = B [[(x' Wy,x,)) and B (x""Wy ( [[x" (y""Wyy,,)e Mx Wy.z,]- Then let
Wyxx)]
(x Wyya,)(x' Wyxx)) [e,, i
and B B[xx Wy:ns C3.
(22. i
where (c1, (C1,0,,0)) (2W) cWW).. We get
aL(w)
2oww 1, 2(c2 c1) 203.
dw,
If C1 = 0, then
aLw = w0 1, 22g for all n and any E R achieves the same performance.
C2 - C1 # O, setting
(Liw) Aug
0 returns
G3 C--6 2ッ Wj C
w
for all
C2 Ci
Therefore the optimal loss achieved via setting w Cl Without loss of generality, we can update W + CW. Then w = 1,, and we obtain
min B Was
(x (WTT(000)=)) -) min E
((( (x WX y- y>²
W
which completes the proof of (23).
A.2 Proof of Lemma
Proof. Recap the loss (Fw)(W) in (16a) and prediction (Z) in (15a), have
(W) = B((y - 8PGD
B (x - WX (XB+ (XB+
= E (x x TT XB) 2(x B- --x WX XXX x WX + -6-x E)
= B B xx WXT xB) + x 5) + + B- WXT XB)(E - x WXT 5)
= B [x -- x "'X XB)² + +6--x"XX 5)2
(25)
= E (x' T W xB) + (x WX 5) -2BB xx WX Xp+ +x+ WX g BI B) + 1
((W
fe(W)
constant
where (25) follows Assumption 2. Since f2(W) is convex, ( is strongly-convex if and only (W) is strongly-come -convxxx, which completes the proof of strong convexity.
Next, (20) and (21) the proof of Lemma demonstrate that the optimal loss is achievable and achieved at & = 0. Subsequently, (17) indicates that SÄTT has the same form as gyop Under the strong convexity assumption, & is unique, which leads to the conclusion that goco = EATT
20

---

A.3 Proof of Lemma
According to Lemma 1, (W) is strongly-conve as long as either E[(x" T WX XB)* or W[xx"WXT WXT ] is strongly-convex. Therefore, in this lemma, the two claims correspond to the strong convexity of B[(x" WX g)*] and B[(x" W xB)² terms, respectively.
Proof.
Suppose the decomposition claim holds. Without losing generality, we may assume (I1, S,,x0 X1) are zero-mean because we can allocate the mean component to (x2.B2, X2) without changing the covariance.
Σn B[ and Σx B[X]X X1J. If the first claim holds, using independence, observe that we can write
Claim 1: Let Σ, ]
B[ WX 5) = B[(]] W W B[(x 5 + B[ W B[<x workswxx33 5 where the last three terms of the right hand side are convex and the first term obeys
B[(x] [X]] of W
= of [[xxx WX
= of WEx
= =0 0 V Σx
Since noise level using the full-rankness of covariance matrices Σ, and Σx, we conclude with strong convexity of B[(x" J.
Claim 2: Now recall that Σx x[XTX; X1] and set A X X1 Σx and B X Σx. Observe that E[A]= = O. If the second claim holds, B[ X] = E[A + B). Note that (A. are independent of each other and (2. x2).. Using independence and B[A]
similarly write
( W XB)² B[(x" B[(x"WABX) WAB31 + B((x WBB
Now using BIB, |[x;] 0 and their independence from rest, these terms obeys
B[(x WAB)²] (xx W WABB) B(= B(x
WAB B WAB:)
B[(x WBB [[(x] WBB + WBB B[(x] WBB B[(x) B(x
WBB.)
In both equations, the last three terms the right hand side are convex. To proceed, we focus on the first terms. Using independence and setting Σx = EX X] Ex > 0, we note that
[x W E[(x W BB
= [xx WX XB;)1
where BI, X are independent and full-rank covariance. To proceed, note that
[[x]] W xB = B[xx] W B((] W(X 1x
Observing the convexity of the right hand side and focusing on the first term, we get
= V. |
[[ WEx tr (2, (r
Using the fact that covariance matrices, Σ, Ex. Ep. are full rank concludes the strong convexity proof of TXX X
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) := Lewr(")) and := grep in this section.
21

---

B.1 Supporting Results
We begin by deriving the moments of random variables.
2n'th moment of a normally distributed variable: Let N(O, o²). Then we have
B B[u²"]= = (2n 1)!!.
(26)
4'th moment: Let u N(0, Ia). Then for any W, W E R we have
B
(u Wu)(u W
=B
Wijutar
Wi
Wijuay
lifol
Wane
Σ Wi
E
Σ
Wisua Win
Σ
W راا
Wa W E [+]]
Σ
Www,, BlG
Σ
W.y W. Blu + Σ
w., w", Blu
[n
64
W W
W. Σ Wi Σ
Wo,W',
d Σ W +
W W +
Σ Wi,W'
W=
tr (W) tr (W") + tr ( W' trWw"").
(27)
4'th cross-moment: Let N(O, La) and for any W E Raxd let Aw W O la. Then we have
E
(u Wvr [u)]
d
Σ u/vi
B
WaJur
B
Σ WButv?
Σ
WiyWermuuvyy
+
Σ
Wir
Σ
Wy, W.WWeruueyy Σ 12 + Σ
W
Σ u²
+ Σ
Wi, WWWWWuuuu33
E
Σ>
W
B
Σ
Σ
W
+
Σ
Σ
Σ
Σ
W,
(4)
14)
121
12
9
Σ
W (d (d-1)
Σ
W + 6
Σ
W3 (d - 2) Σ W + >
WryW
1=1
1=1
(+)
l
d
4
d
3 Σ
W + (d 44)
Σ
W2
+
Wi;WWr n
:=1
1
3tr (Ax) + (d 4)tr (WWT) + tr (w3).
(28)
22

---

6'th moment: Let u N(0, Ia). Then for any W. W E R we have
B
(u T Wu)(u W
B
Wijuiu
Wijua
Σ
1=1
d
Σ
Wing
Σ
+ B
Σ
Wiju
Σ
Wijuit
Σ
Wi W E
Wi W E
u
Wij Wi 18
W B
(d + 4)
3
Wi WoW-+ Σ WaW",
Σ
Wi,
Wey Wi
(29)
1=1
121
+d++4)
Σ
WiWWi + Σ Wi, Σ
= (d + (tr tr (W) + tr (W tr (WW'))
(30)
where (29) is obtained by following
d
= B[ud] (d 1) B("") B[u2] 3(d 4).
up
d
up
B(u 2B + 2) Bln Blu [Bu +
8'th moment: Let u ~ N(O, I.). Then for any W. W E Redd, we have
E
T Wu)(u W W
E
W Wepue
Σ
B
WH
Σ Wiie
Σ
Σ
B
Σ
W Wijuuui
Σ
W Wiju
Σ
:21
W3 W B
Σ #
Σ
W W", B
uuf
Σ
Σ
1+f
F=1
W W B
Σ
4 + Σ
W., W E uu
Σ
Σ
47 n
(d+ 4) 6)
Σ
Wi W + Σ
W. +
Σ W
Σ
W.
(31)
[=1
[+]
d
= -d 4)(d + 6)
W Jj
+
Σ
Σ
WoyW's
LF
iF1
(d+ p)(d + 6) (T (W) tr (W')
tr W' trW"))
(32)
23

---

where (31) is obtained by following
E
M
f=1
=B(u"] + (d - 1) 1 (d Did B[ B[u"]
-105 + 9(d 1) + 30 1) 3(d ()dd 2)
=3dd 4) + 6),
2 u
Σ
Σ
E
up +
Pay
=2E[('')Euu> (d 2) B[u 1) + +2E("") 4(d ") (d- 2Xdd -
=30 + 3(d - 2)+ + 18 + 12 2) 2d
=d +d 6)-
B.2 Independent Data with General Covariance
Proof of Theorem 1. Consider general independent linear model as defined in (7) where Mg and Ep are full-rank feature and task convariance matrices and
Ex). B N(O, EB), g N(O,0²), and y= x=x B++ g.
Let
x= [x]
= 19 ... E,1", and y= [y ... yal' = XB + 5. xn
To simplify and without loss of generality, let i = x, x, X = xxx 1/2 B 2. B where we have
x ~ I)..I).
B N(O. / 2
and
= . XB+
Then recap the loss from (5a), and we obtain
L(W = B (u 8(Z))²
= B [xT 1-++--0 xx= x=WX (XB++)))
= = =Txx' xB)² 2(x B - WX XB)E - x WX +) ( - x [) E3
= =B WX xB)²| B 5 +0²
(33)
where the last equality comes from the independence of label noise E.5.
We first consider the following term
B (x W 5) B( / TT = =o² tr (WW))
where we define W /1//2x!// 1/2 Next, focus on the following
E B((x B r WX xB)²| E (I B F WX XB)
= (F (1 WX x |]
tr (=((1 W x 2(1 ( WXTx)")
tr tr (Z(W W) X1) tr ( WWTWE XEXT XI)
tr 2n tr (ΣW) tr tr W XEX XI).
where E =:=E!/20 23/2Ep1!?
24

---

Let F E IR" be the i'th column of X and Eij be the (i. f' ''th entry of I. Then the (i, f) entry of matrix
X is
XEX
Eppk F
km
Then we get
i j:
B XEX'
- E,
i= j:
XEX" x)
=I, [
Σ
-n B +... X
+
زز
jFi
2nn + nnn - 1)) + Σ
jai
En(n + 1)+ Σ
F=1
# tr))))
Therefore
B[X XEX X] n(n 1)E + tr I.
Combining all together results in
L(W) = tr (z) - 2ntr (Σ + nin l)t + +o +o.
+00
= .M - 22tr [W) n(n + 1)tr (ΣW W) +
(FV") (WW)).
(34)
where M
(Σ) o Setting V 0 returns
-2n Σ 2n(n 1)- EW + 2nMW =0 - W. (In 1)I M--
Then we have
W, =,, (11 1)l ME ) 1
and
L. = L(W,) = M ntr ((12 + 1) + MI 2-1)
B.3 Retrieval Augmented Generation with a Correlation
In this section, consider the retrieval augmented generation (RAG) linear model similar to (9), where we first draw the query vector x and task vector B via
NO and ß N(O,
We then draw data (x;)",, be used in-context according the rule corr_ corr_coef(x,ee xx)) 2 a Hence, for i & we sample
x; x ~ N(ax, & ~ NN(,00)) and y; x==xBB B E..
(35)
which results in (9) by setting y² a
Theorem 4 (Extended version of Theorem 2) Consider linear model defined in (35). Recap the objective from (5a) and let W. := arg minw Lxw(W), and L. = LPGO W.). Then W. and L. satisfy cnd(a L. =d=d+ - +))
W, cl and
(36)
where
(dd +
c= a'n(d + 2)d 4) + a2 2)d + 2n 3) y n 1) o y²)
25

---

Suppose a = =0(1/ vd). d/n = (1) and d is sufficiently large. Let K = a² + and y² 1 aaa
Then W. and L. have approximate forms
knd
W
and
L. = =
(37)
d 02
+ o
Proof. Here, for clean notation and without loss of generality, define and rewrite (35) via
gi - D. & - ~(0,0007) and x; = ax 7gi.
(ax + yg.) &c.
Then we obtain
L(W) B (((v-((Z))
= B (x 5 Wx (XB+ (xB+ E))
= = x wx xB)² 2(x B - x' WX Xp)( - W + (g 5)
B wx xB) + B Wx +
(38)
To begin with, let
N1 tr (W) + tr (WW") tr (W2). N2 N2= tr (WW"). and N₂ tr (W).
We first focus on the second term in (38)
B (x Wx x3 B
Ex W(ax y8))
no² E W(ax yg) yg) W
no=² (a² B[x Wxx WT = W Wgg's W'xl)
= =no² (a2 + N2)-
(It follows (27) and independence of x, g.)
Next, the first term in (38) can be decomposed into
B B WX' " B B} + (x WXT XB)²
2E [Bxx'
Bx W XB
In the following, we consider solving (a)-(c) sequentially.
(a): B B((x B) = d.
(b): (((x W XB)²
B
x W
(ax + yg:) ygi) B
B
Σ
( xxx +y88.8 arygg + +ryggx }
a'n Wxx pJ + B Σ W'iBe
+ a Σ x WxgT B
+ Σ gxx B
+202 [x Wxx BB gg" WTx 2a 2n Wxgg fxx Wgx T
=au'i(a + 4)N1 n'd d + n IyNs) + (a2y'ndN, + a nld + 2)N:) (2u2r2n2 2u?y
+ 2)) N₁ (n²>pn(d 2) And (d + n 1)) No
AiN A2N2.
26

---

(:: B Bx Wx XB =B Σ x' Bx W(ax + ygx + yg:) B
Σ x Bx x(xxxx gigi ayxgi + ayygi }
= BBnBBB xx'xx" Bx Wxx B) + + Bx Wgg' B]
= = + 2)tr (W) + (W)
= =aan(d +2 -yn) N₂
A3N33.
Here, (b) utilizes the 4'th and 6'th moment results (27) and (30) and we define
A = a'n'dd a2 d 2)
A2 a + y n
A3 = anddddddnnnnnnnn +) +y
Then combining all together results in
L(W) ANN A2N2 2A,N3 no (a2-N + N2) +d
To find the optimal solution, set V.((W and we obtain
Av A2VN2 VNs (a y VN2 O.
(39)
Note that we have
VN, = ((r(WW) + tr (wWWT) + tr( 2tr (W) / 2W 2W
VN₂ = Vtr(WWT)= 2 2W
VN₂ = Vtr (W) I.
Therefore, (39) returns
2A, (tr (tr(W) W + WT) 2 2A3 + 2n (W) W +T) y 0,
(40)
which implies that the optimal solution W, has the form of cl for some constant c. Then suppose
W. = cl, we have tr (W) cd and (40) returns
(d + 2)el + 2A2c1 - 2A,1 + 2a 2)cl + x2cc) = -0
As
c A,(d + 2) + A2 no (a (d + +2)
a² + 2)
a'n(d + 2) + 4) + (d + 2)(d + 2n + 3) +d + +2) +))
Then the optimal loss is obtained by setting W. cl and
L. L(W.) A +) A 2Azed ndd
= (A ( + 24gcd
=d 0² A3cd.
completes the proof of (36). Now if assuming a dimension d, we have the approximate
d/n d/n= = and sufficiently large
od
at a 2n) (
ad
(a?d+ 1)n + (a (
and
(a + 1)nd
L (a=d + 1)n d 2
27

---

B.4 Task-feature Alignment with Correlation
In this section, we consider the task-feature alignment data model similar (11), where we first draw task vector B via
B N(O,), I).
Then we generate examples (x;, y1) according to the rule coir corr_coef(xx;B) a 2 0 via
xi B N(aB. I). &: N(O, and y1 y B Ei.
(41)
which results in (11) by setting y² = 1/(
Theorem (Extended version of Theorem 3) Consider linear model defined in (41). Recap the objective from (5a) and let W. == arg minw Lowo (W), and L. = Lrco .). Then W and L. satisfy
W. = cl and L. = dy (Aoa + 1) + o² -cn-nndy cndy' 2Dga
(42)
where
Aja² + 2Agar +
Ara Azart Asa² (d + + 1) + (A + 2x² + 1)/y²
and
Ao = d
A (d+ 4
A2 (d+ 2)( 4)(d 6)n
A3 (d 2)(d +))3n+ 4)
A4 = (d + 2)(n + d+ (d+D) 8).
Suppose = (1/ vα). d/n (1) and d is sufficiently large. Let a²d and x² 1/x. Then
W, and L. have upproximate forms
knd
W+ (d 0²)/K
and x (d 02)/K
(43)
Proof. Here, for clean notation and without loss of generality, we define and rewrite (41) via
I), ~~ N(O, and x; x; aB + gi, y1 = yx, = y (aB + gi) B E..
Recap the loss function from (5a), we obtain
L(W) B (y- 8(2))?
B (xx B+ g - W (yXB+ (xxB+)))*
(WX = Bx B x-x wwx xB)² + 2y(x B B- TWX xB)( - x x) + (g- - WX E)
(x - wxx xB) + B
Wx Txx 5 +o?.
(44)
Similar to Appendix B.3, to begin with, let
N1 = tr (w)² + tr (WWT) + tr (W2). N2 = tr (WWT). and N = tr (w), and additionally, given Aw
let
Na = 3tr (A (d 4)tr (WW") tr tw?))
We first focus on the second term in (44)
WX
b (aB + 8) Σ ElaB + gi)
no B B(aB g)_ W(B gi)
no (a'B ( WB) 2a w + (g (
no (Er tr (w3) tr (WW + (2a² + 1)tr (WWT))
no (a + (20² + 1)N₂).
(It follows (27) and independence of B. g')
28

---

Next, the first term of (44) (omitting y²) returns the following decomposition:
B (x B-x WX xB)
((aB + 8) WX wx' xB))
E (aB aB WX XB + g B WX XØ)
a² BIBB B) + a² BlB wx xB) + Bigg B) Bl(g WX XB)
- -2a² BB Wx - gg WX Xß]
= a²d(d + 2) a BI wx xp)*] d Bl(g WX'
B BB WX XB) 2B|B" gg W XB).
(d)
Consider solving (a)-(d) sequentially follows: To begin with, we use the following decomposition for all (a)-(d):
X XB
Σ x B
Σ (aB gi) gg)
Σ a BB B aBg ag. BBB B.g.
Then, have
(a) BIs WX xB)²
Σ B wPp B + oB wBg. a WgB Wggg
d'nn
( WBB B)
a²B Σ B WBg B
a Σ B w WgBB B
+ Σ B W WIETD
+B
B WBB BB WeggB
2 B
B WBg BB WgB
=atn = (B'BBB
WBB B)²
a 3 [ WBg B)
nBn (B Wg'B 8) +
B
B2a WBB BB Wg's B n wBg. BB WgB B)
=d'if(d 4 + 6)N; +ann(d 4)N; + a²n + 2) + 4)N₂
(45)
+n 1)N1 + nNa
(46)
(d 4)N1 + 20 n(d 4)N
(47)
=a2nnd 4) 6) 2n 3) n(n - 1)) Ni ann(d 4)N, nN4
(48)
BiN + B2N + nN4.
where (45) and (47) utilize (30) and (32), and (46) is obtained via
2 B
B
Σ
W
gig
B W W''g g" DJ + n(n D) (T Wg'g' T Wg"g" g
nN4 n(n 1)N1,
which follows (27) and (28).
29

---

(b) [g W XB)²
Σ 8 WBB B ag wwg. ag WgB B 8 8"Wg:g. B
WBB B)
a2 Σ 8 2'BBg// B
a²3 Σ g W Wgg B
+3 Σ 8
w Wgig
+2d'nB
8' 2'WBB" Bg B
+20
WBg Bg WgB
= = 3 WBB B]²
n²nn (8 WBg' B)
n'n (s Wg'B B) +B Σx W Wg.g;
+209//1B8 WBB Bg" Wg'g B + n0 WBg. 8g Wg. B)
= =a'nfd+ + 2) 4)N aan + 2)N₂ + a'nd(d + 2)N₂ nld + 1)N2
(49)
22 (d+ + 2)N2 2a' +
(50)
=aa'nd 2) n(d + 4) 2n + +) n(d + 1)) N2
= ByN2,
where (49) and (50) are obtained using (27), (30) and
= =nB g Wg'g T B) + nin - 1)B Wg'g' BB gg"g"" B
B
Wg.g. B
=nid + 2)N2 n(n 1)N₂ n(n + 1)N2.
(c):
B B BB WX XB]
= =n BB W(aB g (') B)
= f BB wpp B + B B BB Wg 8 B)
=a2n(d + 2) + 4)tr (W)- + n(d + 2)tr (W)
= =aa'ntd 2) 4) n(d +2)) N₁
= B4N3.
(d) B8 gg" W XB}
= nE B8 gg W(aB + g (aB + g)
=an | gg WBp B + nB B gg Wg'g TT B
=d2nld 2)tr (W) ntr (W)
= =(d2n(d + 2) + n)) N3
BsNs3.
Here we define
B ann(d 4)(a2n(d 6) 2n +) +(n n(n-1) 1)
B2 = a' a'nid 2) 4)
B; = n²n(d 2)(a² n(d + 4) 2n + d 3) n(d + n 1)
Ba n2n(d 2)(d 4) n(d 22
B, ann 2)
Then combining all together results in L(W) = x² (22 (a'ddd + 2) + a (B N, B,Ng nN.) + ByNg - 20 BaN; - 2BBNs) + + (2x² 1)Na) o²
BiNi + (a2 B, B. 2(a² B. Bs)Nss d?nN.) (a? + (2a² 1)N+) + 1) +o = x BiN
30

---

and differentiating it results in
VZ(W) >² (a²B,VN, (d''B,,, (a²B2 B,) 2(a² B4 + BB)VN H (
Similar to the proof in Appendix B.3, W. has the form of W. cf and we have
VNi V (wr) + [WW") tr (W2)) 2tr (W) 2W + 2W = 2c(d 2)1
VN2 Vtr Vtr(wW") = 2W 2cl
VN: Vtr (t) =
VN+ = v(3tr (A + (d + 4)tr (WWT) tr (W2))
6 diag (Aw) 2(d 4)W 2W
= =2c(d + 8)I.
Therefore, setting VL(W) returns
y (2cld + B) + 2c(a² B2 B3) 2(a Ba + B5) + 2old 8)a²n) (c =0
a² + Bs
2)a + (a + B3) + (d + 8) + (o 2) 2a 1)/y²
atmid
2n +
abnidd 20d 4)d
2 41(3n + 4) a-mild + 2X3n 3) (d 8)) and ao DA
a'mid + 2xd a 2ndd add 2xd
2)+1
(-( 2x33
(d (d 1)+ 20 11/y²
Then the optimal loss is obtained by setting W cl and
( (d 2) + 1)+ 12 B5)cd.
completes the proof of (42). Now if assuming = 01///))///n=000)) 1/(a2d 1) and sufficiently large dimension d, we have the approximate
A 222
R d 3na + +
1)2
n(ad 1)3
d(a 2)
1
(a2d+ 1)m (d+ (d+)))==+++)))
and
(a + nd
L. / 12+
(a 1)n + (d + +z) 1)
(a 1)nd
d =d
(a + I)n (d (0)
C Analysis of Low-Rank Parameterization
C.1 Proof of Lemma 3
Proof. Recall the loss function from (34)
L(W) M 2ntr ΣW n(n + 1)tr (WW W) nMtr (WWT
where W E. wx 1/2 - !!! > and M tr (E) + o². For any W. let us parameterize
W UEUT where U E d denotes the eigenvectors of W and E E Iprrr is symmetric square
31

---

1.0
0.9
1.0
Awer Averaged test rikk
43
Linsar Al:
Lincor Alt
0.77
0.4
10 10 30
50 59
50
10 20 30 40 50 60
50
30
40
50
60
70
10
20 20 samples
50
samples
samples
Rroi
(a) Linear attention
(b) H3
(c) Averaged risk
Evolving B
Figure 4: Further comparison for linear attention and H3. In (a) and (b), given maximum context lengths Rmax we train linear attention and models to minimize the average loss across all positions from 1 to Mmax. Averaged test risks are presented in (c). In (d), the task vector B evolves gradually over the context positions S n via B, = (i/n)ß + (1 10l-i/n -i/n) both scenarios, H3 outperforms linear attention benefiting from its additional convolutional filter (c.f. f in (2b)). Implementation details are discussed in Section
matrix. We will first treat U fixed and optimize E. We will then optimize U. Fixing U, setting Σ U'EU, we obtain
(E = M 2ntr n(n+ l)tr (EEE) + nMtr (E²).
Differentiating, obtain
0.5n = -E+ (n 1 ME.
Setting V((E) 0 returns
E. = (MI (n 1)
(51)
Let 2; denote the i'th largest eigenvalue of Σ. Plugging in this value, we obtain the optimal risk as function of U is given by
L.(U) M - tr(ÉE.) M-n-tr(iie -M --nn-rr-e (MII+ (n 1)
(52)
X;
M n (n + 1)2; M
M n ++ M
(53)
Now observe that, the right hand side is strictly decreasing function of the eigenvalues 7; of Σ UTU Thus, to minimize L. (U), need to maximize Σ -+1+MAT It follows from Cauchy
interlacing theorem that 2; A where A; is the i'th largest eigenvalue of Σ since Σ is orthogonal projection of E on U. Consequently, we find the desired bound where
M
L. = M
n
x+1+ +
The equality holds by setting U to be the top-r eigenvectors of Σ and E = E.(U) to be the diagonal matrix according to (51).
D Additional Experiments
In this section, 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 equiva- lence 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 mex-tength mmux and loss evaluated as the average loss given samples from mmax Such setting has been wildly studied in the previous ICL work [Garg al., 2022,
32

---

Akyürek et al., 2023, Li et al., 2023]. We generate data according to (7) with Σ Eg la and 0. and train l-layer linear attention (Fig. 4a) and H3 (Fig. 4b) models with different max- -eax-eenghhs Mlmax = 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 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 B. Specifically, consider that each sequence corre- sponds to two individual task parameters ( N(0. Ia) and B ~ N(0, Ld). Then the i'th sample is generated via x; N(O, Iu) and y; B; x; where B, 2 (1 2)) and A; 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 continuously enabling novel applications for LLMs [GeminiTTemm et al., 2023, OpenAI, 2023, Touvron et al., 2023]. Towards this, Xie et al. [2022] explain ICL by language model's 's ability to perform implicit Bayesian inference where, under specific assumptions 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.
the pre-training data distribution, the model infers
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 target function from a given function class, including linear models. A number of works have studied this and related settings to develop 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 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, 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. 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 veta--gaddeens 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-relater parameters. Inspired by the connection between linear attention and GD, they developed 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 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 test prompt, thereby corroborating the findings of [Von Oswald et al., 2023]. They also show that the learned model implements 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 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
33

---

suitable non non-linear target functions, showing that learned Transformer again implements a single step of GD on lineare regression objective. For single-layer linear attention model, Zhang al. [2024] study the optimization dynamics of gradient flow while training such 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 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 distributional shift. Furthermore, we strive to capture the effect 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 demonstrations are provided as in-context examples [Wang et al., 2022, Basu et al., 2023]. Through 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 successful retrieval of query-relevant demonstrations in a practical retrieval augmented setup.
significantly outperform the typical ICL setup where fixed task-specific
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, 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 H3 architecture [Fu et al., 2023), highlight its advantages over linear attention in specific ICL settings.
34

---

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 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 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 other settings.
is fine include aspirational goals as motivation long 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 a theoretical study which (similar to prior studies the field) relies precise but simplified data model to draw quantitatively precise conclusions. All the assumptions on the data model are clearly stated in Section and Section 3. We have also added paragraph after the conclusion to specifically highlight various limitations of our work.
Guidelines:
The answer NA means that the paper has 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 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 few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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, 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 impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
35

---

Question: For each theoretical result, does the paper provide the full set of assumptions and 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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide 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 ex- perimental results the paper to the extent that 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 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 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 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 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
36

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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 the data and code, with sufficient instruc- tions 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 corroborate our theoretical findings. We have provided sufficient details reproduce these experiments in Section 4.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc)
public//uiddss/odeeSuueee sot.eessssissssioooooooooooo ssionPolccy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: htte:/giidee//coeSuneieee
/ //nips.cc...c
stiiibimiiiioooooiccc.... for more details.
The authors should provide instructions 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. only 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 supplemental material (appended 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, hyper- parameters, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
37

---

Justification: This work is a theoretical work studying the optimization landscape of linear attention/H3 under population risk. Then our goal of simulations to find the optimal solution corresponding to the minimal risks. Therefore, we do not report the error bars and we have included the discussion 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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 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 the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: Yes, the authors confirm that the research conducted the paper conform wiht the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
38

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: In its current form, we don't 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 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 no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact 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), faireess considerations (e.g., deploymenmn 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 generic algorithm for optimizing neural networks could enable people 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 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 best faith effort.
12. Licenses for existing assets
39

---

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 URL.
The name of the license (c.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 souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 encouraged to reach out to the asset's '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 release new assets.
Researchers should communicate the details the dataset/cod//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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include 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 sereenshots, 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS 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.
40

---

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 equivalent approval/teview based the requirements of your country 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 crowdsoureing 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.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
41


Reviews:
[{"review_id": "3KNFFhwQQ5", "reviewer": "Reviewer_LuUw", "summary": "This paper studies how transformers can use ICL to solve linear regression problems. It is shown that state space models, transformers are both capable of performing linear regression as well as gradient descent (which implements the least squares solution). There are results about LORA and RAG.", "questions": "What is the significance of the RAG result? It makes sense that giving the learner access to $\\beta$ through $X$ directly improves the error. A comparison to some reasonable baseline might be a nice story if it came about that transformers are able to leverage this kind of side information particularly well. What i mean is to compare this with a loss like $\\Vert y-X\\beta\\Vert^2 -\\Vert X\\beta\\Vert^2 + \\Vert \\beta\\Vert^2$ which is aware of this kind of side information.\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$?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The presentation is clear.", "weaknesses": "The analysis of vanilla regression is not novel, I think, and has been done for instance in Theorem 1 of Ahn et al (transformers learn to do preconditioned gradient descent) where it is even shown that a pretrained transformer (with GD) learns to do this in a similar setting (which is stronger than showing that it can learn it). Maybe for SSMs this result is new but I am not sure how strong of a contribution this is considering that it only establishes the existence of such a solution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed comments and questions on our submission.\n\n**W1: Novelty of contribution and prior art.** Ahn et al. analyzes the loss landscape of linear transformers. Their results only apply to special IID data models (see their Table 1) but they also characterize critical points of multilayer attention. Also the more recent [[Zhang et al.]](https://arxiv.org/pdf/2306.09927) studies gradient flow under the IID data model with isotropic task covariance.\n\nSimilar to Ahn et al., we characterize the **loss landscape**. While our Theorem 1 may appear similar to them, our work makes multiple novel contributions discussed under the Common Response. To recap,\n- We provide the first analysis establishing the optimization-theoretic equivalence of single-layer SSM/H3, linear attention, and optimal PGD under suitable data settings.\n- Comparable works assume IID data whereas our theory allows for correlated data under Assumptions 1 and 2. \n- Comparable works assume full parameterization ($d\\times d$ weights). This is not realistic in practice. Our Lemma 1 characterizes the landscape of low-rank attention for the first time.\n\n **W2: \"The results only show the existence of the solution.\"** As explained above, our study extends beyond showing the existence of a solution and thoroughly investigates the **optimization landscape of ICL** across various architectures and data settings.\n\n**Q1: Significance of RAG.** In a typical RAG setup, given a query, one retrieves relevant demonstrations to create in-context prompts which often leads to improved performance compared to utilizing query-independent demonstrations for ICL. With the objective of providing a theoretical justification to this phenomenon observed in practice, we consider the data model in Eq. (11), where the query token $x$ and the in-context features $(x_i)_{i=1}^n$ are defined to be relevant with $\\alpha$ being the Pearson correlation coefficient among them (motivated by RAG practice). However, it’s important to note that the query and in-context features  are all independent of the task vector $\\beta$. Thus, contrary to the reviewer’s statement, we are not providing the learner access to $\\beta$ through $X$. As a key contribution, under the data model in Eq. (11), we theoretically quantify the improvement in the ICL sample complexity realized via RAG. As shown in Eq. (12), the improvement factor is $(\\alpha^2d+1)$. Higher $\\alpha$ values lead to greater sample efficiency benefits. \n\n**Fair comparison baseline:** In Eq. (11), we sample $(x_i)_{i=1}^n$ via $\\mathcal{N}(\\alpha x,(1−\\alpha^2)I)$ to ensure that the features and labels for different correlation coefficients $\\alpha$ have the same norm, i.e., $E[||x_i||^2]=d$ and $E[y_i^2]=d+\\sigma^2$. With this normalization, $\\alpha=0$ serves as the baseline, and increasing $\\alpha$ benefits in reducing the sample complexity of ICL as validated in Fig 1(b). Finally, under the RAG setting where $X$ and $\\beta$ are independent, the loss $||y-X\\beta||^2-||X\\beta||^2+||\\beta||^2$ suggested by the reviewer will yield similar results, subject to some normalization. \n\n**Q2: Significance of LoRA.** We study low-rank attention and LoRA adaptation because these are typically what is used in real applications. Lemma 1 shows that attention with low-rank parameterization learns to recover the $r$-dimensional data-task eigenspace corresponding to the top eigenvalues to achieve optimal loss. Based on this insight from Lemma 1, we derive the LoRA result, by considering the distance between the old and new covariance eigen-spectrums. Also see Common Response.\n \n **Q3: Questions on Fig.3.** \n> - **Why linear attention performs better for shorter contexts.** In Fig. 3(a)-(c), we train the model to minimize the average risk over all in-context instances, not just the last/query token as in other results. According to Theorem 1, the optimal weight $W_\\star$ varies with context length $n$, so the model must learn to optimize $W$ across different lengths. Thus, a model trained on shorter in-context windows can outperform the one trained on longer windows when tested on shorter contexts.\n> - **How Fig. 3(b) shows that H3 is better for longer context...** Inspired by this question, we tested both linear attention and SSM models (from Fig. 3 (a) and (b)) over unseen context lengths up to $100$. The results are provided in Fig. 2 of the provided pdf file in the Common Response. These results clearly show that, compared to linear attention, SSM generalizes much better to unseen context lengths, supporting our claim. Additionally, these results indicate that while models trained on shorter in-context lengths perform better over shorter ranges, their performance does not compare as well over longer ranges. We will include these new experiments.\n> - **...performance on varying $\\beta$ in Fig. 3(d) better than the IID case…**  Nice question! The reviewer is correct in suggesting an IID baseline. We have identified that the experimental setup led to varying expected norms of the label $E[y_i^2]$ due to the way $\\beta_i$ is defined in the experiment section. To recap, given two random vectors $\\beta_1$ and $\\beta_2$ that follow the standard Gaussian distribution, the task vector at the $i$th position is defined by $\\beta_i=\\alpha_i\\beta_1+(1-\\alpha_i)\\beta_2$. Then the expected norm of the label $E[y_i^2]=\\alpha_i^2E[||x^\\top\\beta_1||]+(1-\\alpha_i)^2E[||x^\\top\\beta_2||^2]=(\\alpha_i^2+(1-\\alpha_i)^2)d$ varies with $i$. Based on this finding, we have updated the black curve of Fig.3(d) in the paper and the revised figure is shown in the Fig.3 in the provided pdf file. The updated curve now provides a reasonable IID baseline. We appreciate the reviewer's helpful comment and will revise the paper accordingly. \n\nWe thank the reviewer for the detailed comments. Although the reviewer noted a lack of novelty and recommended rejection, we hope our responses highlight the novelty and clarify any misunderstandings.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9cIXcDmYBR", "reviewer": "Reviewer_HGXo", "summary": "The authors examine the capabilities of Transformers with linear attention in performing in-context learning (ICL) by implementing a linear estimator through gradient descent. The existing studies mostly consider IID task and feature vectors and fully parameterized attention weights. This work expands on these studies by analyzing:\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.", "questions": "1. How does the distributional alignment in RAG quantitatively affect the sample complexity of ICL? Are there specific metrics or case studies that illustrate these benefits in practical applications?\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?", "limitations": "As mentioned in Weakness: 1. The analysis is limited to single-layer linear attention and linear distribution data, which might not fully capture the complexities of multi-layer architectures. 2. The RAG and LoRA analyses are not precise and fully formal.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's constructive feedback and the recognition of the credibility of our findings. Below, we address the questions and concerns raised by the reviewer.\n\n**W1:** While our study focuses on single-layer architectures similar to previous work, it nonetheless presents novel and insightful conclusions for various practical settings pertaining to in-context learning (ICL). We have briefly summarized the contributions in the Common Response. Although there exist works investigating the ICL performance using multi-layer models [Ahn et al., Von Oswald et al.], analyzing multi-layer architectures still introduces challenges, particularly when exploring the optimization landscape rather than existence or approximation results. We acknowledge this as an important area for future research.\n \n**W2:** Given that this concern appears to be common, we have derived the exact formula for the RAG setting and provided empirical evidence supporting its validity in the Common Response. We have also included additional discussion regarding the LoRA setting. The agreement between our theoretical analysis and empirical results in the figures (Fig. 1 for RAG in the provided pdf file and Fig. 2(c) for LoRA in the submission) further validate our findings.\n \n**Q1:** The benefit of RAG has been empirically studied in existing literature [Lewis et al., Nakano et al., Izacard et al.] which allows for selecting relevant demonstrations from a collection of instances. In our work, we model it within a linear data setting in Eq. (11) such that query token $x$ is correlated with the in-context features $(x_i)_{i=1}^n$. We derive results showing that $\\alpha$-correlated RAG data model achieves a reduction in sample complexity by a factor of $(\\alpha^2d+1)$. This implies that highly relevant demonstrations require significantly fewer in-context samples to achieve comparable performance compared to unrelated/fixed demonstrations.\n \n**Q2:** It is acknowledged that LoRA underperforms fully parameterized models, and our results reflect this. As demonstrated in Eq. (14), setting $r=d$ returns the optimal risk when the model is fully parameterized, highlighting a discrepancy when $r\\neq d$ or when $\\Sigma^{old}\\neq\\Sigma^{new}$. Fig. 2(c) also empirically validates this by varying the rank of LoRA weights; a rank of 20 corresponds to full parameterization, achieving the lowest risk. Reducing the rank $r$ increases the test risk, as expected.\n\n**Q3:** Multilayer architectures will be more challenging due to their non-convex loss landscape. The most related work is by Ahn et al. who characterizes the critical points of ICL with multi-layer linear attention. Extrapolating from our and their results, we suspect that certain critical points of the $L$-layer SSM would correspond to implementing a $L$-step weighted PGD algorithm. Ahn et al. also makes stringent assumptions on data (besides the IID data model, they also require specific covariance structures as shown in their Table 1). It would be interesting to see if our general data assumptions can be imported to the multilayer models.\n\n**Q4:** LoRA [Hu et al.] is a popular technique proposed to adapt large models to downstream tasks in a parameter-efficient fashion. Typically, downstream tasks have different distributions from pretrained tasks, and traditional methods such as fine-tuning the entire model over the new task are expensive and inefficient. LoRA uses fewer parameters to adapt the model to new tasks without modifying the model weights, making it both memory- and parameter-efficient. In practice, it has been widely applied and shown to realize significant improvements (e.g., Llama3). Our focus is on studying its theoretical aspects for ICL settings.\n \n**Q5:**  As noted in our discussion of relevant literature in the submission, Mahankali et al.'s work has demonstrated that even for nonlinear in-context tasks, a linear attention model still implements one step of gradient descent on the *linear regression objective*. This suggests that for nonlinear tasks, a single-layer linear attention model may still achieve the same optimal loss as optimally PGD, and the challenge lies in finding the optimal preconditioning weight. While we have not yet explored nonlinear settings, we recognize this as an important direction for future research.\n \nWe hope this response adequately addresses the reviewer's concerns. We are committed to enhancing our submission based on their feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UXPnQal5ys", "reviewer": "Reviewer_EE6j", "summary": "The paper studies the in-context learning capabilities of linear attention (ATT) and linear state space layers (SSM) on a linear regression task. It shows that for both ATT and SSM there exists a parametrization such that they perform as well as one step of preconditioned gradient descent (PGD) with optimal preconditioning in expectation, and that they cannot perform better.\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.", "questions": "* Under which assumptions are the approximation in (12) valid? What is the nature of the approximations made in the proof of the statement (high probability, asymptotic, ...)?\n\n* Can you elaborate on how claim (A3) follow from (14)?\n\n* Define $\\alpha$-correlated\n* l.342: wildly studied -> widely studied (?)", "limitations": "The authors address the limitations of their work and acknowledge that their analyses are \"not precise and fully formal\".", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "In 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for recognizing the relevance of our work to contemporary ML research and AI applications. Indeed, our focus is developing a theoretical understanding of in-context learning (ICL) that is insightful for practical settings. Below, we respond to the concerns raised by the reviewer point by point.\n \n**W1/W2:** We have clarified these in our Common Response to all reviewers.\n\n**W3:** We acknowledge this concern and will better explain Figure 3. Proposition 1 establishes the equivalence of linear attention (ATT) and SSM under Assumptions 1 and 2, namely, $(i)$ all ICL examples within the prompt share the same task vector, $(ii)$ they are conditioned on the task and query vector, and $(iii)$ the model is trained to predict query at a fixed time $n+1$. The point of Fig. 3 is that, *when these assumptions are not valid, SSM and ATT are not necessarily equivalent and also that SSM can implement a more general predictor, namely weighted PGD*. \n\n- In Fig 3(a,b,c) we train ATT and SSM using the **average loss** where we treat all examples within the prompt as a query. That is, rather than fitting only the last token (index $n+1$), we fit all timesteps where we predict time $i+1$ using examples from $1$ to $i$. This is in line with how LLMs are trained seq2seq in practice as well as ICL experiments of [[Garg et al. NeurIPS’22]](https://proceedings.neurips.cc/paper_files/paper/2022/file/c529dba08a146ea8d6cf715ae8930cbe-Paper-Conference.pdf). Under this setting, Fig 3(a,b) demonstrate that SSM and ATT have different behavior whereas Fig 3(c) demonstrates that SSM achieves a strictly lower risk which is in line with our statement (as weighted PGD is more general than PGD). Here, the intuition is that SSM can intelligently weight different positions in the prompt to better minimize average risk. \n\n- Finally, Fig 3(d) showcases another setting where Proposition 1 does not hold: The task vector $\\beta$ is not shared and is varying along the prompt (i.e., temporally-drifting distribution). In this setting, SSM outperforms ATT for the same reason. The difference becomes visible as $n$ grows and there is more room for weighing the context window via weighted PGD.\n\nWe will revise the current text so that these points come across more clearly.\n\n \n**Q1:** We clarified this in the Common Response: This approximation is valid when $\\alpha=\\mathcal{O}(1/\\sqrt{d}), n/d=\\mathcal{O}(1)$ and $d$ is sufficiently large. We will include a full theorem stating the precise risk in the revised manuscript.\n\n**Q2:** (A3) follows from Lemma 1 and Eq. (14). Recall that $\\Sigma$ and $\\Sigma^{new}$ are the *effective covariances* of the source and target distributions. Here, *effective covariance* refers to the product of the task and feature covariances. Lemma 1 studies the landscape of low-rank parameterized attention: Eq. (13) is a strict generalization of the risk formula $\\cal{L}_*$ of Theorem 1 where the model reduces the risk precisely along the top-$r$ eigendirections. Similarly, Eq. (14) considers the fine-tuning setting and shows that LoRA can improve the risk of an initial model by updating the attention weights along the top-$r$ eigendirections with the maximal gain. Additional discussion is under the Common Response.\n\n**Q3:** The notion of $\\alpha$-correlation is defined in Eq.(11) where $\\mathbb{E}[\\cos(x_i,x)]=\\alpha$. Thus, $\\alpha\\in[0,1]$ is the Pearson correlation coefficient. We will revise the manuscript ensuring that the term \"$\\alpha$-correlated\" is clearly linked to its definition, especially in the introduction section.\n\n**Q4:** Thanks for identifying the typo error and we have corrected it. \n\nWe hope that the clarifications provided in our responses have addressed the concerns highlighted by the reviewer. We would be happy to respond to further feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Common response to all the reviewers
We thank the reviewers for their constructive comments and insightful questions. We are glad that Reviewer HGXo acknowledges the credibility of our conclusions and Reviewer EE6j notes the high relevance of our study to contemporary machine learning research and applications of AI. Here, we would like to restate our key contributions and respond to other shared concerns raised by the reviewers. 

Our main contributions are as follows:

1. We establish the theoretical and empirical equivalence among optimizing single-layer **linear attention**, single-layer **SSM**, and implementing *optimally*-**preconditioned gradient descent (PGD)**. While previous works (e.g., Ahn et al.) have noted the equivalence between linear attention and PGD, to the best of our knowledge, our work is the first to elucidate the equivalence between SSM and PGD. 

2. Our Proposition 1 extends the equivalence among attention, SSM, and PGD to more general and realistic settings, subsuming but going beyond the independent data scenario (as in Ahn et al. and our Theorem 1). Two key contributions are:

- Our contribution on SSM is entirely novel and relies on establishing an optimization-theoretic equivalence between gating (within the SSM) and linear attention.
 - By considering **dependent data** (e.g., in RAG) and **low-rank parameterizations** (e.g., for LoRA adaption) — factors not assumed or analyzed in previous studies — we enhance the understanding of model behavior under more complex yet highly practical settings.

3. The alignments between theoretical predictions and empirical results demonstrate the accuracy and value of our theoretical insights. 

To proceed, we address the shared concerns by the reviewers:

 - **The exact formula of RAG analysis:** We recognize that the analysis of the RAG setting in our submission is not fully precise due to the complexity involved in the high-order (up to 6-order) moments of $x$, $X$, and $\beta$. To address this main concern, we have recalculated the exact formulations for the RAG data setting. In particular, the final solution takes the following exact form:
$$W_\star=cI\quad\text{and}\quad {\mathcal{L}_\star}=d+\sigma^2-cnd(\alpha^2(d+1)+1)$$
where 
$$c=\frac{\alpha^2(d+1)+1}{\alpha^4n(d+2)(d+4)+\alpha^2(1-\alpha^2)(d+2)(d+2n+3)+(1-\alpha^2)^2(d+n+1)+\sigma^2(\alpha^2(d+1)+1)}.$$
Here $\alpha=0$ corresponds to iid setting. Note that Eq. (12) in our submission provided an approximate solution assuming that $\alpha=\mathcal{O}(1/\sqrt{d})$, $d/n=\mathcal{O}(1)$ and large enough $d$.  Additionally, we have also updated the RAG figure (Fig.1(b)) based on the exact formula provided above and the results are shown in Fig.1 of the provided pdf file. The new theoretical predictions now perfectly align with empirical observations. We will incorporate these updates in the final version of the paper.

- **The interpretation of LoRA results:** Our LoRA results in Eq. (14) show that when there is distribution shift and the joint diagonalizability assumption holds, LoRA leads to the model adapting the initial preconditioning weights to a target distribution over the principal $r$-dimensional eigenspace with the maximal gain. Though the analysis presents only an upper bound due to the complexity of arbitrary distribution shifts (e.g., arbitrary $\Sigma^{old}$ and $\Sigma^{new}$ matrices), it marks the first optimization-theoretic exploration of LoRA in ICL. Our empirical results in Figure 2(c) validate the tightness of the prediction of Eq. (14). Notably, Lemma 1 on low-rank attention provides a stronger theoretical guarantee on the landscape, namely, Eq. (13). This can be viewed as a special case of the LoRA setting where the weights of the initial model are set to zero. We will enhance the discussion of our LoRA results. We believe that extending our analysis of LoRA to broader settings is an interesting avenue for future work.

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

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

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