技術的概要：混合専門家（MoE）層におけるハイパーパラメータ転送

1. 問題定義

混合専門家（MoE）層は、フォワードパス中のアクティブなパラメータ数から学習可能なパラメータの総数を分離することで、現代のニューラルネットワークのスケーリングに不可欠なメカニズムとなっています。しかし、スパースな MoE モデルの訓練は、ハイパーパラメータ（HP）の選択に関して大きな複雑さを伴います。密なモデルとは異なり、MoE 構造は、慎重な調整を必要とする新しい学習可能なパラメータ（ルーター重み）と新しいアーキテクチャ次元（専門家の数、専門家のサイズ）を導入します。

大規模なスケールでハイパーパラメータ（学習率、初期化スケール、重み減衰など）を直接調整することは、計算コストの面で非現実的です。密なトランスフォーマーには HP 転送技術が存在し、小規模モデルで見つかった最適なハイパーパラメータを大規模モデルへ外挿可能にしていますが、これらの手法はスパースな MoE 層特有のスケーリングダイナミクスに対して厳密に適応されていません。核心的な課題は、モデルの幅、深さ、専門家の数、専門家のサイズを同時に増大させる際に、各スケールで再調整を行うことなく、訓練の安定性と性能を確保しつつ、ハイパーパラメータをどのようにスケーリングするかを決定することです。

2. 手法

2.1. 提案するパラメータ化

著者らは、MoE 層を備えたトランスフォーマーモデルに対する新しいパラメータ化を提案し、以前に密なトランスフォーマー向けに開発されたCompletePパラメータ化を拡張し、MoE 固有のスケーリング規則を含めるものです。その目的は、モデル次元がスケーリングするにつれて、生ハイパーパラメータ値（初期化標準偏差 $\sigma$ と学習率 $\eta$）がどのように変化すべきかを予測する規則のセットを定義し、訓練ダイナミクスが一貫して維持されるようにすることです。

この導出は、ネットワークコンポーネント（プレアクティベーションと残差更新）が初期化時に $O(1)$ であり、各訓練ステップで $\Theta(1)$ の更新を受けることを要求する**最大更新パラメータ化（$\mu$P）**の原理に依存しています。著者らはこれを MoE に拡張し、最大更新の条件が層の出力だけでなく、個々の専門家コンポーネント（混合係数と専門家の出力）に対しても成り立つことを要求します。

MoE モジュールに対して導出された主要なスケーリング規則には以下が含まれます：

• ルーター重み: 学習率は $\eta \propto n_{embd}^{-1}$ としてスケーリングします。初期化は $n_{embd}^{-\gamma}$（ただし $\gamma \geq 0.5$）としてスケーリングします。
• 専門家のバイアス: 0 で初期化され、スパース性が固定されている場合、専門家の数に依存しない一定の学習率 $\eta \propto 1$ を使用します。
• 専門家の MLP 重み:
  • アップ・プロジェクション（$W_{up}$）: $\sigma_{init} \propto n_{embd}^{-1/2}$、$\eta \propto n_{embd}^{-1}$。
  • ダウン・プロジェクション（$W_{down}$）: $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$、$\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$。
  • ここで、$\alpha_{ffn}$ は埋め込み次元に対する専門家の隠れ層サイズの乗数を表します。

重要なのは、著者らが専門家の総数（$n_{exp}$）とアクティブな専門家の数（$n_{act}$）をスケーリングする際に、スパース性比率 $\kappa = n_{act}/n_{exp}$（アクティブな専門家の割合）を一定に固定している点です。これは、アクティブな専門家の数を固定して総プールを増加させるアプローチとは対照的です。

2.2. 理論的根拠：動的平均場理論（DMFT）

これらのヒューリスティックなスケーリング規則を正当化するために、著者らは**動的平均場理論（DMFT）**を採用します。彼らは、一定の活性化スパース性 $\kappa$ を維持しつつ、無限の幅（$n_{embd}$）、深さ（$L$）、専門家サイズ（$n_{hid}$）、および専門家数（$n_{exp}$）の同時極限における MoE 層を備えた残差ネットワークの訓練ダイナミクスを分析します。

この分析は、新たな 3 段階の平均場階層を明らかにします：

1. 残差ストリーム表現: 専門家出力に対する平均場。
2. 専門家の出力: 個々の専門家ニューロンに対する平均場。
3. 個々のニューロン: 各専門家内の平均場。

DMFT 分析は、提案されたパラメータ化の下では、極限の訓練ダイナミクスが以下であることを示しています：

• FFN 比率（$\alpha_{ffn}$）に依存しない: 結合スケーリング極限が取られる限り、ダイナミクスは埋め込み次元に対する専門家の特定のサイズに依存しません。
• スパース性 $\kappa$ のみによって決まる: スパース性比率が固定されている限り、ダイナミクスはすべてのスケーリングパラメータに対して一貫しています。
• スケール不変性: ネットワーク要約統計量（例えば、層ごとの特徴カーネル）の進化はスケール間で一貫しており、理論的にハイパーパラメータが信頼性を持って転送されることを保証します。

2.3. 実験設定

著者らは、FineWebおよびC4データセット上のデコーダのみのトランスフォーマー言語モデルを用いて、経験的にパラメータ化を検証しました。

