[论文解读] Accelerating Transformer Decoding via a Hybrid of Self-attention and Recurrent Neural Network

Chengyi Wang, Shuangzhi Wu+1 more|arXiv (Cornell University)|Sep 5, 2019

Natural Language Processing Techniques参考文献 13被引用 5

一句话总结

本文提出一种混合Transformer架构，用单层RNN替代自回归解码器，同时保留自注意力编码器，使解码速度相比原始Transformer最高提升4.1倍。通过从完整Transformer教师模型进行知识蒸馏，该混合模型在NIST和WMT2017基准测试中保持了具有竞争力的翻译质量，表现匹配或超过基线模型。

ABSTRACT

Due to the highly parallelizable architecture, Transformer is faster to train than RNN-based models and popularly used in machine translation tasks. However, at inference time, each output word requires all the hidden states of the previously generated words, which limits the parallelization capability, and makes it much slower than RNN-based ones. In this paper, we systematically analyze the time cost of different components of both the Transformer and RNN-based model. Based on it, we propose a hybrid network of self-attention and RNN structures, in which, the highly parallelizable self-attention is utilized as the encoder, and the simpler RNN structure is used as the decoder. Our hybrid network can decode 4-times faster than the Transformer. In addition, with the help of knowledge distillation, our hybrid network achieves comparable translation quality to the original Transformer.

研究动机与目标

• 为解决Transformer模型推理速度慢的问题，其根源在于自回归解码和多头注意力带来的高计算成本。
• 探索将自注意力的并行计算优势与RNN的快速推理能力相结合，是否能构建更高效的解码系统。
• 在用更简单的RNN结构替换完整Transformer解码器的情况下，仍保持高翻译质量。
• 评估知识蒸馏在提升轻量化混合模型性能方面的有效性。
• 证明该混合模型可在标准基准测试中实现显著提速，而不会牺牲翻译质量。

提出的方法

• 模型采用原始Transformer中的标准自注意力编码器，通过多头注意力和残差连接保留捕捉长距离依赖关系的能力。
• 将解码器替换为基于单层GRU的RNN，通过在每一步生成时避免完整序列注意力计算，实现更快的自回归生成。
• 混合模型首先在目标序列上使用最大似然估计（MLE）进行预训练。
• 应用序列级知识蒸馏，其中完整Transformer作为教师模型，在微调过程中指导学生混合模型，从而提升翻译质量。
• 在基线模型和混合模型中均采用关键值对预计算和权重融合技术，以确保速度比较的公平性。
• 在NIST和WMT2017中文-英文翻译任务上进行实验，采用束搜索（beam size=12，长度惩罚=1.0）。

实验结果

研究问题

• RQ1 用轻量级RNN替换Transformer解码器是否能显著加速推理，同时不损失翻译质量？
• RQ2 在相同条件下，混合模型的解码速度与标准RNN模型和完整Transformer模型相比如何？
• RQ3 知识蒸馏在多大程度上能提升轻量化混合模型的性能，使其达到完整Transformer的水平？
• RQ4 自注意力、RNN和知识蒸馏各组件对整体速度与准确率权衡的贡献如何？
• RQ5 该混合架构是否在不同模型深度和不同数据集类型下均保持强性能？

主要发现

• 在NIST数据集上，6层混合模型相比6层Transformer基线模型实现4.1倍加速，解码时间从420秒降至107秒。
• 在WMT2017中文-英文任务上，混合模型相比6层Transformer实现3.9倍加速，解码时间从420秒降至114秒。
• 通过知识蒸馏，混合模型在WMT2017上取得22.50的BLEU分数，与完整Transformer教师模型的23.08分非常接近。
• 即使不使用知识蒸馏，混合模型也优于标准RNMT基线；在使用KD后，其性能超过RNMT超过1个BLEU点。
• 采用单层RNN解码器的混合模型在知识蒸馏后，其BLEU分数与更深的RNN变体（2、4、6层）相当，表明在应用蒸馏时，网络深度的影响较小。
• RNN解码器的时间开销显著低于Transformer解码器——在NIST上分别为138.0秒与434.1秒，证实了用RNN替代计算密集型注意力解码可带来显著效率提升。