mixture of experts training

Mixture of Experts (MoE) Training is the specialized training methodology for sparse conditional computation models where only a subset of parameters (experts) are activated per input — requiring careful handling of expert load balancing, routing stability, communication patterns across devices, and auxiliary losses to prevent expert collapse, with techniques like expert parallelism, top-k gating, and capacity factors enabling models like Mixtral 8x7B, GPT-4 (rumored MoE), and Switch Transformer to achieve dense-model quality at a fraction of the per-token compute cost.

MoE Architecture

Standard Transformer FFN:
  x → [FFN: 4096 → 16384 → 4096] → y
  Every token uses ALL parameters

MoE Layer (8 experts, top-2 routing):
  x → [Router/Gate network] → selects Expert 3 and Expert 7
  x → [Expert 3: 4096 → 16384 → 4096] × w_3
    + [Expert 7: 4096 → 16384 → 4096] × w_7 → y
  Each token uses only 2 of 8 experts (25% of FFN params)

Key Training Challenges

Load Balancing Loss

# Auxiliary loss to encourage balanced expert usage
def load_balance_loss(router_probs, expert_indices, num_experts):
    # f_i = fraction of tokens routed to expert i
    # p_i = average router probability for expert i
    f = torch.zeros(num_experts)
    p = torch.zeros(num_experts)
    for i in range(num_experts):
        mask = (expert_indices == i).float()
        f[i] = mask.mean()
        p[i] = router_probs[:, i].mean()
    # Loss encourages uniform f_i (each expert gets equal tokens)
    return num_experts * (f * p).sum()

Expert Parallelism

8 GPUs, 8 experts, 4-way data parallel:

GPU 0: Expert 0,1  |  Tokens from all GPUs routed to Exp 0,1
GPU 1: Expert 2,3  |  Tokens from all GPUs routed to Exp 2,3
GPU 2: Expert 4,5  |  Tokens from all GPUs routed to Exp 4,5
GPU 3: Expert 6,7  |  Tokens from all GPUs routed to Exp 6,7
GPU 4-7: Duplicate of GPU 0-3 (data parallel)

all-to-all communication: Each GPU sends tokens to correct expert GPU

MoE Model Comparison

Capacity Factor and Token Dropping

• Capacity factor C: Maximum tokens per expert = C × (total_tokens / num_experts).
• C = 1.0: Perfect balance, may drop tokens if routing is uneven.
• C = 1.25: 25% buffer for imbalance (common choice).
• Dropped tokens: Skip the MoE layer, use residual connection only.
• Training: Some dropping is acceptable. Inference: Never drop (use auxiliary buffer).

Training Tips

• Router z-loss: Penalize large logits to stabilize gating → prevents routing oscillation.
• Expert jitter: Add small noise to router inputs during training → prevents collapse.
• Gradient scaling: Scale expert gradients by 1/num_selected_experts.
• Initialization: Initialize router weights small → initially uniform routing → gradual specialization.

MoE training is the methodology that enables trillion-parameter models with affordable compute — by activating only a fraction of parameters per token and carefully managing expert load balancing, routing stability, and communication across devices, MoE architectures achieve the quality of dense models 5-10× larger while requiring only the inference compute of much smaller models, making them the dominant architecture choice for frontier language models.