Sparse routing and gating networks

Overview

Sparse routing is a technique where we activate only a subset of model parameters for each input, rather than using the entire network. A gating network decides which subset to activate. This enables models to scale to massive parameter counts while keeping per-example compute constant.

Core intuition: Why run the entire neural network when different inputs need different expertise? A question about chemistry shouldn't need to activate poetry-writing neurons.

The Core Architecture

1. Mixture of Experts (MoE) Framework

The output is a weighted combination:

$y = \sum_{i=1}^{N} G(x)_i \cdot E_i(x)$

But crucially, we only compute the top-$k$ terms.

2. Deriving the Gating Function

Starting point: We need a function that maps input $x \in \mathbb{R}^d$ to $N$ routing scores.

Step 1 — Raw scores: Apply a learned linear transformation: $\text{logits} = W_g x + b_g$ where $W_g \in \mathbb{R}^{N \times d}$. This gives $N$ scalar scores.

Step 2 — Normalization: We want probabilities (sum to 1). Apply softmax: $G(x)_i = \frac{\exp(\text{logits}_i)}{\sum_{j=1}^{N} \exp(\text{logits}_j)}$

Why softmax? It's differentiable (enables backprop) and ensures valid probability distribution.

Step 3 — Sparsification: Keep only top-$k$ entries, zero out others:

