Mixture of Experts: Gating and Routing Explained

Introduction

In our last lesson, we established the high-level architecture of an MoE layer, which replaces the standard Feed-Forward Network (FFN) in a Transformer block. We learned that it consists of two main components:

1. A set of parallel FFNs called experts.
2. A router or gating network that conditionally directs input tokens to a small subset of these experts.

This conditional computation is the key to scaling a model's parameter count while keeping the computational cost manageable. We also briefly touched on the idea of Top-K routing, where only the "top K" experts are chosen for each token.

Today, we will dive into the mechanics of that process. Our learning goal is to explain the gating/routing mechanism that assigns tokens to a subset of experts. We'll explore how the router calculates its preferences, how tokens are actually assigned to experts, and the critical problem of ensuring that the workload is distributed evenly.

Let's begin by dissecting how the router makes its choices.

1. The Gating Mechanism: From Input to Expert Assignment

The router's job is to look at an incoming token and decide which of the available experts are best suited to process it. This is a multi-step process.

Step 1: Computing Router Logits

The router is a simple, learnable neural network—typically a single linear layer. For each token's hidden state vector, it produces a vector of "logits" (raw, unnormalized scores), one for each expert.

The following video provides an excellent, code-oriented walkthrough of how these logits are generated and used to select the top experts.

As you saw, the router produces a score for every expert. The "Top-K" selection (where K is typically 1, 2, or in DeepSeek-V3's case, 8) ensures sparsity by only considering the experts with the highest scores for the next step.

Step 2: Normalizing Scores into Gating Weights

Once the top K experts are chosen, their logits need to be converted into a set of weights (or probabilities) that will be used to combine their outputs.

A common way to do this is to effectively "mask" the unselected experts and then apply a softmax function. Masking involves setting the logits of all non-top-K experts to negative infinity. When the softmax function is applied, the scores of these masked experts become zero.

This simple image provides a great visualization of the process for a Top-2 routing scenario.

Let's see the code-level view of this normalization and masking process.

This process of calculating logits, selecting the top K, and applying softmax is the core of the gating mechanism. The final output of the MoE layer is then the weighted sum of the outputs from the selected experts, using these probabilities as the weights.

It's worth noting a subtle but important variant used in modern models like DeepSeek-V3. Instead of softmax, they use a sigmoid function on the logits.

• Softmax creates a competitive distribution where the weights must sum to 1. Experts are effectively competing for a token.
• Sigmoid treats each expert's score independently, outputting a value between 0 and 1 for each. This can be useful when multiple experts might be equally relevant.

2. Expert Capacity and Token Assignment

The gating weights tell us which experts to use, but there's a practical constraint we haven't discussed: expert capacity. For hardware efficiency, each expert is designed to handle a fixed number of tokens from a batch.

What happens if the router tries to send too many tokens to one popular expert? This is where the concepts of slots and token dropping come into play.

The HuggingFace video provides a brilliant and detailed walkthrough of how this slot assignment and potential dropping works in practice.

This mechanism is crucial. The router doesn't just pick experts; it works within a system that has fixed capacity limits. This leads directly to our next topic: if experts have limited capacity, the model must be encouraged to distribute the load evenly.

3. The Load Balancing Problem

If left unconstrained, a router might learn to favor a small number of "superstar" experts, sending most tokens their way. This is called routing collapse. It's inefficient because:

1. The favored experts become computational bottlenecks.
2. The other experts are underutilized and their parameters are wasted.

To combat this, MoE models employ load balancing strategies.

Solution 1: Auxiliary Loss (The Traditional Approach)

The most common approach is to add an auxiliary loss term to the model's main training objective. This is an extra penalty that encourages the router to spread its assignments more evenly across all experts.

While effective, this approach has a downside: the model has to spend computational resources on optimizing this secondary objective, which can sometimes detract from its primary goal of language modeling.

Solution 2: Auxiliary-Loss-Free Balancing (The DeepSeek-V3 Innovation)

DeepSeek-V3 pioneers a more elegant solution that avoids a separate loss function. It directly adjusts the routing logits during training to encourage balance.

The following video explains this clever mechanism.

Let's look at the formal description from the DeepSeek-V3 paper. The key is Equation (16), which shows the bias term $b_j$ being added to the affinity score (logit) only for the purpose of routing.

This dynamic adjustment allows the model to maintain a balanced load without the trade-off of an auxiliary loss, letting it focus solely on the primary language modeling task. The result is better performance and more natural expert specialization, as shown in the ablation studies from the paper. Because this strategy is so effective, DeepSeek-V3 does not need to drop tokens due to capacity limits.

Conclusion

In this lesson, we have thoroughly explored the gating and routing mechanism that underpins Mixture-of-Experts architectures.

Key Takeaways:

• Gating Mechanism: A router network computes logits (scores) for each expert. A Top-K selection process identifies the most suitable experts. Their scores are then normalized into weights, often using a softmax or sigmoid function.
• Expert Assignment: For efficiency, experts have a fixed capacity (number of slots). If a popular expert is assigned more tokens than its capacity, overflowed tokens may be "dropped" (passed through a residual connection without expert computation).
• Load Balancing: To prevent routing collapse, MoE models need a load balancing strategy. This can be a traditional auxiliary loss that penalizes imbalance, or a more advanced auxiliary-loss-free method like in DeepSeek-V3, which dynamically adjusts a bias term to guide the router's decisions.

Preview of the Next Lesson:

We now have a complete picture of the MoE layer's architecture and its dynamic routing process. In our final lesson on MoE, we will put it all together. We will examine the specific design of the MoE layers in DeepSeek-V3, including its use of shared vs. routed experts, and clarify the crucial distinction between a model's total parameters and its active parameters per token.