← Back to Table of Contents

Chapter 33 — Mixture of Experts (MoE)

“Why activate all 600 billion parameters for every token when only 100 billion are needed? MoE models are sparse — they think with a fraction of their brain.”

Dense vs Sparse Models

The key insight: MoE gives you the quality of a large model with the inference speed of a small model (if memory allows).

MoE Architecture

In a standard transformer, each layer has one FFN. In MoE, the FFN is replaced by N experts + a router:

Tensor Shapes

The Router

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch
import torch.nn as nn
import torch.nn.functional as F

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.top_k = top_k
    
    def forward(self, x):
        # x: [B, T, d]
        logits = self.gate(x)                    # [B, T, n_experts]
        
        # Select top-K experts per token
        top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)
        top_k_weights = F.softmax(top_k_logits, dim=-1)  # [B, T, K]
        
        return top_k_weights, top_k_indices

Load Balancing

Without intervention, the router collapses — sending most tokens to a few “popular” experts while others are idle. This is solved with an auxiliary load balancing loss:

Where:

• $f_i$ = fraction of tokens routed to expert $i$
• $P_i$ = average router probability for expert $i$
• $\alpha$ = balance loss weight (typically 0.01)

This loss encourages uniform distribution of tokens across experts.

Simple MoE Layer Implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts=8, top_k=2):
        super().__init__()
        self.router = TopKRouter(d_model, n_experts, top_k)
        self.experts = nn.ModuleList([
            SwiGLU_FFN(d_model, d_ff) for _ in range(n_experts)
        ])
        self.n_experts = n_experts
        self.top_k = top_k
    
    def forward(self, x):
        B, T, d = x.shape
        weights, indices = self.router(x)    # [B,T,K], [B,T,K]
        
        # Flatten for routing
        x_flat = x.view(-1, d)               # [B*T, d]
        output = torch.zeros_like(x_flat)     # [B*T, d]
        
        for k in range(self.top_k):
            expert_indices = indices[:, :, k].reshape(-1)    # [B*T]
            expert_weights = weights[:, :, k].reshape(-1, 1) # [B*T, 1]
            
            for e in range(self.n_experts):
                mask = (expert_indices == e)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += expert_weights[mask] * expert_output
        
        return output.view(B, T, d)

In production, this loop is replaced by optimized scatter/gather operations (e.g., Megablocks library).

Notable MoE Models

Fine-Grained Experts (DeepSeek Approach)

DeepSeek splits each expert into smaller “fine-grained” experts and adds a shared expert that always activates:

This design gives better expert utilization and reduces the load balancing problem.

MoE KV-Cache

An important practical point: MoE does not affect the KV-cache. Attention layers are shared (not sparse) — only the FFN is replaced by experts. This means:

• KV-cache size is identical to a dense model with the same attention dimensions
• Memory calculations from Chapter 17 apply directly
• The total model memory = attention params + all expert params (even though only K are active)

This creates a unique memory profile: MoE models need lots of memory for weights but have the same KV-cache footprint as their “active parameter” equivalent.

MoE Challenges

What’s Next

MoE shows that not all model architectures need to be dense transformers. The next chapter explores more radical departures — State Space Models (SSMs) and other alternatives to attention-based architectures.

← Previous: Chapter 32 — Scaling Laws & Emergent Abilities · Next: Chapter 34 — SSMs & Beyond Transformers →

Last updated: April 2026