Mixture-of-Experts (MoE) architecture

WHAT is a Mixture-of-Experts?

The core distinction to burn in:

Quantity	Meaning	Scales with
Total (dense) params	all experts stored	$N$
Active params / token	what actually runs	$k$
Compute (FLOPs)	forward cost per token	$k$ , not $N$

HOW does routing work? (Derivation from scratch)

We want, for each token embedding $x \in \mathbb{R}^d$ , to (1) pick which experts, and (2) blend their outputs.

Step 1 — Score every expert. Give the router a learnable weight matrix $W_g \in \mathbb{R}^{N \times d}$ . The raw score (logit) for expert $i$ is the dot product: $h_i = (W_g\, x)_i$ Why this step? A dot product measures alignment between the token and each expert's learned "preference vector" (row of $W_g$ ). High alignment → this expert is relevant.

Step 2 — Turn scores into probabilities. Apply softmax so scores are comparable and sum to 1: $p_i = \frac{e^{h_i}}{\sum_{j=1}^{N} e^{h_j}}$ Why this step? We need normalized weights to combine expert outputs, and softmax is the smooth, differentiable way to convert "who is best" into weights.

Step 3 — Keep only the Top- $k$ . Let $\mathcal{T}$ be the set of indices of the $k$ largest $p_i$ . Zero out the rest: $g_i = \begin{cases} \dfrac{p_i}{\sum_{j\in\mathcal{T}} p_j} & i \in \mathcal{T}\\[4pt] 0 & \text{otherwise}\end{cases}$ Why this step? Setting non-selected gates to 0 means those experts do zero computation — that's where the FLOP savings come from. Re-normalizing over $\mathcal{T}$ keeps the used weights summing to 1.

Step 4 — Combine expert outputs. $y = \sum_{i \in \mathcal{T}} g_i \, E_i(x)$ Why this step? We only evaluate the $k$ chosen experts $E_i(x)$ and take a gate-weighted sum. This is the layer's output.

Figure — Mixture-of-Experts (MoE) architecture

The load-balancing problem (WHY MoE needs a helper loss)

Fix: an auxiliary load-balancing loss. For a batch, define:

$f_i$ = fraction of tokens routed to expert $i$ (dispatch fraction).
$P_i$ = average router probability assigned to expert $i$ .

$\mathcal{L}_{\text{aux}} = \alpha \, N \sum_{i=1}^{N} f_i \, P_i$

Why this exact form? It is minimized when both $f_i$ and $P_i$ are uniform ( $=1/N$ ), i.e. every expert gets an equal share. Using $P_i$ (soft, differentiable) instead of only $f_i$ (from a hard argmax, non-differentiable) lets gradients actually flow to fix imbalance. The factor $N$ keeps the loss scale independent of expert count; $\alpha$ (e.g. $0.01$ ) keeps it a gentle nudge, not the main objective.

Worked Examples

Worked example Example 2 — Routing a single token

$N=4$ experts. Router logits for token $x$ : $h = [2.0,\; 1.0,\; 0.5,\; 3.0]$ , Top- $k$ with $k=2$ .

Softmax: $e^h = [7.39, 2.72, 1.65, 20.09]$ , sum $=31.85$ . $p = [0.232, 0.085, 0.052, 0.631]$ . Why? Normalizes scores into a probability distribution over experts.

Top-2 experts: indices $4$ ( $p=0.631$ ) and $1$ ( $p=0.232$ ). Why? Two highest probabilities; the other two get gate $0$ and do nothing.

Re-normalize: denominator $=0.631+0.232=0.863$ . $g_4 = 0.731,\ g_1 = 0.269$ . Why? Used gates must sum to 1 so the weighted output is a proper convex blend.

Output: $y = 0.731\,E_4(x) + 0.269\,E_1(x)$ . Experts 2 and 3 never evaluated.

Common Mistakes (Steel-manned)

Recall Feynman: explain to a 12-year-old

Imagine a hospital with 64 specialist doctors. A dense model is like making every patient see all 64 doctors — accurate but insanely slow and expensive. MoE puts a smart receptionist (the router) at the door who looks at your problem and sends you to just the 2 best-matched doctors. The hospital still has all that expertise, but each patient only uses a tiny bit of it. To keep it fair, there's a rule so one popular doctor doesn't get all the patients while others sit idle (load balancing), and each doctor can only see so many patients per hour (capacity).

