[论文解读] Sharing Attention Weights for Fast Transformer

Tong Xiao, Yinqiao Li+3 more|arXiv (Cornell University)|Jun 26, 2019

Topic Modeling参考文献 20被引用 7

一句话总结

本文提出共享注意力网络（San），通过在自注意力和编码器-解码器注意力机制中跨相邻层共享注意力权重，加速Transformer推理。通过实现隐藏状态的垂直复用并联合学习共享策略，San在缓存基线基础上实现1.3倍加速，在无缓存情况下最高实现16倍加速，BLEU分数损失极小。

ABSTRACT

Recently, the Transformer machine translation system has shown strong results by stacking attention layers on both the source and target-language sides. But the inference of this model is slow due to the heavy use of dot-product attention in auto-regressive decoding. In this paper we speed up Transformer via a fast and lightweight attention model. More specifically, we share attention weights in adjacent layers and enable the efficient re-use of hidden states in a vertical manner. Moreover, the sharing policy can be jointly learned with the MT model. We test our approach on ten WMT and NIST OpenMT tasks. Experimental results show that it yields an average of 1.3X speed-up (with almost no decrease in BLEU) on top of a state-of-the-art implementation that has already adopted a cache for fast inference. Also, our approach obtains a 1.8X speed-up when it works with the extsc{Aan} model. This is even 16 times faster than the baseline with no use of the attention cache.

研究动机与目标

• 解决由自回归解码和重复的点积注意力计算引起的Transformer模型高推理延迟问题。
• 通过利用相邻层间注意力权重分布的结构相似性，减少堆叠注意力层中的冗余计算。
• 开发一种联合学习框架，端到端优化共享策略与机器翻译模型，而非使用启发式共享规则。
• 通过权重共享复用隐藏状态，减少内存占用。
• 在不显著损失准确率的前提下实现快速推理，尤其在与现有加速技术（如注意力缓存）结合时表现更优。

提出的方法

• 提出一种共享注意力网络（San），在自注意力和编码器-解码器注意力子层中，跨相邻层垂直共享注意力权重。
• 将注意力机制设计为在多个堆叠层中使用相同的注意力权重矩阵，从而减少冗余计算。
• 引入一种联合训练策略，使共享策略与机器翻译模型端到端联合学习，使网络能自动发现最优共享模式。
• 通过将隐藏状态存储在共享内存位置，实现高效复用，减少内存占用并加速推理。
• 将San与现有加速技术（如注意力缓存和Aan模型）集成，证明其改进具有正交性和互补性。
• 以缩放点积注意力机制为基础，跨层应用共享权重，在保持模型结构的同时提升效率。

实验结果

研究问题

• RQ1 在不降低翻译质量的前提下，跨相邻层共享注意力权重是否能减少Transformer的推理时间？
• RQ2 将共享策略与机器翻译模型联合学习，是否能带来比启发式共享策略更好的性能和更快的收敛速度？
• RQ3 当与现有加速方法（如注意力缓存或Aan模型）结合时，San的效率如何？
• RQ4 训练后的Transformer中，相邻层在注意力权重分布上表现出多大程度的相似性，从而支持权重共享的可行性？
• RQ5 San是否能在保持或提升推理速度的同时减少内存使用？

主要发现

• 与已使用注意力缓存的最先进Transformer实现相比，San实现了平均1.3倍的加速，BLEU分数几乎无下降。
• 与Aan模型结合时，San实现了1.8倍的加速，表明不同加速技术之间具有互补优势。
• San与Aan联合使用的系统，相比无任何缓存机制的基线，最快可提升16倍。
• 训练过程中，相邻层间注意力权重分布的Jensen-Shannon散度显著降低，表明层自然收敛至相似的注意力模式。
• San的训练似然略高于基线，表明在共享机制下模型仍保持有效性和稳定性。
• 该方法实现简单，与现有技术正交，可直接应用于多种Transformer变体。