[论文解读] Faster Transformer Decoding: N-gram Masked Self-Attention

Ciprian Chelba, Mia Xu Chen+2 more|arXiv (Cornell University)|Jan 14, 2020

Topic Modeling参考文献 8被引用 13

一句话总结

本文提出N-gram掩码自注意力机制，通过仅将自注意力上下文限制在前N−1个目标词元内，从而加速Transformer解码，将计算复杂度从O(T²)降低至O(N·T)。在WMT EnDe和EnFr数据集上的实验表明，当N=8时，BLEU得分接近基线水平（低至0.3–0.4 BLEU），同时实现高达2–3倍的加速并减少内存带宽。

ABSTRACT

Motivated by the fact that most of the information relevant to the prediction of target tokens is drawn from the source sentence $S=s_1, \ldots, s_S$, we propose truncating the target-side window used for computing self-attention by making an $N$-gram assumption. Experiments on WMT EnDe and EnFr data sets show that the $N$-gram masked self-attention model loses very little in BLEU score for $N$ values in the range $4, \ldots, 8$, depending on the task.

研究动机与目标

• 在不显著降低性能的前提下，减少自回归Transformer解码的计算与内存开销。
• 探究将自注意力上下文限制在固定N-gram窗口内是否能保持神经机器翻译中的翻译质量。
• 评估序列生成中推理速度、内存效率与模型准确率之间的权衡。
• 探索在束搜索过程中仅存储最近N−1个词元的固定大小缓冲区的可行性，以减少内存带宽。

提出的方法

• 引入一种N-gram掩码自注意力机制，将解码器中的注意力计算限制在前N−1个目标词元内。
• 通过应用仅排除目标序列中前N−1个位置以外的词元的掩码，对标准因果自注意力机制进行修改。
• 在Lingvo框架中实现N-gram注意力作为可配置选项，替换解码器层中的完整因果注意力。
• 使用大小为N−1的滑动窗口缓冲区，在解码过程中逐次存储和更新上下文词元，以减少内存访问开销。
• 保持完整的编码器-解码器注意力机制，仅将解码器的自注意力限制为N-gram形式。
• 在WMT EnDe和EnFr数据集上，使用标准Transformer架构、子词分词和固定超参数进行模型训练与评估。

实验结果

研究问题

• RQ1 将自注意力上下文限制在前N−1个词元内，对BLEU得分的翻译质量有何影响？
• RQ2 在机器翻译任务中，最优的N-gram阶数（N）是什么，能够实现速度与性能的最佳平衡？
• RQ3 N-gram自注意力在推理过程中能在多大程度上降低计算复杂度与内存带宽？
• RQ4 尽管存在上下文截断，N-gram机制是否仍能保持模型捕捉长距离依赖关系的能力？
• RQ5 N-gram方法是否能通过固定大小的缓冲区高效实现，从而在束搜索过程中减少内存访问？

主要发现

• 在WMT EnDe数据集上，当N=8时，模型在测试集上的BLEU得分为22.5，仅比基线（22.4）低0.1 BLEU，且与最佳基线模型相比仅低0.3–0.4 BLEU。
• 在WMT EnFr数据集上，N=8在测试数据上的BLEU得分为32.7，与基线（32.9）相差仅0.2 BLEU，且优于更小的N-gram模型。
• N-gram模型将计算复杂度从O(T²)降低至O(N·T)，对于长度T≈16–25的序列，理论加速比约为2–3倍。
• 在WMT EnDe测试数据上，N=8模型的对数困惑度为2.66，仅略高于基线的2.64，表明语言建模能力几乎无退化。
• N-gram机制允许使用长度为N−1的固定大小内存缓冲区，将内存带宽降低O(T/N)倍，尤其在TPU硬件上优势显著。
• 当N≥3时，性能呈平稳下降趋势，N=4在WMT EnDe上已达到22.2 BLEU，表明即使较小的N值也适用于速度-精度权衡。