Handling Dropped Tokens

In the architecture of sparse Mixture of Experts models, particularly those employing Top-K routing, each expert is typically assigned a fixed processing capacity, denoted as $C$. This capacity represents the maximum number of tokens that an expert can handle within a given computational step (often a micro-batch in distributed training). This fixed capacity is essential for maintaining predictable computational load and memory usage across devices, especially in large-scale distributed settings where experts are parallelized across different processors.

However, the routing mechanism, driven by the gating network, doesn't inherently guarantee that the number of tokens assigned to any single expert will strictly adhere to this capacity $C$. The gating network selects the "best" expert(s) for each token based on learned affinities. Due to variations in input data and the dynamic nature of router learning, it's common for some experts to be assigned more tokens than their designated capacity allows, particularly if the load balancing isn't perfectly achieved.

Tokens assigned to an expert that has already reached its capacity $C$ are referred to as "dropped tokens."

The Consequence of Dropping Tokens

When a token is dropped by its assigned expert due to capacity limits, it effectively bypasses the specialized computation that the MoE layer was intended to provide for it. The most common strategy for handling these dropped tokens is simple: their input representation is passed through the MoE layer unchanged.

Consider an MoE layer performing a residual update:

$$y = x + \text{MoE}(x)$$

For a token $x_i$ that is successfully routed to and processed by expert $j$, the output is:

$$y_i = x_i + g_i \cdot \text{Expert}_j(x_i)$$

where $g_i$ is the gating weight.

However, if token $x_k$ is assigned to expert $j$, but expert $j$ has already reached its capacity $C$, then token $x_k$ is dropped. In this standard pass-through approach, its output calculation effectively becomes:

$$y_k = x_k + 0$$

The MoE contribution is zero because the token never entered the expert computation pipeline for that layer.

This has several negative implications:

1. Loss of Specialization: The token misses out on the potential benefits of processing by a specialized expert, potentially degrading model quality. The model doesn't learn as effectively for inputs that frequently lead to dropped tokens.
2. Reduced Learning Signal: Since the dropped token doesn't pass through an expert, it provides no gradient signal to update the weights of the expert it was assigned to. This can hinder the expert's specialization process. It also provides a less informative signal back to the router about the quality of its routing decision for that token.
3. Indicator of Imbalance: A high percentage of dropped tokens is a direct symptom of poor load balancing. It indicates that the auxiliary loss term ($L_{aux}$) discussed previously is not sufficiently penalizing imbalanced assignments, or that the router is struggling to distribute tokens effectively.

Managing Capacity and Dropping

Handling dropped tokens primarily involves strategies to prevent excessive dropping, rather than complex mechanisms to process them after the fact.

Setting Expert Capacity ($C$)

The expert capacity $C$ is a critical hyperparameter. It needs to be large enough to accommodate a reasonable degree of imbalance inherent in the routing process, but not so large that it leads to excessive computational waste (padding) when experts are underutilized.

Capacity is typically set based on the ideal uniform load per expert, scaled by a capacity_factor. If $N$ is the number of tokens being processed (e.g., per device in expert parallelism) and $E$ is the number of experts available locally, the ideal load is $N/E$. The capacity is then set as:

$$C = \lceil \text{capacity\_factor} \times \frac{N}{E} \rceil$$

Common values for capacity_factor range from 1.0 to 2.0.

• A factor closer to 1.0 (e.g., 1.1, 1.25) minimizes computational overhead but increases the risk of dropping tokens if the load is even slightly imbalanced.
• A larger factor (e.g., 1.5, 2.0) reduces the likelihood of dropping tokens but potentially leads to more computation being allocated than strictly necessary, as experts might often process fewer than $C$ tokens.

The choice involves a trade-off between computational efficiency and minimizing information loss due to dropped tokens. It often requires tuning based on observed drop rates during initial training runs.

The Role of Auxiliary Load Balancing Loss ($L_{aux}$)

As detailed in the previous section, the auxiliary loss term $L_{aux}$ is the primary mechanism for encouraging the router to distribute tokens evenly, thereby minimizing the conditions that lead to dropped tokens. While $L_{aux}$ doesn't directly handle a token once it's dropped, effective tuning of its coefficient $\alpha$ in $L_{total} = L_{task} + \alpha L_{aux}$ is essential for keeping the token drop rate low. A higher $\alpha$ generally pushes the router towards more balanced assignments, reducing drops but potentially interfering with optimal specialization if set too high.

Visualizing the Dropping Mechanism

The following diagram illustrates the process where tokens are routed, capacity is checked, and some tokens are dropped.

Flow of tokens through an MoE layer with limited expert capacity. Tokens are routed to experts. If an expert's capacity (here, C=2) is full, subsequent tokens assigned to it are dropped and typically bypass expert computation via a pass-through mechanism.

Monitoring Dropped Tokens

During training, it is essential to monitor the percentage of dropped tokens per MoE layer or averaged across the model. This metric serves as an important health check:

• Consistently high drop rates (e.g., > 1-5%) often indicate insufficient expert capacity (capacity_factor too low), ineffective load balancing ( $L_{aux}$ coefficient $\alpha$ needs tuning or router isn't learning well), or potential issues with data distribution.
• Sudden spikes in drop rates might signal training instability or shifts in data characteristics that the router is struggling to adapt to.

Logging this metric helps in diagnosing training problems and tuning the relevant hyperparameters (capacity_factor, $\alpha$).

Alternative Strategies (Less Common)

While pass-through is the standard, research has explored other ways to handle capacity overflows, such as:

• Routing to Second Choice: If using Top-1 routing, a dropped token could potentially be sent to its second-choice expert if that expert has available capacity. This adds routing complexity.
• Overflow Experts: Designating specific experts to handle overflow tokens. This requires careful design to ensure these overflow experts are utilized effectively.

These methods are not widely adopted in standard large-scale MoE implementations due to added complexity in routing logic and communication patterns. The focus remains on preventing drops through capacity management and load balancing losses.

In summary, handling dropped tokens in MoE training is less about actively processing them and more about implementing preventative measures. Setting an appropriate expert capacity and effectively tuning the load balancing auxiliary loss are the primary strategies to minimize token dropping and ensure that most tokens benefit from specialized expert computation. Monitoring the drop rate is important for maintaining training stability and model performance.