Mixture of Experts (MoE)

Architecture, routing, training, and inference optimization for sparse expert models.

MoE Fundamentals

Architecture

Input → Router → Top-k Experts → Combine → Output
         │
         ├── Expert 1 (FFN)
         ├── Expert 2 (FFN)
         ├── ...
         └── Expert N (FFN)

Key Equations

# Router logits
router_logits = x @ W_router^T

# Top-k selection
top_k_indices, top_k_weights = top_k(router_logits, k)

# Expert outputs
expert_outputs = [expert_i(x) for i in top_k_indices]

# Weighted combination
output = Σ(top_k_weights[i] * expert_outputs[i])

Sparsity Benefits

Routing Mechanisms

Top-k Gating

class TopKRouter:
    def __init__(self, d_model: int, num_experts: int, k: int = 2):
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.k = k

    def forward(self, x):
        logits = self.router(x)  # [batch, seq, num_experts]
        weights, indices = torch.topk(logits, self.k, dim=-1)
        weights = F.softmax(weights, dim=-1)
        return weights, indices

Load Balancing (Critical)

# Auxiliary loss: encourage uniform expert usage
def load_balancing_loss(router_logits, top_k_indices):
    # Fraction of tokens routed to each expert
    expert_mask = F.one_hot(top_k_indices, num_classes=num_experts).float()
    tokens_per_expert = expert_mask.mean(dim=[0, 1])

    # Router probabilities
    router_probs = F.softmax(router_logits, dim=-1).mean(dim=[0, 1])

    # Loss: minimize coefficient of variation
    return num_experts * (tokens_per_expert * router_probs).sum()

# Z-loss (logit normalization)
def z_loss(router_logits):
    return torch.logsumexp(router_logits, dim=-1).mean() ** 2

Router Variants

Expert Architectures

Standard (FFN Experts)

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.gate = nn.Linear(d_model, d_ff, bias=False)  # SwiGLU

    def forward(self, x):
        return self.w2(F.silu(self.gate(x)) * self.w1(x))

Shared + Specialized

class SharedExpertMoE(nn.Module):
    def __init__(self, d_model, num_experts, shared_experts=1):
        self.shared = nn.ModuleList([Expert(d_model, d_ff) for _ in range(shared_experts)])
        self.routed = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])
        self.router = Router(d_model, num_experts)

    def forward(self, x):
        # Shared always active
        shared_out = sum(e(x) for e in self.shared)
        # Routed top-k
        routed_out = self.moe_forward(x)
        return shared_out + routed_out

DeepSeek MoE (Fine-Grained)

# DeepSeek-V2/V3: 160-256 experts, 6-8 active
# Key innovations:
# 1. Fine-grained experts (smaller d_ff)
# 2. Shared experts (always active, capture common knowledge)
# 3. Auxiliary-loss-free load balancing (device-level)
# 4. MLA (Multi-head Latent Attention) for KV compression

class DeepSeekMoE(nn.Module):
    def __init__(self, config):
        self.num_experts = 160  # Fine-grained
        self.num_shared = 2     # Always active
        self.k = 6              # Top-6 routed
        # ... implementation

Training MoE

Parallelism Strategy

# Expert Parallelism (EP)
# Each expert on different device
# Data parallel within expert

# Tensor Parallelism (TP) within expert
# Sequence Parallelism (SP)

# Pipeline Parallelism (PP) across layers

# 3D Parallelism: DP × TP × EP
# DeepSeek-V3: 2048 GPUs, 3D parallelism

Training Stability

# Gradient clipping per expert
def clip_expert_gradients(model, max_norm=1.0):
    for expert in model.experts:
        nn.utils.clip_grad_norm_(expert.parameters(), max_norm)

# Router gradient scaling
router_grad_scale = 1.0 / math.sqrt(num_experts)

Communication Optimization

# All-to-all for token dispatch
# Overlap computation with communication
async def moe_forward(x):
    # 1. Compute router logits
    logits = router(x)

    # 2. Dispatch tokens (async all-to-all)
    dispatch_future = async_all_to_all(x, indices)

    # 3. Compute shared experts while waiting
    shared_out = shared_experts(x)

    # 4. Wait for routed tokens
    routed_x = await dispatch_future

    # 5. Compute routed experts
    routed_out = routed_experts(routed_x)

    # 6. Combine
    return shared_out + combine(routed_out, weights)

Inference Optimization

Expert Caching

class ExpertCache:
    def __init__(self, max_experts_in_memory=8):
        self.cache = {}  # expert_id → (weights, device)
        self.max_size = max_experts_in_memory
        self.access_order = []

    def get(self, expert_id):
        if expert_id in self.cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.cache[expert_id]

        # Load from disk
        weights = load_expert(expert_id)
        self.evict_if_needed()
        self.cache[expert_id] = weights
        self.access_order.append(expert_id)
        return weights

Speculative Expert Activation

# Predict which experts needed next
class SpeculativeRouter:
    def __init__(self, router, predictor):
        self.router = router
        self.predictor = predictor  # Small model predicting next experts

    def forward(self, x, prev_experts=None):
        # Use predictor to pre-load experts
        if prev_experts:
            predicted = self.predictor(prev_experts)
            preload_experts(predicted)

        return self.router(x)

Quantization

# Expert-specific quantization
def quantize_moe(model):
    for i, expert in enumerate(model.experts):
        # High-traffic experts: higher precision
        if expert.traffic > threshold:
            quantize(expert, bits=8)  # INT8
        else:
            quantize(expert, bits=4)  # INT4

    # Router: keep FP16/BF16
    quantize(model.router, bits=16)

MoE vs Dense: When to Use

Notable MoE Models

Related Notes

• Model Routing — Routing between models
• Inference Optimization — MoE inference
• Small Language Models — Distillation from MoE
• Reasoning Models — MoE for reasoning
• Cost Optimization — MoE cost benefits
• Scaling Strategies — MoE parallelism
• Research Directions — MoE research frontiers

Sources: "GShard" (Lepikhin et al.), "Switch Transformers" (Fedus et al.), "Mixtral" (Mistral AI), "DeepSeek-V2/V3" (DeepSeek), "GLaM" (Google), "MoE" surveys, "Expert Choice Routing", "Soft MoE", "Scaling MoE" papers, vLLM/SGLang MoE support.