Mixture of Experts: Conditional Computation

Walk into the OPD block of a large multi-specialty hospital — the kind with forty-plus departments, from cardiology to dermatology to oncology. Every doctor on staff represents years of specialized training; collectively the hospital has enormous medical capacity. But no patient sees forty doctors. A triage desk reads the patient's symptoms and sends them to one, occasionally two, relevant specialists. A person with a fractured wrist never occupies an oncologist's time, and the oncology department's capacity sits idle for that patient without costing the hospital anything for that visit. The hospital's total expertise can keep growing — add a new department, hire more specialists — without making every single patient visit slower, because each visit only ever touches a small, relevant slice of that expertise.

That is precisely the bet a Mixture-of-Experts (MoE) layer makes inside a neural network, and it is the central idea behind "conditional computation": which parameters get used is decided per input, instead of every input marching through every parameter. This chapter builds that mechanism from the dense feed-forward layer you already know, traces one token through a small MoE layer by hand, and shows why this is the technique that let recent large language models grow to hundreds of billions of parameters without a proportional blowup in inference cost.

From One FFN to a Bank of Experts

Recall the feed-forward sub-layer inside a Transformer block. For a token's hidden vector x of dimension d_model, the standard FFN is

FFN(x) = W2 · ReLU(W1 · x + b1) + b2

with an inner "up-projection" to a wider dimension d_ff (typically 4 × d_model) and a "down-projection" back to d_model. Every token, at every layer, passes through this entire block — all of W1 and W2 get multiplied against it, regardless of whether the token is a comma, a rare technical term, or the start of a new sentence. In a dense Transformer, the FLOPs spent per token scale directly with the total number of FFN parameters. Want a bigger, more knowledgeable model? You pay for it on every single token, forever, at inference time.

A Mixture-of-Experts FFN layer breaks that coupling. Instead of one FFN, the layer holds N independent FFNs called experts, E_1, E_2, …, E_N, each shaped like the FFN above. A small extra component, the router (or gating network), looks at the token and decides which k of the N experts (with k far smaller than N — commonly 1 or 2) should process it. Only those k experts actually run their forward pass; the other N − k experts contribute nothing and cost nothing for that token. Total parameter count scales with N (model capacity), while compute per token scales with only k (inference cost) — the two are decoupled. That decoupling, not any single equation, is the entire point of the technique, and it is why the topic sits under "efficient deep learning": you are buying capacity with parameters (cheap: they just sit in memory) instead of buying it with FLOPs (expensive: they run on every forward pass).

The Router: Learning Where to Send Each Token

The router is a single learned linear layer followed by a softmax. Given token hidden state x (dimension d) and a learned weight matrix W_g of shape d × N, it produces one score per expert:

h = x · W_g          // router logits, one per expert
p = softmax(h)        // gate probabilities, sum to 1

The layer then keeps only the top-k entries of p (call this set TopK) and renormalizes them so the kept weights sum to exactly 1:

g_i = p_i / (sum of p_j for j in TopK)   for i in TopK
g_i = 0                                  for i not in TopK

y = sum over i in TopK of  g_i × E_i(x)

The crucial detail that separates this from ordinary attention-style weighting is that g_i = 0 is not "multiply by zero after computing E_i(x)" — it means E_i(x) is never computed at all for that token. The routing decision is a piece of control flow, not just a set of weights. That is what makes this conditional computation rather than conditional weighting.

Architecture of a Top-2 MoE Layer

The diagram below traces exactly the case worked out numerically in the next section: 4 experts, top-2 routing. Two experts execute (solid green, weighted by their renormalized gate values); two are skipped entirely (dashed grey) — their router score was computed, but their FFN body never ran.

Worked Example: Tracing One Token Through a 4-Expert Layer

Take a token with hidden state x = [1, 2, −1] (using d_model = 3 to keep this hand-traceable) entering a layer with N = 4 experts and top-k = 2 routing.

Step 1 — router logits. With router weight matrix

