📝

ReMoE : Fully Differentiable Mixture-of-Experts with ReLU Routing

に公開

https://arxiv.org/abs/2412.14711

要約

Mixture-of-Experts (MoE) モデルの新しいアーキテクチャーである ReMoE を提案

  1. 主要な革新点

    • 従来の TopK + Softmax ルーティングの代わりに ReLU ルーティングを導入
    • 完全に微分可能な設計によりトレーニングの安定性と性能を向上
    • トークンごとに動的に専門家 (expert) を割り当てる柔軟な仕組み
  2. 技術的な特徴

    • ReLU ゲートを使って各 expert の活性化状態を制御
    • 適応的な L1 正則化により計算コストを管理
    • ロードバランシングを組み込んだ正則化手法を採用
  3. トレーニングの特徴

    • 3 段階の自然なトレーニングプロセス
      • ウォームアップ段階(密な状態)
      • スパース化段階(密からスパースへ移行)
      • 安定段階(スパースな状態を維持)
  4. 実験結果

    • 様々なモデルサイズ (182M-978M parameters)
    • 異なる expert 数 (4-128)
    • 異なる粒度レベル (1-64)
      において従来の MoE を一貫して上回る性能を示した
  5. 主な利点

    • より効果的なドメイン特化
    • expert を動的に割り当て
    • expert の数に対するより良いスケーリング特性
    • トレーニングとインファレンスの速度を維持

Abstract

  1. 問題設定

    • スパース活性化を用いる MoE モデルは計算予算を増やさずモデル容量を拡大出来る
    • しかし従来の TopK ルーターは連続的でなく微分不可能な訓練を必要とし、性能とスケーラビリティーが制限される
  2. 提案手法

    • ReLU をルーターとして使用する完全微分可能な MoE アーキテクチャー ReMoE を提案
    • 従来の TopK + Softmax ルーティングに代わるシンプルかつ効果的な代替手法
    • ルーターのスパース性を制御しながらエキスパート間の負荷分散を行う手法も提案
  3. 利点

    • トークンとレイヤー間で計算を動的に割り当て可能
    • ドメイン特化を示す
    • 様々なモデルサイズ、エキスパート数、粒度レベルで従来の TopK ルーティング MoE を一貫して上回る性能を達成
    • エキスパート数に関するスケーラビリティーも向上
  4. 実装

1 Introduction

  1. 背景

    • Transformer モデルはパラメーター数増加に伴い性能が向上
    • しかし計算リソースによる制約が存在
    • スパース活性化 MoE は訓練と推論時に一部のパラメーターのみを活性化する事でこの課題に対処
    • 計算コストを増やさずにモデル容量を拡大可能
  2. ルーティングの課題

    • MoE の核となるのはルーティングネットワーク
    • TopK ルーティングが最も一般的
    • しかし離散的で微分不可能な訓練目的関数を導入する事が問題
  3. 既存の微分可能 MoE 研究

    • Soft MoE : トークンマージングを導入
    • SMEAR : エキスパートマージングを提案
    • しかしこれらはトークンの因果性を破壊するので自己回帰モデルに適さない
    • Lory : SMEAR を改良し自己回帰モデルに適用可能だが TopK ルーティングより性能が劣る
  4. 提案手法 ReMoE

    • ReLU ルーティングを TopK ルーティングの代替として導入
    • Softmax と TopK の代わりに ReLU ゲートを使用
    • 活性エキスパート数は ReLU 関数のスパース性により決定
    • スパース性を維持するため、適応的に調整された係数を持つ負荷分散 L1 正則化を導入
  5. 利点

    • 連続的で完全に微分可能
    • エキスパートの活性状態を独立に制御可能
    • トークンとレイヤー間で動的な計算リソース割り当てが可能
    • ドメイン特化を示す
    • 様々な設定で従来手法を上回る性能を達成

Figure 1

2 Preliminaries

