Training Dynamics of In-Context Learning in Linear Attention

Thank you very much for your interest in our work and for raising clear, constructive questions. We're glad to hear that you like our results. Below, we present some new experimental results and respond to questions. The numbering follows the reviewer's "Questions for Authors" list.

1. Nonzero $\boldsymbol v_i,\boldsymbol u_i$

   Thank you for this insightful question. The validity of fixed points is not affect by $\boldsymbol v_i,\boldsymbol u_i$. In other words, setting $\boldsymbol v_i=\boldsymbol u_i = \boldsymbol 0$ and the rest of the weights to Equation (9) or (15) results in valid fixed points. This follows from our proof that (i) $\boldsymbol v_i,\boldsymbol u_i$ have zero gradient when they are zero; (ii) the rest of the weights have zero gradient when they satisfy Equation (9) or (15). Therefore, all weights in the full model have zero gradient.

   As for dynamics, we agree that in practice, $\boldsymbol v_i,\boldsymbol u_i$ can be nonzero and they evolve in training, affecting the dynamics in interesting ways. We actually discussed this issue in our Appendix F "Training Dynamics of In-Context and In-Weight Learning". In our setup, nonzero $\boldsymbol v_i,\boldsymbol u_i$ can be interpreted as a form of in-weight learning (IWL). In our main setup of sampling the task vector $\boldsymbol w$ iid from $\mathcal N (0,I)$, $\boldsymbol v_i,\boldsymbol u_i$ remain near zero throughout training, and the model develops only ICL, as shown in Figure 9a. However, when a portion of $\boldsymbol w$ is fixed, the evolution of $\boldsymbol v_i,\boldsymbol u_i$ is non-negligible. In this case, the model first develops IWL, then loses IWL while simultaneously developing ICL, as shown in Figure 9b-e. This competition between ICL and IWL is indeed intriguing. While a full theoretical treatment is beyond our scope, we hope future studies can build on our analysis of the purely ICL dynamics and explore its interaction with IWL dynamics. If the reviewer finds it beneficial, we're happy to move some of Appendix F to main text.

   We also note that setting $\boldsymbol v_i=\boldsymbol u_i = \boldsymbol 0$ is common in theory literature on ICL in linear [1,2] and softmax [3,4] attention.

2. Cumulative Loss

   We have added Figure 14 (see URL) to demonstrate that using the cumulative loss doesn't change the training dynamics, other than the loss values being different. Consistent with using the last-token loss, training $\mathsf{ATTN_M}$ with cumulative loss exhibits a single abrupt loss drop, and training $\mathsf{ATTN_S}$ exhibits multiple loss drops. In our in-context linear regression setup, training with cumulative loss is effectively equivalent to training with the last-token loss on sequences of varying lengths. Because the linear attention model includes a $1/N$ scaling factor, it can be trained and tested with varying sequence lengths $N$. For a detailed discussion on the $1/N$ scaling, please refer to our rebuttal to reviewer etiN.

   It may be worth noting that computing loss for only the last token is also common in the literature, such as in Von Oswald et al. [5] and Zhang et al. [1].

3. Softmax Attention

   We have added Figure 10 (see URL) to show that the different training dynamics of linear $\mathsf{ATTN_M}$ and linear $\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.

We've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.

[1] Ruiqi Zhang, Spencer Frei, Peter Bartlett. Trained transformers learn linear models in-context. JMLR 2024.

