Top-K Routing: Expert Selection in Mixture of Experts Models

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.

Top-1 routing is the most extreme form of sparsity available in a Mixture of Experts architecture: each token is sent to exactly one expert, the one with the highest routing score, and all other experts are completely bypassed. This approach is also called hard routing because the decision is binary: a token either goes to an expert or it does not, with no intermediate values allowed. The simplicity of this approach is appealing because it maximizes the efficiency advantage of MoE. When each token visits only one expert, we can scale the number of experts arbitrarily without changing the computational cost for any individual token.

To understand why top-1 routing is so powerful, consider the arithmetic. If you have 64 experts and each expert's feed-forward network contains 1 billion parameters, the model has 64 billion expert parameters in total. But when routing each token to exactly one expert, only 1 billion of those parameters participate in computing the output for any given token. The other 63 billion parameters sit idle for that token, available to be activated for other tokens that the router deems appropriate. This separation between total model capacity (determined by N experts) and per-token computational cost (determined by K activations) is the basic economic advantage of sparse MoE architectures.

Think of top-1 routing as a binary tournament. Each expert submits a "bid" for each token in the form of a routing score. The highest bidder wins the token entirely. The losers receive nothing from that token, no gradient signal, no opportunity to refine their representation of that kind of input. This creates strong selective pressure: to win routing competitions, an expert must learn to be distinctively better than all other experts at some identifiable class of tokens. The result, when training works well, is a collection of specialists with sharply differentiated capabilities.

To understand how top-1 routing works, we will walk through the process step by step. The computation begins by projecting the token's hidden representation through the router's learnable weight matrix to obtain a score for each expert. Specifically, we compute the routing scores (logits) for all experts by projecting the token representation $\mathbf{x}$ through the router's learnable weights $W_r$:

where:

• $\mathbf{h}$: the vector of routing logits (scores) for all experts, one score per expert
• $W_r$: the router weight matrix of shape $(\text{num\_experts} \times \text{hidden\_dim})$
• $\mathbf{x}$: the input token representation vector of dimension $\text{hidden\_dim}$

This linear transformation is a compatibility function. Each row of the weight matrix $W_r$ can be thought of as a "query vector" for one expert, and the dot product between this query and the token representation measures how well-suited that expert is for processing this particular token. A high dot product means the token's representation is well-aligned with what this expert has learned to handle. Higher scores indicate stronger affinity between the token and the expert.

Why a linear transformation? A more complex routing function, perhaps a small neural network, might produce better routing decisions. The linear router has a practical advantage: it is computationally cheap. Routing must happen before any expert computation, so the overhead of the router should be minimal relative to the cost of the experts themselves. A single matrix multiplication accomplishes this efficiently. Because it operates on the same hidden representation that the model uses for everything else, the linear router can use that representation's full semantic content to make routing decisions.

Then, top-1 routing selects the single expert with the highest score:

where:

• $k^*$: the index of the selected expert
• $\arg\max_i$: the operation that finds the index $i$ maximizing the value
• $h_i$: the routing logit (score) for the $i$-th expert

The argmax operation identifies the expert with the highest score. This expert then processes the token. The output for that token becomes simply the output of the selected expert:

where:

• $\mathbf{y}$: the output vector for the token
• $E_{k^*}$: the feed-forward network function of the selected expert $k^*$
• $\mathbf{x}$: the input token representation vector

This formula is strikingly clean. The entire computation of the expert layer for this token reduces to a single feed-forward pass through one expert network. There is no mixing, no weighting, no combination: just a direct pass through the winner.

Top-1 routing introduces a significant challenge that is easy to overlook at first: the $\arg\max$ operation is not differentiable. Backpropagation requires gradients, and gradients require the operation to be continuous, but argmax produces a discrete integer index that jumps discontinuously as inputs change. How do gradients flow back through a hard selection?

The standard solution uses the straight-through estimator. It uses different computations for the forward and backward passes, accepting the inconsistency as a pragmatic compromise that enables training. During the forward pass, the token goes through only one expert via true hard $\arg\max$ selection. During the backward pass, the operation is treated as a softmax, allowing gradients to flow back to all experts in proportion to their routing probabilities.