2.1 MoE for Decoder-Only Transformer

  1. Decoder-Only Transformer の構造

    • L 個のレイヤーで構成
    • 各レイヤーは Self-Attention モジュールと Feed-Forward Network (FFN) モジュールを含む
  2. MoE による変更点

    • 各 FFN モジュールを MoE モジュールに置き換え
    • MoE モジュールは小規模なルーターと E 個のエキスパート FFN で構成
    • 各エキスパートは元の FFN と同等の構造を持つ
  3. 数式による定義

    • 入力 : \bm{x}^l = (\bm{x}^l_t)^T_{t=1} \in \mathbb{R}^{T \times d}
      • T はバッチ内のトークン数
      • d は隠れ層のサイズ
      • l はレイヤーを示す
    • 出力 : \bm{y}^l = (\bm{y}^l_t)^T_{t=1} は以下の式で計算
    \bm{y}^l_t = \sum^E_{e=1} R(\bm{x}^l_t)_e\text{FFN}_e(\bm{x}^l_t; d_{ffn})
    • ここで
      R(\cdot) はルーティング関数
      d_{ffn} は FFN の中間層サイズ
      通常 d_{ffn} = 4d に設定

この構造により MoE は元のモデルの FFN をスパース化された複数のエキスパートで置き換える事が可能となる

2.2 TopK Routing

  1. TopK ルーティングの定義

    数式で表すと

    R(\bm{x}^l_t) = \text{TopK}(\text{Softmax}(\bm{x}^l_t\bm{W}_l), k)

    ここで

    • \bm{W}_l \in \mathbb{R}^{d \times E} はルーターの重み行列
    • \text{TopK}(\cdot ·, k) は上位 k 個の最大値を保持し、残りを 0 に設定する関数
  2. 計算削減の仕組み

    • R(\bm{x}^l_t)_e0 となるエキスパート e について
      • 順伝播時の \text{FFN}_e の計算をスキップ
      • 逆伝播時も同様に計算をスキップ
  3. 特徴

    • 最も一般的に使用されるルーティング手法
    • Softmax で確率分布を計算
    • TopK で上位 k 個のエキスパートのみを選択
    • 残りのエキスパートの出力を 0 に設定
  4. 課題

    • 離散的な選択プロセスを導入
    • 微分不可能な操作を含む
    • これらが性能とスケーラビリティーを制限する要因となる

この手法は高い計算効率を実現する一方で離散的な性質による最適化の課題を抱えている

3 Our Method : ReMoE

3.1 Motivation : From TopK to ReLU

  1. TopK の問題点

    • Softmax 後の入力 \bm{x} = (x_e)^E_{e=1} に対し
    • k 番目に大きい値 x_{[k]} で不連続な飛びが発生
    • 数式で表すと
      \text{TopK}(\bm{x},k)_e = x_e \cdot\bm{1}\{x_e \geq t(\bm{x},k)\}, t(\bm{x},k) = x_{[k]}
    • ここで \bm{1}\{\cdot\} は条件が満たされる時 1、それ以外で 0 を返す指示関数
  2. ReLU による解決

    • 不連続性を排除するため閾値 t(\bm{x},k)0 に固定
    • これは ReLU 関数に対応
      \text{ReLU}(\bm{x})_e = x_e \cdot 1\{x_e \geq 0\}
  3. 改善点

    • 全ての入力の閾値を 0 に揃える事で不連続点を整理
    • エキスパートの活性・非活性の遷移点である 0 で出力が連続に
    • これにより訓練パイプラインが完全に微分可能に
  4. 概念的な利点

    • TopK の離散的な損失関数を排除
    • より安定した訓練が可能
    • エキスパート選択の学習をより滑らかに実現

この変更により離散的な TopK の問題を解決しながらスパース性を維持する事が可能となる

Figure 2

3.2 Differentiable ReLU Routing

  1. ReLU ルーティング関数の定義

    &&R(\bm{x}^l_t) = \text{ReLU}(\bm{x}^l_t\bm{W}_l)&&

    • スパース率の目標を (1 - \dfrac{k}{E}) に設定
    • k は活性エキスパート数
    • E は全エキスパート数
    • これにより TopK ルーティングと同等の計算コストを実現
  2. 従来の TopK ルーティングとの違い

    • TopK : Softmax 出力は合計が 1 となる確率分布を表現
    • ReLU : Softmax を使わず ReLU の非負出力を直接使用
    • エキスパートへの重み付けとして 0 を含む値を許容
    • 不連続な TopK による選択を排除
  3. 柔軟性の向上

    • TopK では各トークンが正確に k 個のエキスパートに割り当てられる
    • ReLU ではルーティング決定が独立
    • トークンごとに異なる数のエキスパートを使用可能
    • 難しいトークンにより多くの計算リソースを割り当て可能
  4. 連続性の保証

    • TopK による離散的な損失関数を回避
    • 完全に微分可能な処理を実現
    • ルーティング出力が 0 付近の場合のみエキスパートの活性状態が変化
    • これにより安定した訓練が可能

