[論文レビュー] BASE Layers: Simplifying Training of Large, Sparse Models

Mike Lewis, Shruti Bhosale+3 more|arXiv (Cornell University)|Mar 30, 2021

Topic Modeling参考文献 29被引用数 63

ひとこと要約

BASEレイヤーは最適でバランスの取れた token-to-expert アサインメントを線形割り当て問題を用いて実現し、トークンごとに単一エキスパートを割り当てた疎なルーティングを追加損失やハイパーパラメータなしで可能にし、計算効率を向上させる。

ABSTRACT

We introduce a new balanced assignment of experts (BASE) layer for large language models that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

研究の動機と目的

• 言語モデルをスケールさせつつトレーニングコストを削減するために、疎なエキスパートモデルの利用を促す。
• バランスの取れた線形割り当て BASEレイヤーを提案し、Balancing lossesと capacity factorsを排除する。
• バランスの取れた割り当てによるトークンあたり単一エキスパートのルーティングが、大規模モデル全体で強い計算効率を示すことを示す。

提案手法

• トークンとエキスパートの割り当てを線形割り当て問題として定式化し、各エキスパートが等しい数のトークンを扱うようにする。
• アー Auction algorithm を使用して、複数のワーカー間で parallelに割り当てを解決し、頑健性のためにトークンシャッフルを含むエッジケースにも対応。
• 各トークンを1つのエキスパートにルーティングする BASEレイヤーを実装し、続いて残差接続を介してエキスパート出力をソフトミックスする。
• バランスの取れたトークン割当で学習しスループットを最大化し、追加の balancing losses を回避する一方、推論時にはグリーディエキスパート選択を用いてテストする。
• 固定GPUランタイムを前提にパープレキシティとして計算効率を測定し、密結合および従来の疎なアプローチと比較する。

実験結果

リサーチクエスチョン

• RQ1 線形割り当てベースのルーティングスキームは、追加の損失項なしでエキスパートの使用をバランスさせることができるか。
• RQ2 単一エキスパートのBASEレイヤーは、既存のDenseおよびSparse法の効率と同等かそれを上回るか、 大規模モデル全体で。
• RQ3 BASEの配置とサイズは性能と計算効率にどのように影響するか。
• RQ4 トレーニング中のエキスパートの専門化と負荷分散に対するバランスの取れたルーティングの影響は何か。

主な発見

• BASEレイヤーは Balancing losses や capacity factors なしで線形割り当てによるトークン-to-エキスパートのバランスの取れたルーティングを達成する。
• 単一のBASEレイヤーは、より高い計算予算において密データ並列およびモデル並列のベースラインを大幅に上回る。
• BASEは計算効率でSparsely Gated MoEおよびSwitch transformersと有利に比較され、類似の予算で彼らの性能に匹敵または上回ることがある。
• 複数のBASEレイヤーを交互に配置することは、パラメータ総数を概ね一致させつつ性能を段階的に向上させる。
• BASEレイヤーは配置やBASEサブレイヤーの数が異なっても堅牢な性能を示し、アーキテクチャの頑健性を示す。