Mixture of Experts（MoE）アーキテクチャ論文レビューと実践スケーリング 2026

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

class Expert(nn.Module):
    """単一エキスパートFFNモジュール"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class MoELayer(nn.Module):
    """Mixture of Expertsレイヤー"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        top_k: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        # ゲーティングネットワーク: 入力 -> エキスパートスコア
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # (B*S, d_model)

        # ゲーティングスコア計算
        gate_logits = self.gate(x_flat)  # (B*S, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-Kエキスパート選択
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )  # (B*S, top_k)
        # 選択されたエキスパート重みの再正規化
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # エキスパート別出力の合算
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]  # (B*S,)
            weight = top_k_probs[:, k].unsqueeze(-1)  # (B*S, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output.view(batch_size, seq_len, d_model)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TopKRouter(nn.Module):
    """Top-Kルーティング with ノイズ注入"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 2,
        noise_std: float = 1.0,
    ):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch * seq_len, d_model)
        Returns:
            top_k_indices: (batch * seq_len, top_k)
            top_k_weights: (batch * seq_len, top_k)
            gate_logits: (batch * seq_len, num_experts)
        """
        gate_logits = self.gate(x)

        # 学習時にノイズ注入で探索を促進
        if self.training:
            noise = torch.randn_like(gate_logits) * self.noise_std
            gate_logits = gate_logits + noise

        gate_probs = F.softmax(gate_logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )
        # 再正規化
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_indices, top_k_weights, gate_logits

def load_balancing_loss(
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int = 2,
) -> torch.Tensor:
    """
    Switch Transformerスタイルのロードバランシング損失計算

    各エキスパートに割り当てられたトークンの比率(f_i)と
    ゲート確率の平均(p_i)の内積を最小化する。
    均等分布の時に最小値を持つ。

    Args:
        gate_logits: (num_tokens, num_experts)
        top_k_indices: (num_tokens, top_k)
        num_experts: エキスパート数
        top_k: 活性エキスパート数
    Returns:
        スカラーロードバランシング損失
    """
    num_tokens = gate_logits.shape[0]
    gate_probs = F.softmax(gate_logits, dim=-1)

    # f_i: 各エキスパートに割り当てられたトークン比率
    # one-hotエンコーディングで割り当てカウント
    expert_mask = F.one_hot(
        top_k_indices, num_classes=num_experts
    ).float()  # (num_tokens, top_k, num_experts)
    expert_mask = expert_mask.sum(dim=1)  # (num_tokens, num_experts)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # p_i: 各エキスパートに対する平均ゲート確率
    p = gate_probs.mean(dim=0)  # (num_experts,)

    # バランシング損失: num_experts * sum(f_i * p_i)
    balance_loss = num_experts * (f * p).sum()

    return balance_loss

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """
    ST-MoE Router Z-Loss実装

    ルーターロジットのlog-sum-exp値の二乗平均を最小化。
    ロジットが過度に大きくなることを防止して学習安定性を確保。

    Args:
        gate_logits: (num_tokens, num_experts)
    Returns:
        スカラーz-loss値
    """
    # log(sum(exp(x)))の二乗
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss


# 学習ループでの統合
def compute_moe_loss(
    task_loss: torch.Tensor,
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001,
) -> torch.Tensor:
    """
    全体MoE損失 = タスク損失 + バランシング損失 + Z-Loss
    """
    bal_loss = load_balancing_loss(
        gate_logits, top_k_indices, num_experts
    )
    z_loss = router_z_loss(gate_logits)

    total_loss = task_loss + alpha_balance * bal_loss + alpha_z * z_loss
    return total_loss

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3スタイル Auxiliary-Loss-Freeルーター"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 8,
        bias_update_speed: float = 0.001,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = bias_update_speed
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # バイアスはグラディエント更新ではなくルールベースで調整
        self.register_buffer(
            "expert_bias", torch.zeros(num_experts)
        )
        self.register_buffer(
            "target_load", torch.ones(num_experts) / num_experts
        )

    def forward(self, x: torch.Tensor):
        gate_logits = self.gate(x)
        # バイアスを加えてルーティング決定に反映
        biased_logits = gate_logits + self.expert_bias.unsqueeze(0)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K選択はバイアス含有ロジットで実行
        _, top_k_indices = torch.topk(
            biased_logits, self.top_k, dim=-1
        )
        # 重みは元のゲート確率から抽出
        top_k_probs = torch.gather(gate_probs, 1, top_k_indices)
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 学習時にバイアス更新（グラディエント不要）
        if self.training:
            with torch.no_grad():
                expert_load = F.one_hot(
                    top_k_indices.view(-1), self.num_experts
                ).float().sum(0)
                expert_load = expert_load / expert_load.sum()
                # 過負荷エキスパートはバイアス減少、低負荷は増加
                self.expert_bias += self.gamma * (
                    self.target_load - expert_load
                )

        return top_k_indices, top_k_weights, gate_logits

from dataclasses import dataclass
from typing import Dict, Optional
import torch

@dataclass
class ExpertCacheConfig:
    """エキスパートキャッシュ設定"""
    gpu_cache_size: int = 4   # GPUに常駐するエキスパート数
    prefetch_count: int = 2    # 先行ロードするエキスパート数
    device: str = "cuda:0"

class ExpertOffloadManager:
    """
    エキスパートオフローディング管理者。
    LRUベースでGPUキャッシュを管理し、
    次レイヤーのエキスパートを先行ロードする。
    """
    def __init__(
        self,
        experts: Dict[int, torch.nn.Module],
        config: ExpertCacheConfig,
    ):
        self.experts = experts
        self.config = config
        self.gpu_cache: Dict[int, torch.nn.Module] = {}
        self.access_order: list = []

    def get_expert(self, expert_id: int) -> torch.nn.Module:
        """エキスパートをGPUにロード（LRUキャッシュ）"""
        if expert_id in self.gpu_cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.gpu_cache[expert_id]

        # GPUキャッシュが満杯の場合、LRUエキスパートを退避
        if len(self.gpu_cache) >= self.config.gpu_cache_size:
            evict_id = self.access_order.pop(0)
            self.gpu_cache[evict_id].cpu()
            del self.gpu_cache[evict_id]

        # エキスパートをGPUに移動
        expert = self.experts[expert_id].to(self.config.device)
        self.gpu_cache[expert_id] = expert
        self.access_order.append(expert_id)
        return expert

    def prefetch(self, predicted_expert_ids: list):
        """次レイヤーで使用されるエキスパートを非同期で先行ロード"""
        for eid in predicted_expert_ids[:self.config.prefetch_count]:
            if eid not in self.gpu_cache:
                self.get_expert(eid)