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

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.

W1/W2: We have clarified these in our Common Response to all reviewers.

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.

• 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]. 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.

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

We will revise the current text so that these points come across more clearly.

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.

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.

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.

Q4: Thanks for identifying the typo error and we have corrected it.

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

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.

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.

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.

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.

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.

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.

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.

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.

We hope this response adequately addresses the reviewer's concerns. We are committed to enhancing our submission based on their feedback.

We thank the reviewer for the detailed comments and questions on our submission.

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.] studies gradient flow under the IID data model with isotropic task covariance.

Similar 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,

• We provide the first analysis establishing the optimization-theoretic equivalence of single-layer SSM/H3, linear attention, and optimal PGD under suitable data settings.
• Comparable works assume IID data whereas our theory allows for correlated data under Assumptions 1 and 2.
• 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.

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.

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.

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.

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.

Q3: Questions on Fig.3.

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

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