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

• URL: http://arxiv.org/abs/2503.12734v1
• Date: Mon, 17 Mar 2025 02:00:49 GMT
• Title: In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
• Authors: Jianliang He, Xintian Pan, Siyu Chen, Zhuoran Yang,
• Abstract summary: We study how multi-head softmax attention models are trained to perform in-context learning on linear data.<n>Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution.
• Score: 52.159541540613915
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We study how multi-head softmax attention models are trained to perform in-context learning on linear data. Through extensive empirical experiments and rigorous theoretical analysis, we demystify the emergence of elegant attention patterns: a diagonal and homogeneous pattern in the key-query (KQ) weights, and a last-entry-only and zero-sum pattern in the output-value (OV) weights. Remarkably, these patterns consistently appear from gradient-based training starting from random initialization. Our analysis reveals that such emergent structures enable multi-head attention to approximately implement a debiased gradient descent predictor -- one that outperforms single-head attention and nearly achieves Bayesian optimality up to proportional factor. Furthermore, compared to linear transformers, the softmax attention readily generalizes to sequences longer than those seen during training. We also extend our study to scenarios with non-isotropic covariates and multi-task linear regression. In the former, multi-head attention learns to implement a form of pre-conditioned gradient descent. In the latter, we uncover an intriguing regime where the interplay between head number and task number triggers a superposition phenomenon that efficiently resolves multi-task in-context learning. Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution, paving the way for deeper understanding and broader applications of in-context learning.

Related papers

• Training Dynamics of In-Context Learning in Linear Attention [6.663503238373593]
  We study descent gradient dynamics of multi-head linear self-attention trained for in-context linear regression.<n>We characterize how in-context learning abilities evolve during descent training of linear attention.
  arXiv  Detail & Related papers  (2025-01-27T18:03:00Z)
• How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression [19.64743851296488]
  In this study, we consider a sparse linear regression problem and investigate how a trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heads exhibits different patterns across layers. We demonstrate that such a preprocess-then-optimize algorithm can significantly outperform naive gradient descent and ridge regression algorithms.
  arXiv  Detail & Related papers  (2024-08-08T15:33:02Z)
• Understanding Forgetting in Continual Learning with Linear Regression [21.8755265936716]
  Continual learning, focused on sequentially learning multiple tasks, has gained significant attention recently. We provide a general theoretical analysis of forgetting in the linear regression model via Gradient Descent. We demonstrate that, given a sufficiently large data size, the arrangement of tasks in a sequence, where tasks with larger eigenvalues in their population data covariance matrices are trained later, tends to result in increased forgetting.
  arXiv  Detail & Related papers  (2024-05-27T18:33:37Z)
• Asymptotic theory of in-context learning by linear attention [33.53106537972063]
  In-context learning is a cornerstone of Transformers' success.<n>Questions about the necessary sample complexity, pretraining task diversity, and context length for successful ICL remain unresolved.
  arXiv  Detail & Related papers  (2024-05-20T03:24:24Z)
• Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality [54.20763128054692]
  We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics.
  arXiv  Detail & Related papers  (2024-02-29T18:43:52Z)
• How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? [92.90857135952231]
  Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities. We study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear attention model for linear regression.
  arXiv  Detail & Related papers  (2023-10-12T15:01:43Z)
• In-Context Convergence of Transformers [63.04956160537308]
  We study the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent. For data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process.
  arXiv  Detail & Related papers  (2023-10-08T17:55:33Z)
• Multiplicative noise and heavy tails in stochastic optimization [62.993432503309485]
  empirical optimization is central to modern machine learning, but its role in its success is still unclear. We show that it commonly arises in parameters of discrete multiplicative noise due to variance. A detailed analysis is conducted in which we describe on key factors, including recent step size, and data, all exhibit similar results on state-of-the-art neural network models.
  arXiv  Detail & Related papers  (2020-06-11T09:58:01Z)
• Kernel and Rich Regimes in Overparametrized Models [69.40899443842443]
  We show that gradient descent on overparametrized multilayer networks can induce rich implicit biases that are not RKHS norms. We also demonstrate this transition empirically for more complex matrix factorization models and multilayer non-linear networks.
  arXiv  Detail & Related papers  (2020-02-20T15:43:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.