[論文レビュー] Scalable Transformers for Neural Machine Translation

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

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

ひとこと要約

本稿では、共有可能で切り出し可能な重みを用いて、FLOPsおよびパラメータ数の異なる複数のスケールに動的にスケーリング可能な、統合型の1つのモデルであるスケーラブルトランスフォーマー（ST）を提案する。オンラインの単語レベルおよびオフラインのシーケンスレベルの自己蒸留を用いた3段階の訓練スキームにより、再訓練を伴わずに多様なモデルサイズで最先端の性能を達成し、サブモデルは独立に訓練されたものよりも優れており、一部のケースでは元の最も広いトランスフォーマーをも上回る。

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）からクラスタ（高精度）まで多様なハードウェア制約に適合させるために、再訓練を伴わずトランスフォーマーを展開する課題に対処すること。
• 異なるスケール用に別々に訓練・保存する複数のモデルの冗長性を排除すること。
• モデル幅の柔軟かつ即時のスケーリングを可能にする統合アーキテクチャの開発を目的とすること。
• 共有パラメータを用いた一括訓練における、異なるスケールのサブトランスフォーマー間の訓練干渉を知識蒸留により軽減すること。

提案手法

• すべてのサブトランスフォーマーが層の最も広いモデルの重みを共有し、その一部を切り出すことで得られる、統合型のスケーラブルトランスフォーマーのアーキテクチャを提案する。
• 3段階の訓練スキームを採用：まず最も広いモデルを事前訓練し、次に訓練中に小さなサブトランスフォーマーを監視するためにオンラインの単語レベル自己蒸留を適用する。
• 最も広いモデルからのソフトラベルを用いたオフラインのシーケンスレベル自己蒸留を導入し、より小さなサブトランスフォーマーの精錬を促進する。
• 訓練中に異なるサブトランスフォーマー幅をランダムにサンプリングすることで、スケール全体にわたる頑健性と一般化性能を確保する。
• 最も広いモデルを切り出すことで、性能の低下を最小限に抑えた動的推論を可能にする。
• パラメータ共有を活用することで、複数の独立したモデルを訓練するのと比較して、メモリおよび計算コストを大幅に削減する。

実験結果

リサーチクエスチョン

• RQ1 再訓練を伴わず、FLOPsおよびパラメータ数の異なる複数の推論スケール（性能）をサポートできる1つのトランスフォーマー・モデルを訓練できるか？
• RQ2 共有パラメータを用いて一括訓練する際、異なるサイズのサブトランスフォーマー間の訓練干渉をどのように軽減できるか？
• RQ3 自己蒸留技術を用いることで、小さなサブトランスフォーマーの性能を向上させるとともに、最大サイズのモデルの性能を維持または向上させられるか？
• RQ4 訓練後、最適なサブトランスフォーマーを探索できる能力が、スタンドアロンの最大サイズモデルを訓練するのと比較して、より優れた性能をもたらすか？
• RQ5 提案手法は、総合的な訓練および保存コストを削減しながら、最先端のモデルと同等またはそれ以上の性能を達成できるか？

主な発見

• タイプ-2スケーラブルトランスフォーマーからランダムに探索された最良のサブトランスフォーマーは、WMT En-Deテストセットで29.7 BLEUを達成し、元の最も広いトランスフォーマー（29.3 BLEU）を上回った。
• タイプ-2モデルの上位10個のサブトランスフォーマーは、テストセットで平均29.60 ± 0.061 BLEUを達成し、すべてスタンドアロンの最も広いトランスフォーマーの性能を上回った。
• スケーラブルトランスフォーマーの訓練には、8台のV100 GPUで212時間、209Mパラメータを要したが、幅の異なる12個の独立モデルを訓練するには712時間および1.29Bパラメータが必要だった。
• タイプ-1 STモデルは、En-Deで29.3 BLEU、En-Frで43.1 BLEUを達成し、ベースラインのトランスフォーマーと同等の性能を示したが、13個の異なるサイズのサブモデルが同じパラメータを共有していた。
• タイプ-2 STの性能向上は、チェックポイントの平均化によるものではない。スタンドアロンの最も広いモデルの1,000個のチェックポイントから上位10個を平均化しても29.4 BLEUにとどまり、タイプ-2 STの最良サブモデルを下回った。
• 本手法により柔軟な展開が可能である：再訓練を経ずに、最も広いモデルを切り出すことで任意のサブトランスフォーマーを取得でき、訓練およびメモリのオーバーヘッドを顕著に削減できる。