Apollo Hospitals in Chennai runs one of India's largest multi-specialty networks. A patient walking into the OPD does not get examined by every doctor on the payroll — the cardiologist, the neurologist, the orthopedist, the dermatologist, all forty of them. A triage step reads the symptoms and routes the patient to one or two relevant specialists. The hospital still employs all forty doctors — that total capacity is what lets it handle a heart attack on Monday and a fracture on Tuesday — but any single patient's visit only pays for the two doctors who actually see them. Total specialist capacity and per-patient cost are two different numbers, and the triage step is what decouples them.
This is exactly the design problem a Mixture-of-Experts (MoE) layer solves inside a neural network, and it is worth being precise about why it is a problem at all. In a standard Transformer block, every token that flows through the network is processed by the same feed-forward network (FFN) — one fixed set of weights, applied identically to every token, every time. If you want the model to know more — to have more parameters, more representational capacity — the conventional route is to make that one FFN wider. But a wider FFN means more multiply-adds for every single token, whether that token needed the extra capacity or not. Parameter count and compute cost are welded together. MoE unwelds them: replace the single FFN with a bank of N smaller FFNs ("experts"), add a lightweight router that looks at each token and decides which k of the N experts should process it, and run only those k. Total parameter capacity scales with N. Compute cost per token scales with k. You can grow N — the hospital's specialist roster — without growing the bill for any individual patient.
From dense computation to conditional computation
Recall the standard Transformer FFN sublayer: a token's hidden vector x (dimension d_model) is projected up to a wider dimension d_ff, passed through a nonlinearity, then projected back down to d_model. Two weight matrices, sizes d_model × d_ff and d_ff × d_model, so the parameter count of one dense FFN (ignoring biases) is 2 × d_model × d_ff, and the compute cost of one forward pass through it is likewise O(d_model × d_ff) multiply-adds. Every token in every sequence pays this cost. This is dense computation: fixed work per token, no exceptions.
Mixture-of-Experts replaces that single FFN with N independent FFNs of the same shape, plus a small router network — typically just one linear layer W_g of shape d_model × N — that produces a score for each expert given the token. A gating function (almost always softmax followed by a top-k selection) converts those scores into a sparse set of active experts and their combination weights. Only the chosen k experts actually run their FFN on that token; the other N − k contribute nothing to that token's forward pass and receive no gradient from it. This is conditional computation: the network decides, per input, how much and which part of itself to use. The word "mixture" refers to the weighted combination of the chosen experts' outputs; "of experts" refers to the specialised sub-networks; the entire mechanism is often just called a "sparse MoE layer" or "sparsely-gated layer" in the research literature, distinguishing it from a dense model where the "mixture" would trivially include every expert with equal weight.
The mechanism, formally
Given a token vector x ∈ R^d_model and N experts E_1, …, E_N, each a full FFN sub-network, an MoE layer computes:
logits = x · W_g # shape (N,), one score per expert
p = softmax(logits) # shape (N,), a probability per expert
S = top_k_indices(p, k) # the k highest-scoring experts
g_i = p_i / Σ_{j in S} p_j for i in S # renormalize the chosen weights to sum to 1
y = Σ_{i in S} g_i · E_i(x) # weighted sum of only the chosen experts' outputs
Two design choices are doing all the work here. First, top-k selection happens before renormalization — the softmax is computed over all N experts, but only the k largest probabilities survive, and those survivors are rescaled to sum to 1 so the final combination is still a proper weighted average. Second, experts that are not selected are not merely "weighted by zero" in the sense of being computed and then discarded — they are never computed at all for that token. This is the entire source of the compute saving; a zero weight multiplied by an already-computed expert output saves nothing, but skipping the computation entirely saves the full O(d_model × d_ff) cost of that expert's forward pass.
Worked example 1: counting parameters versus counting compute
Take a Transformer layer with d_model = 512 and the conventional 4× expansion, d_ff = 2048. A single dense FFN sublayer has:
dense_params = 2 × d_model × d_ff
= 2 × 512 × 2048
= 2,097,152 ≈ 2.10M parameters
Now replace it with an MoE layer of N = 8 experts, each the same shape as the dense FFN above, with a router using top-k = 2 (top-2 routing, the setting used by Mixtral and by the original Shazeer et al. design). The gating matrix contributes d_model × N = 512 × 8 = 4,096 parameters — small enough to fold into a rounding error against the experts themselves.
moe_total_params = N × dense_params + gating_params
= 8 × 2,097,152 + 4,096
= 16,777,216 + 4,096
= 16,781,312 ≈ 16.78M parameters
moe_active_params = k × dense_params (per token, ignoring the negligible gating cost)
= 2 × 2,097,152
= 4,194,304 ≈ 4.19M parameters
Two ratios matter here, and they are different numbers on purpose. The layer's total parameter capacity is 16.78M / 2.10M ≈ 8× the dense layer — exactly N, since capacity scales with the number of experts you choose to build. But the compute a single token actually triggers is 4.19M / 2.10M = 2× the dense layer — exactly k, since only the chosen experts run. The ratio of total capacity to per-token active compute is N / k = 8 / 2 = 4: this layer holds four times more knowledge than any single token's forward pass touches. That gap is the entire value proposition of MoE — you can keep growing N (more specialists, more total knowledge, better model quality) with only a linear-in-k increase in the FLOPs and latency any given request actually pays for. This is precisely how the largest production language models reach hundreds of billions of total parameters while their inference cost tracks a much smaller "active" parameter count.
Worked example 2: tracing the router and the combination step by hand
The parameter count above explains why MoE is attractive; it says nothing about how the routing decision is actually made. Work through one token completely, with numbers small enough to check by hand.
Let d_model = 4, N = 4 experts, top-k = 2 routing. Take a token embedding x = [1, 0, 1, 0] and a gating matrix
W_g = [[ 1, 0, 2, -1],
[ 0, 1, 0, 1],
[-1, 2, 0, 1],
[ 1, 1, 1, 0]]
Because x has zeros in positions 1 and 3, only rows 0 and 2 of W_g contribute to the matrix–vector product:
logits = x · W_g = row0 + row2
= [1, 0, 2, -1] + [-1, 2, 0, 1]
= [0, 2, 2, 0]
Apply softmax. With exp(0) = 1 and exp(2) ≈ 7.389, the denominator is 1 + 7.389 + 7.389 + 1 = 16.778, giving
p ≈ [0.0596, 0.4404, 0.4404, 0.0596]
Experts 1 and 2 (zero-indexed) tie for the two highest scores, so top-2 selects S = {1, 2}. Renormalizing over just those two: 0.4404 + 0.4404 = 0.8808, and each becomes 0.4404 / 0.8808 = 0.5. The tie makes the renormalized weights come out exactly equal — a useful sanity check, not a coincidence induced by rounding.
Now compute the two selected experts' outputs. Treat each expert here as a single linear map (a stand-in for the two-layer ReLU FFN from the previous section, chosen only so the arithmetic stays checkable by hand):
E1(x) = x · A1, where A1's nonzero rows are row0=[2,1,0,0], row2=[1,1,2,0]
= [3, 2, 2, 0]
E2(x) = x · A2, where A2's nonzero rows are row0=[0,1,1,1], row2=[2,0,1,1]
= [2, 1, 2, 2]
The layer's output is the renormalized weighted sum of only these two — experts 0 and 3 never ran:
y = 0.5 × [3, 2, 2, 0] + 0.5 × [2, 1, 2, 2]
= [1.5, 1.0, 1.0, 0.0] + [1.0, 0.5, 1.0, 1.0]
= [2.5, 1.5, 2.0, 1.0]
Here is the same computation as runnable code, using NumPy directly so every array is defined before use and the printed values can be checked against the hand trace above:
import numpy as np
def softmax(v):
e = np.exp(v - np.max(v))
return e / e.sum()
x = np.array([1, 0, 1, 0], dtype=float)
Wg = np.array([
[ 1, 0, 2, -1],
[ 0, 1, 0, 1],
[-1, 2, 0, 1],
[ 1, 1, 1, 0],
], dtype=float)
logits = x @ Wg # [0. 2. 2. 0.]
print(logits)
p = softmax(logits) # [0.0596 0.4404 0.4404 0.0596]
print(np.round(p, 4))
k = 2
top_idx = np.argsort(p)[-k:] # indices {1, 2} (order between the tie is arbitrary)
top_p = p[top_idx]
g = top_p / top_p.sum() # [0.5 0.5]
print(g)
A1 = np.array([[2,1,0,0],[0,0,0,0],[1,1,2,0],[0,0,0,0]], dtype=float)
A2 = np.array([[0,1,1,1],[0,0,0,0],[2,0,1,1],[0,0,0,0]], dtype=float)
experts = {1: A1, 2: A2} # only the routed experts are ever evaluated
y = sum(g[j] * (x @ experts[idx]) for j, idx in enumerate(top_idx))
print(np.round(y, 4)) # [2.5 1.5 2. 1. ]
Running this produces logits = [0. 2. 2. 0.], p ≈ [0.0596, 0.4404, 0.4404, 0.0596], g = [0.5, 0.5], and a final output of [2.5, 1.5, 2.0, 1.0] — matching the hand derivation exactly. Notice the dictionary experts holds only the two matrices that were actually selected; a real implementation never materializes A0 or A3 for this token at all, which is the whole point.
How expensive is the router itself? A complexity check
It is fair to ask whether the router's own overhead eats into the savings. It does not, and the reason is a clean complexity argument. Computing the N logits costs O(d_model × N) multiply-adds — one matrix–vector product. Selecting the top-k of N scores can be done by a full sort in O(N log N), or faster with a size-k min-heap in O(N log k), or, using an introselect/quickselect-style partial-selection algorithm, in expected O(N). Compare this against the cost of running a single expert's FFN, O(d_model × d_ff).
Take realistic production-scale numbers: d_model = 4096, d_ff = 16384, N = 128 experts. The router's logit computation costs 4096 × 128 = 524,288 multiply-adds; the top-k selection over 128 scores costs at most a few hundred comparisons — utterly negligible by comparison. A single expert's FFN forward pass costs 2 × 4096 × 16384 = 134,217,728 multiply-adds — about 256 times larger than the router's entire logit computation. The router is, computationally, almost free next to the thing it is deciding whether to run. This is precisely why conditional computation pays off: you can afford a cheap decision step in exchange for skipping N − k expensive ones, and the decision step's own cost does not eat meaningfully into the savings even as N grows into the hundreds.
The failure mode nobody's diagram shows: load imbalance
A common misconception is that because the router is trained end-to-end by gradient descent, it will naturally learn to spread tokens evenly across experts, the way a competent human triage nurse would balance patient load across doctors. It does not, and left alone it tends to do close to the opposite. Early in training, gating scores are only mildly non-uniform, but whichever experts happen to receive slightly more tokens also receive proportionally more gradient updates — an expert's weights only change from tokens actually routed to it. Those experts improve faster, which lowers the loss when the router sends them more tokens, which the router's own gradient signal then reinforces by routing them even more traffic. This is a self-reinforcing, rich-get-richer dynamic: a handful of "winner" experts can end up absorbing nearly all tokens while the rest of the bank goes almost untrained, silently throwing away most of the N-fold capacity increase MoE was built to provide.
The fix used across the major systems is an explicit auxiliary load-balancing loss added to the training objective, penalizing the model whenever routing probability mass concentrates on too few experts across a batch (Shazeer et al.'s original importance-and-load losses; Switch Transformer's simplified single load-balancing term). Production systems additionally enforce a hard expert capacity: each expert may accept at most capacity_factor × (tokens_per_batch / N) tokens per batch. A token whose chosen expert is already full is simply dropped from that layer — its representation passes through via the residual connection untransformed, rather than the system dynamically reassigning it. This is a genuine, deliberate engineering tradeoff: a larger capacity factor drops fewer tokens but wastes more memory and compute on padding for underfull experts, while a smaller one is more efficient but degrades more tokens per batch. There is no free setting; every real MoE deployment picks a point on that curve.
From research idea to production systems
Sparsely-gated Mixture-of-Experts layers for neural networks were introduced by Shazeer, Mirhoseini, Maziarz, Davis, Le, Hinton, and Dean in "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer" (ICLR 2017), which inserted an MoE layer between LSTM layers in a language model and scaled to 137 billion parameters — enormous for 2017 — while keeping per-token compute close to that of a much smaller dense model, using noisy top-k gating with an explicit load-balancing loss.
Fedus, Zoph, and Shazeer's Switch Transformer (arXiv 2021; Journal of Machine Learning Research, 2022) simplified the routing to top-1 — each token goes to exactly one expert, cutting router complexity and communication cost — and formalized the expert-capacity-with-token-dropping mechanism described above. Switch Transformer scaled a T5-style model to 1.6 trillion total parameters (the "Switch-C" configuration) while each token's actual FFN computation stayed comparable to a much smaller dense model, demonstrating that total parameter count and per-token compute really could be decoupled at that scale.
More recently, Mistral AI's Mixtral 8x7B (Jiang et al., 2024) brought the idea into widely-used open-weight language models: each Transformer block's FFN sublayer is an MoE layer with 8 experts and top-2 routing, matching this chapter's worked examples exactly in structure. The model's attention layers, embeddings, and layer norms are dense and shared by every token — they are not duplicated per expert — while only the FFN sublayers are split into experts. The result is roughly 46.7 billion total parameters but only about 12.9 billion active parameters per forward pass. Note that 12.9B is not simply 46.7B × (2/8): that naive scaling would only be correct if the entire model were made of MoE layers, but the shared dense components (attention, embeddings) count fully toward both the total and the active figure regardless of routing, so the true active count sits above the naive fraction of the total.
Active recall
Attempt each question before reading its answer.
Q1. In the parameter-counting example (d_model = 512, d_ff = 2048, N = 8, top-2 routing), compute the exact ratio of total layer parameters to active parameters per token, ignoring the gating network.
Q2. Starting from the same example, suppose top-k is increased from 2 to 4 while N = 8 stays fixed. What happens to (a) total parameter count, (b) active compute per token, and (c) the total-to-active ratio?
Q3. In worked example 2 (x = [1,0,1,0], N = 4, logits [0, 2, 2, 0]), suppose the layer instead used top-1 routing. What is the new output, and what practical instability does the exact tie in the logits expose?
Q4. For N = 128 experts, d_model = 4096, d_ff = 16384, compare the cost of computing the router's logits to the cost of one expert's FFN forward pass. Which dominates, and by roughly what factor?
Q5. Explain, mechanistically, why an MoE router trained with no auxiliary load-balancing loss tends to collapse onto a small subset of experts rather than spreading tokens evenly.
Q6. Mixtral 8x7B has 8 experts per layer with top-2 routing, about 46.7B total parameters and about 12.9B active parameters per token. Why is the active count not simply 46.7B × (2/8) = 11.675B?
A1. Total parameters = N × dense_params = 8 × 2,097,152 = 16,777,216. Active parameters per token = k × dense_params = 2 × 2,097,152 = 4,194,304. Ratio = 16,777,216 / 4,194,304 = 4, which is exactly N/k = 8/2 — the ratio depends only on the routing sparsity, not on the absolute size of each expert.
A2. (a) Total parameter count is unchanged at 16,777,216 — it depends only on N, the number of experts you built, not on how many you route to. (b) Active compute per token doubles to 4 × 2,097,152 = 8,388,608, since active compute scales with k. (c) The ratio falls to N/k = 8/4 = 2: the layer is now less sparse, each token pays for more of the total capacity, and the gap between "parameters you have" and "parameters a token uses" — the entire advantage MoE offers over a dense layer of the same total size — shrinks by half.
A3. The softmax probabilities are unchanged, p ≈ [0.0596, 0.4404, 0.4404, 0.0596], because the gating scores are computed before any top-k truncation. But experts 1 and 2 are in an exact tie for the maximum. A real implementation's argmax breaks ties by taking the first occurrence in index order, so top-1 deterministically selects expert 1 (index 1). Since only one expert is chosen, renormalization gives it weight g = 1.0, so the output is that expert's raw output in full: y = E1(x) = [3, 2, 2, 0] — not the blended [2.5, 1.5, 2.0, 1.0] that top-2 produced. The instability this exposes: with real-valued, continuously trained weights, exact ties are rare, but near-ties are common, and top-1 routing means an arbitrarily small perturbation to the input or the weights (a slightly different token, one gradient step) can flip which single expert handles a token, discontinuously changing the output. Top-k ≥ 2 routing is less brittle precisely because it blends across a small set rather than committing to one winner.
A4. Router logits: d_model × N = 4096 × 128 = 524,288 multiply-adds (top-k selection over 128 scores adds at most a few hundred more operations, negligible by comparison). One expert's FFN: 2 × d_model × d_ff = 2 × 4096 × 16384 = 134,217,728 multiply-adds. The expert FFN dominates, by a factor of 134,217,728 / 524,288 ≈ 256. The router's cost is more than two orders of magnitude smaller than a single expert it might decide to skip, which is why the routing decision itself is essentially free relative to the compute it saves.
A5. An expert's parameters only receive gradient updates from tokens actually routed to it. If, early in training, an expert receives even a slightly larger share of tokens (from initialization noise or the data distribution), it gets more gradient signal and improves faster than its peers. A better-performing expert produces lower loss when the router sends it more tokens, and the router's own parameters are trained to minimize loss — so gradient descent on the router itself increases the probability of routing to that already-ahead expert. This closes a positive feedback loop: more tokens → more training → better performance → even more tokens, with no natural force pulling traffic back toward the underused experts. Without an explicit penalty on uneven routing, this loop can collapse most of the traffic onto a small handful of experts, leaving the rest of the parameter bank essentially untrained.
A6. The naive calculation 46.7B × (2/8) would only be valid if every parameter in the model belonged to a mixture-of-experts layer scaled by the routing fraction k/N. But attention layers, token embeddings, and layer norms are dense components shared identically by all tokens — they are not split into experts and are not subject to routing at all, so 100% of their parameters count toward both the total and the active figure. Only the FFN sublayers are structured as 8-expert, top-2 MoE layers. The true active parameter count is therefore the full weight of the shared dense components plus 2/8 of the MoE FFN parameters, which works out to roughly 12.9B — higher than the naive uniform-scaling estimate because the dense components don't get the routing discount at all.
Think About It
Think about this: How would you explain mixture of experts: conditional computation and scaling 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: conditional computation and scaling 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: conditional computation and scaling 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: conditional computation and scaling, 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.