W_g =
[  1.0   0.0  -1.0   0.5 ]
[  0.0   1.0   0.5  -1.0 ]
[ -0.5   0.5   1.0   0.0 ]

compute h = x · W_g column by column:

h1 = (1)(1.0) + (2)(0.0) + (-1)(-0.5) = 1.0 + 0.0 + 0.5 = 1.5
h2 = (1)(0.0) + (2)(1.0) + (-1)(0.5)  = 0.0 + 2.0 - 0.5 = 1.5
h3 = (1)(-1.0) + (2)(0.5) + (-1)(1.0) = -1.0 + 1.0 - 1.0 = -1.0
h4 = (1)(0.5) + (2)(-1.0) + (-1)(0.0) = 0.5 - 2.0 + 0.0 = -1.5

h = [1.5, 1.5, -1.0, -1.5]

Step 2 — softmax. exp(1.5) ≈ 4.4817, exp(-1.0) ≈ 0.3679, exp(-1.5) ≈ 0.2231; the sum is 4.4817 + 4.4817 + 0.3679 + 0.2231 ≈ 9.5544, giving

p = [0.4691, 0.4691, 0.0385, 0.0233]   (sums to 1.0000)

Step 3 — top-2 and renormalize. Experts 1 and 2 tie for the top two scores. Their raw probabilities sum to 0.4691 + 0.4691 = 0.9382, so the renormalized gate weights are g_1 = 0.4691 / 0.9382 = 0.50 and g_2 = 0.50. Experts 3 and 4 get g_3 = g_4 = 0 and are never evaluated.

Step 4 — run only the selected experts. Each expert is a 2-layer FFN with a hidden width of 2 (kept tiny for hand-tracing; real experts use d_ff = 4 × d_model or more, but the arithmetic is identical in kind). Expert 1's parameters:

W1_1 = [[1,0],[0,1],[1,-1]],  b1_1 = [0,0]
W2_1 = [[1,0,1],[0.5,0.5,-1]], b2_1 = [0,0,0]

hidden = ReLU(x · W1_1 + b1_1)
       = ReLU([(1)(1)+(2)(0)+(-1)(1), (1)(0)+(2)(1)+(-1)(-1)])
       = ReLU([0, 3]) = [0, 3]

E1(x) = hidden · W2_1 + b2_1
      = [0(1)+3(0.5), 0(0)+3(0.5), 0(1)+3(-1)]
      = [1.5, 1.5, -3.0]

Expert 2's parameters:

W1_2 = [[-1,1],[1,0],[0,1]],  b1_2 = [0,0]
W2_2 = [[2,-1,0],[0,1,1]],    b2_2 = [0,0,0]

hidden = ReLU(x · W1_2 + b1_2)
       = ReLU([(1)(-1)+(2)(1)+(-1)(0), (1)(1)+(2)(0)+(-1)(1)])
       = ReLU([1, 0]) = [1, 0]

E2(x) = hidden · W2_2 + b2_2
      = [1(2)+0(0), 1(-1)+0(1), 1(0)+0(1)]
      = [2.0, -1.0, 0.0]

Step 5 — combine. y = g_1 · E1(x) + g_2 · E2(x) = 0.5×[1.5,1.5,-3.0] + 0.5×[2.0,-1.0,0.0] = [1.75, 0.25, -1.5]. Experts 3 and 4's parameters — potentially millions of them in a real model — never entered a single multiply for this token.

Here is that entire trace as runnable NumPy, reproducing every number above exactly:

import numpy as np

x = np.array([1.0, 2.0, -1.0])
Wg = np.array([[1.0, 0.0, -1.0, 0.5],
               [0.0, 1.0, 0.5, -1.0],
               [-0.5, 0.5, 1.0, 0.0]])

def softmax(z):
    e = np.exp(z - np.max(z))
    return e / e.sum()

def relu(v):
    return np.maximum(0, v)

