Mixture of Experts (MoE)

Standard Transformer Layer:
Input → Self-Attention (always on)
      → Feed-Forward Network (always on)
      → Output

MoE Layer:
Input → Self-Attention (always on)
      → Gate/Router (decides which experts)
        ├─ Expert 1 (only if needed)
        ├─ Expert 2 (only if needed)
        ├─ Expert 3 (only if needed)
        └─ Expert N (only if needed)
      → Combine outputs
      → Output

Input: [batch, seq, d_model]

Router: Linear layer
scores = Linear(input) @ expert_weights

Gate: Softmax
probs = softmax(scores)

Top-K Selection (e.g., K=2):
top_2_probs, top_2_indices = topk(probs, k=2)

Weighted Combination:
output = Σ top_2_probs[i] * Expert[top_2_indices[i]](input)

Standard 70B Model:
- All parameters → all operations
- 70B multiply-accumulate per token

MoE 56B Model (Mixtral 8×7B):
- Attention: 56B operations (same)
- FFN: Only 2 of 8 experts (2/8 × 56B ≈ 14B)
- Total: ~70B but more efficient!

Inference Speed:
- Standard 70B: 60 tokens/sec
- MoE 56B: 120 tokens/sec (2× faster!)
- Same memory footprint!

Idea: Different experts specialize

Expert 1: Grammar
Expert 2: Medical knowledge
Expert 3: Code generation
Expert 4: Math reasoning
...
Expert 8: Common sense

Router learns: Medical question → activate Expert 2

Empirically: Hard to prove, but plausible

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    torch_dtype=torch.float16,
)

# Architecture:
# - 8 Experts (7B each)
# - Top-2 routing (2 of 8 activate)
# - Quasi-Parameters: 56B
# - Effective: 14B active per token

DeepSeek-MoE-16B:
- 16 Experts (2B each)
- Top-2 routing
- Only 2.8B parameters active per token
- Performance: ≈ 7B standard models
- Speed: 3× faster!

Flagship: DeepSeek-MoE-145B (very good quality)

Speculation from OpenAI papers + hints:
- Likely MoE architecture
- Size: ~1.8T parameters
- Active per token: ~100-300B

Evidence:
- Very fast inference (suggests sparse activation)
- Costs consistent despite 10× larger than GPT-3.5
- OpenAI published MoE scaling papers

Conclusion: Likely true but unconfirmed

Problem: All tokens activate same experts
         Others never used

Result: Ineffective → just 2-expert model

Solution: Add balance loss
Loss_total = Main_Loss + λ × Balance_Loss
             where Balance_Loss = Std_Dev(expert_usage)

Forces fair distribution

Operation: Gather/scatter outputs from experts

GPU Overhead:
- Custom CUDA kernels needed
- vLLM: Has these → 95% theoretical speedup
- Others: Only 50-70% speedup

Router can "jump" (suddenly activate different experts)
→ Gradient explosions

Solutions:
- Gradient clipping
- Learning rate scheduling
- Auxiliary losses

MoE models need more training care than dense

1. Inference-speed critical (chatbots <100ms latency)
   → Mixtral 8×7B 2× faster

2. Broad knowledge (search, general assistant)
   → Experts can specialize

3. Budget-limited inference
   → 70% compute savings
   → 70% infrastructure cost reduction

4. Hardware bottleneck
   → Smaller footprint with sparse activation

1. Your own training (custom MoE too complex)
2. Latency must be consistent (variable routing)
3. Specialized narrow tasks (single expert better)
4. Low-precision deployment (quantization trickier)

Why:
1. Inference efficiency (2-3× speedup)
2. Training efficiency (not much harder)
3. Scaling better to huge models
4. Hardware trends (more parallelism)

Predictions:
- 2026: More open-source MoE models
- 2027: MoE standard for >100B models
- 2028: Hybrid dense-MoE (best of both)

For production: Mixtral 8×7B ready now.
For training: Still early stage.

import torch
import torch.nn as nn

