arXiv:2503.12734cs.LGstat.ML2025-03ICML被引 8

多头注意力通过训练自发形成特定权重模式,实现高效线性回归预测。

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

论文配图:In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
图 1 · 摘自论文原文
  • 训练中自动生成关键-查询的对角同质模式与输出-值的末项非零零和模式。
  • 模型性能接近贝叶斯最优,且可外推到更长序列,优于单头注意力。
  • 适用于多任务学习与非各向同性数据,揭示了架构与数据分布的协同作用。

我们研究多头softmax注意力模型在线性数据上进行上下文学习的训练机制。通过大量实验与严格理论分析,揭示了其自发形成的优雅权重模式:键-查询(KQ)权重呈现对角且同质结构,输出-值(OV)权重为仅末项非零且总和为零。这些模式在随机初始化下由梯度训练稳定出现。分析表明,此类结构使多头注意力近似实现无偏梯度下降预测器,性能超越单头注意力,接近贝叶斯最优(比例因子内)。相较于线性Transformer,该注意力能自然泛化至训练时未见的更长序列。研究还扩展至各向异性协变量与多任务线性回归场景:前者中模型学习预条件梯度下降;后者中头数与任务数的相互作用引发叠加现象,高效解决多任务上下文学习。结果表明,上下文学习能力是模型架构与数据分布共同作用下的涌现特性,为深入理解与广泛应用奠定了基础。

原文摘要 · Abstract (English)

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 anisotropic 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.

注意力机制在线学习线性回归可解释性

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。