h = x @ Wg                                  # [1.5, 1.5, -1.0, -1.5]
p = softmax(h)                              # [0.4691, 0.4691, 0.0385, 0.0233]
top_idx = np.argsort(-p)[:2]                # [0, 1] -> experts 1 and 2
g = p[top_idx] / p[top_idx].sum()           # [0.5, 0.5]

W1_1 = np.array([[1.0,0.0],[0.0,1.0],[1.0,-1.0]])
W2_1 = np.array([[1.0,0.0,1.0],[0.5,0.5,-1.0]])
E1 = relu(x @ W1_1) @ W2_1                  # [1.5, 1.5, -3.0]

W1_2 = np.array([[-1.0,1.0],[1.0,0.0],[0.0,1.0]])
W2_2 = np.array([[2.0,-1.0,0.0],[0.0,1.0,1.0]])
E2 = relu(x @ W1_2) @ W2_2                  # [2.0, -1.0, 0.0]

y = g[0] * E1 + g[1] * E2
print(y)                                    # [1.75, 0.25, -1.5]

Why This Is Efficient: Decoupling Parameters from FLOPs

Generalize the arithmetic above. If every expert has the same size as a normal dense FFN block, a layer with N experts holds N × the parameters of one dense FFN, but a token only ever triggers k of them, so it costs k × the FLOPs of one dense FFN. Compare that to the only alternative for reaching the same total parameter count densely: a dense model with N × the parameters must run all of them, every token, so it costs N × the FLOPs. The MoE layer reaches the same capacity for a fraction k/N of the compute a dense model would need — equivalently, it buys N/k times more parameter capacity than a dense model at matched compute. In the worked example, N/k = 4/2 = 2×. Mixtral 8x7B, a widely studied open model, uses N = 8, k = 2 per FFN layer (attention layers are shared, not duplicated), giving N/k = 4×: its published parameter count is about 46.7 billion total, but only about 12.9 billion parameters are active per token — roughly the FLOPs of a 13B dense model with the standing knowledge of a much larger one.

Two systems details make this actually work on hardware rather than just on paper. First, routing must be balanced. Nothing in the softmax-plus-top-k math above stops the router from learning to send almost every token to the same one or two experts — a failure called expert collapse, where the favored experts get most of the gradient signal, become better, and attract even more traffic, while the rest sit undertrained and wasted. Training adds an auxiliary load-balancing loss, in the style introduced by the Switch Transformer: for a batch, let f_i be the fraction of tokens actually dispatched to expert i, and P_i the average softmax probability the router assigned to expert i across the batch. The auxiliary loss L_aux = N × sum_i(f_i × P_i), added to the main training loss with a small weight, is minimized when routing is uniform across experts; it works around the fact that the hard top-k choice itself isn't differentiable, by pushing on the differentiable probability P_i instead. Second, each expert is given a capacity factor: a hard cap on how many tokens it will accept in one training or inference batch, so that every expert's matrix multiply has a fixed, predictable shape for efficient batched execution on accelerators (this matters especially when experts live on different devices, a setup called expert parallelism). Tokens that route to an already-full expert beyond its capacity are dropped for that expert, which is a real cost of running MoE at scale, not a theoretical footnote. Later designs such as DeepSeek-MoE push the idea further by using many small, fine-grained experts plus a few "shared" experts that every token always uses, splitting general-purpose computation from the sparsely-routed specialized computation.

Common Misconception

Students who first meet MoE almost always assume: "Mixture-of-Experts makes the model smaller." It does the opposite. An MoE layer increases total parameter count — that is the entire mechanism by which it adds capacity. What shrinks is the compute (FLOPs) and, correspondingly, the latency per token, because only k of the N experts run. The two are not the same thing, and the gap between them has a very concrete consequence: GPU memory. Because any token could route to any expert, every expert's parameters must generally sit resident in accelerator memory (or be fetched with real latency cost if they don't), even though only a small fraction actually gets used per token. Mixtral 8x7B needs to hold all 46.7B parameters in memory to serve traffic, even though it only computes with about 12.9B of them for any single token — it needs more VRAM than a dense 13B model with matched compute, not less. The efficiency MoE buys is in FLOPs and therefore inference speed at a given parameter budget; it is not a reduction in the model's footprint on disk or in memory. If a claim about MoE doesn't specify whether it's talking about total parameters, active parameters, or memory footprint, treat it as underspecified.