class MoELayer(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Create N experts (each is a small FFN)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_size, hidden_size * 4),
                nn.ReLU(),
                nn.Linear(hidden_size * 4, hidden_size),
            )
            for _ in range(num_experts)
        ])

        # Gating network (router)
        self.gate = nn.Linear(hidden_size, num_experts)

    def forward(self, x):
        batch_size, seq_len, hidden_size = x.shape

        # Reshape for batch processing
        x_flat = x.view(-1, hidden_size)  # [batch*seq, hidden]

        # Compute routing scores
        scores = self.gate(x_flat)  # [batch*seq, num_experts]

        # Select top-k experts
        top_k_scores, top_k_indices = torch.topk(scores, self.top_k, dim=-1)
        # top_k_scores: [batch*seq, top_k]
        # top_k_indices: [batch*seq, top_k]  indices of which experts

        # Normalize scores (softmax) for weighting
        weights = torch.softmax(top_k_scores, dim=-1)  # [batch*seq, top_k]

        # Gather expert outputs
        output = torch.zeros_like(x_flat)  # [batch*seq, hidden]

        for i, expert in enumerate(self.experts):
            mask = (top_k_indices == i).any(dim=-1)  # Which tokens use expert i

            if mask.any():
                expert_output = expert(x_flat[mask])  # Only process selected tokens
                weight_i = (weights * (top_k_indices == i).float()).sum(dim=-1, keepdim=True)
                output[mask] += weight_i[mask] * expert_output

        # Reshape back
        return output.view(batch_size, seq_len, hidden_size)

Scenario:
8 experts, but tokens only route to experts 1 and 2.
Experts 3-8 unused → wasted parameters.

Solution: Add auxiliary loss to encourage balanced routing

def moe_loss_with_balance(gate_scores, num_experts, top_k):
    """
    gate_scores: [batch_size, num_experts]
    Maximize performance AND balance across experts
    """

    # Main loss (standard cross-entropy or similar)
    # loss_main = ...

    # Balance loss: Encourage uniform expert selection
    expert_selection = torch.softmax(gate_scores, dim=-1)  # [batch, experts]

    # Expected selection per expert
    mean_selection = expert_selection.mean(dim=0)  # [experts]

    # Variance (low variance = balanced, high variance = collapse)
    balance_loss = torch.var(mean_selection)

    # Combined loss
    total_loss = loss_main + 0.01 * balance_loss
    # 0.01 weight (tune this)

    return total_loss

# During training:
# - If balance_loss increases → experts not balanced → increase weight to 0.02
# - If performance drops too much → decrease weight to 0.005

class CapacityMoE(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k, capacity_factor=1.25):
        super().__init__()
        self.capacity_factor = capacity_factor
        self.num_experts = num_experts
        self.top_k = top_k

    def forward(self, x, max_tokens=None):
        """
        capacity_factor=1.25 means each expert can handle 125% of avg load
        If exceeded → overflow handling (drop or reassign)
        """

        batch_size, seq_len, hidden_size = x.shape
        num_tokens = batch_size * seq_len

        # Capacity per expert
        capacity_per_expert = int((num_tokens * self.capacity_factor) / self.num_experts)

        scores = self.gate(x)
        top_k_scores, top_k_indices = torch.topk(scores, self.top_k, dim=-1)

        # Check if capacity exceeded
        expert_load = torch.bincount(top_k_indices.flatten(), minlength=self.num_experts)
        overloaded = expert_load > capacity_per_expert

        if overloaded.any():
            # Option 1: Drop tokens (simple but loses info)
            # Option 2: Reassign to less-loaded expert (better)
            # Option 3: Use secondary expert set (complex)
            print(f"Warning: {overloaded.sum()} experts overloaded")

        # Continue with standard routing...

Do you have inference latency constraints (< 100ms)?
  ├─ YES → Use MoE (Mixtral)
  │   └─ Do you have >1M tokens/sec to process?
  │       ├─ YES → Deploy Mixtral locally
  │       └─ NO → Cloud API fine
  │
  └─ NO → Use dense model
      └─ Do you need maximum quality?
          ├─ YES → Claude/GPT-4 (proprietary)
          └─ NO → Llama 3 70B (dense, cheaper)

Scenario: 1 million API calls/month, 300 tokens per call

Option 1: Mixtral 8×7B (Local)
  - Hardware: RTX 4090 (€3K), storage (€500)
  - Electricity: €200/month
  - Cost per 1M calls: €0.30 (amortized)

Option 2: Claude 3.5 (API)
  - €0.003 per 1K input tokens
  - 300 tokens × 1M = 300B tokens
  - Cost: €900/month

