A Routing Testbed for MoEs, Part 2: Visualizing Balancing Methods

In Part 1 I argued that two properties of the learned routing decide how much of an MoE’s total capacity is actually used. The load should be balanced across the experts, so that no expert stays undertrained and wastes the capacity it holds, and the experts should specialize, so that the additional parameters learn distinct things instead of duplicating each other. Load balance can be read off the routing statistics of any training run, while specialization cannot, which is what motivated our MoE Routing Testbed: A data mix of Wikipedia in eight languages with distinct scripts provides a strong prior for how tokens should be routed, a reference routing derived from that prior which we verified to be better than what our models learn on their own and the two metrics of expert utilization and routing purity that quantify the two goals separately.

With that setup, the routing behavior of a training run becomes directly visible instead of having to be inferred from a validation loss. This second part of the post uses it on the mechanism that shapes routing most directly, namely the token balancing that essentially every MoE training run applies in one form or another. We trained hundreds of configurations covering four balancing methods and the ways in which they can be applied. Via the expert utilization and routing purity metrics that our testbed can measure, the differences between these balancing approaches can be nicely visualized, as I show in this post.

The token balancing methods

We compared four approaches from the literature that have been used in MoE work over the past years:

The load balancing loss (LBL), introduced with GShard and popularized by the Switch Transformer, is the oldest of them and still the most widely used one. For each MoE layer it computes two quantities per expert: The share $f_i$ of the token-expert assignments in the batch that went to expert $i$, and the average routing probability $P_i$ that the router assigned to expert $i$ over the same tokens. The auxiliary loss added to the language modeling loss is then $\alpha E \sum_{i=1}^{E} f_i P_i$. While $f$ is the distribution we want to be uniform, the surrogate $P$ is used additionally to make the loss differentiable, which is not the case for $f$ due to top-$k$ selection. $\alpha$ is a hyperparameter that weighs this auxiliary loss against the main LM loss.

Balanced assignment (BA), introduced with BASE layers, decouples balancing from the router instead of penalizing it. Given the score $s_{ti}$ that the router assigns to each token $t$ in a batch for each expert $i$, it looks for the binary assignment that maximizes the total score of the selected token-expert pairs, subject to every token receiving $k$ experts and every expert receiving the same number of tokens. This is a linear assignment problem and can be solved with an auction algorithm. BASE layers itself assigns a single expert per token, the $k>1$ case used here is the straightforward generalization of that. Balance is perfect by construction rather than encouraged, and there is no strength left to tune, but as a consequence also no flexibility in that regard.

Sinkhorn routing (SH), as used in the scaling study of Clark et al. (2022), keeps the same assignment view as BA but relaxes it. Rather than solving the assignment exactly, it adds an entropy term to the objective and the solution of that relaxed problem is found by the Sinkhorn algorithm. A handful of iterations are sufficient and they are cheap matrix operations, so the cost stays well below an exact solver, at the price of balance that is only approximate. Like BA, it acts on the routing itself rather than on the loss modifying the router.

The expert bias (EB) method, introduced by Wang et al. (2024) and popularized by DeepSeek V3, often called auxiliary-loss-free balancing, takes a third route and manipulates the selection. Each expert gets a scalar bias $b_i$ that enters the top-$k$ selection only, so the selected experts are those with the highest $s_{ti} + b_i$, while the gating weights that scale their outputs are still computed from the unbiased $s_{ti}$. The bias hence steers which experts are chosen without distorting how their outputs are combined. After every step, the biases are updated by a fixed increment $\gamma$ based on the observed load, as $b_i \leftarrow b_i + \gamma \operatorname{sign}(\bar{c} - c_i)$, with $c_i$ the number of tokens expert $i$ received in that step and $\bar{c}$ the average across experts, so the bias decreases for experts above their share and increases for those below it. Balance is therefore approached over the course of many steps instead of being enforced within a single one, and the appeal of the method is that no additional gradient competes with the LM objective.

Three dimensions, not one

These four methods are usually compared as alternatives, i.e. the question is posed as which balancing method to use. But the method is only one part of the configuration. Applying it also requires deciding how hard the balance is enforced, and, before that, over which set of tokens it is applied. We therefore treat balancing as three separate dimensions: The method (how balance is encouraged), the strength (how hard it is enforced) and the scope (over which tokens it is computed). Only LBL and EB come with a strength dial at all, and for EB it turns out to be a rather poor one.¹ The scope, on the other hand, applies to all four methods equally, and it is the dimension that is almost never stated, since it follows from how the balancing statistics happen to be aggregated in a given implementation.

