Training Dynamics of In-Context Learning in Linear Attention

TL;DR

The paper investigates how gradient-descent training induces in-context learning (ICL) in linear attention architectures, contrasting a merged key-query parametrization with a separate key-query parametrization. It reveals two distinct dynamical regimes: ATTN_M features two fixed points and a single abrupt loss drop, while ATTN_S exhibits exponentially many fixed points and a cascade of saddle-to-saddle transitions, with ICL materializing as principal component regression in context when early-stopped. A key contribution is the mapping of ATTN_M to a two-layer linear network with a cubic feature map, and the interpretation of ATTN_S as a sum of three-layer convolutional linear networks, enabling precise analysis of fixed points, plateaus, and convergence behavior. The results show that parametrization critically shapes the evolution of ICL during training and provide a theoretical foothold for understanding how ICL abilities evolve in practice, including potential extensions to softmax attention and more complex data regimes.

Abstract

While attention-based models have demonstrated the remarkable ability of in-context learning (ICL), the theoretical understanding of how these models acquired this ability through gradient descent training is still preliminary. Towards answering this question, we study the gradient descent dynamics of multi-head linear self-attention trained for in-context linear regression. We examine two parametrizations of linear self-attention: one with the key and query weights merged as a single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings). For the merged parametrization, we show that the training dynamics has two fixed points and the loss trajectory exhibits a single, abrupt drop. We derive an analytical time-course solution for a certain class of datasets and initialization. For the separate parametrization, we show that the training dynamics has exponentially many fixed points and the loss exhibits saddle-to-saddle dynamics, which we reduce to scalar ordinary differential equations. During training, the model implements principal component regression in context with the number of principal components increasing over training time. Overall, we provide a theoretical description of how ICL abilities evolve during gradient descent training of linear attention, revealing abrupt acquisition or progressive improvements depending on how the key and query are parametrized.

Figures (14)

• Figure 1: Multi-head linear attention with merged key and query, $\mathsf{ATTN}_{\text{M}}({\bm{X}})_{D+1,N+1}$, is equivalent to a two-layer fully-connected linear network with cubic feature input, $\mathsf{MLP} ({\bm{z}})$. Left: Schematic of the equivalence. Right: Loss trajectories of linear attention and the fully-connected linear network match well. The two models are trained with the same data and initialization. Both exhibit the characteristic abrupt loss drop documented by prior work on the ICL dynamics in linear oswald23GD and softmax attention aaditya24induction. Here $D=4,N=31,H=8$.
• Figure 2: Multi-head linear attention with separate rank-one key and query $\mathsf{ATTN}_{\text{S}}({\bm{X}})_{D+1,N+1}$ is a sum of $H$ (number of heads) three-layer convolutional linear networks with the cubic feature ${\bm{z}}$ as input. Here we take $D=3$ to avoid clutter. Entries in the vectors are denoted as ${\bm{x}}_q = \left[x_q^1, x_q^2, x_q^3 \right]^\top, {\bm{\beta}} = \left[\beta^1, \beta^2, \beta^3 \right]^\top$.
• Figure 3: Multi-head linear attention with separate rank-one key and query exhibits saddle-to-saddle dynamics. (a) The loss curve has $D$ abrupt drops, separated by plateaus (six runs from different random initialization are plotted). The loss at each plateau matches our theoretical prediction in \ref{['eq:attnS-loss']} (dashed gray lines). (b) The value weight $v_i$ in each head for one of the runs in (a) is plotted in solid blue curves. The numerical solutions of $v_i$ from \ref{['eq:1Ddynamics']} are plotted in dashed blue curves and match the simulations well. The shades of blue distinguish different heads. (c) The key weights during the loss plateau are plotted in color. When the model moves from one fixed point to the next, the key weight in a head, ${\bm{k}}_i$, aligns with a new eigenvector of the input token covariance ${\bm{\Lambda}}$. The key weights ${\bm{k}}_{1:4}$ and the eigenvectors ${\bm{e}}_{1:4}$ are rows in the heatmaps. A video of the dynamics is provided at https://yedizhang.github.io/img/attnS.gif. Here $D=4,N=31,H=4$, and ${\bm{\Lambda}}$ has eigenvalues $0.4,0.3,0.2,0.1$ and eigenvectors as plotted in (c).
• Figure 4: Multi-head linear attention with separate low-rank key and query exhibits saddle-to-saddle dynamics, with the duration of plateaus depending on the rank $R$. Solid black curves are loss trajectories from six random initializations. Dashed gray lines mark the loss values predicted by \ref{['eq:attnS-loss']} at nine fixed points, which are $\mathcal{L}({\mathcal{M}}_0),\mathcal{L}({\mathcal{M}}_1),\cdots,\mathcal{L}({\mathcal{M}}_8)$ from top to bottom. The four panels differ only in the rank of the key and query weights. Here $D=8,N=31,H=9$, ${\bm{\Lambda}}$ has trace $1$ and eigenvalues $\lambda_d\propto d^{-1}$.
• Figure 5: Loss trajectories of softmax attention with merged or separate key and query. Six runs from different random initialization are plotted. Similar to the linear attention case, softmax $\mathsf{ATTN}_{\text{M}}$ exhibits one abrupt loss drop, while softmax $\mathsf{ATTN}_{\text{S}}$ exhibits multiple loss drops. The dataset and model setup are the same as \ref{['fig:attnM', 'fig:attnS']}, except for adding the softmax activation function.

Theorems & Definitions (8)

• proof
• proof
• proof
• proof
• proof
• proof
• proof
• proof