MoEアーキテクチャ大解剖

🔀 MoEとは何か — 疎活性化の革命

MoE（Mixture of Experts, 専門家の混合）とは、「モデル全体のごく一部のパラメータだけを、各トークンで活性化させる」アーキテクチャだ。従来のDenseモデルが全パラメータを毎回使うのに対し、MoEは入力ごとに最適な数個の「専門家（Expert）」だけを選んで走らせる。この「疎活性化（Sparse Activation）」こそが、MoEの核心だ。

💾 なぜMoEがメモリ問題を解くのか

UIAPduinoページで語った「500Bを4GBで」の実現に直結する重要セクションだ。

Denseモデルでは全重みを毎トークンで使うため、SSDオフロードは推論速度が壊滅的に遅くなる。MoEは「このトークンではこのExpertしか使わない」と事前に分かるため、使わないExpertは物理メモリから外して良い。これがSSDストリーミング推論と相性が良い本質的理由だ。

🧠 MoEの5つの部品

📐 Dense vs MoE の構造比較

class DenseBlock(nn.Module):
    def __init__(self, dim, ffn_dim):
        super().__init__()
        self.attn = Attention(dim)
        self.ffn = FeedForward(dim, ffn_dim)  # 全トークン共通

    def forward(self, x):
        x = x + self.attn(x)
        x = x + self.ffn(x)  # 毎回全パラメータ使用
        return x

class MoEBlock(nn.Module):
    def __init__(self, dim, ffn_dim, num_experts=8, top_k=2):
        super().__init__()
        self.attn = Attention(dim)
        self.moe = MixtureOfExperts(
            dim, ffn_dim,
            num_experts=num_experts,
            top_k=top_k
        )

    def forward(self, x):
        x = x + self.attn(x)
        x = x + self.moe(x)  # Top-K個のExpertだけ使用
        return x

🔍 ソースコード解剖：Mixtral 8x7B

Mixtralは、Mistral AI が公開した最もシンプルで参照しやすいMoEの公式実装だ。公式リポジトリ mistralai/mistral-src は約1500行。以下はその実装の骨格を、初学者向けに整理したもの。

import torch
import torch.nn as nn
import torch.nn.functional as F

class MixtralMoE(nn.Module):
    def __init__(self, hidden_dim=4096, ffn_dim=14336,
                 num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # ① Router: 軽量な1層線形層（4096→8 で32KBのみ）
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)

        # ② Experts: 独立したFFNを8個保持（ここがメモリの大部分）
        self.experts = nn.ModuleList([
            FeedForward(hidden_dim, ffn_dim)
            for _ in range(num_experts)
        ])

    def forward(self, x):
        # x: [batch, seq_len, hidden_dim]
        batch, seq, dim = x.shape
        x_flat = x.view(-1, dim)  # [batch*seq, dim]

        # ③ Router: 各トークンの「どのExpertに行くか」のスコア計算
        router_logits = self.gate(x_flat)  # [batch*seq, num_experts]

        # ④ Top-K選択: 上位2個のExpertだけ残す
        routing_weights, selected = torch.topk(
            router_logits, self.top_k, dim=-1
        )
        routing_weights = F.softmax(routing_weights, dim=-1)

        # ⑤ 選ばれたExpertだけ実行
        output = torch.zeros_like(x_flat)
        for expert_id in range(self.num_experts):
            # このExpertが選ばれたトークンのインデックス
            mask = (selected == expert_id).any(dim=-1)
            if not mask.any():
                continue  # このExpertは使われない → メモリから外せる

            tokens = x_flat[mask]
            expert_out = self.experts[expert_id](tokens)

            # ⑥ Routerの重みで加重
            weight = routing_weights[mask].sum(dim=-1, keepdim=True)
            output[mask] += weight * expert_out

        return output.view(batch, seq, dim)

🧮 Routerの数学とLoad Balance

