Mixture of Experts: Sparse Gating Networks

The parameter wall

Every dense transformer feed-forward sublayer you have studied so far pays for its full parameter count on every single token, every single forward pass, every single backward pass. Double the hidden width of the FFN and you double the FLOPs per token, whether that token needed the extra capacity or not. This coupling between "how much the model knows" and "how much compute each token consumes" is the central economic constraint of dense scaling: GPT-3-class dense models became expensive to train and, worse, expensive to serve, because inference cost tracks parameter count one-to-one. Sparse Mixture-of-Experts (MoE) layers break that coupling. An MoE layer can store far more parameters than any single token actually touches, because a learned router looks at each token and activates only a small subset of the available sub-networks — the "experts" — for that token. Capacity grows with total parameters; compute per token grows only with the size of the active subset. This chapter derives exactly how that routing works, traces one token through a real gating computation by hand, and shows why the parameter/compute decoupling is not a footnote but the entire point of the architecture.

Grounding the mechanism: IRCTC's grievance triage

IRCTC's customer grievance system receives complaints that span very different domains: catering quality, PNR and booking errors, refund delays, station security. No single desk officer is equally good at all four, and it would be wasteful to route every complaint to every desk. Instead, a triage step reads the complaint and scores it against each desk's specialization, then forwards the complaint to the top two most relevant desks, weighting their combined response by how confident the triage was in each match. A "seat not allotted despite confirmed PNR" complaint might score high on both the booking desk and the refund desk, and receive a resolution that blends both — mostly booking guidance, partially refund guidance. Each desk also has a daily capacity; if the booking desk is already saturated with today's complaints, an overflowing complaint either waits or gets handled by a fallback path, and if triage always dumps complaints onto one popular desk while others sit idle, the whole system needs a correction to keep every desk usefully staffed.

Translate this directly into MoE vocabulary. The complaint's feature representation is the token vector x. The triage classifier is the router. Each desk is an expert — a small neural network specialized (through training, not hand-design) in some region of the input space. The triage's confidence score per desk is the router logit. Selecting only the top few desks per complaint is sparse top-k gating. How much each selected desk's answer counts toward the final resolution is the softmax gate weight. The daily desk quota is the expert capacity factor, and the correction that keeps desks evenly loaded is the load-balancing auxiliary loss. Every one of these has an exact mathematical counterpart in a transformer MoE layer, which we now define formally.

The formal sparsely-gated MoE layer

A transformer block ordinarily alternates a self-attention sublayer with a dense feed-forward sublayer. An MoE transformer keeps attention dense and shared across all tokens, but replaces the FFN sublayer with N independent expert networks E₁, …, E_N, each structured like a normal FFN (Linear → activation → Linear) mapping ℝ^d → ℝ^d. A router — a single learned weight matrix W_g ∈ ℝ^N×d — produces one logit per expert for each token: h(x) = W_gx. Shazeer, Mirhoseini, Maziarz, Davis, Le, Hinton, and Dean (2017), in the paper that introduced this architecture for language modeling ("Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer," ICLR 2017), define a KeepTopK function that sets every logit outside the top k to −∞, and the gate as G(x) = softmax(KeepTopK(h(x), k)). Because softmax of −∞ is exactly 0, the layer's output

y(x) = sum_{i=1}^{N} G(x)_i * E_i(x)

only needs the k experts with nonzero gate weight actually computed — the rest contribute an exact zero and can be skipped entirely. This is what makes the layer sparse: compute per token scales with k, not with N, while the parameter count stored on disk and in accelerator memory scales with N. The 2017 paper also proposes noisy top-k gating, which adds tunable Gaussian noise to the logits before the KeepTopK step during training, encouraging the router to explore experts it would otherwise never select and reducing the risk of a small clique of experts dominating early in training.

