Mixture of Experts (Sparse MoE)

1. The Mixture of Experts Framework

A Mixture of Experts (MoE) model combines multiple expert networks with a gating mechanism:

where is the number of experts and is the gating weight for expert .

1.1 Sparse vs. Dense MoE

1.2 The Efficiency Advantage

For experts, active:

• Parameters: single expert
• Compute: single expert
• Memory: single expert

2. MoE Architecture with Router

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  {/* Title */}
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Mixture of Experts Architecture</text>
  
  {/* Input */}
  <g transform="translate(50, 150)">
    <rect width="80" height="100" rx="8" fill="#f8fafc" stroke="#3b82f6" stroke-width="2"/>
    <text x="40" y="45" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10" font-weight="bold">Input</text>
    <text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">x ∈ ℝᵈ</text>
  </g>
  
  {/* Router/Gate */}
  <g transform="translate(170, 150)">
    <rect width="100" height="100" rx="8" fill="#f8fafc" stroke="#a855f7" stroke-width="2"/>
    <text x="50" y="35" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10" font-weight="bold">Router</text>
    <text x="50" y="55" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">G(x) = softmax</text>
    <text x="50" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">(Wg · x)</text>
    <text x="50" y="90" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="8">Top-k selection</text>
  </g>
  
  {/* Experts */}
  <g transform="translate(320, 50)">
    {/* Expert 1 */}
    <rect width="100" height="60" rx="5" fill="#f8fafc" stroke="#22c55e" stroke-width="2"/>
    <text x="50" y="25" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">Expert 1</text>
    <text x="50" y="45" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">FFN(x)</text>
    
    {/* Expert 2 */}
    <g transform="translate(0, 80)">
      <rect width="100" height="60" rx="5" fill="#f8fafc" stroke="#f59e0b" stroke-width="2"/>
      <text x="50" y="25" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10">Expert 2</text>
      <text x="50" y="45" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">FFN(x)</text>
    </g>
    
    {/* Expert 3 */}
    <g transform="translate(0, 160)">
      <rect width="100" height="60" rx="5" fill="#f8fafc" stroke="#3b82f6" stroke-width="2"/>
      <text x="50" y="25" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10">Expert 3</text>
      <text x="50" y="45" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">FFN(x)</text>
    </g>
    
    {/* Expert 4 */}
    <g transform="translate(0, 240)">
      <rect width="100" height="60" rx="5" fill="#f8fafc" stroke="#ec4899" stroke-width="2"/>
      <text x="50" y="25" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="10">Expert 4</text>
      <text x="50" y="45" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">FFN(x)</text>
    </g>
    
    {/* ... */}
    <text x="50" y="325" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="10">⋮</text>
    
    {/* Expert E */}
    <g transform="translate(0, 340)">
      <rect width="100" height="60" rx="5" fill="#f8fafc" stroke="#a855f7" stroke-width="2"/>
      <text x="50" y="25" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10">Expert E</text>
      <text x="50" y="45" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">FFN(x)</text>
    </g>
  </g>
  
  {/* Output */}
  <g transform="translate(620, 180)">
    <rect width="80" height="100" rx="8" fill="#f8fafc" stroke="#f59e0b" stroke-width="2"/>
    <text x="40" y="45" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10" font-weight="bold">Output</text>
    <text x="40" y="70" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">y ∈ ℝᵈ</text>
  </g>
  
  {/* Arrows */}
  <defs>
    <marker id="f6b5_arrowBlue5" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#60a5fa"/>
    </marker>
    <marker id="f6b5_arrowPurple5" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#c084fc"/>
    </marker>
    <marker id="f6b5_arrowGreen5" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#22c55e"/>
    </marker>
  </defs>
  
  <line x1="130" y1="200" x2="170" y2="200" stroke="#60a5fa" stroke-width="2" marker-end="url(#f6b5_arrowBlue5)"/>
  
  {/* Router to experts */}
  <line x1="270" y1="180" x2="320" y2="80" stroke="#c084fc" stroke-width="2" marker-end="url(#f6b5_arrowPurple5)"/>
  <line x1="270" y1="200" x2="320" y2="200" stroke="#c084fc" stroke-width="2" marker-end="url(#f6b5_arrowPurple5)"/>
  
  {/* Experts to output */}
  <line x1="420" y1="80" x2="620" y2="220" stroke="#22c55e" stroke-width="2" marker-end="url(#f6b5_arrowGreen5)"/>
  <line x1="420" y1="200" x2="620" y2="220" stroke="#22c55e" stroke-width="2" marker-end="url(#f6b5_arrowGreen5)"/>
  
  {/* Active indicator */}
  <circle cx="315" cy="80" r="4" fill="#22c55e"/>
  <circle cx="315" cy="200" r="4" fill="#f59e0b"/>
  <circle cx="315" cy="280" r="4" fill="#3b82f6" opacity="0.3"/>
  <circle cx="315" cy="360" r="4" fill="#ec4899" opacity="0.3"/>
  
  {/* Legend */}
  <g transform="translate(50, 350)">
    <rect width="200" height="40" rx="5" fill="#f8fafc" stroke="#334155"/>
    <circle cx="20" cy="20" r="5" fill="#22c55e"/>
    <text x="35" y="24" fill="#94a3b8" font-family="monospace" font-size="8">Active expert (selected)</text>
    <circle cx="130" cy="20" r="5" fill="#3b82f6" opacity="0.3"/>
    <text x="145" y="24" fill="#94a3b8" font-family="monospace" font-size="8">Inactive</text>
  </g>
</svg>

3. Gating Function

3.1 Top-k Routing

The gating function selects experts per token:

where is the gating weight matrix.

3.2 Gating Weights

3.3 Switch Transformer (Fedus et al., 2022)

