技術要約: DTOP-p MoE

問題提起

疎なMixture-of-Experts (MoE) アーキテクチャは、パラメータ数と計算コストを切り離すことで、基盤モデルを効率的にスケールアップするために極めて重要である。しかし、標準的な Top-k ルーティング は、トークンの難易度やレイヤーごとのニーズに関わらず、固定された数のエキスパートを選択するという硬直したスパース性パターンを強いる。これは、入力の複雑性に内在する分散を無視し、適応的なリソース割り当てを制限してしまう。

一方、Top-p ルーティング（核サンプリング）は、累積確率が閾値 $p$ に達するまでエキスパートを選択することで、より適応的な代替案を提供するが、既存の実装には2つの決定的な限界がある：

1. 固定されたグローバル閾値: 素朴な Top-p 実装は静的なグローバル確率閾値を使用しており、進化する学習ダイナミクスに適応できない。
2. 制御不能なスパース性: これらの手法は、予測不可能な計算コストをもたらす。活性化されるエキスパートの数は学習中に大きく変動するため、大規模な事前学習における厳格な計算予算を遵守することが困難になる。さらに、固定された閾値はハイパーパラメータの選択に対して非常に敏感であり、計算コストの高い広範なチューニングを必要とすることが多い。

手法: DTOP-p MoE

著者らは、Top-p の適応性と大規模な事前学習における厳格な計算制約を両立させるために設計された、スパース性制御可能な動的ルーティングメカニズムである DTOP-p を提案している。この手法は、主に2つのコンポーネントで構成される：

1. 比例・積分 (PI) コントローラ

閾値選択の非微分性と明示的なリソース制約の欠如に対処するため、DTOP-p は古典的な制御理論からのフィードバック制御ループを採用している。

• メカニズム: システムは、目標とする活性化エキスパート数 ($T$) を設定値（セットポイント）とし、グローバルな確率閾値 ($p$) を制御変数として扱う。
• 更新則: コントローラは、現在のバッチにおけるトークンあたりの平均活性化エキスパート数 ($a_t$) を監視する。そして、離散 PI 制御則を用いて、次のバッチのための閾値 ($p_{t+1}$) を次のように更新する：
  $$p_{t+1} = p_0 + K_{pro} \cdot e_t + K_{int} \cdot \sum_{i=1}^{t} e_i$$
  ここで、$e_t = (T - a_t)/N$ である。
• 効果: これにより、正のエラー（活性化不足）が発生すると $p$ を増加させてより多くのエキスパートを選択し、負のエラーが発生すると $p$ を減少させるという負のフィードバックループが形成される。これにより、ロジットのエントロピーの変化に関わらず、モデルが目標とするスパース性レベルに収束することが保証される。

2. 動的ルーティング正規化 (DRN)

すべてのレイヤーに対して単一のグローバルな閾値を適用することは、各レイヤーが異なる計算ニーズ（例：一般的な特徴のための浅いレイヤー vs 特化した推論のための深いレイヤー）を持つという知見に反し、エキスパートの確信度分布が均一であることを仮定している。

• メカニズム: 生のロジットに対して直接ルーティングを行う代わりに、DTOP-p は各レイヤー $l$ 内で学習可能なスカラー $\theta_l$ を用いてロジットを正規化する：
  $$P(x) = \text{softmax}\left(\theta_l \cdot \frac{z - \mu(z)}{\sigma(z)}\right)$$
  ここで、$z$ は生のロジットであり、$\mu(z)$ と $\sigma(z)$ はそれらの平均と標準偏差である。
• 効果: これにより、グローバルな閾値をローカルな統計的特性から分離できる。大きな $\theta_l$ は分布を鋭くし（より少ないエキスパートを活性化）、小さな $\theta_l$ は分布を平坦にし（より多くのエキスパートを活性化）させる。これにより、グローバルな計算予算を維持しながら、レイヤーごとに異なるスパース性パターンを学習することが可能になる。

レイヤー単位のバリアント

厳格なレイヤーごとの予測可能性を必要とするデプロイメントシナリオのために、著者らは、グローバルなクロスレイヤー予算共有を、ほぼ決定的（near-deterministic）なレイヤー単位のトラッキングに置き換える、各レイヤーが独自の PI コントローラと閾値を保持するレイヤー単位のバリアントも提案している。

主な貢献

1. 固定閾値 Top-p の分析: 本論文は、固定閾値の Top-p MoE が Top-k に対してわずかな利点しか提供せず、閾値 $p$ に対する感受性の高さと不安定なスパース性に苦しむため、制御メカニズムなしでは大規模な事前学習には不適切であることを示している。
2. DTOP-p フレームワーク: スパース性制御のための PI コントローラと、レイヤー単位の適応性のための動的ルーティング正規化を統合した、動的な Top-p MoE の提案。
3. 包括的な実証的研究: 自然言語処理 (NLP) およびコンピュータビジョン (CV) ドメインにおける広範な検証を行い、DTOP-p が Top-k の平均 FLOPs と一致しながら、Top-k および固定 Top-p のベースラインを凌駕することを実証した。

実験結果

著者らは、DCLM-Baseline データセット (NLP) および高品質な内部画像データセット (CV) を使用して、LLM および Diffusion Transformer (DiT) 上で DTOP-p を評価した。

• 性能: NLP ベンチマーク（MMLU、HellaSwag、ARC を含む 13 個のデータセット）において、DTOP-p は目標とする活性化エキスパート数（例：8 エキスパート）を厳密に遵守しながら、Top-k MoE ベースラインに対して平均 1.9% の向上を達成した。
• 安定性: 活性化されるエキスパート数が激しく変動する固定 Top-p とは異なり、DTOP-p は目標とするスパース性に急速に収束し（最初の 1B トークン以内）、低い標準偏差 ($\approx 1$) を維持する。
• スケーリング特性: DTOP-p は以下の項目において堅牢なスケーリングを示した：
  • エキスパートの粒度: 粒度が増すにつれて（例：32 エキスパートから 128 エキスパートへ）、性能向上の幅が広がった。
  • エキスパート容量: 総エキスパート容量が増加するにつれ、性能は単調に向上した。
  • モデルおよびデータセットのサイズ: モデルサイズ（最大 2.4B アクティブパラメータ）およびデータセットサイズ（最大 300B トークン）が増大しても、DTOP-p はベースラインに対する優位性を維持した。
• レイヤー解析: 活性化分布の分析により、DTOP-p が階層的なスパース性パターンを学習し、浅いレイヤーではより少ないエキスパートを、深いレイヤーではより多くのエキスパートを活性化させ、異なる深さで処理される特徴の多様な複雑さに適合していることが確認された。

意義と主張

本論文は、DTOP-p が基盤モデルの事前学習のための堅牢かつ効率的な MoE フレームワークを提供すると主張している。その主な意義は、動的な適応性と厳格な計算制約をうまく両立させたことにある。制御理論を通じて確率閾値を実質的に学習可能にし、レイヤーごとの柔軟性を可能にすることで、DTOP-p は定義済みの FLOP バジェットを超えることなく、トークンの難易度やレイヤーの深さに基づいて計算リソースを割り当てることを可能にする。

著者らは、この研究を持続可能な Green AI への一歩として位置づけており、疎なモデルの計算効率を向上させている。彼らは、この手法が効率を改善する一方で、学習データの固有のバイアスを軽減するものではなく、厳格な安全性アライメントと組み合わせて使用されるべきであることも明記している。本研究は、固定 Top-k があまりに硬直的であり、制御不能な Top-p があまりに予測不可能であるような、大規模な基盤モデルをスケールアップするための実用的なソリューションとして提示されている。