Two routing choices dominate production systems, and they trade smoothness for throughput. Fedus, Zoph, and Shazeer's Switch Transformer (2021, published in JMLR 2022) simplifies routing to k = 1 — a single expert per token — which maximizes sparsity and lets them scale to 1.6 trillion parameters using 2048 experts per MoE layer, but it means each token's fate rests on one router decision with no fallback. Lepikhin et al.'s GShard (2020) and Jiang et al.'s Mixtral of Experts (2024) both use k = 2: a second expert blends in, giving the router room to hedge and giving training a smoother gradient signal, at roughly twice the FLOPs of top-1 for that layer. Neither choice is free; the tradeoff is architectural, not accidental.

Worked example: routing one token through a 4-expert layer

Take a toy MoE layer with N = 4 experts, hidden dimension d = 4, and top-k = 2 gating. The token embedding is x = [1.0, 0.5, −0.5, 2.0]. The router weight matrix, one row per expert, is:

W_g = [[ 0.10,  0.20,  0.00, 0.05],   # Expert 0
       [-0.10,  0.30,  0.20, 0.10],   # Expert 1
       [ 0.05, -0.05,  0.10, 0.20],   # Expert 2
       [ 0.20,  0.10, -0.10, 0.02]]   # Expert 3

Step 1 — router logits. Each logit is the dot product of a router row with x:

h₀ = 0.10(1.0) + 0.20(0.5) + 0.00(−0.5) + 0.05(2.0) = 0.10 + 0.10 + 0 + 0.10 = 0.30
h₁ = −0.10(1.0) + 0.30(0.5) + 0.20(−0.5) + 0.10(2.0) = −0.10 + 0.15 − 0.10 + 0.20 = 0.15
h₂ = 0.05(1.0) − 0.05(0.5) + 0.10(−0.5) + 0.20(2.0) = 0.05 − 0.025 − 0.05 + 0.40 = 0.375
h₃ = 0.20(1.0) + 0.10(0.5) − 0.10(−0.5) + 0.02(2.0) = 0.20 + 0.05 + 0.05 + 0.04 = 0.34

Step 2 — KeepTop2. Sorted descending: h₂ = 0.375, h₃ = 0.34, h₀ = 0.30, h₁ = 0.15. The top two are Expert 2 and Expert 3; Expert 0 and Expert 1 are set to −∞ and drop out entirely.

Step 3 — softmax over the kept two only. Subtract the max (0.375) for numerical stability: exponents are e⁰ = 1 for Expert 2 and e^−0.035 ≈ 0.96561 for Expert 3. Sum = 1.96561.

g₂ = 1 / 1.96561 ≈ 0.5088, g₃ = 0.96561 / 1.96561 ≈ 0.4912 (these sum to 1.0000, confirming the softmax is well-formed over exactly the two active experts).

