一文要約: MoEは「モデルがどれだけ知っているか」と「トークンごとにどれだけ計算を使うか」を切り離します。そのギャップこそが、面白いエンジニアリングがすべて起きる場所です。
30.1 核心的なアイデア
30.1.1 反直感的な結果
2023年12月にMistral AIがMixtral 8x7Bをリリースしたとき、名前が混乱を招きました。8つのモデル?56Bパラメータ?「8x7B」とはどういう意味?
数字はこうです:
- 総パラメータ数: 46.7B
- トークンごとのアクティブパラメータ数: 12.9B
モデルは約47Bパラメータ分の知識を保存していますが、各トークンに対して約13Bの計算量しか使いません。ほとんどのベンチマークで、すべてのトークンに70BパラメータをすべてアクティベートするメガモデルLLaMA 2 70Bに匹敵するか、上回ります。
Mixtral 8x7B vs LLaMA 2 70B:
Mixtral 8x7B LLaMA 2 70B
─────────────────────────────────────────────────────
総パラメータ数 46.7B 70B
アクティブパラメータ数 12.9B 70B
推論速度 約6倍速い ベースライン
MMLU 70.6% 68.9%
この結果は魔法ではありません。シンプルなアーキテクチャ上の決定から生じています: 密なFFNをスパースなMixture of Expertsで置き換える。
30.1.2 密なアクティベーション vs スパースなアクティベーション
標準的なTransformerでは、すべてのトークンがFFN全体を通過します。FFNに4Bパラメータがある場合、文末のピリオドであれ複雑な数学用語であれ、すべてのトークンに対してその4Bパラメータが動きます。
密なモデル:
入力トークン → [すべてのパラメータが参加] → 出力
70Bパラメータ → 70Bのアクティベーション(トークンごと)
MoEモデル:
入力トークン → [ルーターが8つ中2つのエキスパートを選択] → [2つのエキスパートが計算] → 出力
46.7Bパラメータ → 12.9Bのアクティベーション(トークンごと)
ここに埋め込まれた洞察: すべての知識が同時に必要なわけではありません。PythonのFunction呼び出しの一部であるトークンは、散文作成の専門家をアクティブにする必要はありません。英語の物語のトークンは線形代数の専門家をアクティブにする必要はありません。
MoEはこの直感をルーティングメカニズムとして実装します。
30.1.3 ルーティングのアナロジー
専門化されたエンジニアリングチームを考えてみてください。PRを提出するとき:
- セキュリティに特化したレビュアーが認証の変更を確認する
- パフォーマンスのレビュアーがホットパスのコードを確認する
- ドキュメントのレビュアーがAPIの変更を確認する
すべてのレビュアーがすべてのPRを読むわけではありません。チームリード(ルーター)がPRの説明を読んで適切なレビュアーを割り当てます。
ルーティングチームは個々のジェネラリストレビュアーよりも広い専門知識をカバーしますが、PRあたりのコストは実際に参加するレビュアーの数で制限されます。
MoEも同じです: Nエキスパートがいますが、トークンごとにアクティベートされるのはK個だけ。トークンはスペシャリストの扱いを受け、計算量は制限されます。
30.1.4 簡単な歴史
MoEは1991年にJacobsらによって提案されました。フロンティアに到達するまでに3十年かかりました:
タイムライン:
1991 Jacobs et al. — 元のMoEコンセプト
2017 Shazeer et al. — Sparsely-Gated MoE、大規模NLPに適用
2021 Google Switch Transformer — 1.6兆パラメータのMoEモデル
2022 Google GLaM — 1.2兆パラメータ、GPT-3と競合
2023 Mixtral 8x7B — オープンソースMoE、実用的なデプロイ
2024 DeepSeek-V3 — 671B総数、37Bアクティブ、550万ドルの学習コスト
2021〜2024年の加速は2つの力によって推進されています: デプロイされたシステムで推論コストが支配的な費用になったこと、そしてハードウェアがルーティングのオーバーヘッドを無視できるほど速くなったこと。
30.2 MoEアーキテクチャ
30.2.1 スタック内でのMoEの位置
TransformerからMoEへの変更は手術的です。ブロックごとに変更されるコンポーネントは1つだけです:
標準的なTransformerブロック:
入力
↓
Self-Attention
↓
FFN(密) ← これがMoEになる
↓
出力
MoE Transformerブロック:
入力
↓
Self-Attention
↓
MoE層 ← ルーター + NエキスパートなFFN
↓
出力
その他すべて — Attentionメカニズム、残差接続、LayerNorm、位置エンコーディング — はそのままです。
30.2.2 MoE層の内部
MoE層には2つのコンポーネントがあります:
1. ルーター: トークンの隠れ表現をエキスパートの確率分布にマッピングする小さな線形層。
2. エキスパートネットワーク: 各自のウェイトを持つN個の独立したFFN。
MoE層の内部:
x (hidden_size)
↓
┌──────────────┐
│ ルーター │ Linear(hidden_size → num_experts) + Softmax
└──────┬───────┘
↓
┌──────────────┐
│ Top-Kゲート │ K個の最高確率エキスパートのみ保持
└──────┬───────┘
↓
選択されたKエキスパートがxを受け取る:
┌────┬────┬────┬────┬────┬────┬────┬────┐
│ E0 │ E1 │ E2 │ E3 │ E4 │ E5 │ E6 │ E7 │ (8エキスパート合計)
└────┴────┴────┴────┴────┴────┴────┴────┘
↓ ↓
選択されたエキスパートのみが計算する。
出力 = 選択されたエキスパート出力の加重和
30.2.3 ルーターのメカニズム
ルーターはソフトマックスを続けた単純な線形層です:
# 入力: x、shape = (batch_size, seq_len, hidden_size)
# ステップ1: 各エキスパートのスコアを計算
router_logits = Linear(hidden_size, num_experts)(x)
# shape: (batch_size, seq_len, num_experts)
# ステップ2: エキスパートに対してソフトマックス
router_probs = softmax(router_logits, dim=-1)
# 例、1トークンの場合: [0.40, 0.30, 0.10, 0.05, 0.05, 0.03, 0.04, 0.03]
# ステップ3: Top-Kを選択
top_k_probs, top_k_indices = topk(router_probs, k=2)
# top_k_probs: [0.40, 0.30]
# top_k_indices: [0, 1] → エキスパート0とエキスパート1
# ステップ4: ウェイトを再正規化
weights = top_k_probs / top_k_probs.sum()
# weights: [0.57, 0.43]
ルーターはトークンの種類を認識してスペシャリストにマッピングすることを学習します。この専門化はプログラムされたものではなく、学習から生まれます。実際にMixtralのルーターは測定可能な専門化を示します:
観察されたルーティングの傾向(Mixtralの分析):
Pythonキーワード → エキスパート3、エキスパート7(コードスペシャリスト)
数学的用語 → エキスパート1、エキスパート5(数学スペシャリスト)
機能語 → エキスパート0、エキスパート4(構文スペシャリスト)
30.2.4 エキスパートネットワーク
各エキスパートは標準的なFFNです:
class Expert(nn.Module):
def __init__(self, hidden_size: int, intermediate_size: int):
super().__init__()
self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False)
self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False)
self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False)
self.act = nn.SiLU()
def forward(self, x):
# SwiGLUゲーティング: w3からの要素ごとのゲートがw1の出力を制御
return self.w2(self.act(self.w1(x)) * self.w3(x))
各エキスパートは独立しています — 別々のウェイト、別々の勾配。8つのエキスパートはFFNパラメータの8倍を意味します。これが余分な容量の出所です。
30.2.5 Top-K選択: なぜK=2か?
Top-1(トークンごとに1エキスパート):
- 最小の計算量
- 勾配は1つのエキスパートにのみ流れる → 学習の不安定性
- ルーティングが間違った場合に冗長性なし
Top-2(トークンごとに2エキスパート):
- 2つのエキスパートが補完し合える
- 勾配が2つのエキスパートに届く → より安定
- Mixtralとほとんどのプロダクションシステムの標準的な慣行
K > 2のTop-K:
- エキスパートが増えるごとにスパース性が減少
- K = Nのとき、モデルは密なFFNに退化する
- 計算量はKに比例して増加
実証的な勝者はK = 2です。安定性、冗長性、効率性のスイートスポットに当たっています。
30.2.6 負荷分散
制約なしでは、ルーターは崩壊します。少数の「安全な」エキスパートを見つけて、ほとんどすべてをそこにルーティングします:
崩壊したルーティング(病的なケース):
エキスパート0: トークンの85% ← ボトルネック、すべての勾配を受け取る
エキスパート1: 9%
エキスパート2: 3%
...
エキスパート7: 0.1% ← ほぼ未使用、パラメータの無駄遣い
これは2つの問題を引き起こします: エキスパート0が計算上のボトルネックになり、エキスパート2〜7がほとんど学習されません。
解決策: 補助的な負荷分散損失
ルーティングが不均等な場合にペナルティを追加します:
def load_balancing_loss(router_probs, expert_indices, num_experts):
# expert_fraction: 各エキスパートが選択される頻度
expert_mask = F.one_hot(expert_indices, num_experts).float()
expert_fraction = expert_mask.mean(dim=(0, 1)) # エキスパートごとの選択率
# router_fraction: エキスパートごとの平均ルーティング確率
router_fraction = router_probs.mean(dim=(0, 1))
# ペナルティ = num_experts × sum(fraction × probability)
# 分布が均一なときに最小
aux_loss = num_experts * (expert_fraction * router_fraction).sum()
return aux_loss
直感: エキスパート0がトークンの85%(高いexpert_fraction)を選択され、かつルーターが高い確率(高いrouter_fraction)を割り当てた場合、積が大きくなりペナルティが強くなります。これがルーターを均一な分布に向けて押し進めます。
損失は係数(通常aux_loss_coef = 0.01)でスケールされ、メインの学習損失に追加されます。
30.3 Mixtral 8x7B
30.3.1 設定
| パラメータ | 値 | 備考 |
|---|---|---|
| 総パラメータ数 | 46.7B | すべてのエキスパートのウェイトを含む |
| アクティブパラメータ数 | 12.9B | トークンごとに8つ中2つのエキスパートのみアクティブ |
| 層ごとのエキスパート数 | 8 | それぞれが完全なSwiGLU FFN |
| アクティブエキスパート数 | 2 | Top-2ルーティング |
| 隠れ次元数 | 4096 | LLaMA 2と同じ |
| 層数 | 32 | 32 Transformerブロック |
| Attention | GQA、32 Qヘッド / 8 KVヘッド | 効率化のためのグループクエリ |
| コンテキスト長 | 32K | RoPEスライディングウィンドウ付き |
30.3.2 パラメータ数の導出
46.7Bパラメータはどこから来るのでしょうか?
埋め込み層:
32000 × 4096 ≈ 1.31億
層ごと:
Self-Attention(GQA):
Q: 4096 × 4096 = 1,680万
K: 4096 × 1024 = 420万 (8 KVヘッド × 128 head_dim)
V: 4096 × 1024 = 420万
O: 4096 × 4096 = 1,680万
小計: ≈ 4,200万
MoE層(8エキスパート、SwiGLU):
各エキスパート:
w1: 4096 × 14336 = 5,870万
w2: 14336 × 4096 = 5,870万
w3: 4096 × 14336 = 5,870万 (ゲート)
エキスパートごとの小計: ≈ 1.76億
8エキスパート: ≈ 14.08億 = 1.4B
ルーター: 4096 × 8 ≈ 3.3万(無視できる)
層合計: 4,200万 + 14.08億 ≈ 14.5億
32層: 14.5億 × 32 ≈ 464億
埋め込みとLMヘッド含む: ≈ 46.7B総計
アクティブパラメータの計算:
トークンごとに8つ中2つのエキスパートのみが動く:
MoE以外の部分(Attention × 32層): ≈ 13億
MoE部分(8つ中2つのエキスパート × 32層): ≈ 112億
埋め込み: ≈ 4億
合計アクティブ: ≈ 12.9B
これがMixtralが密な13Bモデルよりトークンごとの計算量が少ないのに、47Bモデルの知識容量を持つ理由です。
30.3.3 Mixtral vs LLaMA 2 70B
| 指標 | Mixtral 8x7B | LLaMA 2 70B |
|---|---|---|
| 総パラメータ数 | 46.7B | 70B |
| アクティブパラメータ数 | 12.9B | 70B |
| 推論FLOP | 約13B相当 | 70B |
| トークン/秒 | 約6倍速い | ベースライン |
| VRAM(FP16) | 約90 GB | 約140 GB |
| MMLU | 70.6% | 68.9% |
| HumanEval(コード) | 60.7% | 約30% |
| 多言語 | 強い | 中程度 |
効率の差は顕著です。1日に1,000万トークンを処理するサービングデプロイメントでは、Mixtralは同等かそれ以上の品質でLLaMA 2 70Bの約1/6の計算量しか使いません。
30.3.4 観察されたルーターの挙動
Mistralの学習済みルーターの分析で、実際の専門化が確認されています:
エキスパートは独自のドメインを発展させます。ラベルが割り当てられていないにもかかわらず、学習信号だけを通してルーターはどのエキスパートがどのトークンタイプに最適かを学習します。
位置が重要です: 文の先頭の「The」は文中の「the」とは異なるルーティングをされる可能性があります。ルーターはトークンのアイデンティティだけでなく、構文的コンテキストに敏感です。
隣接するトークンは多様化します: シーケンス内の隣接するトークンは異なるエキスパートのサブセットを選択する傾向があり、モデルがシーケンス全体で暗黙的な分業形式を学習したことを示唆しています。
30.4 DeepSeek-V3: MoEをさらに押し進める
30.4.1 コストの話
DeepSeek-V3(2024年12月)は6,710億パラメータのモデルを 550万ドル で学習しました。GPT-4の推定学習コストは1億ドルを超えます。ほぼ同等の能力で18倍安い。このギャップはアーキテクチャ効率から生まれています: Multi-head Latent Attention(MLA) と 細粒度MoE。
30.4.2 設定
| パラメータ | DeepSeek-V3 | 備考 |
|---|---|---|
| 総パラメータ数 | 671B | 非常に大きな総容量 |
| アクティブパラメータ数 | 37B | トークンごとの計算量は管理可能 |
| ルーティングされるエキスパート数 | 256 | 細粒度の専門化 |
| 共有エキスパート | 1 | 常にアクティブ、ユニバーサルバックボーン |
| アクティブなルーティングエキスパート数 | 8 | 256からTop-8 |
| 層数 | 61 | Mixtralより深い |
| 隠れ次元数 | 7168 | |
| コンテキスト長 | 128K |
30.4.3 Multi-head Latent Attention(MLA)
128Kコンテキストウィンドウでは、KVキャッシュが支配的な制約になります:
標準MHAのKVキャッシュ:
サイズ ∝ num_heads × head_dim × seq_len × num_layers
128Kでは: 巨大で、GPUメモリを埋め尽くす
MLAのKVキャッシュ:
KとVを低次元の潜在ベクトル c_KV に圧縮
フルのKとVの代わりに c_KV をキャッシュ
Attention計算時に展開
MLAは低ランク射影を適用します:
標準MHAのパス:
x → W_K → K (フルKをキャッシュ)
x → W_V → V (フルVをキャッシュ)
KVキャッシュ ∝ num_heads × head_dim
MLAのパス:
x → W_DKV → c_KV (圧縮)
c_KV をキャッシュ (はるかに小さい)
c_KV → W_UK → K (計算時に展開)
c_KV → W_UV → V
KVキャッシュ ∝ latent_dim (latent_dim << num_heads × head_dim)
latent_dim = 0.25 × (num_heads × head_dim) なら、KVキャッシュは75%縮小します。128Kコンテキストでは、これはシングルノードに収まるか否かの違いです。
30.4.4 細粒度MoE
Mixtralは8つの大きなエキスパートを使います。DeepSeek-V3は256の小さなエキスパートを使います。この違いは重要です:
粗粒度(8つの大きなエキスパート、Top-2):
- 各エキスパートはフルサイズのFFN
- 2/8 = 25%のアクティベーション率
- ルーティングの決定は粗い
細粒度(256の小さなエキスパート、Top-8):
- 各エキスパートはFFNの一部
- 8/256 ≈ 3%のアクティベーション率
- はるかに精密なルーティング
- より多くのエキスパートでより良い負荷分散
トークンごとの総計算量は似たようなものです(8つの小さなエキスパートは2つの大きなエキスパートとFLOP的に等しくなれます)が、ルーティングの粒度は32倍細かくなります。これはモデルがどのスペシャリストを使うかについてはるかに精密な決定を下せることを意味します。
30.4.5 共有エキスパート
DeepSeek-V3は すべてのトークンに対して常にアクティブな 1つのエキスパートを追加します:
DeepSeek-V3 MoE:
x
├── ルーター → 256のルーティングエキスパートから8つを選択
│ ↓
│ routed_output
│
└── 共有エキスパート(常にアクティブ)
↓
shared_output
final_output = routed_output + shared_output
共有エキスパートはユニバーサルパターンを処理します — 一般的な文法、標準的な推論ステップ、頻繁なサブワードなど、どのドメインのトークンにも必要なもの。ルーティングエキスパートは差別化されたドメイン固有の計算を処理します。
これにより、ルーティングエキスパートが一般的なパターンに容量を無駄遣いするのを防ぎます。
30.4.6 550万ドルで何が買えたか
学習コストの内訳:
- FP8混合精度: 演算ごとのメモリ帯域幅を半分に
- MLA: より小さなKVキャッシュによる大きなバッチサイズ
- 計算-通信オーバーラップ: 勾配all-reduce中に計算が進む
- エキスパート並列性: 256エキスパートが2,048GPUクラスターにきれいにシャードされる
- 14.8兆の学習トークン: 高品質データ、マルチステージカリキュラム
この組み合わせがGPT-4クラスのベンチマーク結果を、推定学習コストの約1/18で達成します。アーキテクチャの選択が複合的に積み重なります。
30.5 MoEの課題
30.5.1 学習の不安定性
MoEモデルは同等の計算量の密なモデルよりも学習が難しいです:
ルーターの崩壊: ルーターが少数のエキスパートにトラフィックを集中させ、それらのエキスパートがすべての勾配を受け取り、他のエキスパートの学習が止まり、問題が悪化します。防御策: 負荷分散損失、初期化ノイズ、エキスパートドロップアウト。
損失スパイク: ルーティングの決定がバッチ間で急激に変化し、勾配の不連続性を引き起こします。防御策: 勾配クリッピング、小さな学習率、大きなバッチサイズ。
エキスパートの飢餓: 一部のエキスパートが適切に学習するのに十分なトークンを受け取れません。防御策: 使用率の低いエキスパートへの強制的な再ルーティングを行うキャパシティファクター。
30.5.2 実際の負荷不均衡
補助損失があっても、完全なバランスは保証されません:
現実的なルーティング分布(学習後):
エキスパート0: 18% ← 中程度に人気
エキスパート1: 14%
エキスパート2: 13%
エキスパート3: 12%
エキスパート4: 11%
エキスパート5: 11%
エキスパート6: 10%
エキスパート7: 11%
vs 理想的な均一分布:
各エキスパート: 12.5%
これは許容範囲です。しかし分散設定でエキスパート0がGPU 0にあり、エキスパート7がGPU 7にある場合、この不均衡は直接計算のレイテンシに変換されます。
キャパシティファクター: バッチごとに各エキスパートが処理できるトークン数のハードキャップ。オーバーフローするトークンはドロップされるか次のベストエキスパートにリダイレクトされます。一般的な値: 1.0〜1.5。
capacity = (total_tokens / num_experts) * capacity_factor
# expert_queue > capacity の場合: 超過トークンは次のベストエキスパートが処理
30.5.3 All-to-All通信
分散学習では、エキスパートはGPUをまたいでシャードされています。トークンのルーティングはGPUの境界を越えます:
セットアップ: 4 GPU、各2エキスパート
GPU 0: エキスパート0、1
GPU 1: エキスパート2、3
GPU 2: エキスパート4、5
GPU 3: エキスパート6、7
GPU 0のバッチはエキスパート5(GPU 2)とエキスパート7(GPU 3)にトークンをルーティングするかもしれない:
GPU 0 → GPU 2: トークンのアクティベーションを送信
GPU 2 → GPU 0: 計算結果を返す
(すべてのGPUが同時にこれを行う → all-to-allパターン)
all-to-allは の通信コストを持ち、MoE層ごとに2回発生します(送信に1回、受信に1回)。注意深く処理しないと、この通信が実際の計算時間を支配する可能性があります。
緩和策: 計算-通信オーバーラップ、ノード間トラフィックを最小化するエキスパートグループ配置、ディスパッチ前にトークンをバッチにまとめる。
30.5.4 サービングの複雑さ
動的バッチングがより難しい: 密なモデルでは、バッチ内のすべてのトークンが同じ計算パスを辿ります。MoEモデルでは、異なるトークンが異なるエキスパートをアクティベートします。密なモデルに機能するバッチング戦略は、MoEルーティングの下では悲惨に断片化する可能性があります。
メモリプロファイル: トークンごとにアクティブなのは2〜8エキスパートだけでも、すべてのエキスパートウェイトがメモリに存在しなければなりません。Mixtralはアクティブパラメータが12.9Bにもかかわらず、FP16推論に約90 GB VRAMが必要です。「軽い計算」の恩恵はVRAMの比例削減には変換されません。
30.6 MoEの実装
30.6.1 コアMoE層
import torch
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
def __init__(self, hidden_size: int, intermediate_size: int):
super().__init__()
self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False)
self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False)
self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False)
self.act = nn.SiLU()
def forward(self, x):
return self.w2(self.act(self.w1(x)) * self.w3(x))
class MoELayer(nn.Module):
def __init__(
self,
hidden_size: int = 4096,
intermediate_size: int = 14336,
num_experts: int = 8,
top_k: int = 2,
aux_loss_coef: float = 0.01,
):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.aux_loss_coef = aux_loss_coef
self.router = nn.Linear(hidden_size, num_experts, bias=False)
self.experts = nn.ModuleList([
Expert(hidden_size, intermediate_size) for _ in range(num_experts)
])
def forward(self, x):
batch, seq_len, hidden_size = x.shape
# ルーター: スコアリングと選択
router_logits = self.router(x)
router_probs = F.softmax(router_logits, dim=-1)
top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
# ディスパッチ: 選択されたエキスパートにトークンを送る
x_flat = x.view(-1, hidden_size)
output = torch.zeros_like(x_flat)
for expert_idx in range(self.num_experts):
# このエキスパートにルーティングされるトークンはどれか?
expert_mask = (top_k_indices == expert_idx).any(dim=-1).view(-1)
if not expert_mask.any():
continue
expert_input = x_flat[expert_mask]
expert_output = self.experts[expert_idx](expert_input)
# ウェイト付けして蓄積
weights = torch.where(
top_k_indices == expert_idx, top_k_weights,
torch.zeros_like(top_k_weights),
).sum(dim=-1).view(-1)[expert_mask]
output[expert_mask] += expert_output * weights.unsqueeze(-1)
output = output.view(batch, seq_len, hidden_size)
aux_loss = self._load_balance_loss(router_probs, top_k_indices)
return output, aux_loss
def _load_balance_loss(self, router_probs, expert_indices):
expert_mask = F.one_hot(expert_indices, self.num_experts).float()
expert_fraction = expert_mask.sum(dim=2).mean(dim=(0, 1))
router_fraction = router_probs.mean(dim=(0, 1))
aux_loss = self.num_experts * (expert_fraction * router_fraction).sum()
return aux_loss * self.aux_loss_coef
30.6.2 ノイジールーター(学習安定性)
学習中にノイズを加えることで、ルーターが学習の早期にすべてのエキスパートを探索するよう促します:
class NoisyTopKRouter(nn.Module):
def __init__(self, hidden_size, num_experts, top_k, noise_std=0.1):
super().__init__()
self.top_k = top_k
self.noise_std = noise_std
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
def forward(self, x, training=True):
logits = self.gate(x)
if training and self.noise_std > 0:
logits = logits + torch.randn_like(logits) * self.noise_std
probs = F.softmax(logits, dim=-1)
top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
return weights, top_k_indices, probs
30.6.3 HuggingFaceでMixtralを読み込む
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map="auto", # 利用可能なGPU全体に分散
load_in_4bit=True, # 4ビットでVRAMを約25 GBに削減
)
# MoE構造を確認する
moe_layer = model.model.layers[0].block_sparse_moe
print(f"ルーター: {moe_layer.gate}")
# ルーター: Linear(in_features=4096, out_features=8, bias=False)
print(f"エキスパート数: {len(moe_layer.experts)}")
# エキスパート数: 8
print(f"エキスパート0: {moe_layer.experts[0]}")
# MixtralBlockSparseTop2MLP(
# (w1): Linear(4096 → 14336, bias=False)
# (w2): Linear(14336 → 4096, bias=False)
# (w3): Linear(4096 → 14336, bias=False)
# )
30.7 MoE vs 密なモデル: どちらを使うか
30.7.1 パラメータ数とアクティベーション数
モデル 総パラメータ数 アクティブパラメータ数 アクティベーション率
────────────────────────────────────────────────────────────
LLaMA 2 70B 70B 70B 100%
Mixtral 8x7B 46.7B 12.9B 27.6%
DeepSeek-V3 671B 37B 5.5%
GPT-4(推測) ~1.8T ~110B ~6%
総パラメータ数が増えるにつれ、効率的なフロンティアはますますMoEを有利にします。
30.7.2 学習コスト
| モデル | 推定学習コスト | トークン数 | ハードウェア |
|---|---|---|---|
| LLaMA 2 70B | 約500万ドル | 2T | A100 |
| Mixtral 8x7B | 約200万ドル(推定) | 非公開 | 非公開 |
| DeepSeek-V3 | 550万ドル | 14.8T | H800(2,048 GPU) |
| GPT-4 | >1億ドル(推測) | 13T+ | A100/H100 |
MoEは2つのメカニズムで学習効率を達成します: トークンごとのFLOPが少ない(アクティブエキスパートのみが計算する)、そして計算バジェットをより良く使える(より多くのパラメータ = 同じ計算量でより多くの容量)。
30.7.3 推論効率
| 指標 | 密な70B | MoE 8x7B(12.9Bアクティブ) |
|---|---|---|
| 最初のトークンまでの時間 | ベースライン | 約0.2× |
| スループット | ベースライン | 約3〜4× |
| VRAM(FP16) | 約140 GB | 約90 GB |
| トークン/秒 | ベースライン | 約6× |
スループットの優位性は本物です。レイテンシの優位性も存在しますが、より小さいです。VRAMの優位性も本物ですが、アクティブパラメータ数に比例してスケールしません — すべてのエキスパートをロードしなければなりません。
30.7.4 密なモデルを選ぶとき
- シーケンス長が4K以下(Attentionのオーバーヘッドは管理可能)
- メモリ制約のデプロイメント(推論VRAMバジェットが制約)
- 単一タスクの微調整(MoEの多ドメイン知識が無駄になる)
- より単純なサービングスタックの価値が効率の恩恵より大きい
30.7.5 MoEを選ぶとき
- 高スループット要件(APIサービング、検索拡張)
- 多言語またはマルチドメインタスク
- 利用可能なVRAMが密なモデルが必要とする量を超えている
- 学習バジェットが制約されているが、より多くの総容量が欲しい
30.8 章のまとめ
30.8.1 主要概念
| 概念 | 意味 |
|---|---|
| MoE | Mixture of Experts — 効率的な大型モデルのためのスパースアクティベーション |
| スパースアクティベーション | 各トークンに対してパラメータのサブセットのみが計算する |
| ルーター | Top-K選択を通じてトークンをエキスパートに割り当てる線形層 |
| エキスパート | 独自のパラメータを持つ独立したFFNネットワーク |
| Top-K | トークンごとにK個の最高スコアのエキスパートを選択(通常K=2) |
| 負荷分散 | 均一なエキスパート利用を促進する補助損失 |
| MLA | Multi-head Latent Attention — 低ランク射影によるKVキャッシュの圧縮 |
| 細粒度MoE | 少数の大きなエキスパートの代わりに多数の小さなエキスパート; 低いアクティベーション率 |
| 共有エキスパート | 常にアクティブな1つのエキスパート; ユニバーサルなトークンパターンを処理 |
30.8.2 重要な数値
Mixtral 8x7B:
総パラメータ数: 46.7B | アクティブ: 12.9B (27.6%)
エキスパート数: 8 | トークンごとのアクティブ: 2
結果: LLaMA 2 70Bに匹敵し、推論速度は約6倍
DeepSeek-V3:
総パラメータ数: 671B | アクティブ: 37B (5.5%)
エキスパート数: 256 + 1 | トークンごとのアクティブ: 8 + 1
学習コスト: 550万ドル(GPT-4: >1億ドル)
30.8.3 コアな数式
ルーターの計算:
router_logits = Linear(x) # hidden_size → num_experts
router_probs = softmax(router_logits)
top_k_weights, top_k_indices = topk(router_probs, k)
MoEの出力:
負荷分散損失:
ここで はエキスパートの選択頻度、 は平均ルーティング確率です。
30.8.4 私の見解
MoEは「フロンティアAIはシステムエンジニアリングである」という最も明確な例です。アルゴリズム — 各トークンをKエキスパートにルーティングし、負荷分散ペナルティで学習する — は複雑ではありません。難しいのは大規模でそれを機能させることです: all-to-all通信がボトルネックにならずに数百のGPUをまたいでトークンをルーティングし、671Bモデルでのルーターの崩壊をデバッグし、動的バッチングの病理を処理するサービングスタックを構築すること。
DeepSeek-V3が重要なのは推論コストが安いからではなく、550万ドルの学習コストという数字がフロンティア能力がもはや学習バジェットの関数だけではないことを証明するからです。アーキテクチャ効率は複合的に積み重なります。
章のチェックリスト
この章を終えたら、次のことができるようになっているはずです:
- スパースアクティベーションと、なぜそれが総容量とトークンごとの計算量を切り離すのかを説明できる。
- MoE層の構造を説明できる: ルーター、Top-Kゲーティング、エキスパートFFN。
- なぜK=2がTop-K選択の標準的な選択なのかを説明できる。
- 負荷分散損失と、それなしに何が起こるかを説明できる。
- Mixtral 8x7Bのアクティブパラメータおよびトータルパラメータを計算できる。
- MLAと、長コンテキストMoEモデルにとって何が重要かを説明できる。
- MoEの失敗モードを少なくとも2つ挙げ、その緩和策を説明できる。
次の章へ
MoEは学習と推論の両方で計算を賢く使うことについてです。まったく別の次元もあります: より良い答えを得るために推論時により多くの計算を使うこと。第31章では推論モデルの革命を解説します — GPT-4oのAIME 2024での12%からo3の96.7%まで、そしてDeepSeek-R1がオープンソースで可能にしたストーリーをたどりましょう。