これらの特徴により、より柔軟で効率的なエキスパート割り当てを実現している

3.3 Controlling Sparsity via Adaptive L_1 Regularization

  1. 基本的な課題

    • ReMoE は計算コストを ReLU 出力のスパース性で制御
    • 目標スパース率は (1 - \dfrac{k}{E})
    • しかし直接訓練すると低いスパース性となりがち
    • モデルが容量増加のため多くのエキスパートを活性化する傾向
  2. 正則化の導入

    • 言語モデルの損失 L_{lm} に正則化項 L_{reg} を追加
    L = L_{lm} + \lambda_iL_{reg}
    • \lambda_i は現在の訓練ステップ i に基づく適応係数
  3. 適応的な係数更新

    \lambda_{i+1} = \lambda_i \cdot \alpha^{\text{sign}((1-\frac{k}{E})-S_i)}
    • \alpha\gt1 は更新乗数
    • S_i は全ルーター出力の平均スパース率
    S_i = 1 - \frac{1}{LTE}\sum^L_{l=1}\sum^T_{t=1}\sum^E_{e=1}\bm{1}\{R(\bm{x}^l_t)_e > 0\}
  4. L1 正則化項

    L_{reg} = \frac{1}{LT}\sum^L_{l=1}\sum^T_{t=1}\|R(\bm{x}^l_t)\|_1
    • ReLU 出力が非負のため L1 ノルムは要素の合計と一致
    • \dfrac{\lambda_i}{LT} の勾配を各非ゼロ出力に追加
    • 出力を 0 に向けてスパース化を促進
  5. 効果

    • 目標スパース率の周りで制御が可能
    • わずかな変動のみで安定
    • 平均的に k 個のエキスパートを活性化
    • TopK ルーティング MoE と同等の計算量を実現

このアプローチによりスパース性を効果的に制御しながら訓練速度を維持する事が可能となる

Figure 3

3.4 Integrate Load Balancing into L_1 Regularization

  1. 負荷分散が必要

    • MoE 設計における重要な課題
    • ルーティングが崩壊する可能性
    • 複数デバイス間で計算分布が不均衡に
  2. 負荷分散を考慮した L1 正則化を導入

    L_{reg,lb} = \frac{1}{LT}\sum^L_{l=1}\sum^T_{t=1}\sum^E_{e=1}f_{l,e}R(\bm{x}^l_t)_e
    f_{l,e} = \frac{E}{kT}\sum^T_{t=1}\bm{1}\{R(\bm{x}^l_t)_e > 0\}
  3. 特徴と動作

    • f_{l,e} は微分不可能
    • レイヤー l におけるエキスパート e の平均活性化率を表現
    • 目標比率 \dfrac{k}{E} に対する相対的な重みとして機能
    • ルーター出力の勾配を \dfrac{f_{l,e}\lambda_i}{LT} に修正
    • より多くのトークンを受け取るエキスパートにペナルティーを課す
  4. TopK ルーティングとの比較

    • TopK では Softmax 出力の合計が 1
    • 損失の下限が 1 となる
    • ReLU では出力が任意に小さくなり得る
    • L_{reg},lb は 0 で下界付け
  5. 実装上の工夫

    • ReMoE では係数 \lambda_i を固定出来ない
    • 適応的な更新によりスパース性制御と負荷分散のバランスを取る

この手法により計算負荷の均衡を保ちながらスパース性を制御する事が可能となる

