[论文解读] Fast Transformer Decoding: One Write-Head is All You Need

Noam Shazeer|arXiv (Cornell University)|Nov 6, 2019

Topic Modeling参考文献 5被引用 26

一句话总结

本文提出多查询注意力（multi-query attention），一种多头注意力的变体，通过在所有注意力头之间共享键（keys）和值（values），显著降低了增量解码过程中的内存带宽消耗。该方法在仅造成轻微性能下降的情况下，将推理速度提升至多12倍，非常适合对延迟敏感的应用。

ABSTRACT

Multi-head attention layers, as used in the Transformer neural sequence model, are a powerful alternative to RNNs for moving information across and between sequences. While training these layers is generally fast and simple, due to parallelizability across the length of the sequence, incremental inference (where such paralleization is impossible) is often slow, due to the memory-bandwidth cost of repeatedly loading the large "keys" and "values" tensors. We propose a variant called multi-query attention, where the keys and values are shared across all of the different attention "heads", greatly reducing the size of these tensors and hence the memory bandwidth requirements of incremental decoding. We verify experimentally that the resulting models can indeed be much faster to decode, and incur only minor quality degradation from the baseline.

研究动机与目标

• 为解决现代硬件上Transformer模型增量解码过程中高内存带宽消耗的问题，从而限制推理速度。
• 在不显著降低模型性能的前提下，减小多头注意力中键和值张量的大小。
• 通过最小化自回归生成过程中的重复内存访问，实现在延迟敏感型应用中的更快推理。
• 评估在所有注意力头之间共享键和值是否能保持与标准多头注意力相当的性能。

提出的方法

• 提出多查询注意力，其中所有注意力头共享相同的键和值投影，从而减少参数数量和内存带宽。
• 通过将每个头独立的键和值投影替换为所有头共享的单一投影，修改标准多头注意力机制。
• 每个头使用相同的查询投影，但所有头共享键和值投影，使K和V张量的大小从O(hd)降低到O(d)。
• 采用标准的自注意力计算，使用共享的K和V矩阵，保持相同的注意力机制但内存占用更小。
• 将该方法应用于标准注意力和局部注意力变体，在不同设置下评估性能。
• 在TPUv2硬件上使用批处理推理和增量解码，测量训练和推理成本。

实验结果

研究问题

• RQ1 在增量解码过程中，通过在所有注意力头之间共享键和值，是否能显著降低内存带宽而不会造成显著性能损失？
• RQ2 所提出的多查询注意力在性能上与标准多头注意力及其他减少注意力参数的变体相比如何？
• RQ3 多查询注意力在自回归生成中能将推理速度提升多少？
• RQ4 该方法在机器翻译和语言建模范例上的性能是否保持竞争力？
• RQ5 该加速效果是否在不增加训练时间或模型大小的前提下实现？

主要发现

• 在TPUv2上，多查询模型将解码器每标记的推理时间从47μs降低至3.8μs，增量解码阶段提速12.4倍。
• 每标记的摊销推理成本从基线的46μs降至3.8μs，编码器成本也从1.7μs降至1.5μs。
• 在WMT14 EN-DE翻译任务中，多查询模型在束搜索（beam size=4）下取得28.5的BLEU分数，略高于基线（28.4）。
• 在十亿词语言建模范例中，多查询模型的困惑度为30.2，仅略差于基线（29.9）。
• 尽管架构有所改变，模型的训练时间保持相似（13.0μs/token vs. 13.2μs/token），表明无训练开销。
• 该方法在质量和速度上均优于所有通过减少h、dk或dv的替代方案，证明其作为内存带宽优化方案的有效性。