Why does this work? Intuitively, the straight-through estimator provides a smooth surrogate for the discrete selection. Experts with higher routing probabilities receive larger gradient signals, which encourages them to specialize further for the types of tokens that currently prefer them. Experts with lower routing probabilities receive smaller gradients, but they still receive some signal, which prevents them from becoming completely unused. The inconsistency between forward and backward passes is a known theoretical weakness, but empirically it works well enough for MoE training to succeed.

Top-1 routing has another characteristic that can be either a feature or a bug depending on your perspective: it forces complete specialization. Each token must commit fully to one expert's computation, with no blending of perspectives from multiple experts. This creates strong incentives for experts to carve out distinct semantic territories: an expert that specializes narrowly and deeply in a particular kind of input will consistently win the routing competition for that kind of input, while a generalist expert that is merely adequate at everything may rarely be selected.

The flip side of this forced specialization is that it makes the model brittle at boundaries. Consider a token that sits semantically between two categories. With top-1 routing, the model must make a definitive choice about which category's expert to use, potentially losing the fine-grained perspective that a second expert could have contributed. A technical legal term, for instance, might benefit from both a "technical vocabulary" expert and a "legal phrasing" expert, but top-1 routing must choose just one. The model learns to handle these boundary cases as best it can given the hard constraint, but the constraint itself limits expressiveness in a way that top-2 routing avoids.

The most widely used configuration in modern MoE architectures is top-2 routing, which selects the two experts with the highest routing scores and combines their outputs as a weighted sum. This single change from selecting one expert to selecting two has a surprisingly large impact on model behavior. Mixtral and GShard use top-2 routing, as do several other production systems. This reflects a community-wide finding that the benefits of selecting two experts consistently outweigh the costs.

The philosophical shift from top-1 to top-2 is worth pausing on before we get into the mechanics. Top-1 routing treats expert selection as a categorical decision: a token belongs to the domain of exactly one expert, and the model must commit fully. Top-2 routing treats it as a soft assignment: a token likely belongs primarily to one expert's domain, but acknowledges that a second expert may have something useful to contribute. This mirrors how human concepts work. The word "photosynthesis" is primarily a biological term, but it has chemical and physical dimensions too. Routing it to both a "biology" expert and a "chemistry" expert may produce a richer representation than routing it to either expert alone.

At the same time, top-2 routing doubles the per-token compute cost relative to top-1. This is not free. For a fixed total compute budget, the choice of top-2 over top-1 means either using smaller experts, using fewer experts, or accepting a higher overall cost. The empirical finding that top-2 consistently outperforms top-1 in head-to-head comparisons suggests that the representational benefits outweigh the compute costs, but this is a design decision worth making consciously.

The computation for top-2 routing starts identically to top-1: we compute the same routing logits using the same linear transformation $\mathbf{h} = W_r \mathbf{x}$. The difference lies entirely in what we do next. Instead of selecting just the maximum, we identify both the highest and second-highest scoring experts, calling their indices $k_1$ and $k_2$ respectively.

Before combining the outputs of these two experts, we need to determine how much weight to give each one. Simply using their raw logit values would be problematic: the logits are unconstrained real numbers that don't have a natural interpretation as combination weights. We need the weights to be non-negative and to sum to 1, so that the combined output is a proper convex combination of the two expert outputs. The solution is to apply softmax, but with a necessary design choice: we apply softmax only over the two selected experts, not over all experts.

The gating weights for these experts are computed by applying softmax only over the selected experts:

where:

• $g_1, g_2$: the normalized gating weights for the first and second selected experts, both in the range $(0, 1)$
• $h_{k_1}, h_{k_2}$: the raw routing logits for the first and second selected experts
• $\exp(\cdot)$: the exponential function, which ensures all weights are positive and amplifies differences between values
• $\exp(h_{k_1}) + \exp(h_{k_2})$: the sum of exponentials over the selected pair, serving as a normalizing constant that ensures weights sum to 1