Active Recall Flashcards

#flashcards/ai-ml

What problem does MoE decouple that dense models tie together?

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

In an MoE layer, what does the router output before Top-k selection?

A softmax distribution

p_i

over the

N

experts from logits

W_g x

Why re-normalize the gates over the Top-k set?

So the selected gate weights sum to 1, giving a proper weighted blend of expert outputs.

FLOPs per token in MoE scale with which quantity —

N

k

k

(active experts), not

N

(total experts).

What is the "expert collapse" failure and its fix?

Router sends most tokens to a few experts (rich-get-richer); fixed by an auxiliary load-balancing loss.

Write the load-balancing loss and say when it's minimized.

\mathcal{L}_{aux}=\alpha N\sum_i f_i P_i

; minimized when

f_i=P_i=1/N

(uniform usage).

What is expert capacity

C

and what happens when it's exceeded?

C=\text{cap\_factor}\times(\text{tokens}/N)

; excess tokens are dropped and pass through the residual.

Why not use a hard argmax router instead of softmax?

Argmax is non-differentiable; softmax gives smooth gradients so the router can learn.

For

k=2

N=64

, 7M-param experts, how many active vs total expert params?

Active

=2\times7\text{M}=14

M; total

=64\times7\text{M}=448

Where does the compute saving in MoE physically come from?

Non-selected experts get gate 0 and are never evaluated (zero FLOPs).

Connections

Feed-Forward Networks (Transformer FFN) — each expert is an FFN block.
Softmax and Gating Functions — the router's normalization step.
Scaling Laws — MoE moves along the capacity axis cheaply.
Sparse vs Dense Models — the fundamental design trade-off.
Load Balancing and Auxiliary Losses — keeping experts utilized.
Model Parallelism & Expert Parallelism — how experts are sharded across devices.
Conditional Computation — the general principle MoE instantiates.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, dense model mein har token ko poore network se guzarna padta hai — matlab saare parameters har token ke liye compute hote hain. Yeh accurate to hai par bahut mehenga. MoE ka jugaad simple hai: hum bahut saare chhote "expert" networks rakhte hain, par har token ke liye sirf 2 (ya 1) best experts ko chalate hain. Ek chhota sa "router" decide karta hai ki yeh token kis expert ko jaana chahiye — bilkul receptionist ki tarah jo patient ko sahi doctor ke paas bhejta hai.

Router kaam kaise karta hai? Pehle woh har expert ke liye ek score nikalta hai ( $W_g x$ ka dot product), phir softmax se usko probability banata hai, phir Top-k experts choose karta hai aur baaki ko gate = 0 de deta hai. Gate 0 ka matlab woh expert compute hi nahi hota — yahi se FLOPs ki bachat aati hai. Selected experts ke outputs ko gate weights se multiply karke jod dete hain, bas ho gaya output.

Ek important point: total parameters badhte hain $N$ (experts ki sankhya) ke saath, par compute badhta hai $k$ (active experts) ke saath. Isliye MoE mein aap knowledge/capacity ko sasta mein badha sakte ho bina compute badhaye. Lekin ek problem hai — "rich get richer": ek popular expert ko saare tokens milne lagte hain aur baaki experts bekaar ho jaate hain. Iske liye ek load-balancing loss ( $\alpha N \sum f_i P_i$ ) lagate hain jo sabko barabar kaam dilata hai. Aur har expert ki ek capacity hoti hai — usse zyada tokens aayein to extra tokens drop ho jaate hain.

Exam/interview ke liye yaad rakho: Route → Softmax → Top-k → Combine, aur "compute = k, capacity = N". Bas itna clear ho gaya to MoE tumhara pakka concept hai.

WHAT is a Mixture-of-Experts?

HOW does routing work? (Derivation from scratch)

The load-balancing problem (WHY MoE needs a helper loss)

Worked Examples

Common Mistakes (Steel-manned)

Active Recall Flashcards

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Scaling & Efficient Architectures

Connections