3.5 Natural Three-Stage Training in ReMoE

  1. Stage I : ウォームアップ段階(密な段階)

    • \lambda_i が小さく L_{lm} が大きく急速に減少
    • 半数以上のトークンを各エキスパートが処理
    • ランダム初期化からエキスパートが多様化
    • 実質的に同じパラメーター数を持つ密なモデルの訓練と同等
  2. Stage II : スパース化段階(密からスパースへ)

    • \lambda_iL_{reg} が重要になる
    • ReLU ルーターが活性化するエキスパート数を削減
    • L_{lm} を増加させることなくエキスパートの多様性を強制
    • 密な状態からスパースな状態への移行期間
  3. Stage III : 安定段階(スパースな段階)

    • スパース率 S_i が目標値で安定
    • スパース部分空間で L_{lm} を最適化
    • L_{reg} と $\lambda_i} は緩やかに変化
    • 全体の正則化項 \lambda_i\mathcal L_{reg} は比較的一定
  4. 計算コストに関する考察

    • Stage I と II で追加の計算コストとメモリー消費が発生
    • しかし約 100 イテレーション(全体の約0.17%)のみで必要
    • メモリーオーバーヘッド対策 :
      • マイクロバッチサイズを一時的に削減
      • 活性化チェックポインティング技術を使う
      • 逆伝播時に必要な中間結果を再計算

この自然な 3 段階訓練プロセスにより効率的なスパース化と性能向上を実現している

Figure 4

4 Experiments

4.1 Setup

  1. インフラストラクチャー

    • Megatron-LM をコードベースとして使用
    • ReLU ルーティングを TopK ルーティングの代替として実装
    • サポートする並列化手法 :
      • データ並列
      • テンソル並列
      • パイプライン並列
      • エキスパート並列
  2. モデルアーキテクチャー

    • LLaMA アーキテクチャーを採用
    • 特徴 :
      • GQA (Grouped Query Attention)
      • SwiGLU 活性化関数
      • RoPE 位置エンコーディング
      • RMSNorm
    • コンテキスト長 : 1024
    • バッチサイズ : 512
    • vanilla MoE では重み 0.01 の負荷分散損失を使用
    • ReMoE では式 10 の適応的負荷分散 L1 正則化を使用
  3. 訓練設定

    • データセット : The Pile(800GB の多様なコーパス)
    • 訓練ステップ : 60k(約 30B トークン)
    • トークナイザー : BPE
    • オプティマイザー : AdamW
      • \beta_1 = 0.9
      • \beta_2 = 0.999
      • ZeRO 最適化
    • 学習率 : 5e-4(コサインスケジューラー)
    • 訓練環境 : 8 NVIDIA A100 GPU
  4. 密なバックボーン設定

    • Small: 182M パラメーター
    • Medium: 469M パラメーター
    • Large: 978M パラメーター
    • FLOPs は単一シーケンスに基づき計算

Table 1

4.2 Comparison with Other Routing Methods

  1. 比較対象とした手法

    • dMoE : Token-choice dropless TopK ルーティング
    • EC: Expert-choice TopK ルーティング
    • Hash : 決定論的ハッシュルーティング
    • Lory : 完全微分可能なエキスパートマージングルーティング
    • SparseMixer-v2 : 改良された勾配推定を持つ TopK ルーティング
  2. 実験設定

    • アクティブパラメーター数 N = 182M
    • エキスパート数 E = 8
    • アクティブエキスパート数 k = 1
    • Hash 手法では mod E ハッシュ関数を使用
    • Lory ではセグメント長を 256 に設定
  3. 評価結果

    • 訓練曲線の比較から :
      • すべての MoE モデルが密なモデルを上回る性能
      • 決定論的ハッシュルーティングは学習型ルーティング手法より性能が低い
      • Token-choice dMoE が評価において Expert-choice MoE と SparseMixer-v2 を上回る
  4. ダウンストリームタスクでの評価

    • 以下のタスクで評価 :
      • ARC
      • BoolQ
      • HellaSwag
      • LAMBADA
      • PIQA
      • RACE
    • Lory は訓練ではハッシュルーティングを上回るがダウンストリームタスクでは標準的な TopK ルーティングに劣る
    • ReMoE は全ての手法を上回り、完全微分可能という利点も持つ

Figure 5
Table 2

