Mixture of Experts（MoE）

1. MoEの基本概念

1.1 条件付き計算の動機

通常のニューラルネットワーク：全パラメータが全入力に対して活性化。

スケーリングの法則：性能向上にはパラメータ数の増加が必要。しかし計算コストも増加。

MoEの解決策：パラメータは増やすが、各入力に対しては一部のみ使用。

1.2 アーキテクチャ概要

構成要素：

Expert：複数の独立したサブネットワーク（通常FFN）
Router（Gating Network）：入力をどのExpertに送るか決定

TransformerではFFN層をMoE層に置き換える。Attention層は共有。

1.3 歴史

1991：Jacobs et al. による最初のMoE提案
2017：Shazeer et al. "Outrageously Large Neural Networks"（スパースMoE）
2022：Switch Transformer（1 Expert選択）
2023：Mixtral（実用的MoE LLM）
2024：DeepSeek-V2/V3（細粒度MoE）

2. 数学的定式化

2.1 基本形式

$N$ 個のExpert $E_1, ..., E_N$ と Router $G$ に対し：

$$y = \sum_{i=1}^{N} G(x)_i \cdot E_i(x)$$

$G(x)_i$：入力 $x$ に対するExpert $i$ のゲート重み

2.2 スパースゲーティング

全Expertを使うと計算量削減にならない。Top-K選択：

$$G(x) = \text{Softmax}(\text{TopK}(H(x), k))$$

$H(x) = W_g \cdot x$：ルーター出力（各Expertのスコア）

Top-K以外は0にマスク。通常 $k = 1$ または $k = 2$。

2.3 パラメータ vs 活性化パラメータ

例：Mixtral 8x7B

総パラメータ：約47B（8つの7B Expert）
活性化パラメータ：約13B（Top-2で2つのExpert）

推論時の計算量は活性化パラメータに比例。

2.4 計算量の比較

Dense 70Bモデル vs MoE 8x7B（Top-2）：

Dense：70B FLOPs/トークン
MoE：約13B FLOPs/トークン（+ルーティングオーバーヘッド）

性能は同等だがMoEは約5倍効率的。

3. ルーティング機構

3.1 基本的なルーター

線形変換 + Softmax：

$$s = W_g \cdot x$$ $$G(x) = \text{Softmax}(\text{TopK}(s, k))$$

$W_g \in \mathbb{R}^{N \times d}$：学習可能なルーター重み

3.2 ノイズ付きTop-K

Shazeer et al. (2017)。探索促進のためノイズを追加。

$$s = W_g \cdot x + \epsilon \cdot \text{Softplus}(W_{noise} \cdot x)$$

$\epsilon \sim \mathcal{N}(0, 1)$

3.3 Expert Choice Routing

Zhou et al. (2022)。トークンがExpertを選ぶのではなく、ExpertがTop-Kトークンを選ぶ。

利点：完全な負荷分散、ドロップなし

欠点：トークン間の依存が生じる

3.4 DeepSeekのルーティング

DeepSeek-V2/V3：

共有Expert（全トークンが使用）+ ルーティングExpert
細粒度Expert（より多くの小さなExpert）
Expert並列とデータ並列の組み合わせ最適化

4. 訓練技術

4.1 負荷分散損失（Auxiliary Loss）

問題：特定のExpertに負荷が集中（Expert崩壊）

解決：負荷分散を促す補助損失を追加

Load Balancing Loss：

$$L_{balance} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i$$

$f_i$：Expert $i$ に送られたトークンの割合
$P_i$：Expert $i$ のルーター確率の平均
$\alpha$：バランス係数（通常0.01-0.1）

均等分散で最小化。

4.2 Router Z-Loss

ST-MoE (2022)。ルーター出力の大きさを制限。

$$L_z = \frac{1}{B} \sum_{x \in B} \left(\log \sum_i e^{s_i(x)}\right)^2$$

訓練安定性の向上。

4.3 Capacity Factor

各Expertが処理できるトークン数の上限。

$$\text{Capacity} = \frac{\text{tokens}}{\text{num\_experts}} \times \text{capacity\_factor}$$

通常 capacity_factor = 1.25〜2.0。

超過トークンはドロップまたは残差接続でバイパス。

4.4 訓練の安定性

MoEはDenseより訓練不安定になりやすい
より小さな学習率
長いWarmup
適切な補助損失の重み調整

5. 実装と分散処理

5.1 Expert並列（EP）

各Expertを異なるデバイスに配置。

通信：All-to-All通信でトークンをExpertに送信。

計算と通信のオーバーラップが重要。

5.2 データ並列 + Expert並列

実際のシステムでは複数の並列化を組み合わせ：

データ並列（DP）：バッチを分割
テンソル並列（TP）：行列を分割
Expert並列（EP）：Expertを分割
パイプライン並列（PP）：層を分割

5.3 推論時の考慮

メモリ：全Expert をロードする必要
バッチ処理：異なるトークンが異なるExpertへ→効率低下
量子化：Expert別に量子化可能

5.4 実装フレームワーク

Megatron-LM（NVIDIA）
DeepSpeed-MoE（Microsoft）
Fairseq（Meta）
Megablocks（Databricks）

6. 主要MoEモデル

6.1 Switch Transformer (2022)

Google。Top-1ルーティング。

シンプルなルーティング
T5ベース
最大1.6Tパラメータ

主要論文：

Fedus et al. (2022) "Switch Transformers: Scaling to Trillion Parameter Models", JMLR

6.2 Mixtral 8x7B (2023)

Mistral AI。実用的なオープンソースMoE。

8 Expert、各7B
Top-2ルーティング
総パラメータ47B、活性化13B
LLaMA 2 70Bに匹敵
Sliding Window Attention併用