Why does applying softmax only over the selected pair matter? If we had applied softmax over all $N$ experts first and then taken the top-2 probabilities, those two probabilities might sum to something like 0.45, not 1.0, because the remaining probability mass would be distributed across the other N-2 experts. The combined output would then be a fraction of the full representation magnitude, not a proper combination. By computing softmax over just the selected pair, we guarantee that $g_1 + g_2 = 1$, making the combination a true weighted average of the two expert outputs.

Why does this formula make sense? Notice that when $h_{k_1} = h_{k_2}$, the two exponentials are equal and the weights work out to exactly $g_1 = g_2 = 0.5$: perfect equal weighting. As $h_{k_1}$ grows larger relative to $h_{k_2}$, $g_1$ grows toward 1 and $g_2$ shrinks toward 0. This is exactly the behavior we want: strongly preferred experts get more weight, weakly preferred experts get less, and the transition is smooth rather than discrete.

The final output combines both expert outputs weighted by these normalized scores:

where:

• $\mathbf{y}$: the combined output vector for the token, with the same dimension as the input
• $g_1, g_2$: the normalized scalar weights for the two experts
• $E_{k_1}, E_{k_2}$: the functions computed by the selected expert networks applied to the input $\mathbf{x}$
• $\mathbf{x}$: the input token representation

This weighted combination means the final representation is a blend of two expert perspectives. If expert $k_1$ had a much higher routing score than expert $k_2$, then $g_1$ will be close to 1 and $g_2$ will be close to 0, making the output dominated by the first expert. Conversely, if both experts had similar scores, the weights will be close to 0.5 each, giving both experts roughly equal influence. This adaptive weighting allows the model to smoothly interpolate between relying on a single dominant expert and equally blending two perspectives.

Why does top-2 work better than top-1 in many settings? Several factors contribute to the performance difference.

The first is gradient flow. With two active experts per token, gradients reach twice as many expert parameters during each training step. This improves training efficiency and helps experts learn faster, particularly in the early stages of training when routing decisions are noisy. During early training, the router has not yet learned which experts are best for which tokens, so routing decisions are essentially random. With top-2 routing, more experts receive gradient signal during this necessary period, helping the entire expert ensemble learn meaningful specializations more quickly. With top-1 routing, an expert that happens to lose the routing competition for every token in an early batch may receive no gradient signal at all, causing it to fall further behind and creating a self-reinforcing spiral toward expert collapse.

The second is representation flexibility. Tokens often don't fit neatly into single semantic categories. A token representing a technical term in a legal document might benefit from both a "technical/scientific" expert and a "legal language" expert. Top-2 routing allows this blending. Language is inherently ambiguous and contextual, so forcing tokens into exactly one conceptual category can be unnecessarily restrictive. The ability to combine two expert perspectives provides a richer representational palette that better captures the semantic complexity of natural language.

The third is training stability. When only one expert processes each token, small perturbations in the router weights can cause dramatic shifts in which expert learns from which data. A small weight update might flip a routing decision for a large fraction of tokens, causing an abrupt redistribution of training signal that can destabilize learning. With two experts active, there is inherent smoothing: even if the router changes its mind about which expert should be the primary handler of a particular kind of token, that token likely still visits at least one of the same experts it did before. This provides continuity in which experts receive which training signal.

The fourth benefit is that top-2 routing provides better gradient estimates for the router itself. With top-1 routing, the straight-through estimator must approximate the gradient through a complete binary selection. With top-2 routing, the normalized softmax weights $g_1$ and $g_2$ provide a continuous signal: the model can gradually shift weight from the second expert toward the first (or vice versa) by adjusting the router's weights, giving the optimizer a smoother objective surface. This continuity in the gating weights makes the router's parameters easier to optimize.

The computational cost of top-2 routing doubles compared to top-1, since each token now passes through two expert networks instead of one. For a model with 64 experts where each expert's feed-forward network has 1 billion parameters, top-2 routing activates 2 billion parameters per token instead of 1 billion. This is still a dramatic reduction from the 64 billion total expert parameters. Doubling the computation is usually worth the gains in stability and model quality, which is why top-2 has become the de facto standard.

