Top-K Routing › Top-K Routing

In the previous chapter, we explored gating networks that compute a probability distribution over all experts for each input token. The full distribution contains a score for every expert, but using every expert to process every token would eliminate the computational advantages that make Mixture of Experts attractive in the first place. The key to making MoE practical is a deceptively simple idea: throw away most of those scores and only forward each token to the handful of experts with the highest routing scores. This selective activation is called top-K routing, and it is the mechanism that allows MoE models to scale to trillions of parameters while keeping per-token computation costs manageable.

Think of top-K routing as a highly specialized staffing system inside a large organization. Imagine a company with hundreds of specialized departments, each with deep expertise in a narrow area: tax law, structural engineering, biomedical research, literary translation, and so on. When a task arrives, rather than routing it through every department in the company (which would be impossibly slow), a dispatcher quickly assesses the nature of the task and forwards it to just the two or three most relevant departments. Those departments work on the task and their outputs are combined into a final recommendation. The dispatcher's judgment about which departments to involve is learned over time based on which assignments produced good outcomes. Top-K routing works the same way: the router learns which experts are relevant for which kinds of tokens, and at inference time it makes rapid, sparse routing decisions that activate only those relevant experts.

Top-K routing makes sparsity useful. By forcing each token to use only K out of N experts, computation scales with K rather than N, so doubling the number of experts does not double the cost of processing a token. The competition to be selected also pressures experts to specialize, since a broadly mediocre expert will rarely win against experts that excel on particular token types.

Top-K routing answers a basic question: given a token and its routing scores across all experts, which experts process that token, and how should we combine their outputs? The choice of K affects model capacity, training stability, inference efficiency, and the degree of specialization that emerges in the expert networks. This chapter explains how top-K routing works, why different K values lead to different model behaviors, how to implement routing and combination correctly and efficiently, and what limitations arise when deploying MoE architectures.

The journey through this chapter will take us from the simplest possible case, routing each token to exactly one expert, through the standard modern choice of two experts, and then to a principled analysis of what happens as K grows. Along the way we will build a complete, working PyTorch implementation and trace through a concrete numerical example that makes the abstract mechanics concrete.