Mixture of Experts（MoE）の理論と実装 — 条件付き計算でLLMを効率化する

x → Attention → Add & Norm → FFN → Add & Norm → y

x → Attention → Add & Norm → MoE(FFN₁, FFN₂, ..., FFNₑ) → Add & Norm → y

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

class Expert:
    """1つのFFNエキスパート。"""

    def __init__(self, d_model, d_ff):
        # SwiGLU風の2つの線形層
        self.W1 = np.random.randn(d_model, d_ff) * 0.02
        self.W2 = np.random.randn(d_ff, d_model) * 0.02
        self.W_gate = np.random.randn(d_model, d_ff) * 0.02

    def forward(self, x):
        """SwiGLU: gate(xW_gate) * (xW1) → W2"""
        gate = x @ self.W_gate
        gate = gate * (1 / (1 + np.exp(-gate)))  # SiLU activation
        hidden = gate * (x @ self.W1)
        return hidden @ self.W2


class MoELayer:
    """Mixture of Experts層。"""

    def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model

        # エキスパート
        self.experts = [Expert(d_model, d_ff) for _ in range(num_experts)]

        # ゲーティングネットワーク
        self.W_gate = np.random.randn(d_model, num_experts) * 0.1

    def route(self, x):
        """トークンをエキスパートにルーティングする。

        Parameters
        ----------
        x : np.ndarray, shape (batch, d_model)

        Returns
        -------
        indices : np.ndarray, shape (batch, top_k) — 選択されたエキスパートの番号
        weights : np.ndarray, shape (batch, top_k) — 正規化された重み
        gate_probs : np.ndarray, shape (batch, num_experts) — 全ゲート確率
        """
        # ゲートロジット
        logits = x @ self.W_gate  # (batch, num_experts)

        # ソフトマックス
        logits_max = np.max(logits, axis=1, keepdims=True)
        exp_logits = np.exp(logits - logits_max)
        gate_probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)

        # Top-k選択
        indices = np.argsort(-gate_probs, axis=1)[:, :self.top_k]

        # 選択されたエキスパートの重みを正規化
        weights = np.zeros((x.shape[0], self.top_k))
        for i in range(x.shape[0]):
            selected_probs = gate_probs[i, indices[i]]
            weights[i] = selected_probs / np.sum(selected_probs)

        return indices, weights, gate_probs

    def forward(self, x):
        """MoE層の前方計算。"""
        batch_size = x.shape[0]
        indices, weights, gate_probs = self.route(x)

        output = np.zeros_like(x)
        for i in range(batch_size):
            for j in range(self.top_k):
                expert_idx = indices[i, j]
                expert_output = self.experts[expert_idx].forward(x[i:i+1])
                output[i] += weights[i, j] * expert_output[0]

        return output, indices, gate_probs

    def load_balancing_loss(self, gate_probs, indices):
        """負荷分散損失を計算する。"""
        batch_size = gate_probs.shape[0]

        # ディスパッチ率 f_i
        f = np.zeros(self.num_experts)
        for i in range(batch_size):
            for j in range(self.top_k):
                f[indices[i, j]] += 1
        f /= batch_size

        # ルーティング確率の平均 p_i
        p = np.mean(gate_probs, axis=0)

        # 負荷分散損失
        loss = self.num_experts * np.sum(f * p)
        return loss, f, p


# MoE層の動作確認
d_model = 64
d_ff = 128
num_experts = 8
top_k = 2
batch_size = 256

moe = MoELayer(d_model, d_ff, num_experts, top_k)
x = np.random.randn(batch_size, d_model).astype(np.float64) * 0.5

output, indices, gate_probs = moe.forward(x)
lb_loss, f, p = moe.load_balancing_loss(gate_probs, indices)

print(f"入力形状: {x.shape}")
print(f"出力形状: {output.shape}")
print(f"負荷分散損失: {lb_loss:.4f} (理想値: {top_k:.4f})")
print(f"エキスパートごとのディスパッチ率: {f}")