Step 4 — expert outputs. For this hand trace, each expert is simplified to a single affine map E_i(x) = v_i·x + b_i (a real transformer expert is a two-layer FFN with a nonlinearity; the weighted-sum combination rule below is identical regardless of how complex each expert's internals are). Let v₂ = [0.5, −0.2, 0.3, 0.1], b₂ = 0.05, and v₃ = [0.1, 0.4, −0.3, 0.2], b₃ = −0.10:

E₂(x) = 0.5(1.0) − 0.2(0.5) + 0.3(−0.5) + 0.1(2.0) + 0.05 = 0.5 − 0.1 − 0.15 + 0.2 + 0.05 = 0.50
E₃(x) = 0.1(1.0) + 0.4(0.5) − 0.3(−0.5) + 0.2(2.0) − 0.10 = 0.1 + 0.2 + 0.15 + 0.4 − 0.10 = 0.75

Step 5 — combine. y(x) = g₂·E₂(x) + g₃·E₃(x) = 0.5088(0.50) + 0.4912(0.75) = 0.2544 + 0.3684 = 0.6228. Note that Expert 0 and Expert 1 never had to be evaluated at all — the layer only computed two of its four expert networks for this token.

The same computation as executable, traceable code:

import numpy as np

def moe_forward(x, W_g, expert_v, expert_b, k=2):
    logits = W_g @ x                              # h(x), shape (N,)
    top_idx = np.argsort(logits)[-k:]             # ascending sort; last k = top-k
    top_logits = logits[top_idx]
    gate = np.exp(top_logits - top_logits.max())  # softmax over the k kept logits only
    gate = gate / gate.sum()
    y = 0.0
    for idx, g in zip(top_idx, gate):
        y += g * (expert_v[idx] @ x + expert_b[idx])
    return logits, dict(zip(top_idx.tolist(), gate.tolist())), y

x = np.array([1.0, 0.5, -0.5, 2.0])
W_g = np.array([
    [ 0.10,  0.20,  0.00, 0.05],
    [-0.10,  0.30,  0.20, 0.10],
    [ 0.05, -0.05,  0.10, 0.20],
    [ 0.20,  0.10, -0.10, 0.02],
])
expert_v = np.array([
    [0.0, 0.0, 0.0, 0.0],   # Expert 0, unused this trace
    [0.0, 0.0, 0.0, 0.0],   # Expert 1, unused this trace
    [0.5, -0.2, 0.3, 0.1],  # Expert 2
    [0.1, 0.4, -0.3, 0.2],  # Expert 3
])
expert_b = np.array([0.0, 0.0, 0.05, -0.10])

logits, gate_weights, y = moe_forward(x, W_g, expert_v, expert_b, k=2)
print(logits)
print(gate_weights)
print(round(y, 4))

Tracing this by hand: logits (a NumPy array) holds the values 0.3, 0.15, 0.375, 0.34 — printed by NumPy in its own array format, [0.3 0.15 0.375 0.34 ], not Python list notation. np.argsort returns indices in ascending value order, [1, 0, 3, 2], and slicing [-2:] keeps the last two, [3, 2] — Expert 3 then Expert 2. The softmax over their logits [0.34, 0.375] gives gate = [0.4912..., 0.5088...], so gate_weights holds (rounded to 4 decimals for readability) {3: 0.4912, 2: 0.5088} — the unrounded float64 values the code actually prints are {3: 0.49125089311975967, 2: 0.5087491068802403}, since the dict is built directly from gate.tolist() with no rounding step — and round(y, 4) prints 0.6228 — matching the hand derivation exactly, including which two experts fired.

The gating mechanism, end to end

Why sparsity pays: parameters and FLOPs decouple at scale

The toy example moves two of four experts; production systems move a small fraction of a much larger pool, and the savings compound. Fedus, Zoph, and Shazeer's Switch Transformer (2021/2022) trains a 1.6-trillion-parameter model, Switch-C, with 2048 experts per MoE layer and top-1 routing — each token activates roughly 1/2048th of the expert pool at that layer, while the model's storage footprint reflects the full 1.6T parameters. Lepikhin et al.'s GShard (2020) scaled a top-2 multilingual translation Transformer to 600 billion parameters using automatic sharding across TPU pods, demonstrating that the routing idea composes with distributed training infrastructure, not just with a single accelerator. Jiang et al.'s Mixtral of Experts (2024) gives the cleanest recent numbers to reason about: 8 experts per MoE layer, top-2 routing, 46.7 billion total parameters, but only about 12.9 billion parameters active per forward pass, because attention layers are shared and only two of the eight FFN experts per layer fire for any given token.

That 46.7B-to-12.9B gap is the entire economic argument for the architecture, made concrete: 46.7 / 12.9 ≈ 3.62. A dense model matched to Mixtral's active compute budget would need roughly 3.6 times fewer stored parameters and, correspondingly, far less representational capacity — the knowledge and specialization spread across all eight experts, of which any given token draws on only two. Mixtral is reported to match or exceed dense Llama-2-70B on most benchmarks while running inference at roughly the compute cost of a 13B dense model. The FLOPs a forward pass consumes are governed by k and the size of each expert, full stop — not by N. This is precisely why doubling the number of experts in an MoE layer, unlike doubling the width of a dense FFN, does not double the per-token compute cost; it only requires more memory to store the additional experts.

Keeping every expert busy: the load-balancing auxiliary loss

A router trained purely to minimize task loss has no built-in incentive to spread tokens evenly across experts. Left unchecked, it can collapse onto a small clique of "favorite" experts early in training — those experts get more gradient updates, become better at whatever they already handle, and are selected even more often, while the rest stay undertrained. This is analogous to the IRCTC triage system dumping every complaint on the refund desk while the catering desk sits idle: not because refunds are the right answer more often, but because the triage system got stuck in that habit. Shazeer et al. (2017) and later Fedus et al. (2021/2022) address this with an auxiliary load-balancing loss added to the training objective:

L_aux = alpha * N * sum_i ( f_i * P_i )

where N is the number of experts, f_i is the fraction of tokens in the batch whose top-1 choice is expert i, P_i is the mean router probability mass assigned to expert i across the batch, and α is a small coefficient (commonly around 0.01) that keeps this term from dominating the main task loss.

Suppose a batch of 4 tokens is routed top-1 with counts [1, 0, 3, 0] across four experts, so f = [0.25, 0, 0.75, 0], and the mean softmax probabilities across the batch are P = [0.20, 0.15, 0.55, 0.10]. Then Σf_iP_i = 0.25(0.20) + 0(0.15) + 0.75(0.55) + 0(0.10) = 0.05 + 0.4125 = 0.4625, and with α = 0.01, N = 4: L_aux = 0.01 × 4 × 0.4625 = 0.0185. Compare a perfectly balanced batch where every expert gets exactly one token and equal probability mass, f_i = P_i = 0.25 for all i: Σf_iP_i = 4 × (0.25 × 0.25) = 0.25, and L_aux = 0.01 × 4 × 0.25 = 0.01, the minimum value this loss can take. The imbalanced routing above (0.0185) is penalized nearly twice as heavily as the balanced case (0.01), which is exactly the corrective pressure needed to push the router away from expert collapse.

Balancing has a second, purely computational side: each expert typically has a fixed capacity — a maximum number of tokens it will process in a batch, set by a capacity factor to bound memory and keep distributed training synchronized (GShard, Lepikhin et al., 2020). If more tokens are routed to an expert than its capacity allows, the overflow tokens are dropped for that expert, regardless of how confident the router was. This is the direct computational counterpart of the IRCTC desk hitting its daily quota — some complaints simply cannot be seen by their first-choice desk today.

The misconception

Every architecture a student has studied before MoE — deeper CNNs, wider dense FFNs, more attention heads — ties "more parameters" directly to "more compute per input." It is natural, and wrong, to carry that assumption into Mixture-of-Experts and conclude that an MoE layer with more experts is proportionally more expensive to run per token. The opposite is true by construction: the KeepTopK gate guarantees exactly k experts execute per token no matter how large N grows. Growing N grows the model's total parameter count, its storage footprint, and its representational breadth — but not the FLOPs any single token incurs, because the router still only wakes up k of them. Mixtral's 46.7B stored parameters against roughly 12.9B active per token is the concrete proof: quadrupling the number of experts in a layer while holding k fixed leaves per-token inference cost essentially unchanged, while quadrupling the width of a dense FFN would quadruple it. The parameter count is a statement about memory and capacity; the active-expert count, not the total, is the statement about compute.

Active recall

Attempt each question before reading its answer.

Q1. Why does an MoE layer reduce inference compute relative to a dense model with the same total parameter count?
Q2. Using the router matrix W_g from the worked example, a new token arrives: x′ = [0, 1, 1, 0]. Compute all four router logits and the top-2 gate weights.
Q3. If this layer used top-1 routing (as in Switch Transformer) instead of top-2, holding the router weights W_g fixed, what is the gate weight and output y for the original token x = [1.0, 0.5, -0.5, 2.0]?
Q4. Suppose Expert 2 is at capacity when the original token x arrives, so it must be dropped for that expert and only Expert 3 fires. Give the output y under (a) gate renormalization to 1.0 for the sole surviving expert, and (b) no renormalization, where Expert 3 simply keeps its original softmax weight.
Q5. A batch of 4 tokens is routed top-1 with counts [1, 0, 3, 0] across 4 experts, and mean router probabilities are P = [0.20, 0.15, 0.55, 0.10]. Compute L_aux with α = 0.01, and compare it to the theoretical minimum for a perfectly balanced batch.
Q6. True or false: an 8-expert MoE layer with top-2 routing has 8 times the inference FLOPs of a single dense expert-sized network. Justify your answer using Mixtral's published figures.

A1. Compute is governed by how many experts actually execute per token (k), not by how many exist (N), because the KeepTopK gate zeroes out every non-selected expert before the softmax; those zero-weighted experts are never evaluated. Total parameters — and therefore memory footprint — scale with N, but FLOPs per token scale with k, so a model can carry far more capacity than it spends per token.

A2. h₀ = 0.10(0)+0.20(1)+0.00(1)+0.05(0) = 0.20. h₁ = -0.10(0)+0.30(1)+0.20(1)+0.10(0) = 0.50. h₂ = 0.05(0)-0.05(1)+0.10(1)+0.20(0) = 0.05. h₃ = 0.20(0)+0.10(1)-0.10(1)+0.02(0) = 0.00. Logits: [0.20, 0.50, 0.05, 0.00]. Top-2 are Expert 1 (0.50) and Expert 0 (0.20). Softmax: exponents e⁰=1 (Expert 1) and e^-0.30≈0.7408 (Expert 0), sum≈1.7408. g₁ = 1/1.7408 ≈ 0.5744, g₀ = 0.7408/1.7408 ≈ 0.4256. Note the routing decision itself flipped entirely: this token activates Experts 0 and 1, not Experts 2 and 3 as before — a different token, a different pair of specialists.

A3. KeepTop1 keeps only the single highest logit, h₂ = 0.375 (Expert 2). Softmax over one value is trivially 1.0, so g₂ = 1.0 and every other gate is exactly 0. Output y = 1.0 × E₂(x) = 0.50 — identical to Expert 2's raw output, with no blending from Expert 3 at all. This exposes the brittleness top-1 trades for its extra sparsity: the entire output now rests on one router decision, and the smoothing effect the second expert provided in the top-2 case (which pulled y up to 0.6228) disappears.

A4. Original top-2 gates were g₂ ≈ 0.5088, g₃ ≈ 0.4912. With Expert 2 dropped: (a) renormalizing the sole survivor to gate 1.0 gives y = 1.0 × E₃(x) = 0.75. (b) Leaving Expert 3's original softmax weight unchanged (not renormalizing) gives y = 0.4912 × 0.75 = 0.3684, noticeably smaller — the token effectively loses the portion of the layer's contribution that Expert 2 would have supplied, typically compensated elsewhere by the transformer's residual connection. Real systems differ on which convention they use; GShard-style capacity-factor implementations (Lepikhin et al., 2020) make this an explicit design choice, not an incidental detail.

A5. f = [0.25, 0, 0.75, 0]. Σf_iP_i = 0.25(0.20) + 0 + 0.75(0.55) + 0 = 0.05 + 0.4125 = 0.4625. L_aux = 0.01 × 4 × 0.4625 = 0.0185. The perfectly balanced baseline (f_i=P_i=0.25 for all four) gives Σ = 4(0.0625) = 0.25 and L_aux = 0.01 × 4 × 0.25 = 0.01, the global minimum. 0.0185 is nearly double the minimum, correctly flagging that 3 of 4 tokens concentrated on one expert.

A6. False. Mixtral's own published figures show 46.7B total parameters against roughly 12.9B active parameters per token with 8 experts and top-2 routing — active compute corresponds to about 2 experts' worth of FFN work (plus shared attention), not 8. If the claim were true, active parameters would equal or exceed total parameters, which contradicts the 12.9B < 46.7B relationship Mixtral reports. Total expert count N sets capacity and memory footprint; only k, the number of experts actually gated on, sets per-token FLOPs.

Think About It

Think about this: How would you explain mixture of experts: sparse gating networks to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where mixture of experts: sparse gating networks is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting mixture of experts: sparse gating networks to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind mixture of experts: sparse gating networks, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.