Top-K Routing › Summary

Top-K routing makes Mixture of Experts computationally practical by selecting only the K experts with the highest routing scores for each token, completely bypassing all other experts. This selective activation is what separates MoE's theoretical scaling advantage from a practical engineering achievement: by keeping K small relative to the total number of experts, MoE models can have far more parameters than dense models while maintaining comparable per-token compute costs.

The routing computation has two steps. First, the router applies a linear transformation to the token's hidden representation to obtain a score for each expert, where higher scores indicate stronger affinity between the token and that expert. Second, the top-K scores are identified using an argmax selection, and the selected experts' scores are normalized via softmax over just the selected pair to produce proper gating weights that sum to 1.

The choice of K involves basic trade-offs. Top-1 routing maximizes efficiency by requiring each token to visit only one expert, but it is susceptible to training instability and can lead to expert collapse without careful auxiliary losses. Top-2 routing, used by models like Mixtral and GShard, doubles the per-token compute but provides substantially smoother training dynamics, better gradient flow to experts, and the ability to blend multiple expert perspectives for tokens with ambiguous routing. Values of K beyond 2 show diminishing returns in most settings and approach the compute cost of dense architectures.

The implementation involves two key steps beyond the routing formula itself: efficient token dispatch that groups tokens by assigned expert for batched GPU processing, and careful preservation of the full routing probability distribution (even for top-1 routing) so that auxiliary load balancing losses can be computed. Naive sequential implementations are functionally correct but impractically slow; efficient implementations process each expert's assigned tokens as a single batched operation.

The main limitation of top-K routing is load imbalance: without additional mechanisms, some experts receive far more tokens than others, wasting capacity, creating computational bottlenecks, and potentially leading to expert collapse. The next chapter addresses this directly with load balancing techniques and auxiliary losses that encourage more uniform token distribution across experts while preserving the learned specialization that makes MoE valuable.