Active Recall

Attempt these before reading the answers below.

1. In an MoE layer with N = 8 experts and top-k = 2 routing, expert 5 gets the highest router probability for a token, but expert 5 has already reached its capacity-factor limit for this batch. What typically happens to that token's routing to expert 5?
2. A dense Transformer FFN block has d_model = 1024, d_ff = 4096, giving roughly 2 × 1024 × 4096 ≈ 8.39M parameters (ignoring biases). If this block is replaced by 8 experts of that same size with top-2 routing: (a) what is the new total parameter count for the layer, and (b) by what factor does the per-token FLOPs of this MoE layer compare to a dense layer holding the same total (67.1M) parameters?
3. True or False, with correction: "A model with MoE layers uses less GPU memory than a dense model with matched active compute."
4. Why is an auxiliary load-balancing loss necessary during MoE training, and what does it act on given that the top-k routing choice itself is not differentiable?
5. In this chapter's worked example, experts 1 and 2 tied at raw softmax probability 0.4691 each. Why must these be renormalized after top-k selection rather than used directly as the combination weights?
6. Why does a systems engineer, not only the model architect, need to care about the capacity-factor hyperparameter?

Answers.

1. The token overflows expert 5's capacity and is dropped for that expert — in a Switch-Transformer-style implementation it typically passes through with no expert contribution for that layer (effectively just the residual connection), or in some implementations spills to its next-highest-probability expert if that expert still has room. Either way, capacity overflow is a real quality cost, which is exactly why load balancing during training matters: a well-balanced router produces far fewer overflow events at inference.

2. (a) 8 × 8.39M ≈ 67.1M total parameters in the layer. (b) The MoE layer computes with only 2 of the 8 experts per token, i.e. 2 × 8.39M ≈ 16.8M-parameters'-worth of FLOPs, versus a dense layer holding all 67.1M parameters, which must compute with every one of them. That's a 67.1 / 16.8 = N/k = 8/2 = 4× FLOPs reduction for identical total parameter count — the same ratio, and the same reasoning, used above for Mixtral 8x7B.

3. False. MoE models generally need more GPU memory than a dense model with the same active compute, because all N experts' parameters must stay reachable even though only k are used per token. MoE trades memory footprint for FLOPs savings; it does not reduce memory.

4. Without it, the router can collapse onto a small favored subset of experts, which then absorb most of the training signal and improve further while the rest stay undertrained — wasted parameters and, at inference, overloaded hot experts hitting their capacity limits. Because the hard top-k selection isn't differentiable, the auxiliary loss L_aux = N × Σ_i (f_i × P_i) acts on the differentiable soft probability P_i (multiplied by the actual dispatch fraction f_i) so gradient descent can still push routing toward balance.

5. Before renormalization the two kept probabilities sum to only 0.9382, not 1. Using raw probabilities as weights would silently shrink the FFN sub-layer's contribution to the residual stream by a variable amount every time (depending on how much mass the discarded experts happened to hold for that specific token), which is inconsistent output scaling the rest of the network wasn't trained to expect. Renormalizing to make the kept weights sum to exactly 1 keeps the combination a proper weighted average regardless of how routing turned out.

6. Capacity factor fixes the maximum tokens each expert processes per step, which is what gives every expert's matrix multiply a static, predictable shape — essential for efficient batched execution on accelerators, and doubly so when different experts live on different devices (expert parallelism) and communication has to be scheduled around fixed buffer sizes. Set it too low and quality suffers from dropped tokens; set it too high and you pay for padded, unused compute slots. That trade-off is a hardware-utilization and latency-budget decision as much as a modeling one.

Think About It

Think about this: How would you explain mixture of experts: conditional computation 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 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 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, 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.