4.3 Scalability of ReMoE

  1. パラメーター数 N に関するスケーリング

    • N を 182M から 978M まで変化させて評価
    • E = 8, k = 1 に固定
    • 総パラメーター数は 777M から 5.73B
    • 結果 :
      • ReMoE は全モデルサイズで MoE を上回る
      • モデルサイズが増加しても性能差は維持
      • より大規模なモデルでも優位性を保持
  2. エキスパート数 E に関するスケーリング

    • N = 182M, k = 1 に固定
    • E を 4 から 128 まで変化
    • 結果 :
      • ReMoE は全設定で標準 MoE を上回る
      • E の増加に対する性能向上が MoE より急峻
      • より大きなエキスパートプールを効果的に活用
      • 表現力と汎化性能の向上を実現
  3. 粒度 G に関するスケーリング

    • Fine-grained MoE の評価
    • エキスパートを G 個の小さなエキスパートに分割
    • 計算式 :
    \bm{y}^l_t = \sum^{EG}_{e=1}R(\bm{x}^l_t)_e\text{FFN}_e(\bm{x}^l_t; d_{ffn}/G)
    R(\bm{x}^l_t) = \text{TopK}(\text{Softmax}(\bm{x}^l_t\bm{W}_l), kG)
    • 結果 :
      • Fine-grained ReMoE が Fine-grained MoE を一貫して上回る
      • G = 32, 64 で理論上限 (Dense×8) に匹敵する性能を達成
      • より効率的な計算リソース利用を実現

Figure 6

5 Discussion

5.1 Dynamic Expert Allocation in ReMoE

  1. 動的割り当ての仕組み

    • 各トークンが動的にエキスパートの部分集合を活性化
    • モデルが適応的にリソースを割り当て可能
    • N = 182M, E = 8, k = 1 の ReMoE で分析を実施
  2. トークン頻度との関係

    • 稀少トークンに対する割り当て :
      • '©', 'OTAL', '@#' などのレアトークン
      • より多くのエキスパートを割り当て
    • 頻出トークンに対する割り当て :
      • スペース、改行、'the' などの一般的なトークン
      • より少ないエキスパートを割り当て
  3. ハフマン木との類似性

    • 頻出シンボルには短いコード
    • 稀少シンボルには長いコード
    • ReMoE も同様の原理 :
      • 頻出トークンは少数のエキスパートで「クラスター化」
      • 頻出トークンの「表現」を圧縮
      • 稀少トークンにはより多様なエキスパートの線形結合を使用
  4. 利点

    • リソース使用とモデル容量のバランスを最適化
    • エキスパートの制限下で効率的な性能を実現
    • ドメインレベルでも動的割り当てが機能

この動的割り当て機構により、計算リソースの効率的な活用が可能となっている

Figure 7

5.2 The Role of Load Balancing in ReMoE

  1. 負荷分散の重要性と課題

    • 従来の TopK ルーティング MoE ではルーティング崩壊の可能性
    • 全入力を同一エキスパートに割り当てる傾向
    • 訓練目的関数が連続的かつ完全微分可能になる問題
    • 負荷分散ありとなしで大きな性能差が存在
  2. ReMoE における負荷分散の効果

    • 微分可能性により L1 正則化のみでも十分な性能
    • しかし負荷分散なしの場合 :
      • 一部のエキスパートが非活性に
      • モデル容量を制限する可能性
    • 負荷分散を導入した効果 :
      • エキスパートへのトークン割り当てが均一化
      • 全エキスパートが活用される
      • 最終的な損失が減少
  3. レイヤー間のスパース性分布

    • 負荷分散によりレイヤー間でよりスムーズなスパース性分布を実現
    • f_{l,e} が割り当てられたトークン数に基づき計算される事で、より密なレイヤーにより強いペナルティーを課す
  4. 完全な均一性は不要

    • 負荷分散を導入しても完全に均一な分布は達成されない
    • トレードオフを調整可能 :
      • L1 正則化を修正(例 : f_{l,e}f_{l,e}^2 に変更)
      • デバイスレベルで負荷分散技術を活用
    • ReMoE ではルーティング崩壊が深刻でないため主にハードウェア利用率の問題として扱える

これらの特徴により ReMoE は効果的な負荷分散を実現しながら高い性能を維持している

Figure 8

