Skip to content
Published on

Mixture of Experts（MoE）アーキテクチャ完全分析

🇰🇷KO🇺🇸EN🇯🇵JA
Authors
  • Name
    Twitter
Mixture of Experts Architecture

1. MoEとは何か？

**Mixture of Experts（MoE）は、モデルの全パラメータの一部のみを活性化して演算効率を高めるアーキテクチャです。すべての入力に対してモデル全体を使用するDenseモデルとは異なり、MoEは入力に応じて最適なエキスパート（Expert）**のみを選択して処理します。

Dense vs Sparseモデル

  • Denseモデル: すべてのパラメータが毎回活性化（例：LLaMA、GPT-4）
  • Sparse MoE: 全パラメータの一部のみ活性化（例：Mixtral、DeepSeek-V3）

核心的な利点は、パラメータ数は大きいが演算量は少ないということです。Mixtral 8x7Bは合計46.7Bのパラメータを持ちますが、推論時に活性化されるパラメータは約12.9Bに過ぎません。

2. MoEアーキテクチャの核心構成要素

Expert Network

各Expertは独立したFFN（Feed-Forward Network）です：

import torch
import torch.nn as nn

class Expert(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)  # SwiGLU gate

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU activation
        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))

Router（Gating Network）

Routerは各トークンをどのExpertに送るかを決定します：

class TopKRouter(nn.Module):
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.top_k = top_k

    def forward(self, x: torch.Tensor):
        # x shape: (batch, seq_len, d_model)
        logits = self.gate(x)  # (batch, seq_len, num_experts)
        top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)
        top_k_weights = torch.softmax(top_k_logits, dim=-1)
        return top_k_weights, top_k_indices

MoEレイヤー全体の実装

class MoELayer(nn.Module):
    def __init__(self, d_model: int, d_ff: int,
                 num_experts: int = 8, top_k: int = 2):
        super().__init__()
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])
        self.router = TopKRouter(d_model, num_experts, top_k)
        self.num_experts = num_experts

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, d_model = x.shape
        weights, indices = self.router(x)

        # Reshape for expert processing
        flat_x = x.view(-1, d_model)
        flat_weights = weights.view(-1, weights.shape[-1])
        flat_indices = indices.view(-1, indices.shape[-1])

        output = torch.zeros_like(flat_x)
        for i, expert in enumerate(self.experts):
            # Find tokens routed to this expert
            mask = (flat_indices == i).any(dim=-1)
            if mask.any():
                expert_input = flat_x[mask]
                expert_output = expert(expert_input)
                # Weight by router probability
                idx = (flat_indices[mask] == i).float()
                w = (flat_weights[mask] * idx).sum(dim=-1, keepdim=True)
                output[mask] += w * expert_output

        return output.view(batch_size, seq_len, d_model)

3. 主要MoEモデル分析

Mixtral 8x7B（Mistral AI）

  • 8つのExpert、Top-2ルーティング
  • 合計46.7Bパラメータ、活性12.9B
  • Attentionレイヤーは共有、FFNのみExpertに分離

DeepSeek-V3 MoE

DeepSeek-V3はより精巧なMoE設計を採用しています：

