Sparse routing and gating networks

6.1.5 · AI-ML › Scaling & Efficient Architectures

Overview

Sparse routing ek aisi technique hai jisme hum har input ke liye sirf kuch model parameters activate karte hain, poore network ko use karne ki jagah. Gating network decide karti hai ki kaun sa subset activate karna hai. Isse models bahut bade parameter counts tak scale ho sakte hain jabki per-example compute constant rehta hai.

Core intuition: Jab alag-alag inputs ko alag expertise chahiye, toh poora neural network kyun run karein? Chemistry ke ek question ko poetry-writing neurons activate karne ki zarurat nahi honi chahiye.

The Core Architecture

1. Mixture of Experts (MoE) Framework

Output ek weighted combination hota hai:

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

Lekin sabse important baat, hum sirf top-$k$ terms compute karte hain.

2. Deriving the Gating Function

Starting point: Humein ek aisi function chahiye jo input $x\in {\mathbb{R}}^d$ ko $N$ routing scores mein map kare.

Step 1 — Raw scores: Ek learned linear transformation apply karo: $\text{logits}=W_{g}x+b_g$ jahan $W_g\in {\mathbb{R}}^{N\times d}$. Isse $N$ scalar scores milte hain.

Step 2 — Normalization: Humein probabilities chahiye (jinка sum 1 ho). Softmax apply karo: $G(x{)}_i=\frac{\exp ({\text{logits}}_i)}{{\sum }_{j=1}^N\exp ({\text{logits}}_j)}$

Softmax kyun? Yeh differentiable hai (backprop enable karta hai) aur valid probability distribution ensure karta hai.

Step 3 — Sparsification: Sirf top-$k$ entries rakho, baaki zero kar do:

G(x)_i & \text{if } i \in \text{TopK}(G(x)) \\ 0 & \text{otherwise} \end{cases}$$ **Yeh kyun matter karta hai**: $N=128$ experts aur $k=2$ ke saath, hum forward pass mein **128 mein se sirf 2 experts run karte hain**, jabki poora parameter store (100% capacity) available rehta hai. > [!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) > ``` **Computational savings ki derivation**: - Dense FFN sublayer jo har token ko full hidden width se route karta hai: $\mathcal{O}(d \cdot d_{\text{ffn}})$ per token - Har expert khud ek FFN hai width $d_{\text{ffn}}$ ka, jis token ko process karta hai uske liye $\mathcal{O}(d \cdot d_{\text{ffn}})$ cost aata hai - Top-$k$ routing wala MoE layer: ek token exactly $k$ experts se guzarta hai, isliye active compute hai $\mathcal{O}(k \cdot d \cdot d_{\text{ffn}})$ per token - **Key point**: active compute $k$ pe depend karta hai (ek chota constant jaise 1 ya 2), $N$ pe **nahi**. Gating cost $\mathcal{O}(N \cdot d)$ expert FFNs ke mukable mein negligible hai. Toh hum $N$ (aur total parameters) badha sakte hain almost free mein — yahi hai **sublinear scaling** of active compute with total parameters ka source! ### 3. Load Balancing Problem > [!intuition] Load Balancing Kyun Zaroori Hai > Socho agar gating network seekh le "Expert 3 sab kuch ke liye best hai!" Toh sab examples Expert 3 ko route ho jayenge, aur baaki 127 experts waste ho jayenge. Humein **gating ko encourage karna hoga ki woh sabhi experts use kare**. **Mathematically**: Expert $i$ ke liye ==load== define karo: $$\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}$ sabhi $i$ ke liye. **Intuitive diagnostic**: Hum load ka coefficient of variation $\text{CV} = \frac{\sigma}{\mu}$ se imbalance measure kar sakte hain. Yeh scale-invariant hai aur exactly 0 hota hai jab sabhi loads equal hon. Yeh ek achha *monitoring metric* hai, lekin yeh discrete top-k ke through differentiable nahi hai, isliye hum actually ise optimize nahi karte. **Switch Transformer auxiliary loss ki derivation**: Humein ek **differentiable** surrogate chahiye. Batch ke upar har expert ke liye do quantities define karo: - $f_i = \frac{1}{|\text{batch}|} \sum_{x} \mathbb{1}[i \in \text{TopK}(G(x))]$ — expert $i$ ko dispatch kiye gaye **tokens ka fraction** (hard, non-differentiable part) - $P_i = \frac{1}{|\text{batch}|} \sum_{x} G(x)_i$ — expert $i$ ko assign ki gayi **mean gate probability** (soft, differentiable) Switch Transformer auxiliary loss hai: $$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i$$ **Yeh specific form kyun?** Gradients sirf differentiable $P_i$ se flow karte hain (kyunki backprop ke dauran $f_i$ ek detached constant hai). $\sum_i f_i P_i$ minimize karna, $\sum_i P_i = 1$ ke constraint ke saath, gate ko push karta hai ki woh *already-overloaded experts par probability kam kare* (bada $f_i$) aur *underloaded walon par badhaye*. Loss minimized hota hai (value $= 1$) exactly jab $f_i = 1/N$ aur $P_i = 1/N$ sabhi $i$ ke liye ho — yaani uniform routing. Factor $N$ loss ko $\mathcal{O}(1)$ rakhta hai expert count chahe kuch bhi ho. **Historical note (Shazeer et al. 2017)**: Original sparsely-gated MoE ne batch-summed *importance* $\text{Imp}_i = \sum_x G(x)_i$ aur *load* $\text{Load}_i$ par **do alag squared terms** use kiye the: $$\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}$$ Dono squared-CV terms minimize hote hain jab respective quantities uniform hon. Switch ka single $f_i P_i$ product usi idea ka ek simpler, baad mein aaya refinement hai. ## Key Variants ### Switch Transformer (Google, 2021) **Key insight**: Maximum sparsity ke liye $k=1$ (har token ke liye single expert) use karo. > [!formula] Switch Routing > $$i^* = \arg\max_i G(x)_i$$ > $$y = E_{i^*}(x)$$ **Yeh kyun kaam karta hai**: - Training faster hai (har example ke liye sirf 1 expert ke gradients) - Load balancing aasan hai (clear winner-take-all) - Inference faster hai (weighted combination nahi) **Trade-off**: Har token ke liye kam capacity utilization, lekin 1000+ experts tak scale ho sakta hai. ### Expert Choice Routing (Google, 2022) **Paradigm shift**: Tokens experts choose karne ki jagah, ==experts tokens choose karte hain==! **Kaise kaam karta hai**: 1. Har expert ki capacity $C = \frac{k \cdot |\text{batch}|}{N}$ tokens hoti hai 2. Har expert $i$ ke liye, $G(x)_i$ score se top-$C$ tokens select karo 3. Sirf selected tokens par expert output compute karo **Yeh better kyun hai**: - **Perfect load balancing**: Har expert exactly $C$ tokens process karta hai - **Koi auxiliary loss nahi chahiye**: Balance structural hai, learned nahi - **Better utilization**: High-affinity tokens process hote hain **Derivation**: Agar batch mein $B$ tokens hain, $N$ experts hain, aur hum chahte hain ki har token average $k$ experts dekhe: - Total expert slots chahiye: $B \cdot k$ - Slots per expert: $C = \frac{B \cdot k}{N}$ - Har expert affinity score $G(x)_i$ se apne top-$C$ tokens pick karta hai ### Soft MoE (Meta, 2023) **Key innovation**: Koi discrete routing nahi! Experts sabhi input tokens ka **weighted combination** process karte hain. > [!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)) ``` **Slots N·S columns kyun span karte hain**: har expert $M = N\cdot S$ slots mein se $S$ own karta hai. Assignment matrix $\Phi$ ki shape isliye $B \times (N\cdot S)$ hoti hai, $B \times S$ nahi — har expert ko apna dedicated block of slots milna chahiye. Isse routing fully differentiable ho jaata hai (koi top-$k$ discontinuity nahi) aur construction se perfectly balanced hota hai, $M$ slot mixtures compute karne ki cost par. ## Training Considerations ### 1. Gradient Flow > [!mistake] Top-K Se Broken Gradients > **Galat soch**: "Top-$k$ selection theek hai, gradients selected experts se flow karte hain." **Yeh sahi kyun lagta hai**: Selected experts ko unke outputs se gradients milte hain. **Problem**: Non-selected experts ko **zero gradient** milta hai. Gating function $G(x)$ ko koi signal nahi milta ki usne non-selected experts ke liye achhe choices kiye ya nahi. Isse experts initially jo randomly achhe the, wahi cheezein specialize karne lage hain. **Fix**: - **Straight-through estimators** use karo: Forward pass top-$k$ use karta hai, backward pass softmax - **Auxiliary losses** add karo jo exploration encourage karein - Training ke dauran zyada tokens per expert allow karne ke liye **expert capacity factor** $> 1$ use karo ### 2. Data Parallelism with Expert Parallelism **Challenge**: Har expert alag GPU par ho sakta hai. GPU 3 par Expert 5 ko route hone wale token ke liye humein: 1. Token GPU 3 par bhejna hoga 2. Expert 5 ka output compute karna hoga 3. Result wapas original GPU par bhejna hoga **Communication cost ki derivation**: $B$ tokens, $N$ experts, $D$ devices par, top-$k$ routing ke saath: - Tokens per device: $\approx \frac{B \cdot k}{N}$ - Agar experts evenly distributed hain: communication volume $\mathcal{O}(B \cdot k \cdot d)$ - **Key bottleneck**: All-to-all scatter/gather operations **Yeh kyun matter karta hai**: $N=512$ experts on 512 GPUs ke saath, communication compute ko dominate kar sakta hai! Solutions: expert replication, hierarchical routing, buffer capacity limits. > [!example] Load Balance Calculate Karna > Maano hamare paas $N=8$ experts hain, 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 ``` **Imbalance calculate karo**: - 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$ **Yeh step kyun?** CV > 0.3 significant imbalance indicate karta hai. Hum $\mathcal{L}_{\text{aux}}$ mein $\alpha$ badhayenge taaki ise zyada penalize kiya ja sake. > [!example] Expert Choice Routing > 16 tokens ka batch, 4 experts, har token $k=2$ experts chahta hai. **Token-to-expert affinities** (rows = tokens, columns = sabhi 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 ka perspective**: Column 1 dekho, affinity se top-8 tokens select karo. Agar T1 (0.80), T3 (0.60), T5 (0.55), ... highest hain, toh unhe 8 process karo. **Yeh step kyun?** Har expert independently apni capacity best-matching tokens se bharta hai. Guaranteed perfect balance (har expert exactly 8 tokens process karta hai). ## Practical Impact ### Switch Transformer Results (2021) - Paper mein roughly ~**100B-parameter dense model** (T5-XXL scale) ke compute cost par train kiya gaya 1.6 trillion parameter model - Perplexity-vs-compute efficiency mein bade speedups (dense T5 baseline se ~7x reported) - Parameters ki near-linear scaling: experts add karne se parameters multiply hote hain jabki active compute almost flat rehta hai ### GLaM (Google, 2021) - Sabse bada GLaM MoE variant **~340B parameters** ke order par hai (ek chota variant ~1.2B hai); yeh GPT-3 175B ko kai benchmarks par outperform karta hai - Train karne ke liye GPT-3 ke mukable mein sirf **1/3 energy** use karta hai - 64 experts per layer ke saath $k=2$ top-expert routing ### Kyun Kaam Karta Hai **Capacity vs. Computation trade-off**: - Dense model: $N$ parameters → $\mathcal{O}(N)$ active compute per token - MoE model: same total parameters, lekin active compute per token $N$ ki jagah $k$ (constant) ke saath scale hota hai $k=2$ aur 128 experts ke saath, hum single expert ke worth se ~64x zyada parameters rakhte hain jabki sirf 2 experts per token run karte hain! > [!recall]- Ek 12-Saal-Ke Bacche Ko Explain Karo > Socho tum ek video game khel rahe ho jahan tum 100 alag power-ups unlock kar sakte ho, lekin tumhara character sirf 2 ek saath carry kar sakta hai. Sparse routing ek smart backpack jaisa hai jo automatically best 2 power-ups pick karta hai jo bhi level tum par ho. Agar tum fire level mein ho, toh woh water blast aur fire shield pick karta hai. Agar tum puzzle level mein ho, toh hint magnifier aur time slower pick karta hai. Tum phir bhi "own" karte ho sabhi 100 power-ups (huge capacity!), lekin sirf 2 ek waqt mein "use" karte ho (low cost!). Gating network woh smart backpack ki AI hai jo decide karti hai ki kaun se 2 load karne hain. Experts woh 100 alag power-ups hain. Aur load balancing ensure karta hai ki tum sabhi power-ups kabhi kabhi use karo, taaki koi waste na ho. Isse game designers aisi characters bana sakte hain jisme zyada abilities hon jo normally active memory mein fit nahi hogi, kyunki zyaatar abilities kisi bhi waqt "so rahi" hoti hain. > [!mnemonic] MoE Core Concept > **"GATE the SPARSE EXPERTS"** > - **G**ating function routes karta hai > - **A**ctivate sirf top-k karo > - **T**raining ko balance chahiye > - **E**xperts specialize karte hain **SPARSE = Subset of Parameters Activate per Route, Saving Enormous compute** ## Common Pitfalls > [!mistake] Capacity Constraints Ignore Karna > **Galat approach**: Experts pe unlimited tokens allow karo, ummeed rakhte hue ki load balance naturally emerge ho jayega. **Yeh sahi kyun lagta hai**: Humne auxiliary loss add ki, toh gating seekh legi balance karna... sahi? **Reality**: Training spikes ke dauran, hundreds of tokens ek expert ko route ho sakte hain. Isse hota hai: - OOM errors (expert ka activation buffer overflow ho jaata hai) - Gradient spikes (ek expert ko massive batch milta hai) - Instability **Fix**: **Expert capacity** set karo $C = \text{capacity\_factor} \times \frac{B \cdot k}{N}$. Capacity se zyada tokens residual connection se pass kiye jaate hain ya drop kar diye jaate hain. Typical capacity_factor = 1.25. **Kyun kaam karta hai**: Har expert ke liye memory aur compute par hard upper bound provide karta hai. Gating ko better balance seekhne par majboor karta hai ek "favorite" par rely karne ki jagah. > [!mistake] Uniform Expert Initialization > **Galat approach**: Sabhi experts ko same weights se initialize karo. **Yeh sahi kyun lagta hai**: Hum fairness chahte hain, toh sabko barabar se shuru karo. **Problem**: Identical experts ke saath, gating network ke paas unhe differentiate karne ka koi signal nahi hota. Sabhi experts identically evolve hote hain, aur effectively tumhare paas 1 expert hai jo $N$ baar copy kiya gaya hai. **Fix**: - Har expert ke liye alag random seeds - Ya: Har expert ki initialization mein small random noise add karo - Ya: Experts ko alag data subsets par pre-train karo **Kyun kaam karta hai**: Symmetry todta hai, experts ko allow karta hai ki woh early specialize karein iss basis par ki woh randomly kisme achhe hain. ## Connections - [[Transformer Architecture]] — MoE layers FFN sublayers replace karte hain - [[Model Parallelism]] — Expert parallelism model sharding ka ek form hai - [[Conditional Computation]] — Input-dependent parameter selection ki broader family - [[Knowledge Distillation]] — Dense student models sparse teachers se distill kar sakte hain - [[Neural Architecture Search]] — Gating ko learned architecture selection ke roop mein dekha ja sakta hai - [[Attention Mechanisms]] — Soft MoE attention-jaisa weighted aggregation use karta hai - [[Batch Normalization]] — Expert-specific norms specialization improve karte hain - [[Long-tail Distribution]] — MoE data diversity ko dense models se better handle karta hai --- #flashcards/ai-ml Top-k gating ke saath sparse routing ka key computational advantage kya hai? ::: N experts mein se sirf k token ke liye run hote hain, isliye active compute $\mathcal{O}(k \cdot d \cdot d_{\text{ffn}})$ hai — yeh k ke saath scale karta hai (ek chota constant), N ke saath nahi. Isse total parameters almost free mein badh sakte hain. Gating function G(x) routing weights kaise produce karta hai? ::: Logits paane ke liye learned linear transformation $W_g \cdot x + b_g$ apply karta hai, phir normalization ke liye softmax, phir sparsify karne ke liye top-k selection, sirf highest-scoring experts rakhta hai. MoE training mein auxiliary load balancing losses kyun zaroori hain? ::: Explicit balancing ke bina, gating collapse ho sakti hai aur sabhi examples kuch hi experts ko route kar sakti hai, baaki experts ki capacity waste ho jaati hai aur model expressiveness kam ho jaati hai. Switch Transformer load-balancing auxiliary loss kya hai? ::: $\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i$ jahan $f_i$ expert $i$ ko dispatch kiye gaye tokens ka fraction hai (detached) aur $P_i$ expert $i$ ke liye mean gate probability hai; gradients $P_i$ se flow karte hain aur overloaded experts se probability door dhakelte hain. Original Shazeer 2017 MoE ne kaun sa load-balancing loss use kiya? ::: Do alag squared-coefficient-of-variation terms — ek importance $\text{Imp}_i = \sum_x G(x)_i$ par aur ek load par — yaani $\sum_i \text{Imp}_i^2$ aur $\sum_i \text{Load}_i^2$ ke proportional, har ek uniform distribution par minimize hota hai. Expert Choice routing perfect load balancing kaise achieve karta hai? ::: Tokens experts choose karne ki jagah experts affinity score se apne top-$C$ tokens choose karte hain, structurally guarantee karta hai ki har expert exactly $C = \frac{k \cdot B}{N}$ tokens process kare. Switch Transformer aur standard MoE mein main difference kya hai? ::: Switch $k=2+$ ki jagah $k=1$ (har token ke liye single expert) use karta hai, zyada experts ke through faster training aur inference ke liye sparsity maximize karta hai. Training ke dauran expert capacity factor > 1.0 kyun zaroori hai? ::: Perfect balance ($C = B \cdot k / N$) se pare buffer room allow karta hai taaki routing spikes handle ho sakein, OOM aur gradient instability prevent ho sake jab load temporarily concentrate ho. Soft MoE dispatch matrix $\Phi$ ki shape kya hai aur kyun? ::: $B \times (N \cdot S)$: $N$ experts mein se har ek $S$ dedicated slots own karta hai, isliye total slots $M = N \cdot S$. $\Phi$ tokens ko sabhi $M$ slots assign karta hai, har expert ko apna dedicated block of inputs deta hai. Soft MoE discrete routing problem kaise eliminate karta hai? ::: Per-expert slots mein sabhi tokens ke softmax-weighted combinations use karta hai, routing fully differentiable banata hai bina gradients mein kisi top-$k$ discontinuity ke. Discrete top-k routing mein broken gradients kya cause karta hai? ::: Non-selected experts ko zero gradient milta hai kyunki woh forward pass mein participate nahi karte, gating ko yeh seekhne se rokta hai ki unchosen experts better hote ya nahi. Switch Transformer ke 1.6T-parameter model se roughly kitna dense-model cost match kiya? ::: Approximately ~100B-parameter dense model (T5-XXL scale) ka compute cost, 10B wala nahi. GLaM ka sabse bada parameter count kitna hai? ::: Approximately ~340B parameters ke order par (ek chota ~1.2B variant ke saath), 64 experts per layer par $k=2$ routing use karta hai — 1.2 trillion nahi. Expert parallelism mein communication bottleneck kya hai? ::: All-to-all scatter/gather operations taaki tokens alag GPUs par experts ko route ho sakein, volume $\mathcal{O}(B \cdot k \cdot d)$ ke saath potentially compute time dominate karte hue. ## 🖼️ 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 ```