G(x)_i & \text{if } i \in \text{TopK}(G(x)) \\ 0 & \text{otherwise} \end{cases}$$ **Why this matters**: With $N=128$ experts and $k=2$, we only **run 2 of 128 experts** per forward pass while retaining the full parameter store (100% capacity available). > [!formula] MoE Layer Forward Pass > ``` > Input: x ∈ ℝ^d > 1. Compute gating scores: s = Wg·x + bg (s ∈ ℝ^N) > 2. Apply softmax: G(x) = softmax(s) > 3. Select top-k: I = TopK(G(x), k) > 4. Renormalize (optional): G'(x)i = G(x)i / Σ(j∈I) G(x)j > 5. Compute: y = Σ(i∈I) G'(x)i · Ei(x) > ``` **Derivation of computational savings**: - Dense FFN sublayer that routes every token through the full hidden width: $\mathcal{O}(d \cdot d_{\text{ffn}})$ per token - Each expert is itself an FFN of width $d_{\text{ffn}}$, costing $\mathcal{O}(d \cdot d_{\text{ffn}})$ per token it processes - MoE layer with top-$k$ routing: a token runs through exactly $k$ experts, so active compute is $\mathcal{O}(k \cdot d \cdot d_{\text{ffn}})$ per token - **Key point**: the active compute depends on $k$ (a small constant like 1 or 2), **not** on $N$. The gating cost $\mathcal{O}(N \cdot d)$ is tiny compared to expert FFNs. So we can grow $N$ (and total parameters) almost for free — this is the source of the **sublinear scaling** of active compute with total parameters! ### 3. Load Balancing Problem > [!intuition] Why We Need Load Balancing > Imagine the gating network learns "Expert 3 is best for everything!" Now all examples route to Expert 3, and the other 127 experts are wasted. We need to **encourage the gating to use all experts**. **Mathematically**: Define ==load== for expert $i$ as: $$\text{Load}_i = \sum_{x \in \text{batch}} \mathbb{1}[i \in \text{TopK}(G(x))]$$ Ideally, $\text{Load}_i \approx \frac{|\text{batch}|}{N}$ for all $i$. **Intuitive diagnostic**: We can measure imbalance with the coefficient of variation of the load, $\text{CV} = \frac{\sigma}{\mu}$. This is scale-invariant and equals 0 exactly when all loads are equal. This is a great *monitoring metric*, but it is not differentiable through the discrete top-k, so it is not what we actually optimize. **Derivation of the Switch Transformer auxiliary loss**: We want a **differentiable** surrogate. Define two per-expert quantities over the batch: - $f_i = \frac{1}{|\text{batch}|} \sum_{x} \mathbb{1}[i \in \text{TopK}(G(x))]$ — the **fraction of tokens** dispatched to expert $i$ (the hard, non-differentiable part) - $P_i = \frac{1}{|\text{batch}|} \sum_{x} G(x)_i$ — the **mean gate probability** assigned to expert $i$ (soft, differentiable) The Switch Transformer auxiliary loss is: $$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i$$ **Why this specific form?** Gradients flow only through the differentiable $P_i$ (since $f_i$ is a detached constant during backprop). Minimizing $\sum_i f_i P_i$ subject to $\sum_i P_i = 1$ pushes the gate to *lower probability on already-overloaded experts* (large $f_i$) and *raise it on underloaded ones*. The loss is minimized (value $= 1$) exactly when both $f_i = 1/N$ and $P_i = 1/N$ for all $i$ — i.e., uniform routing. The factor $N$ keeps the loss $\mathcal{O}(1)$ regardless of expert count. **Historical note (Shazeer et al. 2017)**: The original sparsely-gated MoE used **two separate squared terms** on the batch-summed *importance* $\text{Imp}_i = \sum_x G(x)_i$ and the *load* $\text{Load}_i$: $$\mathcal{L}_{\text{aux}} = \alpha\,\text{CV}(\text{Imp})^2 + \alpha\,\text{CV}(\text{Load})^2 \;\propto\; \frac{\sum_i \text{Imp}_i^2}{(\sum_i \text{Imp}_i)^2} + \frac{\sum_i \text{Load}_i^2}{(\sum_i \text{Load}_i)^2}$$ Both squared-CV terms are minimized when the respective quantities are uniform. Switch's single $f_i P_i$ product is a simpler, later refinement of the same idea. ## Key Variants ### Switch Transformer (Google, 2021) **Key insight**: Use $k=1$ (single expert per token) for maximum sparsity. > [!formula] Switch Routing > $$i^* = \arg\max_i G(x)_i$$ > $$y = E_{i^*}(x)$$ **Why this works**: - Training is faster (only 1 expert's gradients per example) - Load balancing is easier (clear winner-take-all) - Inference is faster (no weighted combination) **Trade-off**: Less capacity utilization per token, but can scale to 1000+ experts. ### Expert Choice Routing (Google, 2022) **Paradigm shift**: Instead of tokens choosing experts, ==experts choose tokens==! **How it works**: 1. Each expert has capacity $C = \frac{k \cdot |\text{batch}|}{N}$ tokens 2. For each expert $i$, select top-$C$ tokens by $G(x)_i$ score 3. Compute expert output only on selected tokens **Why this is better**: - **Perfect load balancing**: Every expert processes exactly $C$ tokens - **No auxiliary loss needed**: Balance is structural, not learned - **Better utilization**: High-affinity tokens get processed **Derivation**: If batch has $B$ tokens, $N$ experts, and we want each token to see $k$ experts on average: - Total expert slots needed: $B \cdot k$ - Slots per expert: $C = \frac{B \cdot k}{N}$ - Each expert picks its top-$C$ tokens by affinity score $G(x)_i$ ### Soft MoE (Meta, 2023) **Key innovation**: No discrete routing! Experts process **weighted combinations** of all input tokens. > [!formula] Soft MoE Forward Pass > ``` > Input: X ∈ ℝ^(B×d) (batch of B tokens) > Config: N experts, S slots per expert → total slots M = N·S 1. Dispatch weights: Φ = softmax_over_tokens(X·Wφ) (Wφ ∈ ℝ^(d×M), Φ ∈ ℝ^(B×M)) Each of the M=N·S slot columns is a softmax over the B tokens. 2. Slot inputs: X̃ = Φᵀ·X (X̃ ∈ ℝ^(M×d)) → rows 1..S go to expert 1, S+1..2S to expert 2, ... 3. Per-expert compute: for expert n, process its S slot rows X̃[(n-1)S+1 : nS] → Ỹ (Ỹ ∈ ℝ^(M×d)) 4. Combine weights: Ψ = softmax_over_slots(X·Wφ) (Ψ ∈ ℝ^(B×M)) 5. Output tokens: Y = Ψ·Ỹ (Y ∈ ℝ^(B×d)) ``` **Why the slots span N·S columns**: each expert owns $S$ of the $M = N\cdot S$ slots. The assignment matrix $\Phi$ therefore has shape $B \times (N\cdot S)$, not $B \times S$ — every expert must receive its own dedicated block of slots. This makes routing fully differentiable (no top-$k$ discontinuity) and perfectly balanced by construction, at the cost of computing $M$ slot mixtures. ## Training Considerations ### 1. Gradient Flow > [!mistake] Broken Gradients from Top-K > **Wrong thinking**: "Top-$k$ selection is fine, gradients flow through selected experts." **Why it feels right**: The selected experts do get gradients from their outputs. **The problem**: Non-selected experts get **zero gradient**. The gating function $G(x)$ receives no signal about whether it made good choices for non-selected experts. This can cause experts to specialize on what they were randomly good at initially. **The fix**: - Use **straight-through estimators**: Forward pass uses top-$k$, backward pass uses softmax - Add **auxiliary losses** that encourage exploration - Use **expert capacity factor** $> 1$ to allow more tokens per expert during training ### 2. Data Parallelism with Expert Parallelism **The challenge**: Each expert might live on a different GPU. For a token routed to Expert 5 on GPU 3, we need: 1. Send token to GPU 3 2. Compute Expert 5's output 3. Send result back to original GPU **Derivation of communication cost**: With $B$ tokens, $N$ experts on $D$ devices, top-$k$ routing: - Tokens per device: $\approx \frac{B \cdot k}{N}$ - If experts evenly distributed: communication volume $\mathcal{O}(B \cdot k \cdot d)$ - **Key bottleneck**: All-to-all scatter/gather operations **Why this matters**: For $N=512$ experts on 512 GPUs, communication can dominate compute! Solutions: expert replication, hierarchical routing, buffer capacity limits. > [!example] Calculating Load Balance > Suppose we have $N=8$ experts, batch size $B=32$, $k=2$. **Ideal load per expert**: $\frac{32 \cdot 2}{8} = 8$ tokens **Observed routing**: ``` Expert 1: 12 tokens Expert 2: 6 tokens Expert 3: 3 tokens Expert 4: 10 tokens Expert 5: 7 tokens Expert 6: 2 tokens Expert 7: 9 tokens Expert 8: 15 tokens ``` **Calculate imbalance**: - Mean: $\mu = 8$ - Variance: $\sigma^2 = \frac{1}{8}\sum (x_i - 8)^2 = \frac{16+4+25+4+1+36+1+49}{8} = 17$ - Standard deviation: $\sigma = 4.12$ - Coefficient of variation: $\text{CV} = \frac{4.12}{8} = 0.515$ **Why this step?** CV > 0.3 indicates significant imbalance. We'd increase $\alpha$ in $\mathcal{L}_{\text{aux}}$ to penalize this more. > [!example] Expert Choice Routing > Batch of 16 tokens, 4 experts, each token wants $k=2$ experts. **Token-to-expert affinities** (rows = tokens, columns = all 4 experts): ``` E1 E2 E3 E4 T1: 0.80 0.10 0.05 0.05 T2: 0.10 0.70 0.15 0.05 T3: 0.60 0.20 0.10 0.10 ... ``` **Expert capacity**: $C = \frac{16 \cdot 2}{4} = 8$ tokens per expert **Expert 1's perspective**: Look at column 1, select top-8 tokens by affinity. If T1 (0.80), T3 (0.60), T5 (0.55), ... are highest, process those 8. **Why this step?** Each expert independently fills its capacity with best-matching tokens. Guaranteed perfect balance (every expert processes exactly 8 tokens). ## Practical Impact ### Switch Transformer Results (2021) - 1.6 trillion parameter model trained at roughly the compute cost of a $\sim$**100B-parameter dense model** (T5-XXL scale) in the paper - Large speedups in the perplexity-vs-compute efficiency (reported ~7x over the dense T5 baseline) - Near-linear scaling of parameters: adding experts multiplies parameters while keeping active compute nearly flat ### GLaM (Google, 2021) - Largest GLaM MoE variant is on the order of **~340B parameters** (a smaller variant is ~1.2B); it outperforms 175B GPT-3 on many benchmarks - Uses only about **1/3 the energy** to train relative to GPT-3 - $k=2$ top-expert routing with 64 experts per layer ### Why It Works **Capacity vs. Computation trade-off**: - Dense model: $N$ parameters → $\mathcal{O}(N)$ active compute per token - MoE model: same total parameters, but active compute per token scales with $k$ (constant) rather than $N$ For $k=2$ with 128 experts, we hold ~64x more parameters than a single expert's worth while running only 2 experts per token! > [!recall]- Explain to a 12-Year-Old > Imagine you're playing a video game where you can unlock 100 different power-ups, but your character can only carry 2 at a time. Sparse routing is like having a smart backpack that automatically picks the best 2 power-ups for whatever level you're on. If you're in a fire level, it picks the water blast and fire shield. If you're in a puzzle level, it picks the hint magnifier and time slower. You still "own" all 100 power-ups (huge capacity!), but you only "use" 2 at a time (low cost!). The gating network is the smart backpack's AI that decides which 2 to load. The experts are the 100 different power-ups. And load balancing is making sure all your power-ups get used sometimes, so you don't waste any. This lets game designers create characters with way more abilities than would normally fit in active memory, because most abilities are "sleeping" at any given time. > [!mnemonic] MoE Core Concept > **"GATE the SPARSE EXPERTS"** > - **G**ating function routes > - **A**ctivate top-k only > - **T**raining needs balance > - **E**xperts specialize **SPARSE = Subset of Parameters Activate per Route, Saving Enormous compute** ## Common Pitfalls > [!mistake] Ignoring Capacity Constraints > **Wrong approach**: Allow unlimited tokens per expert, hoping load balance will naturally emerge. **Why it feels right**: We added the auxiliary loss, so the gating should learn to balance... right? **The reality**: During training spikes, hundreds of tokens might route to one expert. This causes: - OOM errors (expert's activation buffer overflows) - Gradient spikes (one expert gets massive batch) - Instability **The fix**: Set **expert capacity** $C = \text{capacity\_factor} \times \frac{B \cdot k}{N}$. Tokens beyond capacity are passed through a residual connection or dropped. Typical capacity_factor = 1.25. **Why this works**: Provides hard upper bound on memory and compute per expert. Forces gating to learn better balance rather than relying on one "favorite." > [!mistake] Uniform Expert Initialization > **Wrong approach**: Initialize all experts with the same weights. **Why it feels right**: We want fairness, so start everyone equal. **The problem**: With identical experts, the gating network has no signal to differentiate them. All experts evolve identically, and you effectively have 1 expert copied $N$ times. **The fix**: - Different random seeds per expert - Or: Add small random noise to each expert's initialization - Or: Pre-train experts on different data subsets **Why this works**: Breaks symmetry, allows experts to specialize early based on what they randomly excel at. ## Connections - [[Transformer Architecture]] — MoE layers replace FFN sublayers - [[Model Parallelism]] — Expert parallelism is a form of model sharding - [[Conditional Computation]] — Broader family of input-dependent parameter selection - [[Knowledge Distillation]] — Dense student models can distill from sparse teachers - [[Neural Architecture Search]] — Gating can be viewed as learned architecture selection - [[Attention Mechanisms]] — Soft MoE uses attention-like weighted aggregation - [[Batch Normalization]] — Expert-specific norms improve specialization - [[Long-tail Distribution]] — MoE handles data diversity better than dense models --- #flashcards/ai-ml What is the key computational advantage of sparse routing with top-k gating? ::: Only k out of N experts run per token, so active compute is O(k·d·d_ffn) — it scales with k (a small constant), not N. This lets total parameters grow almost for free. How does the gating function G(x) produce routing weights? ::: Applies learned linear transformation Wg·x + bg to get logits, then softmax for normalization, then top-k selection to sparsify, keeping only highest-scoring experts. Why do we need auxiliary load balancing losses in MoE training? ::: Without explicit balancing, gating can collapse to routing all examples to a few experts, wasting the other experts' capacity and reducing model expressiveness. What is the Switch Transformer load-balancing auxiliary loss? ::: L_aux = α·N·Σ(fi·Pi) where fi is the fraction of tokens dispatched to expert i (detached) and Pi is the mean gate probability for expert i; gradients flow through Pi and push probability away from overloaded experts. What load-balancing loss did the original Shazeer 2017 MoE use? ::: Two separate squared-coefficient-of-variation terms — one on importance (Σ_x G(x)i) and one on load — i.e. proportional to ΣImp_i² and ΣLoad_i², each minimized at uniform distribution. How does Expert Choice routing achieve perfect load balancing? ::: Experts choose their top-C tokens by affinity score instead of tokens choosing experts, guaranteeing each expert processes exactly C = (k·B)/N tokens structurally. What is the main difference between Switch Transformer and standard MoE? ::: Switch uses k=1 (single expert per token) instead of k=2+, maximizing sparsity for faster training and inference while maintaining capacity through more experts. Why is expert capacity factor > 1.0 necessary during training? ::: Allows buffer room beyond perfect balance (C = B·k/N) to handle routing spikes, preventing OOM and gradient instability when load temporarily concentrates. What shape is the Soft MoE dispatch matrix Φ and why? ::: B×(N·S): each of the N experts owns S dedicated slots, so total slots M = N·S. Φ assigns tokens to all M slots, giving each expert its own block of inputs. How does Soft MoE eliminate the discrete routing problem? ::: Uses softmax-weighted combinations of all tokens into per-expert slots, making routing fully differentiable with no top-k discontinuity in gradients. What causes broken gradients in discrete top-k routing? ::: Non-selected experts receive zero gradient since they don't participate in forward pass, preventing gating from learning whether unchosen experts would have been better. Roughly what dense-model cost matched Switch Transformer's 1.6T-parameter model? ::: About the compute cost of a ~100B-parameter dense model (T5-XXL scale), not a 10B one. What is GLaM's largest parameter count? ::: On the order of ~340B parameters (with a smaller ~1.2B variant), using k=2 routing over 64 experts per layer — not 1.2 trillion. What is the communication bottleneck in expert parallelism? ::: All-to-all scatter/gather operations to route tokens to experts on different GPUs, with volume O(B·k·d) potentially dominating compute time. ## 🖼️ Concept Map ```mermaid flowchart TD SR[Sparse Routing] GN[Gating Network] MOE[Mixture of Experts Layer] EX[N Experts as FFNs] LOG[Logits Wg x + bg] SM[Softmax Normalization] TK[Top-k Selection] OUT[Weighted Output y] SAVE[Constant Per-Token Compute] SR -->|uses| GN GN -->|decides subset for| MOE MOE -->|contains| EX GN -->|computes| LOG LOG -->|passed to| SM SM -->|sparsified by| TK TK -->|activates k of N| EX TK -->|renormalized weights| OUT EX -->|combined into| OUT TK -->|k not N experts run| SAVE SAVE -->|enables| SR ``` ## 🔊 Hinglish (regional understanding) > [!intuition]- Hinglish mein samjho > Dekho, yahan core idea bahut simple aur powerful hai. Normally jab hum ek neural network chalate hain, toh har input ke liye pura network activate hota hai — matlab saare parameters use hote hain. Lekin sochke dekho, agar tumse chemistry ka sawaal pucha jaaye, toh poetry likhne wale neurons kyu chalane? ==Sparse routing== yahi solve karta hai — ek chota sa **gating network** decide karta hai ki kaunse "experts" (chote sub-networks) ko activate karna hai. Bilkul hospital jaise: reception pe aap aate ho, receptionist aapko sahi specialist ke paas bhej deta hai, aapko har doctor se milne ki zaroorat nahi. Isse hum model mein 100 experts rakh sakte hain (huge capacity), par har input sirf 2-3 se guzarta hai (kam cost). > > Ab technically kaise hota hai — gating network input ko ek linear transformation se scores mein convert karta hai, phir **softmax** laga ke probabilities banata hai, aur uske baad sirf **top-k** experts choose karta hai (baaki ko zero kar deta hai). Ismein sabse important baat samajhna hai: jo actual compute lagta hai woh $k$ pe depend karta hai — ek chota constant jaise 1 ya 2 — **na ki total experts $N$ pe**. Iska matlab hum $N$ badha sakte hain, total parameters badha sakte hain, lekin per-input compute wahi constant rehta hai. Yeh hai woh magic jise ==sublinear scaling== kehte hain — capacity almost free mein grow hoti hai. > > Par ek dikkat aati hai jise **load balancing** kehte hain. Agar gating network seekh le ki "Expert 3 sabke liye best hai," toh saare inputs Expert 3 pe chale jaayenge aur baaki 127 experts bekaar pade rahenge. Isliye hum monitor karte hain ki har expert ka ==load== barabar distribute ho raha hai ya nahi (coefficient of variation jaisi metric se). Yeh matter isliye karta hai kyunki aaj ke bade language models — GPT jaise — inhi MoE techniques pe scale karte hain. Toh yeh samajhna tumhe modern AI ki efficiency ka asli foundation deta hai. ![[audio/6.1.05-Sparse-routing-and-gating-networks.mp3]]