Separating the three makes the comparison a lot more informative, because a difference attributed to a method may well come from the scope at which it was applied.

The role of scope

The effect of the scope is easiest to see in the extreme case of computing the balance over a single training sequence. Sequences are typically packed from a single data source, so such a sequence can be for instance entirely code, and a balancing mechanism applied at this scope pushes the router to spread code tokens uniformly over all experts. That is the opposite of what specialization requires, namely code tokens concentrating on the experts that handle code. At sequence scope, utilization and specialization are directly in conflict by construction.

The other end of the spectrum, and arguably the ideal case, is a scope large enough to have a representative sample of the full training distribution. If the domains are represented at comparable rates, in the testbed exactly so, with eight domains at 1/8 each, then a router that sends each domain to its own group of experts is perfectly specialized and perfectly balanced at the same time. The conflict is hence not inherent to balancing, it is created by computing the balance over token sets that are too small and too homogeneous to reflect the training distribution and hence the distribution that experts should specialize over.

What makes this practically relevant is where real training frameworks sit on that spectrum. As Qiu et al. (2025) pointed out, most public MoE implementations compute the balancing statistics over the tokens that happen to be present on one device. In large-scale training with gradient accumulation or pipeline parallelism, such a per-device micro-batch often contains just one or two sequences, so that training frequently operates at the bad end of that spectrum without anyone having made that decision.

Experimental setup

The testbed lets us measure all of this cheaply. We swept the three dimensions systematically: Four methods (LBL, EB, SH, BA), five scopes (balancing over 1, 2, 8, 16 or all 64 sequences of the global batch, controlled via micro-batch size and cross-device aggregation) and for LBL and EB a range of strengths, resulting in 98 configurations in total. The model is a small single-layer MoE with 50M active parameters (270M total), 32 fine-grained experts under top-4 routing, trained for 20k steps (about 10.5B tokens) at a global batch size of 64 x 8192 tokens. No token dropping is applied in the default setup.

Both our metrics, expert utilization and routing purity, are computed over the full global batch – independent of the scope at which the balancing acts – which is the detail I flagged in Part 1 and which matters here: The token load can be balanced over experts globally without being balanced locally, but not the other way round.

Result 1: Scope decides the trade-off

The following plot shows what I’ve been referring to as the visualization power of the routing testbed. When plotting the expert utilization vs. routing purity trade-offs for different configurations, we can clearly see the effect that scope has:

At a scope of just one sequence, the strength sweep is an almost straight line, from high utilization at low purity for the strongest weight down to low utilization at high purity for the weakest one. Every bit of balance is paid for with purity here, which is exactly the conceptual conflict described above and hence the intuitive result to expect. The run without any balancing, included as the gray diamond, confirms that specialization is what the router does on its own: It lands among the highest purity values of the whole sweep without any help, but its utilization is low. With growing scope, the lines bend towards the top-right corner. At a scope of 16 sequences, the strongest balancing already starts at higher purity and the best configurations reach around 0.9 utilization and 0.6 purity at the same time, resolving the conflict.

The validation loss on the testbed mix follows that trade-off closely, with the lowest loss of every run in the plot coming from the large-scope LBL configurations, while the strongly balanced sequence-scope configurations produce the highest. Coloring the same landscape by that loss makes the correspondence easier to see:

Result 2: Methods are equally good

In the two plots above, the alternative balancing methods were drawn as squares (BA), triangles (SH) and crosses (EB), which unlike LBL have no lines, because there is no strength to vary. BA, true to its auction formulation, pins utilization at exactly 1.0 in every configuration, so its four points sit on a vertical line at the right edge, ordered by scope. Sinkhorn behaves similarly at slightly lower utilization, and EB, applied at global scope as DeepSeek proposed it, reaches 0.90 utilization at 0.56 purity, the same neighborhood as LBL at its largest scopes. The differences between the methods are small compared to the effect of the scope and at a sufficiently large scope all four of them work well.

