In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention

We sincerely thank the reviewer for the positive feedback and appreciation of our work. Here's our response to the questions.

On Heuristic Derivation and Two-Head Gradient Flow. We acknowledge that in Section 4.2, we employ a heuristic derivation to explain how the observed patterns emerge during gradient optimization, using an expansion-based argument. To establish a more rigorous foundation, we analyze the gradient flow in Appendix C to explain why these patterns persist once they emerge and to characterize the trajectory of the parameter evolution.

• (Multi-head Attention Effectively Acts as Two-head Model) Furthermore, our experiments (see Figure 3) demonstrate that a multi-head model essentially implements the same estimator as a two-head model. As detailed in Equation (4.5), on the solution manifold the multi-head model ultimately separates into two distinct head types, determined by the signs of $\mu^{(h)}$ and $\omega^{(h)}$. Within each group, the heads behave equivalently to a single head with $\omega = \omega^{(h)}$ and a $\mu$ equal to the sum of the $\mu^{(h)}$ values within that group. Thus, analyzing the dynamics of a two-head model provides all the essential insights.
• (Dynamics are Similar Under Symmetry) Besides, note that if we initialize a multi-head model with identical absolute values for $\mu$ and $\omega$ (see Definition C.1) and randomly assign each head a positive or negative sign, then as the number of heads increases, each type group is highly likely to contain nearly the same number of heads. By symmetry, tracking the effective KQ and OV parameters within each group is equivalent to analyzing a two-head model.

On Discrepancy between Theoretical Predictions and Experimental Results.

• (Experimental Justification: Small KQ already Behaves as the Limiting Model) As shown by the theoretical results in Section 4.3, when KQ approaches 0, the model exactly implements the GD predictor. Therefore, we compare the performance—measured in terms of mean squared error (MSE)—of the learned model with small KQ against that of the GD predictor. Figure 3(b) demonstrates that the loss curves for the learned transformer model and the GD predictor coincide, indicating identical performance and confirming that the approximation (see Section 4.3) holds well even for a small constant level KQ. Moreover, because the model is trained via gradient optimization, the negligible loss gap between the small KQ multi-head attention and its limiting case suggests that further optimization from a small constant KQ to the limit of 0 would require an exceptionally long time —likely exceeding the number of training steps used in our experiments.
• (Potential Meachanism: Edge of Stability) In practice, the particular magnitudes of KQ parameters learned by the optimization algorithm depends on particular choices of learning rate, training steps, and batch size (if SGD or Adam is applied). This is partly due to the "river valley" loss landscape near the global minimum (see Figure 5(b)). From the theoretical perspective, this phenomenon can be attributed to the edge of stability (e.g., [1]), which is commonly observed in neural network training. In this regime, gradient descent progresses non-monotonically, oscillating between the “valley walls” of the loss surface and failing to fully converge to the minimum. Experimental results indicate that choosing a smaller learning rate, a larger batch size, and sufficiently many training steps can drive the KQ parameters to eventually approach zero.

We would thank the reviewer for the nice questions and we will incorporate these clarifications in the revised version.

[1] Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z. and Talwalkar, A. (2021). Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065.

We sincerely thank the reviewer for their positive feedback and appreciation of our work. Here's our clarifications on theoretical analysis.

On Informal Theoretical Claims. We would like to clarify that most of the theoretical arguments presented in this paper are rigorous, with the exception of the analysis for Stage I in Section 4.2, i.e., theoretical derivation of feature emergence via gradient optimization. While both Stage I and Stage II admit a Taylor expansion-based argument, the analyses are derived from different perspectives. We will provide further clarification as follows.

• (Analysis for Stage I) In Stage I, we perform a gradient-based analysis to identify the driving components in the gradient that lead to the emergence of patterns. We acknowledge that in this stage, we employ an informal argument by focusing on the first-order approximation of $\exp(\cdot)$, leveraging the small-scale initialization that renders higher-order terms negligible. To support this argument, we compare the informal theoretical conjectures with experimental results, which show perfect alignment.
• (Analysis for Stage II) In contrast, during Stage II, we analyze the loss landscape to understand the convergence manifold. Assuming that the patterns identified in the first stage have already emerged, we examine the higher-order terms in the Taylor expansion of the loss function to demonstrate that homogeneous KQ scaling is a necessary condition for loss minimization. We emphasize that the expansion-based argument in this stage is rigorous for the following reasons: (a) the parameter scale remains small during the final stage, allowing the series summation to be analytically tractable, and (b) we derive the optimality condition for each term of the form $\langle\mu_t,\omega_t^{\odot k}\rangle$ for $k\in\mathbb{Z}$, without truncation, ensuring that the result is both rigorous and concise.

