The router: a softmax over experts

The router is a single learned matrix W_g: [d, E]. For a token representation x: [d], it produces one logit per expert and normalizes them into a probability distribution:

h   = x · W_g            # logits, shape [E]
g_i = exp(h_i) / Σ_j exp(h_j)   # gate weights, softmax over E experts
Σ_i g_i = 1,   g_i ≥ 0

So g_i is the router’s confidence that expert i is the right one for this token. The router is minuscule — d × E parameters, e.g. 1024 × 64 ≈ 66K — a rounding error next to the experts themselves. Its whole job is to score, per token, which specialists to consult. Because the softmax is computed independently for each token, two tokens in the same sequence can be routed to entirely different experts; routing is a per-token, not per-sequence, decision.

Advertisement

Top-K selection: where sparsity comes from

A softmax alone still touches all E experts. Sparsity comes from keeping only the top K gate weights and zeroing the rest — almost always K = 1 (Switch-style) or K = 2 (GShard-style), with K << E:

TopK   = indices of the K largest g_i
g'_i  = g_i / Σ_{j ∈ TopK} g_j   for i ∈ TopK,   else 0

The renormalization on the second line matters: after discarding E−K experts, the surviving weights are rescaled to sum to 1 again so the combined output keeps a stable magnitude regardless of how the router split its mass. Top-K is a hard, discrete choice, which is exactly why MoE needs the auxiliary machinery discussed later — the argmax gives no gradient to the experts that were not picked, and left to itself the router tends to keep choosing the same early favorites. Top-K is the lever that sets how sparse the layer is: it fixes how many experts every token pays for.

Advertisement

Combining the chosen experts

Each selected expert runs the full FFN computation on the token, and their outputs are combined as a gate-weighted sum:

y = Σ_{i ∈ TopK} g'_i · E_i(x)
  where  E_i(x) = W2_i · σ(W1_i x)

With K = 1 this is just y = g'_1 · E_1(x) — one expert, scaled by its (renormalized, hence 1.0) gate value; Switch Transformer keeps the raw gate value as the scale so the router still receives a gradient through the chosen expert. With K = 2, two experts run and their outputs blend by confidence. The key structural fact is that y has the same shape [d] as the dense FFN output, so an MoE layer is a drop-in replacement: the rest of the transformer block neither knows nor cares that the FFN became a sparse committee. The gate weight doing double duty — selecting and scaling — is what lets the router learn.

Active vs total parameters

This is the accounting that defines MoE. Split the model into the shared part P_shared (attention, embeddings, layernorms, routers) and the experts, each of size P_exp ≈ 2 · d · d_ff:

Total params  = P_shared + E · P_exp
Active params = P_shared + K · P_exp     (used per token)

Total parameters set the model’s memory footprint — every expert must live in RAM whether or not this token uses it. Active parameters set the per-token compute. The ratio K/E is the sparsity: a model with E = 64, K = 2 stores 32× more FFN parameters than it computes per token. This is why MoE model cards quote two numbers — e.g. ‘46B total, 13B active’ — and why the first tells you what hardware you need to hold the model while the second tells you how fast it runs. Confusing the two is the single most common MoE sizing mistake.

Sparse vs dense FLOPs

A matmul of an activation against a weight matrix of m entries costs about 2m FLOPs (a multiply and an add per entry). One FFN expert’s two matrices hold 2 · d · d_ff entries, so the forward FFN cost per token is:

dense FFN / token  ≈ 2 · (2 · d · d_ff) = 4 · d · d_ff
MoE FFN / token    ≈ K · 4 · d · d_ff     (router cost is negligible)

The compute scales with K, not E. Adding experts grows capacity and memory but leaves per-token FLOPs flat, so long as K is fixed. That is the entire economic argument for MoE: buy extra parameters (cheap — they are just weights sitting in memory) without buying extra compute (expensive — it runs on every token, every step). The router adds only 2 · d · E FLOPs per token, trivial beside the experts. Note this is the theoretical floor; dispatch overhead, load imbalance, and dropped tokens all erode it in practice.