The choice of K in top-K routing is a consequential architectural decision in a Mixture of Experts model. It determines how the model balances computational efficiency against representation richness and how aggressively experts specialize. Understanding these trade-offs supports an informed decision rather than copying prior work.

The core tension is this: smaller K gives better computational efficiency and stronger expert specialization, but at the cost of training instability and representational rigidity. Larger K gives more stable training and richer representations, but approaches the computational cost of a dense model and reduces the pressure for experts to differentiate. The sweet spot for most applications is K=2, but understanding why requires examining each trade-off carefully.

Think of the K selection as analogous to choosing how many consultants to bring in for a difficult decision. With K=1, you pick your best consultant and follow their advice entirely. This is efficient and decisive, but you might miss important perspectives and you're fully exposed to that consultant's blind spots. With K=2, you bring in two consultants and blend their views. You get a richer analysis and protection against any single consultant's weaknesses, at the cost of twice the consulting fees. With K=8, you bring in a panel of consultants and average all their views. You're unlikely to miss anything important, but the panel consensus tends toward mediocrity and the cost is prohibitive. Most real decisions end up somewhere in the K=2 to K=4 range.

Computational cost scales strictly linearly with $K$. If each expert's feed-forward network requires $F$ FLOPs to process a single token, and you select $K$ experts, the expert computation costs exactly $K \times F$ FLOPs per token. This is the most direct and unavoidable trade-off: higher K means more computation per token, with no diminishing returns. For a fixed computational budget, choosing a larger K means either accepting higher inference costs or compensating by using smaller experts (which reduces per-expert capacity) or fewer experts (which reduces total model capacity). The MoE efficiency advantage over dense models comes precisely from keeping K small relative to N.

Model capacity utilization increases with $K$, but in a sublinear fashion. With $K=1$, each token sees $1/N$ of your total expert capacity. With $K=2$, it sees $2/N$. More precisely, the fraction of parameter space that any given token can potentially access is $K/N$. When this ratio approaches 1 (K near N), you have essentially built a very expensive dense model with $N$ times the feed-forward computation. The power of MoE comes from keeping $K \ll N$, typically with $K/N$ in the range of 0.1 to 0.25. This ratio is what determines how "sparse" the architecture truly is.

Training stability generally improves with moderate $K$, but the relationship is not monotone. The original Shazeer et al. MoE work found $K=4$ to work reliably. The GShard paper demonstrated that $K=2$ is sufficient when combined with appropriate auxiliary losses. Switch Transformer pushed to $K=1$ but required careful design of its auxiliary loss and capacity buffer to maintain stable training. Lower $K$ values are more prone to training instability because routing decisions have larger effects. When K is small, the decision of which expert processes a token becomes more consequential and harder to undo. A routing mistake has larger downstream effects on model quality, and the router receives noisier gradient signals because fewer experts are active for each token.

Specialization sharpness decreases as K increases. When K is small, experts must specialize to win the routing competition. An expert that is merely adequate at everything will rarely be in the top K for any particular token, since some other expert is likely better at that specific type of input. This creates strong evolutionary pressure toward expertise in narrow domains. When K is large, experts can afford to be more generalist: even a mediocre expert will frequently appear in the top K simply because there are so few other candidates. This affects what kind of knowledge structure emerges. With aggressive specialization (low K), experts develop distinct "personalities" focused on narrow domains, and the model's routing map becomes a meaningful semantic structure. With milder specialization (higher K), experts overlap considerably in capability and the routing map carries less semantic information.

Here is a summary of how leading MoE models have made the K selection:

The Mixtral model uses $K=2$ with $N=8$ experts, meaning each token activates 2 out of 8 experts, or 25% of expert capacity. This keeps inference costs comparable to a dense model with roughly $2\times$ the feed-forward parameters, while the model has $8\times$ the total expert parameters available to store specialized knowledge. This configuration has proven effective in practice. This provides a significant capacity advantage over dense models at a manageable computational cost.