Thank you for your comments and assessments. Here's our response to the comments.

On Comparision with Existing Work.

• Comparison with [1].
  • ([1] corresponds to single-head case in our paper) [1] considers the multi-head attention for multi-task linear regression, where the number of heads equals the number of tasks. Under a specialized initialization, [1] proves that each head learns to handle one task, corresponding to the single-head case in our paper. In contrast, we consider a more flexible setup in which multiple heads can be freely allocated to a single task, allowing the model architecture to be more expressive and complex. As a result, while [1] shows that the single-head model learns a nonparametric predictor with KQ scaling as $1/\sqrt{d}$, we demonstrate that the multi-head model learns a parametric GD predictor with KQ converging to $0^+$. Hence, the analysis are fundamentally different.
  • ([1] also adopts reparametrization) Note that [1] assumes decomposable weights (Definition 3.1 in [1]), which ensures that the eigenvector space remains fixed. Thus, the analysis in [1] also employs a reparameterization, showing that diagonalizability is preserved along the gradient flow and focusing on the eigenvalues, i.e., the diagonal entries in the isotropic case.
• Comparison with [2].
  • (Empirical Insights) [2] only identifies the diagonal KQ patterns with potentially positive and negative values for two-head attention, as well as identical performance when the number of heads exceeds two. We further quantitively figure out the detailed sign-matching, homogeneous KQ magnitude, and zero-sum OV patterns beyond $H=2$, and reveals that the multi-head softmax attention learns to implement GD predictor, which are not covered in [2]. Thus, we establish a more complete understanding.
  • (Theoretical Insights) We provide a comprehensive explanation by analyzing how patterns emerge through training dynamics, as well as how the trained model operates via a function approximation and optimality analysis. While [2] considers full-model parameterization, it focuses solely on a loss approximation similar to our Proposition 4.1. However, it does not capture the training dynamics and function approximation----that we develop based on the loss approximation results. Moreover, the main result in [2]—the superiority of multi-head over single-head attention—is a natural corollary of our findings, since we show that single-head attention learns a nonparametric kernel regressor, whereas multi-head attention learns a parametric GD. By applying the approximate loss from Proposition 4.1 with $H = 1$ and $H = 2$, we can reproduce the result. Finally, [2] only analyze KQ parameters with OV parameters fixed to the corresponding optimal solution, while we analyze with both KQ and OV parameters.
• Comparison with [3]. [3] focus on the binary classification tasks while we work on the regression task, marking a clear distinction.

On Real-Word Implications. Understanding what transformers learn for specific task serves as a starting point for investigating general ICL meachanism. The theoretical and emppirical insights lay the foundation for understanding how deep models learn language.

On Transformer Architecture. In this paper, we aim to understand the mechanism of the standard multi-head attention architecture and the weighted sum of head is beyond the scope. Also, $\mu^{(h)}$'s naturally act as "weights" for each head, and are expected to learn the "optimal" weight during the training. Thus, we do not anticipate any improvement.

On Reparametrized Model. As shown in Observation 4 of §3, the attention model develops a diagonal-only pattern in the KQ circuit and a last-entry only pattern in the OV circuit during the early stages of training and then continues optimizing within this regime. This indicates that a diagonal parameterization is sufficient to capture the core behavior of model training. Importantly, we do not impose any diagonal constraints on the KQ parameters—the diagonal pattern emerges naturally during full-model training and persists throughout (see Figure 4). Moreover, if the KQ parameters are initialized as a diagonal matrix, this structure is preserved over the course of training.

We thank the reviewer for pointing out the related work, and we will incorporate these comparisons in the revised version.

[1] Chen, S., Sheen, H., Wang, T., & Yang, Z.. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. COLT 2024.

[2] Cui, Y., Ren, J., He, P., Tang, J., & Xing, Y. Superiority of multi-head attention in in-context linear regression. 2024.

[3] Li, H., Wang, M., Lu, S., Cui, X., & Chen, P. Y. How do nonlinear transformers learn and generalize in in-context learning? ICML 2024