Top-K Routing › Visualizing Routing Patterns

Abstract descriptions of routing behavior become much more intuitive when visualized. Let's examine what the routing patterns look like in practice for a sequence of tokens, and what these patterns tell us about how the MoE layer is using its expert capacity.

The heatmap reveals the sparsity pattern of top-K routing directly. Each token (row) activates exactly 2 experts (colored cells), with the color intensity showing the routing weight. White cells represent experts that receive zero weight for that token, and in a model with 8 experts, 75% of the cells are white for any given token. This sparse activation is what makes MoE computationally efficient: instead of all 8 experts processing every token, only 2 do.

Looking at the pattern across tokens, you can also observe how different tokens prefer different expert combinations. Some experts may appear more frequently (showing they specialize in common token types), while others appear less frequently (showing they specialize in rarer but important token types). In a well-trained MoE model, this heatmap would reveal a structured semantic organization: similar tokens would tend to activate similar pairs of experts.

Let's also examine how tokens are distributed across experts, which is a key diagnostic for detecting load imbalance:

This distribution reveals a common challenge with naive top-K routing: without explicit encouragement, tokens do not distribute evenly across experts. Some experts receive many more tokens than others, while some may receive very few. This load imbalance is a necessary issue that the next chapter addresses directly. The red dashed line shows what perfectly uniform distribution would look like: with 16 tokens, K=2, and 8 experts, the ideal distribution would give exactly 4 tokens to each expert. Deviations from this represent wasted capacity: overloaded experts may need to drop tokens, and underloaded experts sit idle.

The weight distribution histogram provides another view of routing behavior. In this example with random router logits, the distribution is roughly symmetric around 0.5: when logits are random, the gap between the top two is equally likely to favor either direction. In a trained model, the distribution would reflect learned routing patterns: tokens where one expert clearly dominates would cluster near 1.0, while ambiguous tokens would cluster near 0.5.