That is worth emphasizing, given how much effort the field spent on the method question. EB was introduced with the claim of being superior to the auxiliary loss, and Qiu et al. (2025) later showed that the two are on par once compared at the same scope, so the original comparison had confounded method with scope. Our experiments confirm this finding, but due to our testbed, we can do so by comprehensively searching the three-dimensional design space and accurately measuring and visualizing the effects, rather than just comparing the loss of two larger models. The same conclusion was recently also reported in Microsoft’s MAI-Thinking-1, which uses a global-batch balancing loss aggregated across data-parallel workers and micro-batches, and the report states that the aggregation strategy mattered much more than the type of loss according to their ablations, matching our observations here.

Result 3: The role of token dropping

Many training setups cap the number of tokens an expert can process, i.e. they use a capacity factor and drop the tokens beyond it, trading model quality for predictable compute and communication. We repeated the grid search experiments under such caps, which inverts the picture:

Two things change: The purity ceiling reduces substantially, so the specialized part of the landscape is no longer reachable at all. And the best run, ringed in both panels, moves from EB at global scope to LBL at a scope of one sequence, while the overall loss level rises. The reason is that the capacity constraint is enforced per device, so only strong local balancing avoids heavy dropping under it. Since the damage from the dropped tokens seems to outweigh what the specialization gains, the best strategy becomes to avoid dropping by strong local balancing. This is also in line with the experience at MAI, who observed that conclusions about balancing methodology can flip depending on the expert capacity and converged on a fully dropless implementation as in our baseline setting.

Result 4: Validation loss hides the effects

One could ask whether the testbed is needed for any of this, or whether a plain validation loss would have shown the same. We checked that by training the same configurations on Dolma3, the open training mix of the OLMo 3 models, and the following plot colors the trade-off landscape by the loss of both side by side:

The two panels hold the same runs at the same coordinates and differ only in which loss the color encodes. On the testbed mix, the color follows the position: The upper right of the landscape is where the low losses sit. On Dolma3, the same runs scatter their colors across the landscape with no usable ordering left, even though the spread of the losses is nearly the same. Measured across all 98 configurations as a rank correlation between the loss and how far up and to the right a run sits, the testbed reaches -0.80 and Dolma3 only -0.28. At 50M active parameters, a loss on a normal pretraining mix simply cannot distinguish routing configurations, and metrics that look at the routing itself are needed instead.

What about expert choice?

Expert choice (EC) routing (Zhou et al., 2022) often comes up in discussions about token balancing and I want to say a few words on why we did not include it in the paper. EC inverts the assignment problem – experts pick their top-$k$ tokens instead of tokens picking experts – and thereby it achieves perfect balance by construction. The catch however is causality: Whether an expert picks a token depends on the other tokens in the batch, including future positions of the same sequence. For an autoregressive decoder this leaks future information during training, and when generating token by token, the assignment cannot be reproduced the way it was trained. Proposed fixes bolt on extra routers and training phases but have rarely been used by production MoE models.

Beyond the causality issue, EC also has another flaw: The balance is not just guaranteed by construction, it is even impossible to remove, which can be problematic at inference time where the token distribution cannot be controlled. If the model has to process a batch of tokens coming purely from one domain, the EC assignment will force them over all experts nonetheless, no matter if they are specialized in that domain or not. LBL, BA, SH and EB do not face this problem because for each of them the balancing can simply be turned off at inference time.

Expert threshold (ET) routing (Sun et al., 2026) is a recent attempt to repair the causality flaw of EC. Each expert maintains a threshold, tracked as an exponential moving average calibrated on the global token distribution, and a token is routed to an expert whenever its score clears that threshold, a decision that depends only on the token itself and hence restores causality. That does indeed solve the problem, but introduces new ones instead: Calibrating the threshold successfully can be challenging. The number of experts per token is no longer fixed and the dynamic allocation might be counterintuitive, as it is driven by the token’s routing score distribution over experts against per-expert (but cross token) thresholds. And the key quality of EC, balance by construction, is given up in this case.

Summary

This concludes my walk through the MoE Routing Testbed paper. Its special data mix with clearly distinguishable domains enables directly measuring expert specialization. Using it, we can clearly measure and visualize how token balancing choices affect the trade-offs between expert specialization and expert utilization. Scope, as noticed by previous work, is the most important factor. Only if it is large enough can the two goals be achieved at the same time instead of conflicting. The balancing method itself, i.e. using a load balancing auxiliary loss or expert biases, plays only a minor role.