[论文解读] Multi-head or Single-head? An Empirical Comparison for Transformer Training

Liyuan Liu, Jialu Liu+1 more|arXiv (Cornell University)|Jun 17, 2021

Natural Language Processing Techniques参考文献 29被引用 23

一句话总结

本文挑战了多头注意力机制优越性源于关注多个位置的普遍认知，表明训练稳定性才是其真正优势。通过用更深的单头变体替代浅层多头Transformer模型（采用自适应初始化以稳定训练），在无需超参数调优的情况下，所有任务均实现一致的性能提升，证明当训练稳定时，性能提升的关键在于模型深度，而非多头结构本身。

ABSTRACT

Multi-head attention plays a crucial role in the recent success of Transformer models, which leads to consistent performance improvements over conventional attention in various applications. The popular belief is that this effectiveness stems from the ability of jointly attending multiple positions. In this paper, we first demonstrate that jointly attending multiple positions is not a unique feature of multi-head attention, as multi-layer single-head attention also attends multiple positions and is more effective. Then, we suggest the main advantage of the multi-head attention is the training stability, since it has less number of layers than the single-head attention, when attending the same number of positions. For example, 24-layer 16-head Transformer (BERT-large) and 384-layer single-head Transformer has the same total attention head number and roughly the same model size, while the multi-head one is significantly shallower. Meanwhile, we show that, with recent advances in deep learning, we can successfully stabilize the training of the 384-layer Transformer. As the training difficulty is no longer a bottleneck, substantially deeper single-head Transformer achieves consistent performance improvements without tuning hyper-parameters.

研究动机与目标

• 探究多头注意力的成功是否源于对多个位置的关注，还是源于训练稳定性。
• 在模型大小和计算复杂度相同的情况下，比较浅层多头Transformer与更深单头变体的性能表现。
• 评估近期的训练稳定化技术是否能使更深的单头Transformer超越标准多头架构。
• 确定深层模型带来的性能提升是源于架构深度本身，还是多头机制的内在特性。
• 评估深层单头与浅层多头Transformer在推理效率和训练收敛速度方面的差异。

提出的方法

• 将标准多头Transformer（如BERT-base、BERT-large）重构为具有相同总注意力头数和模型大小的更深单头变体。
• 采用自适应模型初始化（Admin）技术，以稳定极深单头Transformer（如384层模型）的训练过程。
• 在相同超参数设置下训练浅层多头与深层单头模型，以隔离深度与架构的影响。
• 通过机器翻译和BERT预训练任务衡量性能，以比较泛化能力与鲁棒性。
• 在相同硬件上评估推理速度与训练效率（GPU小时数），以比较实际部署特性。
• 可视化收敛曲线与不同头数下的性能增益，分析架构深度的影响。

实验结果

研究问题

• RQ1 关注多个位置是否唯一提升了多头注意力的效果？还是深层单头注意力也能实现类似效果？
• RQ2 多头注意力的主要优势是其多位置关注能力，还是更好的训练稳定性？
• RQ3 能否通过现代初始化技术成功训练极深的单头Transformer？其性能是否优于浅层多头模型？
• RQ4 在模型大小相同的情况下，深层单头与浅层多头Transformer在性能、推理速度和训练效率方面如何比较？
• RQ5 深度带来的性能增益是否依赖于注意力头的数量？这种改进在不同任务中是否一致？

主要发现

• 在无需任何超参数调优的情况下，384层深层单头Transformer在多个NLP任务中均优于24层浅层多头模型（如BERT-large）。
• 性能提升并非源于多头机制本身，而是当训练稳定时，模型深度带来的增益。
• 自适应初始化（Admin）使384层单头Transformer能够稳定训练，而标准初始化则导致模型发散。
• 尽管模型大小与计算复杂度相同，深层单头Transformer的收敛速度更快，且性能优于浅层多头版本。
• 深层单头与浅层多头Transformer的推理速度几乎相同，具体优势取决于批大小和序列长度。
• 性能增益随头数增加而提升，证实浅层多头与深层单头模型之间的架构差异是性能提升的关键驱动因素。