Top-1 routing: Each token is routed to exactly one expert:

The output:

4. Load Balancing

4.1 The Load Balancing Problem

Without constraints, the router may collapse to sending all tokens to a few experts:

4.2 Auxiliary Load Balancing Loss

where:

• is the fraction of tokens routed to expert
• is the average gating probability for expert
• is the balancing coefficient (typically )

4.3 Expert Capacity

The expert capacity limits the number of tokens each expert can process:

where:

• is the total number of tokens in the batch
• is the number of experts
• is the number of experts per token
• is the capacity factor (typically )

4.4 Token Dropping

When an expert is full, tokens are dropped:

5. Load Balancing Visualization

<svg viewBox="0 0 800 400" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  {/* Title */}
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Load Balancing in Mixture of Experts</text>
  
  {/* Unbalanced */}
  <g transform="translate(50, 60)">
    <text x="175" y="0" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="12" font-weight="bold">Unbalanced (No Loss)</text>
    
    {/* Expert bars */}
    <rect x="30" y="20" width="80" height="200" rx="5" fill="#ec4899" opacity="0.8"/>
    <text x="70" y="230" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="10">75%</text>
    
    <rect x="130" y="150" width="80" height="70" rx="5" fill="#f59e0b" opacity="0.8"/>
    <text x="170" y="230" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10">15%</text>
    
    <rect x="230" y="180" width="80" height="40" rx="5" fill="#3b82f6" opacity="0.8"/>
    <text x="270" y="230" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10">8%</text>
    
    <rect x="330" y="195" width="80" height="25" rx="5" fill="#a855f7" opacity="0.8"/>
    <text x="370" y="230" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10">2%</text>
    
    {/* Labels */}
    <text x="70" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E1</text>
    <text x="170" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E2</text>
    <text x="270" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E3</text>
    <text x="370" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E4</text>
    
    <text x="175" y="250" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="9">⚠ Expert collapse!</text>
    <text x="175" y="270" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Most tokens → Expert 1</text>
  </g>
  
  {/* Balanced */}
  <g transform="translate(450, 60)">
    <text x="175" y="0" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="12" font-weight="bold">Balanced (With Loss)</text>
    
    {/* Expert bars */}
    <rect x="30" y="80" width="80" height="140" rx="5" fill="#22c55e" opacity="0.8"/>
    <text x="70" y="230" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">28%</text>
    
    <rect x="130" y="75" width="80" height="145" rx="5" fill="#22c55e" opacity="0.8"/>
    <text x="170" y="230" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">29%</text>
    
    <rect x="230" y="80" width="80" height="140" rx="5" fill="#22c55e" opacity="0.8"/>
    <text x="270" y="230" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">28%</text>
    
    <rect x="330" y="78" width="80" height="142" rx="5" fill="#22c55e" opacity="0.8"/>
    <text x="370" y="230" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="10">25%</text>
    
    {/* Labels */}
    <text x="70" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E1</text>
    <text x="170" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E2</text>
    <text x="270" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E3</text>
    <text x="370" y="15" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">E4</text>
    
    <text x="175" y="250" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">✓ Uniform distribution</text>
    <text x="175" y="270" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="8">Equal load across experts</text>
  </g>
  
  {/* Balance loss formula */}
  <g transform="translate(50, 320)">
    <rect width="700" height="60" rx="8" fill="#f8fafc" stroke="#334155"/>
    <text x="350" y="20" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="11">Load Balance Loss: L_balance = α · E · Σᵢ fᵢ · pᵢ</text>
    <text x="350" y="40" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">fᵢ = fraction of tokens routed to expert i, pᵢ = mean routing probability for expert i</text>
    <text x="350" y="55" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">α = 0.01 (typical), minimizes when fᵢ = pᵢ = 1/E for all i</text>
  </g>
</svg>

6. Expert Capacity Formula

6.1 Computing Capacity

For a batch of tokens with experts and top- routing:

Example:

• tokens
• experts
• experts per token

6.2 Capacity Factor Impact

6.3 Dynamic Capacity

Adaptive capacity based on routing distribution:

7. Modern MoE Architectures

7.1 Switch Transformer (Fedus et al., 2022)

• Top-1 routing: Each token → 1 expert
• Simpler load balancing: Auxiliary loss only
• Scaling: Up to 1.6T parameters with 2048 experts

7.2 GShard (Lepikhin et al., 2021)

• Top-2 routing: Each token → 2 experts
• Random routing: Second expert selected probabilistically
• Expert capacity: With overflow handling

7.3 Mixtral (Jiang et al., 2024)

• Sliding window attention: Efficient local attention
• Top-2 routing: Standard MoE FFN layers
• 8 experts: 47B total, 12.9B active parameters

7.4 DeepSeek-MoE (DeepSeek-AI, 2024)

• Fine-grained experts: More, smaller experts
• Shared experts: Always activated for common knowledge
• Routed experts: Task-specific experts

8. Routing Strategies

8.1 Token-Choice Routing

Each token selects top- experts (standard approach):

8.2 Expert-Choice Routing (Zhou et al., 2022)

Each expert selects top- tokens:

Advantages: No token dropping, natural load balancing Disadvantages: Some tokens may be selected 0 times

8.3 Hash Routing

Use consistent hashing to assign tokens to experts:

8.4 Expert Routing with Context

where is additional context (e.g., task embedding).

9. Training MoE Models

9.1 Total Loss

9.2 Gradient Computation

For active experts only:

9.3 Expert Parallelism

Distribute experts across devices:

Communication pattern:

1. All-to-all: Route tokens to experts
2. All-to-all: Return outputs

10. MoE vs. Dense Models

MoE models achieve state-of-the-art performance by efficiently scaling parameters while keeping compute manageable through sparse activation.