[論文レビュー] TorchScale: Transformers at Scale

Shuming Ma, Hongyu Wang+9 more|arXiv (Cornell University)|Nov 23, 2022

Topic Modeling被引用数 6

ひとこと要約

TorchScale は、Magneto アーキテクチャ、DeepNorm 初期化、X-MoE を含む先進的なモデリング技術と、新規の勾配クリッピング手法（SparseClip）を統合することで、大規模な Transformer の効率的かつ安定的なスケーリングを可能にするオープンソースの PyTorch ツールキットです。言語モデリングおよびニューラル機械翻訳タスクにおいて、96 層の深さや 256 専門家の MoE 構成ですら安定した学習と性能向上を達成しています。

ABSTRACT

Large Transformers have achieved state-of-the-art performance across many tasks. Most open-source libraries on scaling Transformers focus on improving training or inference with better parallelization. In this work, we present TorchScale, an open-source toolkit that allows researchers and developers to scale up Transformers efficiently and effectively. TorchScale has the implementation of several modeling techniques, which can improve modeling generality and capability, as well as training stability and efficiency. Experimental results on language modeling and neural machine translation demonstrate that TorchScale can successfully scale Transformers to different sizes without tears. The library is available at https://aka.ms/torchscale.

研究の動機と目的

• 標準的なサイズを超えて Transformer をスケーリングする際の不安定性と非効率性の課題に対処すること。
• 特にスパースな MoE アーキテクチャを用いた際の、深く広い Transformer のスケーリングにおける学習安定性の向上。
• 収束性や性能を損なうことなく、モデルの深さ、幅、専門家数の効果的スケーリングを可能にすること。
• 言語モデリング、機械翻訳、マルチモーダル事前学習を含む多様なタスクを統合的にサポートする包括的かつ汎用的なフレームワークの提供。
• スパースな MoE モデルに特化した、新規の勾配クリッピング手法（SparseClip）の導入により、勾配の優位性を軽減し、学習安定性を向上させること。

提案手法

• 汎用的モデリングを目的とした、Sub-LN 初期化と統一されたエンコーダ・デコーダ構造を備えた Magneto アーキテクチャをデフォルトのバックボーンとして採用。
• DeepNet 理論に基づく Sub-LN 初期化を実装し、深層 Transformer における最適化安定性を向上。
• スパarsity を活用したパラメータ効率的スケーリングを実現するため、top-1 および top-2 ルーティングを備えた X-MoE（eXtended Mixture of Experts）を統合。
• MoE の勾配を $ \kappa = 1/\sqrt{E} $ でスケーリングすることで、密度部と MoE 部の勾配ノルムをバランスさせる再重み付け勾配クリッピング手法である SparseClip を提案。
• L2 ノルムに基づく勾配クリッピングを採用し、ノルムの計算を修正：$ ||g||_2 = ||(g_d, \kappa g_s)||_2 $、ここで $ E $ は専門家の数。
• エンコーダ、デコーダ、エンコーダ・デコーダ構造の両方において、密度モデルとスパースな MoE モデルを一貫した学習・推論パイプラインでサポート。

実験結果

リサーチクエスチョン

• RQ1 96 層の深さと 256 専門家の MoE 構成であっても、学習発散を伴わずに Transformer を安定的にスケーリングできるか？
• RQ2 提案手法の SparseClip は、通常の勾配クリッピングと比較して、スパースな MoE モデルにおける学習安定性と収束性をどの程度向上させるか？
• RQ3 top-2 ルーティングを備えた X-MoE アーキテクチャは、パラメータ効率的スケーリングを実現しつつ、モデル性能をどの程度維持できるか？
• RQ4 Sub-LN 初期化を備えた Magneto アーキテクチャは、言語モデリングや機械翻訳などの多様なタスクにおいて、より優れた一般化性能と学習安定性を実現できるか？
• RQ5 TorchScale における先進的なモデリング技術の組み合わせは、異なるモデルサイズやタスクにおいて一貫した性能向上を達成できるか？

主な発見

• TorchScale は、言語モデリングタスクにおいて 96 層のデコーダモデルを安定した収束で学習可能であり、深さが増すにつれて検証 perplexity (PPL) が改善される。
• top-2 ルーティングを用いた 256 専門家の X-MoE モデルは、パラメータ数が 100 倍以上にのぼるにもかかわらず、より良い PPL を達成しており、効果的なスケーリングが実証された。
• SparseClip は、スパースな MoE モデルの学習安定性と収束性を顕著に向上させ、通常の勾配クリッピングと比較して PPL を顕著に低減した。
• SparseClip を用いることで、時間経過に伴い勾配クリッピングのトリガー回数が減少し、より一貫性があり正確な勾配更新が実現された。
• SparseClip ではルーティングのエントロピーが低く抑えられ、学習中の専門家の利用効率が向上し、より安定したルーティング行動が得られた。
• 専門家数が増加しても、トレーニングに要する FLOPs はほぼ一定を維持しており、TorchScale におけるスパースな MoE アプローチの効率性が裏付けられた。