揭示Transformer训练中注意力回路的自发形成机制
A Mechanistic Study of Transformers Training Dynamics
- 在模拟能力任务中，通过梯度下降自然出现专用注意力结构
- 训练呈现两阶段模式，归一化层导致损失突增
- 为大模型预训练提供可解释性洞见，适合研究者参考
大规模Transformer预训练是基础模型成功的关键。然而，模型规模限制了我们对优化过程中机制的理解。本文在可控且可解释的稀疏模块加法任务中，研究Transformer的训练动态。实验表明，通过梯度下降可实现名为聚类头的专用注意力电路来解决该问题，且这些路径在训练中自然涌现。通过可视化沙盒追踪标记演化，我们发现训练存在两个阶段，并观察到由归一化层高曲率引起的损失突增。这些发现为更实际场景（如大语言模型预训练）中的观测模式提供了若干见解。
原文摘要 · Abstract (English)
Large-scale pretraining of transformers has been central to the success of foundation models. However, the scale of those models limits our understanding of the mechanisms at play during optimization. In this work, we study the training dynamics of transformers in a controlled and interpretable setting. On the sparse modular addition task, we demonstrate that specialized attention circuits, called clustering heads, can be implemented during gradient descent to solve the problem. Our experiments show that such pathways naturally emerge during training. By monitoring the evolution of tokens via a visual sandbox, we uncover a two-stage learning and the occurrences of loss spikes due to the high curvature of normalization layers. Our findings provide several insights into patterns observed in more practical settings, such as the pretraining of large language models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。