[論文レビュー] Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training

Zexuan Zhong, Mengzhou Xia+2 more|arXiv (Cornell University)|May 6, 2024

Topic Modeling被引用数 6

ひとこと要約

Lory は、因果セグメントルーティングと類似性ベースのデータバッチ処理を用いた自己回帰型言語モデルの事前学習のための、完全微分可能な MoE フレームワークを導入し、セグメントレベルのエキスパート統合により、強力な perplexity と下流タスクの利得を達成します。

ABSTRACT

Mixture-of-experts (MoE) models facilitate efficient scaling; however, training the router network introduces the challenge of optimizing a non-differentiable, discrete objective. Recently, a fully-differentiable MoE architecture, SMEAR, was proposed (Muqeeth et al., 2023), which softly merges experts in the parameter space; nevertheless, its effectiveness was only demonstrated in downstream fine-tuning on classification tasks. In this paper, we present Lory, the first approach that scales such architectures to autoregressive language model pre-training. Lory introduces two key techniques: (1) a causal segment routing strategy that achieves high efficiency for expert merging operations while preserving the autoregressive nature of language models; (2) a similarity-based data batching method that encourages expert specialization by grouping similar documents in training instances. We pre-train a series of Lory models on 150B tokens from scratch, with up to 32 experts and 30B (1.5B active) parameters. Experimental results show significant performance gains over parameter-matched dense models on both perplexity (+13.9%) and a variety of downstream tasks (+1.5%-11.1%). Despite segment-level routing, Lory models achieve competitive performance compared to state-of-the-art MoE models with token-level routing. We further demonstrate that the trained experts in Lory capture domain-level specialization without supervision. Our work highlights the potential of fully-differentiable MoE architectures for language model pre-training and advocates future research in this area.

研究の動機と目的

• エンコーダ/分類タスクを超える自己回帰LM向けに、完全微分可能なMoEアーキテクチャを動機づけ、実現する。
• 因果セグメントルーティングによる効率的なエキスパート統合を提案し、自己回帰生成を保持しつつ計算量を削減する。
• 事前学習中にドメイン固有のエキスパート専門化を誘導するため、類似性ベースのデータバッチ処理を導入する。
• 150B トークンで最大 30B の全パラメータを持つ Lory モデルを最大 32 エキスパートで事前学習することで、スケーラビリティを示す。

提案手法

• Softmax ルータからの e_i を用いた sum_i e_i * theta_i による o_x = FFN(h_x) のエキスパート統合を採用する。
• 因果セグメントルーティングを用いる：入力を長さ T のセグメントに分割し、前のセグメントから統合済み theta を計算して現在のセグメントに適用し、自己回帰の因果性を保持する。
• 訓練中は前のセグメントの平均隠れ表現に基づいてルーティングウェイトを計算し、漏洩を防ぐために stop-gradient を適用する。
• 意味的に類似する文書を連結するために類似性ベースのデータバッチ処理を採用し、エキスパートの専門化を促進する。
• 推論時にはプロンプトのみのルーティングを使用して、生成全体にわたって各層で単一の統合済み FFN を決定する。
• 150B トークン、0.3B および 1.5B アクティブパラメータ、各層あたり 8/16/32 エキスパート、デンスモデルの初期化とウォームアップスケジュールに従って訓練する。

実験結果

リサーチクエスチョン

• RQ1 完全微分可能な MoE アーキテクチャは、自己回帰型言語モデルの事前学習へ効果的にスケールできるのか？
• RQ2 因果セグメントルーティングは、セグメントレベルでのエキスパート統合を可能にしつつ、自己回帰生成を保持するのか？
• RQ3 類似性ベースのデータバッチ処理はドメインレベルのエキスパート専門化を促進し、下流タスクの性能を改善するのか？
• RQ4 完全微分可能な MoE Lory モデルは、従来のトークンレベル MoE および密なベースラインと比較して perplexity および下流タスクでどうなるのか？

主な発見

• Lory は、評価された領域全体でパラメータ一致密ベースラインに対して顕著な perplexity 改善を達成し、0.3B/32E モデルで Books に対して相対改善 +13.9%、他の領域で arXiv, Books, Wikipedia, C4, Python においても perplexity を低減。
• 選択された構成で、MoE による下流タスクは一貫して改善を示し、+3.7%（常識推論）、+3.3%（読解）、+1.5%（クローズドブ QA）、+11.1%（テキスト分類）。
• 因果統合を伴うセグメントレベルルーティングは、最先端のトークンレベル MoE モデルと競合する性能を示し、同時に完全微分可能な訓練を可能にする。
• エキスパートは監督なしでドメインレベルの専門化を示し、中間層/高層がドメイン信号をより強く示す（例：arXiv, Python）、局所的なトークンパターンを超えた効果的な専門化を示唆。
• 類似性ベースのバッチ処理は、ランダムバッチ処理に対する MoE の性能を大幅に向上させ、エキスパートの訓練と活用にとって重要である。