class DeepSeekMoE(nn.Module):
    """DeepSeek-V3スタイル：共有Expert + Routed Expert"""
    def __init__(self, d_model, d_ff, num_shared=1,
                 num_routed=256, top_k=8):
        super().__init__()
        # すべてのトークンが通過する共有Expert
        self.shared_experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_shared)
        ])
        # トークンごとに選択されるRouted Expert
        self.routed_experts = nn.ModuleList([
            Expert(d_model, d_ff // 4) for _ in range(num_routed)
        ])
        self.router = TopKRouter(d_model, num_routed, top_k)

    def forward(self, x):
        # 共有Expertの結果
        shared_out = sum(e(x) for e in self.shared_experts)
        # Routed Expertの結果
        weights, indices = self.router(x)
        routed_out = self._route_tokens(x, weights, indices)
        return shared_out + routed_out
  • 1つの共有Expert + 256のRouted Expert（Top-8選択）
  • 合計671Bパラメータ、活性37B
  • Auxiliary-loss-freeロードバランシングを導入

モデル比較

モデル総パラメータ活性パラメータExpert数Top-K
Mixtral 8x7B46.7B12.9B82
Mixtral 8x22B141B39B82
DeepSeek-V3671B37B256+18+1
Qwen2.5-MoE14.3B2.7B60+44+4

4. ルーティング戦略

Token Choice vs Expert Choice

# Token Choice: 各トークンがExpertを選択
def token_choice_routing(logits, top_k=2):
    top_k_vals, top_k_idx = logits.topk(top_k, dim=-1)
    weights = torch.softmax(top_k_vals, dim=-1)
    return weights, top_k_idx

# Expert Choice: 各Expertがトークンを選択
def expert_choice_routing(logits, capacity_factor=1.25):
    num_tokens = logits.shape[0]
    num_experts = logits.shape[1]
    capacity = int(num_tokens * capacity_factor / num_experts)

    expert_scores = logits.T  # (num_experts, num_tokens)
    top_k_vals, top_k_idx = expert_scores.topk(capacity, dim=-1)
    return top_k_vals, top_k_idx

5. ロードバランシング

Expert間の負荷不均衡はMoEの核心的な課題です：

def load_balancing_loss(router_logits, top_k_indices, num_experts):
    """Auxiliary load balancing loss (Switch Transformer方式)"""
    # Expert別トークン比率
    mask = torch.zeros_like(router_logits)
    mask.scatter_(-1, top_k_indices, 1.0)
    tokens_per_expert = mask.float().mean(dim=0)  # (num_experts,)

    # Expert別ルーティング確率の平均
    router_probs = torch.softmax(router_logits, dim=-1)
    router_prob_per_expert = router_probs.mean(dim=0)

    # 2つの分布の内積 = 不均衡の尺度
    loss = num_experts * (tokens_per_expert * router_prob_per_expert).sum()
    return loss

DeepSeek-V3のAuxiliary-loss-free方式は、Expert毎のbias termを動的に調整し、トークンを多く受け取るExpertのbiasを下げ、少ないExpertのbiasを上げることで、別途のlossなしに均衡を維持します。

6. 推論最適化

# Expert Parallelism: Expertを複数のGPUに分散
# GPU 0: Expert 0-3, GPU 1: Expert 4-7
class ExpertParallel(nn.Module):
    def __init__(self, experts_per_gpu, rank, world_size):
        super().__init__()
        self.local_experts = nn.ModuleList([
            Expert(d_model, d_ff)
            for _ in range(experts_per_gpu)
        ])
        self.rank = rank
        self.world_size = world_size

    def forward(self, x, indices):
        # All-to-All通信でトークンを再配分
        dispatched = all_to_all(x, indices, self.world_size)
        # ローカルExpertで処理
        output = self._process_local(dispatched)
        # 結果を再結合
        return all_to_all(output, indices, self.world_size)

7. クイズ

Q1: Mixtral 8x7Bで1つのトークンが推論時に使用するパラメータ数は約いくつですか？

約12.9Bパラメータです。8つのExpertのうちTop-2のみが活性化されるため、2つのExpertのFFNパラメータ + 共有されるAttentionパラメータのみが使用されます。全体46.7Bの約28%のみが活性化されることになります。

Q2: DeepSeek-V3が導入したAuxiliary-loss-freeロードバランシングの核心アイデアは？

従来のMoEはロードバランシングのために別途のauxiliary lossを追加しますが、これはモデル性能を低下させる可能性があります。DeepSeek-V3は各Expertに動的なbias termを設け、トークンを多く受け取るExpertのbiasを下げ、少ないExpertのbiasを上げることで自然に均衡を取ります。これにより別途のlossなしで安定したロードバランシングが可能になります。

Q3: Token ChoiceとExpert Choiceルーティングの違いとそれぞれの長所・短所は？
  • Token Choice: 各トークンがTop-K Expertを選択します。実装が簡単ですが、一部のExpertにトークンが集中する不均衡が発生する可能性があります。
  • Expert Choice: 各Expertが処理するトークンを選択します。完全なロードバランシングが保証されますが、一部のトークンがどのExpertにも選択されない可能性があります。

実際にはToken Choice + load balancing lossの組み合わせが最も多く使用されています。

Discuss on Twitter • View on GitHub