[2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[3] Juno Kim, Taiji Suzuki. Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape. ICML 2024.

[4] Yu Huang, Yuan Cheng, Yingbin Liang. In-context Convergence of Transformers. ICML 2024.

[5] Johannes Von Oswald et al. Transformers learn in-context by gradient descent. ICML 2023.

Thank you very much for a very detailed and thoughtful review. We're glad to know that you find our findings instructive, rigorous, and well-presented. We'd like to present some new experimental results and respond to questions. All of our added new figures can be found at URL.

• Rationale for $1/N$ Scaling

  We follow the scaling choice in seminal theory works on linear attention trained for in-context linear regression [1,2]. You're correct that the model can be trained to adapt to the sequence length without the $1/N$ scaling. However, including the scaling factor grants the model the flexibility to be trained and tested with varying context lengths $N$. We've added Figure 14 in which the models are trained on sequences of varying lengths -- this would be unfeasible without the scaling. To explain this point, consider that $\mathsf{ATTN_M}$ computes

  $$\mathsf{ATTN_M}(X) = X +\sum_{i=1}^H W^V_i \frac{XX^\top}N W^{KQ}_i X$$

  The blocks in $XX^\top$ contain the term $\sum_{n=1}^N x_n y_n$, which diverges as $N$ increases. In contrast, the blocks in $\frac{XX^\top}N$ contain $\frac1N \sum_{n=1}^N x_n y_n$, which remains bounded as $N$ increases. Including the $1/N$ scaling allows the model to process sequences of varying lengths without needing to adapt weights specifically for each $N$.

  "we could also instead scale $W_i^{KQ}$ and $W_i^V$ by $1/N$ ..."

  You're correct that we can either scale the $XX^\top$ by $1/N$ or scale $W_i^{KQ}, W_i^V$ by $1/\sqrt N$. However, in both case, we arrive at the same expression in line 246, Equations (18) and (21). In our case, the $1/N$ term is present because of the $1/N$ we build into the model. In your case, the $1/N$ term is also present because $W_i^{KQ}, W_i^V$ are scaled by $1/\sqrt N$. Thus, line 246, Equations (18) and (21) are not affected by whether the scaling is built in the model or adapted through training.

  Thank you for this constructive feedback. We'll incorporate this more detailed explanation for the scaling in revision.

• Why Full-Rank $W^{KQ}_i$ in $\mathsf{ATTN_M}$?

  We agree that $\mathsf{ATTN_M}$ diverges from the established multi-head attention structure in practical transformers. Our intention in studying both $\mathsf{ATTN_M}$ and the low-rank $\mathsf{ATTN_S}$ is that: $\mathsf{ATTN_M}$ is widely used in theoretical literature on transformers (see a list of references in lines 60-65) and low-rank $\mathsf{ATTN_S}$ is closer to the parametrization in practical transformers. Given that a significant portion of the theoretical literature has used $\mathsf{ATTN_M}$, we believe that analyzing it can be of independent interest to transformer theorists, if not to practitioners.

  Your insight regarding that the multiple heads in $\mathsf{ATTN_M}$ are somehow "repetitive" or redundant is also right. In other words, the multi-head $\mathsf{ATTN_M}$ is overparameterized. While overparameterization seems unnecessary in terms of expressivity, its effect on training dynamics and convergence properties is a fundamental, long-standing question in deep learning, and is worth studying. Perhaps surprisingly, we find that the $W^{KQ}_i$ weights in different heads of $\mathsf{ATTN_M}$ stay parallel in training, as shown in Figure 4 and Appendix C.7. This provides a hint at how overparameterized attention models learn and generalize.

• Effect of Initialization Scale

  We've added Figure 11 for this. Please refer to our rebuttal to reviewer Hj7w for details (we reached space limit for this rebuttal).

• Multiplicity of Eigenvalues

  Thank you for this interesting question. We have added Figure 12 in which $\Lambda$ has repeated eigenvalues. Six runs from different random initialization are plotted. Linear $\mathsf{ATTN_S}$ with rank-1 key and query may still exhibit plateaus when learning the repeated eigenvalues due to the different random initial weights in each head. It may also skip the plateau for certain random seeds.

  In the case with distinct eigenvalues, the plateau duration is determined by both the size of the eigenvalue and the random initialization. In the case with equal eigenvalues, the plateau duration is determined by only the random initialization.

• Simulations with Larger $N, D$

  We have added Figure 13 in which we train $\mathsf{ATTN_M}$ and $\mathsf{ATTN_S}$ on a dataset with $N=128,D=64$. The loss trajectories are qualitatively similar to those in lower-dimensional cases, despite being noisier.

• We'll add the reference “Transformers as Statisticians" (Bai et al., 2023). It's indeed an omission -- thanks for your note!

• We'll cut back the introduction for softmax attention in Section 2.2 and clarify the subscript $\mu$ in $X_\mu$ denotes the index of a training sample in Section 2.1. Thanks for your suggestion.

[1] Ahn, Cheng, Daneshmand & Sra. NeurIPS 2023. arXiv:2306.00297

[2] Zhang, Frei & Bartlett. JMLR 2024. arXiv:2306.09927

Thank you very much for your interest in our work and for a thoughtful review. We appreciate how you positioned our contribution at the intersection of three very important and relevant topics, and we’re glad that you found our results provide a clear explanation for progressive ICL acquisition. Below, we present some new experimental results and respond to questions.

• Softmax Attention

  We have added Figure 10 (see URL) to show that the different training dynamics of linear $\mathsf{ATTN_M}$ and linear $\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.

• Effect of Initialization Scale

  We'd like to clarify that many of our results do not rely on initialization: the fixed points, the connection to MLP, and the ICL algorithm we identified for the converged network. Then, we analyze the training dynamics under small initialization.

  We've added Figure 11 (see URL) to demonstrate the effect of initialization scale on dynamics. For $\mathsf{ATTN_M}$, increasing the initialization scale shortens the plateau before the single abrupt loss drop. For $\mathsf{ATTN_S}$, increasing the initialization scale shortens all the plateaus between any two successive abrupt loss drops. For both models, the loss trajectories with the largest initialization exhibit an exponential decay shape, which is the hallmark of lazy learning (NTK regime). In existing theoretical literature, the standard choice of initialization is either the large initialization of exponential-shaped lazy learning, or small initialization of rich feature learning as in our paper. The practical initialization scheme is usually somewhere in between. In Figure 11, we see that dynamics from the intermediate initialization seems like a mix of the exponential-shaped and the abrupt sigmoid-shaped curves. Such mixed curves are often seen in practice, e.g. in induction head emergence in natural language settings (Olsson et al. 2022; Argument 1). Our dynamics analysis focuses on the rich learning regime and provides analytical insight into such phenomena, which we believe is a first step toward understanding dynamics in naturalistic settings.

• Practical Implication

  Given that the scope of our paper is theoretical, we believe our findings and intuition can inform practice by providing a deeper understanding of common phenomena observed in practical settings -- such as abrupt learning and the progressive acquisition of ICL. Though not directly resulting in practical recommendation, we hope our work provides a foundation for future work along the spectrum from theory to practice, ultimately leading to more effective architectural and optimization choices.

  Regarding the implication of the equivalence between $\mathsf{ATTN_M}$ and MLP with cubic features, it can add a perspective to the open question of whether MLPs can learn in-context (Boix-Adsera et al [1]; Tong and Pehlevan [2]). Moreover, we reveal that MLP may perform ICL more comparable to attention models when provided with polynomial features instead of the original sequence. This may explain why Boix-Adsera et al (Figure 25) observed that MLP fails to learn ICL with the original sequence $X$ as input, but succeeds when the input is augmented with $XX^\top$. We'll add a discussion of this point in revision.

  [1] Boix-Adsera, E., Saremi, O., Abbe, E., Bengio, S., Littwin, E., & Susskind, J. M. When can transformers reason with abstract symbols? ICLR 2024.

  [2] Tong, W. L., & Pehlevan, C. MLPs Learn In-Context on Regression and Classification Tasks. ICLR 2025.

We've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.

We thank the reviewer for a detailed review and for raising constructive questions. We're glad to know that you think our claims are backed by clear, precise, and convincing evidence. Below, we present some new experimental results and respond to questions.

• Softmax Attention

  We've added Figure 10 (see URL) to show that the different training dynamics of linear $\mathsf{ATTN_M}$ and linear $\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.

• Effect of Initialization Scale

  We'd like to clarify that many of our results do not rely on initialization: the fixed points, the connection to MLP, and the ICL algorithm we identified for the converged network. Then, we analyze the training dynamics under small initialization.

  We've added Figure 11 (see URL) to demonstrate the effect of initialization scale on dynamics. For $\mathsf{ATTN_M}$, increasing the initialization scale shortens the plateau before the single abrupt loss drop. For $\mathsf{ATTN_S}$, increasing the initialization scale shortens all the plateaus between any two successive abrupt loss drops. For both models, the loss trajectories with the largest initialization exhibit an exponential decay shape, which is the hallmark of lazy learning (NTK regime). In existing theoretical literature, the standard choice of initialization is either the large initialization of exponential-shaped lazy learning, or small initialization of rich feature learning as in our paper. The practical initialization scheme is usually somewhere in between. In Figure 11, we see that dynamics from the intermediate initialization seems like a mix of the exponential-shaped and the abrupt sigmoid-shaped curves. Such mixed curves are often seen in practice, e.g. in induction head emergence in natural language settings (Olsson et al. 2022; Argument 1). Our dynamics analysis focuses on the rich learning regime and provides analytical insight into such phenomena, which we believe is a first step toward understanding dynamics in naturalistic settings.

• Simulations with Larger $N, D$

  Thank you for this suggestion. We've added Figure 13 (see URL) in which we train $\mathsf{ATTN_M}$ and $\mathsf{ATTN_S}$ on a dataset with $N=128,D=64$. The loss trajectories are qualitatively similar to those in lower-dimensional cases, despite being noisier. This suggests that our findings do not break in high-dimensional settings.

We've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.