The key insight from this empirical record is that K=2 has emerged as a stable equilibrium: it is the smallest value that consistently delivers the stability benefits of having more than one active expert, and it is small enough to preserve most of the efficiency advantages of sparse routing. The Switch Transformer's success with K=1 shows that lower values are achievable with sufficient engineering care, but K=2 is the lower bound that most practitioners are comfortable with as a default.

Implementing top-K routing correctly requires handling several details: selecting the top K indices, computing normalized weights, dealing with numerical stability, and letting gradient flow. Each of these sounds simple in isolation, but combining them into a reliable, efficient implementation requires care. Bugs in the routing implementation can silently degrade model quality in ways that are difficult to debug, so it pays to understand what each step is doing and why.

The core routing operation takes router logits and produces both the selected expert indices and their corresponding weights. Let's build this step by step. A reliable implementation must be efficient during both training and inference, and must handle edge cases like ties in routing scores gracefully. PyTorch's torch.topk provides an efficient GPU-accelerated top-K selection that handles most of the heavy lifting.

The routing function itself is remarkably concise once you have the right ingredients:

The output shows that for each token, we get exactly K = 2 expert indices and their corresponding normalized weights that sum to 1. The two selected indices correspond to the two highest-valued entries in the logit vector for that token, and the weights reflect the relative magnitudes of those logits after softmax normalization.

For top-1 routing specifically, we often want to preserve the routing probability distribution even though only one expert is selected. The reason is subtle: load balancing auxiliary losses (which we cover in the next chapter) need the full routing probabilities rather than only the argmax. If we discard the full distribution, we cannot compute the auxiliary losses that keep routing balanced during training.

The implementation needs both discrete forward-pass selection and the full probability distribution for computing auxiliary penalties:

This design pattern, where the routing function returns auxiliary information beyond just the selected experts, is common in practical MoE implementations. The router_logits themselves are also frequently returned from the MoE layer, since auxiliary losses typically operate on the logits directly rather than on the derived probabilities.

Once we have selected experts and computed their weights, tokens must be routed through the selected experts and their outputs combined. Unlike a standard feed-forward layer where all tokens pass through the same parameters simultaneously, MoE dynamically dispatches different tokens to different experts before gathering and combining the results. The dispatch and gather operations are the engineering bottleneck of MoE, and efficient implementation is necessary to achieve the theoretical compute savings.

The basic challenge is that the routing decisions are different for every token. Token A might go to experts 2 and 5, while token B goes to experts 1 and 7, and token C goes to experts 2 and 3. We cannot simply pass all tokens through all experts and then zero out the wrong contributions, because that would defeat the purpose of sparse routing. We need to physically route each token only to its selected experts, and we need to do this in a way that exploits GPU parallelism rather than serializing over individual tokens.

The simplest approach processes tokens individually:

This explicit loop makes the logic clear: for each token, iterate through its K selected experts, compute each expert's output, and accumulate them weighted by the routing weights. However, this approach is extremely slow because it processes tokens sequentially. Modern GPUs are designed to process thousands of operations simultaneously through SIMD parallelism, and this triple-nested loop structure forces strictly sequential execution that underutilizes the hardware. For a sequence of 2048 tokens with K=2, this requires 4096 individual forward passes through expert networks, each processing a single vector. The GPU sits mostly idle, waiting for each sequential operation to complete.

The efficient approach groups tokens by their assigned experts and processes all tokens assigned to the same expert simultaneously:

This version groups all tokens assigned to each expert and processes them in a single batched forward pass through that expert. The key insight is that even though different tokens go to different experts, all tokens assigned to the same expert can be processed together in one operation. The outer loop over K expert positions and the inner loop over experts means we make exactly $K \times N$ forward passes through expert networks, but each pass processes all relevant tokens in parallel. For a model with 8 experts, K=2, and 2048 tokens, this means 16 forward passes through expert networks instead of 4096. Each of those 16 passes processes a batch of tokens rather than a single token, fully exploiting GPU parallelism.

