Mixture-of-Experts (MoE) architecture — Hinglish notes

6.1.4 · AI-ML › Scaling & Efficient Architectures

MoE KYA hota hai?

Yeh core distinction yaad rakho:

Quantity	Meaning	Scales with
Total (dense) params	saare experts stored hain	$N$
Active params / token	jo actually run hota hai	$k$
Compute (FLOPs)	per token forward cost	$k$ , not $N$

Routing KAISE kaam karta hai? (Scratch se Derivation)

Hum chahte hain ki har token embedding $x \in R^{d}$ ke liye, (1) kaunse experts choose karein, aur (2) unke outputs ko blend karein.

Step 1 — Har expert ko score karo. Router ko ek learnable weight matrix $W_{g} \in R^{N \times d}$ do. Expert $i$ ka raw score (logit) dot product hai: $h_{i} = (W_{g} x)_{i}$ Yeh step kyun? Ek dot product token aur har expert ke learned "preference vector" (row of $W_{g}$ ) ke beech alignment measure karta hai. High alignment → yeh expert relevant hai.

Step 2 — Scores ko probabilities mein badlo. Softmax apply karo taaki scores comparable ho jaayein aur 1 mein sum karein: $p_{i} = \frac{e ^{h_{i}}}{\sum _{j = 1}^{N} e ^{h_{j}}}$ Yeh step kyun? Hume expert outputs ko combine karne ke liye normalized weights chahiye, aur softmax "kaun best hai" ko weights mein convert karne ka smooth, differentiable tarika hai.

Step 3 — Sirf Top- $k$ rakho. Maano $T$ $k$ largest $p_{i}$ ke indices ka set hai. Baaki ko zero kar do: $g_{i} = ⎩ ⎨ ⎧ \frac{p _{i}}{\sum _{j \in T} p _{j}} 0 i \in T otherwise$ Yeh step kyun? Non-selected gates ko 0 set karna matlab hai ki woh experts zero computation karte hain — yahi se FLOP savings aati hai. $T$ par re-normalizing rakhta hai ki use hone waale weights 1 mein sum karein.

Step 4 — Expert outputs ko combine karo. $y = \sum_{i \in T} g_{i} E_{i} (x)$ Yeh step kyun? Hum sirf $k$ chosen experts $E_{i} (x)$ evaluate karte hain aur gate-weighted sum lete hain. Yahi layer ka output hai.

Figure — Mixture-of-Experts (MoE) architecture

Load-balancing problem (MoE ko helper loss kyun chahiye)

Fix: ek auxiliary load-balancing loss. Ek batch ke liye define karo:

$f_{i}$ = expert $i$ ko route kiye gaye tokens ka fraction (dispatch fraction).
$P_{i}$ = expert $i$ ko assigned average router probability.

$L_{aux} = α N \sum_{i = 1}^{N} f_{i} P_{i}$

Exactly yeh form kyun? Yeh minimize hota hai jab dono $f_{i}$ aur $P_{i}$ uniform hoon ( $= 1/ N$ ), matlab har expert ko equal share mile. Sirf $f_{i}$ (hard argmax se, non-differentiable) ki jagah $P_{i}$ (soft, differentiable) use karna gradients ko actually flow karne deta hai taaki imbalance theek ho sake. Factor $N$ loss scale ko expert count se independent rakhta hai; $α$ (jaise $0.01$ ) isse ek gentle nudge rakhta hai, main objective nahi.

Worked Examples

Common Mistakes (Steel-manned)

Recall Feynman: 12-saal ke bacche ko explain karo

Ek hospital imagine karo jisme 64 specialist doctors hain. Ek dense model jaisa hai jaise har patient ko saare 64 doctors se milwao — accurate hai lekin insanely slow aur expensive. MoE darwaze par ek samajhdar receptionist (router) rakhta hai jo aapki problem dekh kar aapko sirf 2 best-matched doctors ke paas bhejta hai. Hospital mein abhi bhi woh saari expertise hai, lekin har patient thoda sa hi use karta hai. Fairness ke liye ek rule hai taaki koi ek popular doctor saare patients na le jaaye jabki baaki idle baithe rahein (load balancing), aur har doctor per ghante sirf itne patients dekh sakta hai (capacity).

Active Recall Flashcards

#flashcards/ai-ml

Woh kaunsi problem hai jo MoE decouple karta hai jo dense models baandh dete hain?

Model capacity (total parameters / knowledge) ko per-token compute (FLOPs) se.

MoE layer mein, Top-k selection se pehle router kya output karta hai?

Logits

W_{g} x

N

experts par ek softmax distribution

p_{i}

Top-k set par gates ko re-normalize kyun karein?

Taaki selected gate weights 1 mein sum karein, expert outputs ka ek proper weighted blend milta rahe.

MoE mein FLOPs per token konsi quantity ke saath scale karte hain —

N

k

k

(active experts) ke saath,

N

(total experts) ke saath nahi.

"Expert collapse" failure kya hai aur iska fix kya hai?

Router zyaadatar tokens kuch hi experts ko bhejta hai (rich-get-richer); ek auxiliary load-balancing loss se fix hota hai.

Load-balancing loss likhkar batao yeh kab minimize hoti hai.

L_{a ux} = α N \sum_{i} f_{i} P_{i}

; minimize hoti hai jab

f_{i} = P_{i} = 1/ N

(uniform usage).

Expert capacity

C

kya hai aur jab yeh exceed hoti hai toh kya hota hai?

C = cap_factor \times (tokens / N)

; excess tokens drop ho jaate hain aur residual se pass hote hain.

Softmax ki jagah hard argmax router kyun nahi use karte?

Argmax non-differentiable hai; softmax smooth gradients deta hai taaki router seekh sake.

k = 2

N = 64

, 7M-param experts ke liye, kitne active vs total expert params hain?

Active

= 2 \times 7 M = 14

M; total

= 64 \times 7 M = 448

MoE mein compute saving physically kahan se aati hai?

Non-selected experts ko gate 0 milta hai aur kabhi evaluate nahi hote (zero FLOPs).

Connections

Feed-Forward Networks (Transformer FFN) — har expert ek FFN block hai.
Softmax and Gating Functions — router ka normalization step.
Scaling Laws — MoE capacity axis par sasti tarah move karta hai.
Sparse vs Dense Models — fundamental design trade-off.
Load Balancing and Auxiliary Losses — experts ko utilized rakhna.
Model Parallelism & Expert Parallelism — experts ko devices across kaise shard karte hain.
Conditional Computation — woh general principle jise MoE instantiate karta hai.