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

TL;DR

The paper investigates how in-context learning emerges in 1-layer linear attention and H3 state-space models under realistic data distributions, showing that both architectures implement a $1$-step preconditioned gradient descent. It analyzes three data designs (IID, retrieval-augmented generation, and task-feature alignment) and demonstrates that distributional alignment effectively increases the usable in-context sample size, with a multiplier $oldsymbol{ ext{kappa}= ext{alpha}^2 d+1}$. It also extends the analysis to low-rank parameterizations and LoRA, deriving explicit risk expressions and adaptation bounds that match empirical results. Collectively, the work provides principled insights into the optimization and generalization landscapes of ICL, guiding architecture choice and data-side design to bolster few-shot learning in practical systems.

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.

Figures (4)

• Figure 1: We investigate the optimization landscape of in-context learning from the lens of architecture 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 \ref{['sec:main']}. More experimental details and discussion are deferred to Section \ref{['sec exp']}.
• Figure 2: Empirical evidence validates Theorem \ref{['thm:independent']} and Proposition \ref{['lemma:eqv']}. We train 1-layer linear attention and H3 models with prompts containing independent demonstrations following a linear model, and dotted curves are the theory curves following Eq. \ref{['formula ind']}. (a): We consider noiseless i.i.d. setting where $\boldsymbol{\Sigma}_{\bm{x}}=\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}={\bm{I}}_{d}$ and $\sigma=0$, with results presented in red (attention) and blue (H3) solid curves. (b): We conduct noisy label experiments by choosing $\sigma\neq0$. (c): Consider non-isotropic task by setting $\boldsymbol{\Sigma}_{{\boldsymbol{\beta}}}=\gamma{\mathbf{1}}{\mathbf{1}}^\top+(1-\gamma){\bm{I}}_d$. 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 \ref{['thm:independent']} and Proposition \ref{['lemma:eqv']}. More experimental details are discussed in Section \ref{['sec exp']}.
• Figure 3: Distributional alignment and low-rank parameterization experiments. (a) and (b) show the ICL results using data generated via \ref{['data rag']} and \ref{['data tfa']}, respectively, by changing $\alpha$ from $0$ to $0.6$. In (c), we train low-rank linear attention models by setting $\bm{W}_k,\bm{W}_q\in\mathbb{R}^{(d+1)\times r}$ and in (d), we apply the low-rank LoRA adaptor, $\bm{W}_{lora}:=\bm{W}_{\text{up}}\bm{W}_{\text{down}}^\top$ where $\bm{W}_{\text{up}},\bm{W}_{\text{down}}\in\mathbb{R}^{(d+1)\times r}$, 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 \ref{['sec:main']}), respectively, and the alignments validate our theorems in Section \ref{['sec:main']}. More experimental details are discussed in Section \ref{['sec exp']}.
• Figure 4: Further comparison for linear attention and H3. In (a) and (b), given maximum context lengths $n_{\max}$, we train linear attention and H3 models to minimize the average loss across all positions $n$ from $1$ to $n_{\max}$. Averaged test risks are presented in (c). In (d), the task vector ${\boldsymbol{\beta}}$ evolves gradually over the context positions $i\leq n$ via ${\boldsymbol{\beta}}_i=(i/n){\boldsymbol{\beta}}_1+(1-i/n){\boldsymbol{\beta}}_2$. In both scenarios, H3 outperforms linear attention benefiting from its additional convolutional filter (c.f. $\bm{f}$ in \ref{['ssm']}). More experimental details are discussed in Section \ref{['sec exp']}.

Theorems & Definitions (12)

• Proposition 1
• Lemma 1
• Lemma 2
• Theorem 1
• Corollary 1
• Theorem 2
• Theorem 3
• Lemma 3
• Lemma 4
• Lemma 5