The mask tensor identifies which tokens are assigned to a given expert at a given selection position. The flat_hidden[mask] operation gathers the relevant token representations into a contiguous batch, the expert processes them all simultaneously, and the result is scattered back to the correct positions in the output tensor. This gather-process-scatter pattern is the basic building block of efficient sparse computation on GPUs.

Let's create a complete working implementation combining routing and expert computation:

The output confirms that our MoE layer maintains the expected dimensions while activating only K experts per token. The total expert parameters scale with the number of experts, but each token only uses K experts worth of computation. This separation between total parameter count and per-token computational cost is the central value proposition of MoE architectures: you get the representational power of a large model at the per-token cost of a small one.

To solidify understanding, let's trace through a concrete numerical example of top-2 routing from beginning to end. This example uses a small but fully realistic setup: 4 experts, a hidden dimension of 4, and a single token. Every number is computed explicitly so you can verify each step by hand.

The token has a hidden representation $\mathbf{x}$ with 4 dimensions. The router weight matrix $W_r$ has shape $4 \times 4$ (one row per expert, one column per hidden dimension). We set these to specific values so we can compute routing logits $\mathbf{h} = W_r \mathbf{x}$ by hand:

The routing has decided that expert 3 should contribute roughly 50.2% and expert 0 should contribute roughly 49.8% to this token's output. Experts 1 and 2 are completely bypassed: they compute nothing for this token and receive no gradient signal from it during training.

Because the two selected logits were very close in value (approximately 0.49 and 0.48), the softmax produces nearly equal weights. This near-equal split means both experts contribute roughly equally to the final representation. If we denote expert 3's output as $\mathbf{e}_3$ and expert 0's output as $\mathbf{e}_0$, the final combined output is:

where:

• $\mathbf{y}$: the final combined output vector for this token
• $\mathbf{e}_3, \mathbf{e}_0$: the output vectors computed by expert 3 and expert 0 respectively when given $\mathbf{x}$ as input

This nearly equal weighting means both experts contribute almost equally to the final representation. As a contrast, if the logits had been 0.9 for expert 3 and 0.2 for expert 0, the softmax would produce weights of approximately $\frac{e^{0.9}}{e^{0.9} + e^{0.2}} \approx 0.67$ and $0.33$, giving expert 3 twice as much influence. The further apart the logits, the more the combination resembles a single-expert decision. The closer together, the more it approximates an equal ensemble.

Why does this formula make sense? When two experts have similar routing scores, the model is essentially uncertain which expert is more appropriate. Giving them equal weight is a reasonable hedging strategy: rather than committing fully to one expert despite the uncertainty, the model takes a blend. When one expert strongly dominates, the model correctly places most weight on the clearly superior choice. The softmax over the selected pair implements this hedging automatically, without requiring any explicit uncertainty estimation.

Understanding the key parameters in top-K routing helps you make informed decisions when configuring an MoE layer. These are not simply hyperparameters to tune blindly: each parameter represents a distinct architectural choice that affects capacity and efficiency as well as training dynamics.

The key parameters for the Top-K Routing implementation are:

• num_experts: Total number of experts in the MoE layer. This determines the potential capacity of the MoE layer. More experts means more total parameters and potentially more specialized knowledge, but also more memory requirements and complexity in managing load balance. In practice, values between 8 and 128 are most common. Going beyond 128 experts introduces significant engineering complexity in distributed training without proportional quality gains.
• k: Number of experts to select for each token (typically 1 or 2). This controls the basic trade-off between computational cost and model expressiveness. Lower values of K yield better efficiency and stronger specialization pressure, while higher values provide smoother training and richer token representations. For most applications, K=2 is the right starting point.
• hidden_dim: Dimensionality of the input and output representations. This must match the hidden dimension of the surrounding transformer architecture. The router uses this dimension to compute compatibility scores between tokens and experts. Larger hidden dimensions allow the router to make more fine-grained routing decisions, but the router itself is a single linear layer, so the computational overhead of the routing step is modest.
• intermediate_dim: Dimensionality of the expert feed-forward networks. Larger intermediate dimensions increase the capacity of each individual expert at the cost of higher per-expert computational cost. In transformer architectures, the intermediate dimension is typically 4 times the hidden dimension for dense models, and MoE experts often use the same ratio.

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 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.