5.3 Domain Specialization in ReMoE

  1. ドメインが特化する仕組み

    • 微分可能性と動的割り当て戦略により異なるドメインに特化したエキスパートが発展
    • 様々なエキスパートの専門性を活用してアンサンブル学習を効果的に実行
  2. 実験設定と分析

    • N = 182M, E = 8 のモデルで分析
    • 異なるドメインで評価 :
      • Arxiv
      • Books
      • C4
      • Github
      • Stackexchange
      • Wikipedia
    • レイヤー 0, 5, 11(最初、中間、最後)に注目
  3. MoE vs ReMoE の違い

    • MoE :
      • ほとんどのエキスパートが全ドメインで均一な分布を示す
    • ReMoE :
      • エキスパートが明確なドメイン特化を示す
      • ドメインごとに異なる頻度で活性化
      • 例 : レイヤー 5 のエキスパート 6 は LaTeX や Python などの構造化言語で高頻度に活性化
  4. レイヤー 5 での詳細分析

    • エキスパートが専門化 :
      • エキスパート 1 : 自然言語(husband, wife, lover など)
      • エキスパート 6 : コード関連(variable, env, HEAD など)
    • ドメイン固有の語彙に対する一貫した反応を示す

これらの結果から ReMoE は効果的なドメイン特化を自然に学習し、様々な入力に対して適切なエキスパートを活用している事が示されている

Figure 9

6.1 Mixture-of-Experts

  1. 初期の MoE

    • 1990年代初頭に提案
    • Jacobs, Jordan らが基礎的な概念を確立
    • 局所的なエキスパートの混合モデルとして考案
  2. 大規模ニューラルネットワークへの導入

    • 2017年に Shazeer らが新しい応用を提案
    • スパースサブモジュールとして効率性を重視
    • 計算効率を向上させる手段として活用
  3. 重要な進展

    • GShard
    • Switch Transformer
    • これらにより Transformer モデルへ MoE が統合される
    • 顕著な成果を達成
  4. 最近の商用規模での応用

    • Mixtral-8x7B
    • DeepSeekMoE 16B
    • Snowflake Arctic 17B
    • 大規模言語モデルへの実用的な適用を実証

これらの発展により MoE は理論的な概念から実用的なモデルアーキテクチャーへと進化している

6.2 Routing Mechanisms in MoE

  1. 静的ルーター

    • BASE : 組み合わせ最適化などの事前定義ルール
    • Hash ルーティング : 決定論的ハッシュ関数
    • THOR : 正則化を伴うランダム割り当て
  2. 学習型ルーター

    • 入力トークンに基づき適応的にエキスパートを選択
    • 主なアプローチ :
      • REINFORCE : 強化学習ベース
      • TopK ルーティング : トークンまたはエキスパート選択
        • しかし TopK は勾配推定を妨げる不連続性を導入
  3. 手法の特徴

    • 静的ルーター :
      • 単純で計算効率が良い
      • 柔軟性に欠ける
    • 学習型ルーター :
      • より適応的な振る舞い
      • 最適化が複雑
  4. 課題

    • 離散的最適化は困難
    • 勾配推定問題
    • スケーラビリティーに制限

7 Conclusion

  1. 提案手法 ReMoE の主要な特徴

    • 完全微分可能な MoE アーキテクチャー
    • ReLU ルーティング関数の導入
    • 従来の TopK + Softmax ルーティングの代替として機能
    • 連続性と微分可能性を実現
    • トークンとレイヤー間での動的なエキスパート割り当て
  2. 制御メカニズム

    • 適応的な負荷分散 L1 正則化の導入
    • スパース性と計算効率のバランスを実現
  3. 実験結果

    • 様々なモデルサイズで TopK ルーティング MoE を上回る性能
    • エキスパート数に関してより優れたスケーラビリティー
    • 異なる粒度レベルでの一貫した性能向上
  4. 主要な成果

    • エキスパート数の増加に伴いより急峻な性能向上を示す
    • 従来のアーキテクチャーを超える性能を達成
    • スケーラブルで効率的なアーキテクチャーを実現

Appendix

A. Stability Analysis of TopK and ReLU

B. Insensitivity to \lambda_0 and \alpha

C. Performance for Longer Training

D. Speed Comparison of ReMoE and MoE

E. Downstream Evaluation Results
____E.1 Scaling in Active Parameters N
____E.2 Scaling in Expert Count E
____E.3 Scaling in Granularity G
____E.4 Load Balancing Ablations

F. Detailed Results for Domain Specification

G. Domain-Level Dynamic Expert Allocation in ReMoE

H. Training MoE with Near-Dense Warmup

Discussion