[論文レビュー] Sockeye 3: Fast Neural Machine Translation with PyTorch

Felix Hieber, Michael Denkowski+12 more|arXiv (Cornell University)|Jul 12, 2022

Natural Language Processing Techniques被引用数 11

ひとこと要約

Sockeye 3 は、最適化された混合精度学習、分散データ並列処理、モデル量子化を通じて学習および推論を高速化する PyTorch ベースのニューラル機械翻訳ツールキットです。他の PyTorch 実装と比較して、GPU では最大 126%、CPU では最大 292% 速い推論を達成し、生産および研究用途向けに、ソース/ターゲットのプレフィックス、語彙的ショートリスト、フォーマルリティ/冗長性制御といった高度な機能をサポートしています。

ABSTRACT

Sockeye 3 is the latest version of the Sockeye toolkit for Neural Machine Translation (NMT). Now based on PyTorch, Sockeye 3 provides faster model implementations and more advanced features with a further streamlined codebase. This enables broader experimentation with faster iteration, efficient training of stronger and faster models, and the flexibility to move new ideas quickly from research to production. When running comparable models, Sockeye 3 is up to 126% faster than other PyTorch implementations on GPUs and up to 292% faster on CPUs. Sockeye 3 is open source software released under the Apache 2.0 license.

研究の動機と目的

• PyTorch をベースに、より高速でスケーラブルかつ生産環境対応の NMT ツールキットを開発すること。
• コードベースの単純化と学習・推論パイプラインの最適化を通じて、研究者が迅速にイテレーションを回せるようにすること。
• ソースおよびターゲットのプレフィックス、語彙的ショートリスト、ハイブリッドデコーダーなどの高度な NMT 機能をサポートし、細かく制御可能な翻訳を実現すること。
• 効率的で量子化され、分散処理が可能な学習により、研究のイノベーションと生産環境へのデプロイのギャップを埋めること。
• 翻訳におけるフォーマルリティおよび冗長性制御の実験に柔軟かつオープンソースのフレームワークを提供すること。

提案手法

• Sockeye 3 は PyTorch の JIT コンパイラーを用い、エンコーダー、デコーダー、ビームサーチの静的計算グラフをトレースおよび最適化することで、動的制御フローのオーバーヘッドを低減しています。
• PyTorch の DDP および AMP を用いた分散混合精度学習を実装し、FP16/FP32 の混合精度および完全 FP16 (O2) モードを提供することで、速度とメモリ効率を向上させています。
• データ準備を並列化し、バイナリ形式にシャーディングすることで、固定メモリ使用量で任意に大きなデータセットの学習が可能になっています。
• CPU 推論には PyTorch のダイナミック量子化 (INT8) を、GPU 推論には FP16 キャストを適用することで、最小限の精度損失で推論を高速化しています。
• ソースプレフィックスを用いて、再トレーニングなしでフォーマルリティや冗長性レベルに応じたモデル出力を条件づけています。
• ビームサーチの軽量代替手段としてグリーディサーチを実装し、計算オーバーヘッドを低減しながらも翻訳品質を維持しています。

実験結果

リサーチクエスチョン

• RQ1 PyTorch ベースの NMT ツールキットは、他の実装と比較して、どのようにして著しく高速な学習および推論スルーレットを達成できるか？
• RQ2 混合精度および分散学習は、モデル品質を損なわずに、NMT のスケーラビリティと効率性をどの程度向上できるか？
• RQ3 ソースプレフィックスの条件づけは、ニューラル機械翻訳における出力のフォーマルリティおよび冗長性を効果的に制御できるか？
• RQ4 量子化は、翻訳品質を保持したまま推論速度とメモリフットプリントにどのような影響を与えるか？
• RQ5 グリーディサーチやモデルトレーシングといった最適化された推論コンponent によって、どの程度のパフォーマンス向上が達成できるか？

主な発見

• Sockeye 3 は、他の PyTorch NMT 実装と比較して、GPU で最大 126%、CPU で最大 292% 速い推論を達成しています。
• 8 GPU での混合精度学習では、7.7 倍のスルーレット向上が得られ、GPU 効率は 96.6% に達し、良好なスケーリングが実証されています。
• グリーディサーチは、ビームサイズ 1 のビームサーチと比較して、GPU で 16%、CPU で 6% の翻訳速度向上を実現しています。
• CPU での INT8 量子化により、推論速度が 1 秒あたり 2.6 文から 4.5 文に向上し、BLEU スコアの低下は最小限に抑えられました。
• フォーマルリティ制御モデルは、フォーマル翻訳で 97.8% の正確性を達成し、WMT テストセットでは 45.0 の BLEU スコアを記録し、高い制御忠実度を示しました。
• 長さトークンと N-best 再ランク付けを用いた冗長性制御により、英語-ドイツ語翻訳の最終システムランク指標が、BERTScore × LC で最大 17.4 ポイント向上しました。