Top-K routing enables the core promise of MoE architectures: scaling model capacity without proportionally scaling computation. However, this selective activation introduces several challenges that require attention in any serious deployment.

Load imbalance is the most significant and persistent problem with top-K routing. As our visualization showed, naive top-K routing can lead to severe load imbalance where some experts are overwhelmed with tokens while others sit idle. In the extreme case, a phenomenon called "expert collapse" can occur where the router learns to send all tokens to just one or two experts, effectively wasting the capacity of all other experts. The mechanism of collapse is a positive feedback loop: when an expert receives many tokens, it gets more gradient signal and improves faster, which makes it more attractive to the router, which sends it even more tokens. Without intervention, this feedback loop can concentrate all routing onto a handful of experts.

Expert collapse is particularly problematic during distributed training, where each expert typically resides on a different accelerator (GPU or TPU). If expert 0 receives 50% of tokens while expert 7 receives 5%, you have created a severe computational bottleneck: the accelerator hosting expert 0 is overloaded and becomes the rate-limiting step for the entire batch, while the accelerator hosting expert 7 sits mostly idle. The distributed training efficiency plummets, and the experts that do receive tokens continue to improve while the neglected experts fall further behind. The next chapter introduces auxiliary losses specifically designed to counteract this collapse.

Discrete selection and gradient approximation create a basic tension in the training process. The top-K selection is inherently discrete: a token either goes to an expert or it does not. Backpropagation requires continuous, differentiable operations, so we must use approximations like the straight-through estimator to enable gradient flow through the routing decision. While these approximations work well in practice, they create a mismatch between the forward pass (hard selection) and backward pass (soft gradient flow). This mismatch can occasionally cause training instabilities, particularly with K=1 routing where the approximation is most severe. A router that consistently makes slightly wrong decisions based on the straight-through gradient may not self-correct cleanly, because the gradients it receives describe a different, softer version of the routing function than the one it executes.

Token dropping under capacity constraints introduces a new failure mode in distributed settings. Each expert can only process a limited number of tokens per batch due to memory and compute constraints on the hosting accelerator. This limit is called the "expert capacity." If the router sends more tokens to an expert than its capacity allows, some tokens must be dropped, meaning they skip the expert entirely and pass through unchanged or via a simple residual connection. Dropped tokens miss the expert processing they were routed to, which can degrade model quality. This creates a direct tension: the router wants to send tokens to the best expert for each input (which may create uneven loads), but the capacity constraint punishes any concentration of tokens. Managing this tension is one of the primary engineering challenges in MoE deployment.

Inference complexity adds engineering overhead that does not exist in dense models. During inference, top-K routing requires dynamic dispatch where tokens are grouped by their selected experts before each expert processes its assigned batch. This is fundamentally more complex than dense model inference, where every token passes through the same parameters in a straightforward sequential computation. The dispatch and gather operations add latency that is not reflected in the theoretical FLOP count. When K > 1, you also need to combine outputs from multiple experts per token, adding both computation and memory access overhead. This makes MoE models harder to optimize with inference frameworks that expect uniform computation patterns across tokens.

Despite these challenges, top-K routing has proven remarkably effective at letting models that achieve dramatically better performance for a given computational budget. The Switch Transformer demonstrated that K=1 can work with proper auxiliary losses, establishing that even the most aggressive sparsity is viable. Mixtral showed that K=2 with 8 experts achieves excellent performance relative to its size, competing favorably with much larger dense models. The key insight that practitioners have learned is that the routing mechanism itself is not the whole story: it must be combined with load balancing techniques to realize the full potential of sparse expert architectures. The routing mechanism tells us which experts to use; the balancing mechanism ensures that this selective usage remains distributed enough to avoid collapse and capacity violations. We explore these balancing techniques in the next chapter.

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.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about top-K routing in Mixture of Experts architectures.