Option 3: Llama 3 70B (Local)
  - Hardware: A100 (€8K), storage (€500)
  - Electricity: €500/month
  - Cost per 1M calls: €0.50 (amortized)

Recommendation:
- < 100K calls/month: Use API (setup not worth it)
- 100K-1M calls/month: MoE local (best efficiency)
- > 1M calls/month: Dense local (higher absolute cost, lower per-unit)

Realistic scenarios requiring custom MoE:
1. Specialized domain with millions of training examples
2. Hardware constraints (limited GPUs but lots of tokens)
3. Model size exactly matches your VRAM
4. Need to specialize experts per-domain

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.optim import AdamW

def train_moe_model():
    """Minimal MoE training loop."""

    # Create MoE model (custom, not from transformers)
    model = MoELanguageModel(
        vocab_size=50000,
        hidden_size=4096,
        num_experts=8,
        top_k=2,
    )

    optimizer = AdamW(model.parameters(), lr=1e-4)
    tokenizer = AutoTokenizer.from_pretrained("llama-2-7b")

    # Training loop
    for epoch in range(3):
        for batch in dataloader:
            # Forward pass
            input_ids = batch["input_ids"]
            attention_mask = batch["attention_mask"]

            logits = model(input_ids, attention_mask)

            # Loss
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = input_ids[..., 1:].contiguous()

            loss = torch.nn.functional.cross_entropy(
                shift_logits.view(-1, vocab_size),
                shift_labels.view(-1)
            )

            # Balance loss (critical for MoE)
            balance_loss = model.get_balance_loss()
            total_loss = loss + 0.01 * balance_loss

            # Backward
            optimizer.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            print(f"Loss: {total_loss:.4f}, Balance: {balance_loss:.4f}")

        # Save checkpoint
        torch.save(model.state_dict(), f"moe_model_epoch{epoch}.pt")

# Issues you'll hit:
# - Router gradient explosion (clip_grad_norm helps)
# - Load imbalance (balance_loss helps)
# - Slower training than dense (expert overhead)
# - Harder to debug (which expert failed?)

def diagnose_router_collapse(model, dataloader):
    """Detect if router is collapsed."""

    expert_usage = torch.zeros(model.num_experts)

    with torch.no_grad():
        for batch in dataloader:
            _, routing_indices = model.get_routing(batch["input_ids"])
            # routing_indices shape: [batch_size, seq_len, top_k]

            for expert_id in range(model.num_experts):
                expert_usage[expert_id] += (routing_indices == expert_id).sum().item()

    # Ideal: Each expert used equally (expert_usage / total ≈ 1 / num_experts)
    expected_usage = expert_usage.sum() / model.num_experts

    for i, usage in enumerate(expert_usage):
        percentage = usage / expert_usage.sum() * 100
        if percentage > 20:  # One expert doing 20%+ of work
            print(f"WARNING: Expert {i} is overused ({percentage:.1f}%)")

    return expert_usage

# Fix: Increase balance loss weight during training

Expected speedup: 2×
Actual speedup: 1.2× (disappointing!)

Root cause: Gathering/scattering outputs between experts is slow

Solution:
- Use optimized kernels (vLLM, flash-attn)
- Batch multiple sequences together
- Use GPU-specific optimizations (NVIDIA Apex)

MoE latency varies per request (some tokens hit slow experts)

Problem: Real-time apps need predictable latency

Solutions:
1. Predetermine expert schedule (trade flexibility for consistency)
2. Use only fast experts (subset of 8)
3. Add fixed computational delay (padding) for consistency

1. Hybrid Dense-MoE
   - Dense attention (traditional)
   - MoE feed-forward (experts)
   - Best of both worlds

2. Conditional Computation
   - Experts trained conditionally
   - Token-level decisions more sophisticated
   - Closer to biological brains

3. Grouped Query Attention + MoE
   - Both sparse attention and sparse FFN
   - 3-4× speedups possible

4. Parameter Sharing Across Experts
   - Don't need 8 independent experts
   - Shared base + small differences
   - 50% parameter reduction

❌ Fully differentiable routing (too unstable)
❌ Infinite experts (diminishing returns proven)
❌ MoE for embeddings (too small, overhead kills benefit)
❌ Adaptive expert count (training gets unstable)