Transformer 如何在上下文学习中利用多头注意力？一项关于稀疏线性回归的案例研究
机器学习
2024-08-09 v1
摘要
尽管基于 Transformer 的模型在各种现实任务中取得了显著成功，但其底层机制仍不甚明了。近期研究表明，Transformer 可以作为线性回归问题的上下文学习器实现梯度下降，并据此展开了各种理论分析。然而，这些工作大多通过设计特定的参数构造来关注 Transformer 的表达能力，缺乏对其训练后固有工作机制的全面理解。在本研究中，我们考虑一个稀疏线性回归问题，并探究训练后的多头 Transformer 如何执行上下文学习。我们通过实验发现，多头的利用在不同层间表现出不同的模式：在第一层中，多个头被利用且至关重要，而在后续层中，通常仅需单个头就足够了。我们为这一观察提供了理论解释：第一层对上下文数据进行预处理，后续层则基于预处理后的上下文执行简单的优化步骤。此外，我们证明这种“先预处理后优化”的算法能显著优于朴素的梯度下降和岭回归算法。进一步的实验结果支持了我们的解释。我们的发现为多头注意力的优势提供了洞见，并有助于理解训练后 Transformer 内部隐藏的更复杂机制。
引用
@article{arxiv.2408.04532,
title = {How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression},
author = {Xingwu Chen and Lei Zhao and Difan Zou},
journal= {arXiv preprint arXiv:2408.04532},
year = {2024}
}