# ルーティングパターンの分析
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 左: エキスパートごとのトークン割り当て数
expert_counts = np.zeros(num_experts)
for i in range(batch_size):
    for j in range(top_k):
        expert_counts[indices[i, j]] += 1

colors = plt.cm.Set3(np.linspace(0, 1, num_experts))
bars = axes[0].bar(range(num_experts), expert_counts, color=colors)
axes[0].axhline(y=batch_size * top_k / num_experts, color='red',
                linestyle='--', label=f'Ideal ({batch_size*top_k/num_experts:.0f})')
axes[0].set_xlabel('Expert Index')
axes[0].set_ylabel('Number of Tokens Assigned')
axes[0].set_title('Token Assignment per Expert')
axes[0].legend()
axes[0].grid(True, alpha=0.3, axis='y')

# 中央: ゲーティング確率のヒートマップ（上位50トークン）
im = axes[1].imshow(gate_probs[:50], aspect='auto', cmap='YlOrRd')
axes[1].set_xlabel('Expert Index')
axes[1].set_ylabel('Token Index')
axes[1].set_title('Gating Probabilities (first 50 tokens)')
plt.colorbar(im, ax=axes[1])

# 右: エキスパートペアの共起頻度
pair_counts = np.zeros((num_experts, num_experts))
for i in range(batch_size):
    for j in range(top_k):
        for k_idx in range(j + 1, top_k):
            e1, e2 = sorted([indices[i, j], indices[i, k_idx]])
            pair_counts[e1, e2] += 1

# 対称化
pair_counts = pair_counts + pair_counts.T
im2 = axes[2].imshow(pair_counts, cmap='Blues')
axes[2].set_xlabel('Expert Index')
axes[2].set_ylabel('Expert Index')
axes[2].set_title('Expert Pair Co-occurrence')
plt.colorbar(im2, ax=axes[2])

plt.tight_layout()
plt.savefig('moe_routing.png', dpi=150, bbox_inches='tight')
plt.show()

def simulate_routing_with_balance(n_steps=500, alpha=0.0):
    """負荷分散損失ありなしでルーティングの偏りをシミュレーション。"""
    W_gate = np.random.randn(d_model, num_experts) * 0.1
    lr = 0.01

    imbalance_history = []

    for step in range(n_steps):
        # ランダムバッチ
        x_batch = np.random.randn(64, d_model) * 0.5

        # ゲーティング
        logits = x_batch @ W_gate
        logits_max = np.max(logits, axis=1, keepdims=True)
        probs = np.exp(logits - logits_max) / np.sum(np.exp(logits - logits_max), axis=1, keepdims=True)

        # Top-k
        top_indices = np.argsort(-probs, axis=1)[:, :top_k]

        # 負荷の不均衡度
        counts = np.zeros(num_experts)
        for i in range(64):
            for j in range(top_k):
                counts[top_indices[i, j]] += 1
        counts /= 64
        imbalance = np.std(counts) / np.mean(counts)  # 変動係数
        imbalance_history.append(imbalance)

        # 勾配更新（簡略化: 負荷分散損失の勾配のみ）
        if alpha > 0:
            f = counts
            p = np.mean(probs, axis=0)
            # 負荷分散損失の勾配（近似）
            grad = alpha * num_experts * np.outer(np.mean(x_batch, axis=0), f * p)
            W_gate -= lr * grad

    return imbalance_history

# 比較
imbalance_no_balance = simulate_routing_with_balance(alpha=0.0)
imbalance_with_balance = simulate_routing_with_balance(alpha=0.1)

plt.figure(figsize=(8, 5))
plt.plot(imbalance_no_balance, label='No load balancing (α=0)', alpha=0.8)
plt.plot(imbalance_with_balance, label='With load balancing (α=0.1)', alpha=0.8)
plt.xlabel('Training Step')
plt.ylabel('Routing Imbalance (CV)')
plt.title('Effect of Load Balancing Loss')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('moe_load_balance.png', dpi=150, bbox_inches='tight')
plt.show()