Top-K Gating and Sparse Routing

A mixture-of-experts layer contains many separate sub-networks ("experts"), each capable of specializing in a different kind of input. Running every expert on every input would give you the benefit of all that specialized capacity, but at a cost scaling with the total number of experts — defeating much of the point, since you'd rather have hundreds of experts available without paying for all of them every forward pass. Top-k gating is the routing rule that makes this affordable: score every expert on relevance, then only actually run the handful (typically 1 or 2) that scored highest, leaving the rest of the model's capacity unused — and uncosted — for that input.

The analogy: routing a support ticket to the right specialists

A large company might employ hundreds of specialists, but a single support ticket doesn't get read by all of them — a triage system scores the ticket against each specialist's expertise and routes it to the top two or three most relevant. The company still benefits from hundreds of specialists across its entire ticket volume, but no individual ticket costs more than a couple of people's time. Top-k gating is that triage system, run automatically inside a network: it decides, per input, which small subset of experts is worth consulting.

The mechanics: score, select top-k, weight, combine

Given $N$ experts and an input $x$, a small learned gating network first produces a relevance score for every expert:

$$g = \text{softmax}(W_g x)$$

In words: a linear layer followed by a softmax turns the input into $N$ scores summing to 1, one per expert — a relevance ranking. Top-k gating keeps only the $k$ largest scores (commonly $k=1$ or $2$), zeroes the rest, and typically renormalizes the surviving $k$ to sum to 1. Only experts with a nonzero gate are actually run — a zero-gated expert's output would be wasted work, so implementations skip it entirely. The layer's output is the weighted combination of just those $k$ experts:

$$y = \sum_{i \in \text{top-}k} g_i \cdot E_i(x)$$

In words: only the top-k experts run at all, and their outputs are blended in proportion to how strongly the gate favored each of them.

Worked example 1: scoring and selecting from 6 experts

Suppose a gating network scores 6 experts as $g=(0.05, 0.35, 0.10, 0.30, 0.02, 0.18)$ (softmax-normalized, summing to 1). With $k=2$, keep only the two largest — expert 2 at $0.35$, expert 4 at $0.30$ — and zero the rest. Renormalizing: $0.35/0.65=0.538$ and $0.30/0.65=0.462$. Only experts 2 and 4 run; $y=0.538\,E_2(x)+0.462\,E_4(x)$. The other 4 experts, despite existing and being fully trained, contribute nothing to this input and are never even evaluated.

Worked example 2: the compute savings, concretely

Suppose each expert costs 1 million FLOPs to run, with $N=64$ experts. Running all 64 densely costs $64\text{M}$ FLOPs per input. With top-k gating at $k=2$, only 2 run, costing roughly $2\text{M}$ FLOPs — a $32\times$ reduction, even though the layer's total parameter count (all 64 experts, fully stored) is unchanged. This is the tradeoff sparse routing buys: parameter count scales with $N$, while compute cost per input scales only with $k$.

Picture the curve above as compute cost against $k$ with $N$ held fixed at a large number — cost grows only with the small exponent $k$, not with the much larger base $N$, which is exactly why keeping $k$ small while $N$ grows is the source of the savings.

What this means in practice

Top-k gating is the mechanism behind mixture-of-experts architectures that scale to enormous total parameter counts, many times larger than would be practical to run densely, while keeping the compute cost of any single forward pass close to a much smaller dense network. The gating network must be trained carefully — left unconstrained, it tends to collapse onto favoring a few popular experts, starving the rest of training signal, so real implementations add a load-balancing penalty encouraging routing to spread more evenly across all $N$ experts.

Practice

1. With 8 experts scored $g=(0.02, 0.28, 0.04, 0.20, 0.06, 0.30, 0.05, 0.05)$ and $k=3$, which experts are selected, and what are their renormalized gate weights?
2. If each expert costs 2M FLOPs and $N=128$, $k=4$, compute the FLOPs saved per input versus running all experts densely.
3. Why does an untrained or poorly-regularized gating network risk routing almost every input to the same one or two experts, and why is that a problem?