• ベースモデル: 約 38M のアクティブパラメータを持つ小規模モデルで調整。
• スケーリング: 幅、深さ、専門家の数、専門家のサイズを変化させ、最大 2B の総パラメータを持つモデルまでスケーリング。
• トークン予算: 初期訓練ダイナミクスを分離するために 1B トークン（2000 ステップ）の固定トークン予算で実験を実施し、さらに長い時間枠（最大 7.5B トークン）でも実施。
• オプティマイザ: 標準的な Adam オプティマイザ。
• 負荷分散: 正則化項を損失関数に追加することなく、専門家のバイアスを直接更新して負荷分散を促進する、補助損失不要の戦略を使用。

3. 主要な貢献

1. MoE パラメータ化: 本論文は、CompleteP パラメータ化をスパースな MoE モデルに拡張し、幅、深さ、専門家の数、専門家のサイズ全体にわたるルーター重み、専門家のバイアス、専門家の MLP 重みに対する明示的なスケーリング規則を提供します。
2. DMFT による理論的正当化: 著者らは DMFT を用いて、パラメータ化に対する厳密な理論的基盤を提供します。無限の幅/深さ極限における訓練ダイナミクスの明示的な記述を導出し、ダイナミクスが専門家の次元の特定のスケーリングではなく、スパース性のみによって依存する安定したスケール不変のシステムに収束することを証明します。
3. HP 転送の実証的検証: 本研究は、小規模なベースモデル（38M アクティブパラメータ）で特定された最適なハイパーパラメータ（学習率と初期化スケール）が、さまざまなアーキテクチャ次元において、はるかに大規模なモデル（最大 2B の総パラメータ）へ信頼性を持って転送されることを示しています。
4. アーキテクチャに関する知見: 著者らは経験的に、提案されたパラメータ化の下では、総パラメータ数を固定した状態で専門家のサイズを増やすよりも、専門家の数を増やす方が性能が向上することを検証しました。この知見は最近の文献と一致しますが、ここでは各スケールで高価なハイパーパラメータ掃引を行う必要なく達成されています。

4. 結果

• 信頼性の高い HP 転送: 提案されたスケーリング規則の下では、最適なベース学習率と初期化標準偏差が、51M から 2B の総パラメータを持つモデル間で効果的に転送されます。スケーリングされたモデルの損失曲線は、初期反復ではベースモデルと収束し、その後分岐します（より大きなモデルはより低い損失を達成）。
• 安定性: このパラメータ化は、専門家の数をスケーリングする場合でも、均一な専門家負荷分散を含む安定した訓練ダイナミクスを保証します。著者らは、MoE の事前訓練が一定スケールのハイパーパラメータ（$\Theta(1)$ として扱われる乗数）に対して特に敏感であり、これらを調整することが安定性にとって重要であると指摘しています。
• 性能: 小規模モデルから転送されたゼロショットハイパーパラメータを使用して訓練されたモデルは、アクティブパラメータ数で一致させた場合、密なベースライン（例：GPT-2 small/medium）と競争力のある性能を達成します。
• 専門家の数対サイズ: 実験は、固定されたパラメータ数において、専門家のサイズを増やすよりも専門家の数を増やす方がパラメータ効率が高いことを確認しました。この利点は、より長い訓練時間枠（5B トークン）でも持続します。
• スパース性への感度: 本研究は、スパース性比率 $\kappa$ を固定する場合にのみハイパーパラメータ転送が有効であることを強調しています。アクティブな専門家の数を固定して専門家の数をスケーリングする（$\kappa \to 0$ とする）ことは、最適なハイパーパラメータの転送可能性を破綻させます。

5. 意義と主張

本論文は、MoE モデルのスケーリングのための実用的かつ厳密なフレームワークを提供すると主張しています。ヒューリスティックなパラメータ化と DMFT 分析を組み合わせることで、著者らは以下の方法を提供します：

• 訓練コストの削減: 小規模で安価なベースモデルのみを調整することで、大規模な MoE モデルの最適なハイパーパラメータを選択可能にします。
• 安定性の確保: スパースな MoE 訓練で一般的に見られる訓練の不安定性（例：専門家の崩壊や発散）を防ぐ規則を提供します。
• アーキテクチャ設計の指針: 固定されたスパース性のもとでは、専門家のサイズを増やすよりも専門家の数を増やす方が望ましいという理論的および実証的証拠を提供し、効率的な大規模モデルの設計を支援します。

著者らは限界を認め、現在の作業は学習率と初期化に焦点を当てており、バッチサイズ、重み減衰、LR スケジューリングなどの他のハイパーパラメータは将来の調査に委ねられていると述べています。また、DMFT 解析は無限の幅極限を支持していますが、小規模幅の転送の挙動にはさらなる理論的研究が必要であるとも指摘しています。さらに、本論文は、スパース性によって引き起こされるハードウェア制約のため、MoE における FLOP と性能のトレードオフが密なモデルと大きく異なるため、MoE に対する「計算最適」のスケーリング則（例：Chinchilla 指数）を解決するものではないと主張していません。

要約すると、この研究は、訓練ダイナミクスに関する新たな 3 段階の平均場理論に基づき、小規模から大規模の MoE モデルへのハイパーパラメータの信頼性ある外挿を可能にする基礎的なパラメータ化を確立しています。