[论文解读] Scalable Transformers for Neural Machine Translation

Peng Gao, Shijie Geng+4 more|arXiv (Cornell University)|Jun 4, 2021

Natural Language Processing Techniques参考文献 33被引用 9

一句话总结

本文提出可扩展变换器（Scalable Transformers, ST），一种单一统一模型，通过共享且可裁剪的权重实现在不同计算量（FLOPs）和参数量之间的动态扩展。通过采用包含在线词级与离线序列级自蒸馏的三阶段训练方案，ST 在无需微调的情况下实现了在多种模型规模下的最先进性能，其子模型不仅优于独立训练的对应模型，甚至在某些情况下超越了原始最宽变换器模型。

ABSTRACT

Transformer has been widely adopted in Neural Machine Translation (NMT) because of its large capacity and parallel training of sequence generation. However, the deployment of Transformer is challenging because different scenarios require models of different complexities and scales. Naively training multiple Transformers is redundant in terms of both computation and memory. In this paper, we propose a novel Scalable Transformers, which naturally contains sub-Transformers of different scales and have shared parameters. Each sub-Transformer can be easily obtained by cropping the parameters of the largest Transformer. A three-stage training scheme is proposed to tackle the difficulty of training the Scalable Transformers, which introduces additional supervisions from word-level and sequence-level self-distillation. Extensive experiments were conducted on WMT EN-De and En-Fr to validate our proposed Scalable Transformers.

研究动机与目标

• 为解决在移动设备（低FLOPs）与计算集群（高精度）等不同硬件约束下部署变换器模型的挑战，无需重新训练。
• 消除为不同规模单独训练和存储多个独立模型的冗余。
• 开发一种统一架构，支持灵活的、按需运行时扩展模型宽度，同时保持高翻译精度。
• 通过知识蒸馏缓解联合训练过程中不同规模子变换器之间的训练干扰。

提出的方法

• 提出一种统一的可扩展变换器架构，其中所有子变换器共享参数，并通过裁剪最宽模型的层获得。
• 采用三阶段训练方案：首先预训练最宽模型，然后在训练过程中通过在线词级自蒸馏监督较小的子变换器。
• 引入离线序列级自蒸馏，利用最宽模型生成的软标签进一步优化较小的子变换器。
• 在训练过程中随机采样不同宽度的子变换器，以确保在各规模下均具备鲁棒性与泛化能力。
• 通过裁剪最宽模型实现动态推理，获得性能损失极小的轻量级可部署版本。
• 利用参数共享显著降低内存与计算成本，相比独立训练多个模型具有显著优势。

实验结果

研究问题

• RQ1 是否可以训练单一变换器模型，使其在不重新训练的前提下支持多种推理规模（以FLOPs和参数量衡量）？
• RQ2 当使用共享参数联合训练不同尺寸的子变换器时，如何缓解其间的训练干扰？
• RQ3 自蒸馏技术是否能提升较小子变换器的性能，同时保持或增强最宽模型的性能？
• RQ4 训练后搜索最优子变换器的能力是否优于直接训练一个独立的最宽模型？
• RQ5 所提方法是否能在降低总训练与存储成本的同时，实现优于或相当的性能表现？

主要发现

• 在类型-2可扩展变换器中，通过随机搜索获得的最佳子变换器在WMT En-De测试集上达到29.7 BLEU，优于原始最宽变换器（29.3 BLEU）。
• 类型-2模型中表现最好的前10个子变换器在测试集上的平均BLEU为29.60 ± 0.061，全部超过独立训练的最宽变换器性能。
• 训练可扩展变换器仅需在8块V100 GPU上运行212小时，参数量为209M；而训练12个不同宽度的独立模型则需712小时，参数总量达1.29B。
• 类型-1 ST模型在En-De上达到29.3 BLEU，在En-Fr上达到43.1 BLEU，性能与基线变换器相当，但13个不同尺寸的子模型共享同一套参数。
• 类型-2 ST的性能提升并非源于检查点的模型平均，因为对1000个检查点中表现最好的10个模型进行平均仅得29.4 BLEU，仍低于类型-2 ST中最佳子模型的性能。
• 该方法支持灵活部署：任何子变换器均可通过裁剪最宽模型获得，无需重新训练，显著降低训练与内存开销。