Routerが雑に訓練されると「特定のExpertばかりに仕事が集中する」問題が起きる（死にExpert）。これを防ぐのがLoad Balancing Lossだ。MoEモデルの学習では、通常のCross Entropy Lossに加えてこの補助損失を混ぜる。

def load_balance_loss(router_probs, selected_experts, num_experts):
    """
    router_probs: [batch*seq, num_experts]  Softmax後のRouter確率
    selected_experts: [batch*seq, top_k]     選ばれたExpert ID
    """
    # f_i: Expert iが選ばれた割合（one-hotの平均）
    mask = F.one_hot(selected_experts, num_experts).float()
    f_i = mask.mean(dim=(0, 1))  # [num_experts]

    # P_i: Expert iへのルーティング確率の平均
    P_i = router_probs.mean(dim=0)  # [num_experts]

    # f_i と P_i の両方が均等なほど損失が小さくなる
    return num_experts * (f_i * P_i).sum()

この補助損失により、Routerは「全Expertを平均的に使うように」と逆圧力を受け、Expert間の専門性が自然に分化していく。Mixtralでは router_aux_loss_coef=0.02 程度の重みで混ぜるのが標準。

💾 メモリ割り当ての実際

前ページで掲げた「500Bを4GBで動かす」という挑戦が、MoEによって現実的になる。以下はMixtral 8x7B を例にした、設定別のメモリ要求量。

// エキスパートキャッシュ（LRU方式）の擬似コード
struct ExpertCache {
    std::unordered_map<int, ExpertWeights*> hot;
    std::list<int> lru;
    size_t max_hot_experts = 4;  // 4GB制約内で保持できる数

    ExpertWeights* get(int expert_id) {
        if (hot.count(expert_id)) {
            // ヒット: LRU更新
            lru.remove(expert_id);
            lru.push_front(expert_id);
            return hot[expert_id];
        }
        // ミス: SSDからmmapでロード
        if (hot.size() >= max_hot_experts) {
            int victim = lru.back();
            lru.pop_back();
            munmap(hot[victim]->data, hot[victim]->size);
            hot.erase(victim);
        }
        auto* w = load_from_ssd_mmap(expert_id);
        hot[expert_id] = w;
        lru.push_front(expert_id);
        return w;
    }
};

現実のPCIe 5.0 NVMeは14GB/s の帯域を持つ。Expert 1個（約180MB FP16 / 45MB Q2）のロード時間は3〜13ms。トークン生成が100ms間隔なら、ミス時のペナルティは許容できる。ホットExpertが十分キャッシュされれば、ほとんどのトークンでSSDアクセスは発生しない。

📊 実モデル比較：Mixtral / DeepSeek V3 / Qwen MoE

DeepSeek V3の「Expert 256個中Top-8」は特に興味深い。細分化が進むほどRouterの選択自由度が上がり、Expert間の専門性分化が明確になる。SSDストリーミング視点では、Expertが細かいほど「必要なものだけロード」の粒度が細かくなり有利。

🐧 最初に読むべきLLM — LLM界のUbuntu

「LinuxでいうUbuntuのように、最初に触るのに単純でいい」——この問いへの答えは明確だ。LLMのソースコードを初めて読むなら、Andrej Karpathy の nanoGPT 一択。Ubuntu と同じく、入門者が迷わないよう設計された「デフォルトの選択肢」だ。

📚 候補比較表

💡 Claudeのおすすめ順序

🗺️ 学習ロードマップ

🎯 まとめ：MoEは「計算量を減らすのではなく、メモリを減らす」技術

MoEの本質は、単に「計算量が4分の1になる」ことではない。それ以上に重要なのは、「メモリに何を常駐させるかを選べる」という性質だ。Denseモデルでは全重みが常時必要だが、MoEでは「今このトークンで必要なExpertだけ」を選別できる。この性質が、SSDストリーミング推論・エッジデバイスでの大規模モデル運用を可能にする。

← 関連: UIAPduino — 290円マイコンボードの個人量産