Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts

Abstract: Mixture-of-Experts (MoE) training requires global load balance to prevent expert under-utilization and local balance for efficient expert-parallel execution. Existing distributed Quantile Balancing (QB) uses shard-dependent or approximate global quantiles, while token-independent expert biases cannot ensure microbatch-level balance. We introduce Exact Quantile Balancing (EQB), which computes exact global-batch BF16 quantiles with negligible communication, and Load-Error Injection (LEI), which injects local load errors directly into router-score gradients. On 7.5B-parameter MoEs trained for up to 500B tokens, EQB improves global balance and downstream performance over naive QB, while LEI improves local balance and outperforms the GShard loss at comparable quality.

• The paper presents Exact Quantile Balancing (EQB) and Load-Error Injection (LEI) to separately address global and local load balancing in Mixture-of-Experts (MoE) models, significantly reducing MaxVio values: global from 0.92 to 0.74 (19.4% reduction), and local from 5.96 to 3.49 (41.6%) on a 7.5B-parameter decoder-only MoE.
• EQB precisely calculates the global quantile balance thus optimizing the hardware throughput without requiring to communicate margins over ranks, improving benchmark loss performance (BPB) on all six datasets.
• LEI improves local balance by injecting load errors into the router-score gradients, further reducing local MaxVio by 18.9% over EQB alone, while maintaining high benchmark loss performance, but highlights the need for careful tuning and interaction effects when combined.

Problem formulation and contribution

Mixture-of-Experts (MoE) routing imposes two distinct load-balancing requirements. Global balance, measured over the tokens contributing to an optimizer step, prevents persistent concentration on a subset of experts. Local balance, measured within rank-local expert-parallel microbatches, controls dispatch skew and the resulting variability in expert computation. These objectives are not interchangeable: local imbalances can cancel after aggregation, so a globally balanced optimizer step may still contain severely overloaded local microbatches.

"Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts" (2609.28053) addresses these scales with two complementary mechanisms. Exact Quantile Balancing (EQB) replaces rank-averaged or histogram-approximated quantiles with an exact global-batch BF16 order statistic. Load-Error Injection (LEI) introduces a local balancing signal directly into router-score gradients using the observed hard-routing residual. The central design is therefore a separation between token-independent bias control for global balance and token-dependent gradient correction for local balance.

The paper evaluates the methods on a 7.5B-parameter decoder-only MoE with 256 routed experts, top-6 routing, 0.5B active parameters per token, and 48 MoE layers. The experiments extend to 500B training tokens, providing evidence beyond short ablations, although the empirical scope remains limited to one model scale and one seed.

Exact quantile balancing

Quantile Balancing (QB) computes expert-selection biases from routing margins rather than optimizing an auxiliary loss. For token $t$, let $\tau_t$ be the $(K+1)$st largest adjusted router logit, and define the margin for expert $e$ as $\tau_t-z_{t,e}$. An expert should be selected for approximately a fraction $K/E$ of tokens; consequently, the next bias for that expert is obtained from the corresponding empirical quantile of its margins. Centering the resulting bias vector removes the irrelevant common offset.

The distributed setting creates a nontrivial statistical problem. Averaging rank-local quantiles does not recover the quantile of the union of the rank-local batches, because quantile formation and averaging do not commute. The resulting statistic depends on the partition of tokens across data-parallel ranks. An alternative based on globally aggregated histograms is partition invariant, but its accuracy is limited by histogram resolution.

EQB exploits the BF16 representation of routing margins. It performs two radix-selection passes: a coarse pass over the high byte and a fine pass over the low byte. Each pass constructs 256-bin local histograms and performs a global all-reduce. The selected high-byte range determines which margins participate in the second pass, after which the exact BF16 order statistic is recovered. The method does not gather token-level margins and does not materialize a candidate array.

For $E=256$, EQB communicates $2E \times 256$ int32 counts per layer, equivalent to $0.5$ MiB per layer. Across 48 MoE layers, the stated communication volume is 24 MiB per optimizer step. This is 128 times smaller than an all-reduced histogram with all $2^{16}$ BF16 bins, while retaining exactness with respect to the BF16 empirical quantile. The claim is exact only under the paper’s representation and quantile assumptions: EQB recovers the exact quantile of the BF16-encoded margins, not an exact quantile of higher-precision underlying values.

The method also has a useful optimization interpretation. The paper derives QB from a balanced-assignment formulation in which every token selects $\tau_t$0 experts and every expert receives the same number of assignments. The expert biases appear as dual variables. The quantile update is a coordinate-wise minimizer of the resulting convex dual objective, although the paper explicitly notes that one simultaneous finite-batch update is not generally a complete solution of the coupled dual problem. Ties further matter: an optimal bias may support a balanced assignment without forcing an arbitrary local tie-breaking rule to realize that assignment.

EQB improves both balance and benchmark loss relative to rank-averaged QB in the 100B-token experiment. Global MaxVio falls from $\tau_t$1 to $\tau_t$2, a 19.4% reduction, while Local MaxVio falls from $\tau_t$3 to $\tau_t$4, a 9.8% reduction. BPB improves on all six reported benchmarks. The implication is that correcting the distributed quantile statistic can affect local routing indirectly, even though EQB itself acts through a global, token-independent bias.

Load-Error Injection

Global bias control cannot adapt independently to the routing distributions of different local microbatches. LEI addresses this limitation by injecting the local hard-load error into the router-score gradients.

Let $\tau_t$5 denote the observed fraction of top-$\tau_t$6 assignments received by expert $\tau_t$7 and let $\tau_t$8 be the uniform target. The ideal balancing loss is proportional to $\tau_t$9, but $(K+1)$0 is determined by discrete top-$(K+1)$1 routing and has zero gradient almost everywhere. The paper therefore uses a straight-through estimator in which the forward value is the hard load while the backward Jacobian is supplied by a differentiable surrogate.

The distinction between LEI and the GShard auxiliary loss lies in this surrogate Jacobian. GShard uses the batch-mean normalized routing probability

$(K+1)$2

which couples expert corrections through the normalization denominator. The resulting gradient for an expert depends on the residuals of all experts and is scaled by each token’s total score. Moreover, the probabilities may be poorly aligned with the hard assignments when the selection bias is applied to logits but not to the mixture weights.

LEI instead uses the batch-mean unnormalized score as its surrogate. Its Jacobian is diagonal and constant with respect to the expert index. Consequently, each expert’s observed residual is injected independently and uniformly across the local batch:

$(K+1)$3

where $(K+1)$4 is the relative load error. Overloaded experts receive a positive score-gradient correction and underloaded experts receive a negative correction under the paper’s gradient convention. Forward routing and mixture weights remain unchanged; LEI modifies only the backward signal.

This construction gives LEI a more direct relationship to the hard routing error than GShard. It does not claim to make top-$(K+1)$5 routing differentiable. Rather, it chooses an identity-like surrogate Jacobian so that the correction is determined by the observed discrete load residual without probability normalization or cross-expert coupling.

A stability modification is essential. Unbounded score-side corrections caused runaway softmax log-sum-exp values in sliding-window attention for both GShard and LEI. The paper therefore rescales the residual vector by a common factor based on its maximum absolute component:

$(K+1)$6

This preserves the direction and zero-centering of the residual while bounding its magnitude. The normalization is not an incidental implementation detail: without it, the stronger local balancing signal can destabilize unrelated attention logits.

The 100B-token ablation shows a clear local-balance advantage for LEI. With EQB, normalized LEI obtains Global/Local MaxVio of $(K+1)$7, compared with $(K+1)$8 for normalized GShard and $(K+1)$9 for unnormalized GShard. The strongest contrast is local: normalized LEI reduces Local MaxVio by approximately 18% relative to normalized GShard. Model quality remains comparable, with normalized LEI achieving a mean BPB of $e$0 versus $e$1 for normalized GShard.

Figure 1: 100B-token ablation trajectories. LEI reduces both Global and Local MaxVio more strongly than the GShard loss; normalization prevents the attention-logit instability observed with unbounded score-side corrections.

The figure supports two claims that are not captured by endpoint metrics alone. First, LEI improves the trajectory of both global and local balance relative to GShard. Second, normalization retains most of the balance benefit while preventing the runaway attention-LSE behavior observed with unbounded auxiliary gradients. This establishes that the practical performance of LEI depends on controlling gradient magnitude, not merely on selecting the proposed surrogate.

Empirical results at 100B and 500B tokens

The principal 100B-token comparisons are summarized below.

Method	Global MaxVio	Local MaxVio	Mean BPB
Rank-averaged QB	0.92	5.96	0.6782
EQB	0.74	5.38	0.6570
EQB + GShard	0.71	4.71	0.6853
EQB + LEI	0.60	3.49	0.6758
EQB + normalized GShard	0.74	4.30	0.6758
EQB + normalized LEI	0.60	3.52	0.6720

EQB produces the best mean BPB among the basic global-balancing methods and improves every benchmark reported in the aggregate comparison relative to rank-averaged QB. Adding local score-side balancing substantially improves local MaxVio, but the relationship between balance and quality is not monotonic. For example, unnormalized GShard produces a lower Local MaxVio than EQB alone but a worse mean BPB. Normalized LEI obtains the lowest mean BPB among the six configurations, although the differences are small and no variance estimates are provided.

The longer 500B-token experiment compares EQB with EQB plus normalized LEI. LEI reduces Local MaxVio from $e$2 to $e$3, a reduction of approximately 19.3%, while average accuracy increases from 48.25% to 48.56%. This is a strong result for the paper’s central claim that local balancing can improve without degrading downstream quality. The improvement is not uniform across tasks: ARC rises from 50.05% to 51.49%, whereas TriviaQA decreases from 38.06% to 37.83%. The aggregate increase should therefore be interpreted as a modest net effect across heterogeneous evaluations, not as universal task improvement.

The 500B-token result also exposes a trade-off that is obscured by the shorter ablation. Global MaxVio increases from $e$4 with EQB alone to $e$5 with normalized LEI, even as local balance improves. Thus, LEI is not a uniformly stronger balancing mechanism. Its local gradient corrections can perturb the global behavior established by EQB. This supports the paper’s decomposition of the objectives but also shows that the two controls interact rather than operate independently.

The coefficient sweep reinforces this point. At $e$6, average accuracy reaches 50.53% and Local MaxVio is 4.68, but Global MaxVio rises sharply to 0.94. At $e$7, Local MaxVio falls further to 3.54 and Global MaxVio returns to 0.48, while average accuracy declines to 48.00%. The selected $e$8 is not Pareto-dominant on the reported metrics; it is a compromise with average accuracy of 48.56% and Local MaxVio of 4.15. The sensitivity indicates that LEI’s effect depends materially on coefficient tuning and that local balance alone is an incomplete proxy for training quality.

Limitations and open questions

The empirical conclusions are constrained by one 7.5B-parameter architecture, one training seed, and no reported variance estimates. Consequently, the magnitude and statistical reliability of the improvements cannot be separated from run-to-run variation. The experiments also compare separately tuned GShard and LEI coefficients; the paper correctly notes that equal coefficients would not constitute a fair comparison because the two surrogate Jacobians produce different router-gradient magnitudes.

Local MaxVio is used as a proxy for expert-parallel dispatch efficiency, but actual EP throughput, communication time, padding overhead, and step-time variance are not measured. A lower local maximum violation therefore establishes improved load uniformity, not a demonstrated systems-level speedup. The communication analysis similarly counts EQB metadata but does not provide a full end-to-end latency measurement against rank-averaged quantiles or histogram methods.

The exactness claim is tied to BF16 margins and assumes an integer target rank $e$9. The paper gives a constructive balanced assignment under this divisibility condition, but practical training configurations with nondivisible token counts require a convention not developed in the main algorithm. The treatment of ties also leaves open how the implementation’s tie-breaking behavior affects realized balance at finite precision.

Finally, the 500B-token experiment demonstrates a local/global trade-off rather than complete decoupling. LEI improves Local MaxVio but increases Global MaxVio in the reported comparison. It remains unresolved whether a joint controller, a different residual normalization, or a schedule for $\tau_t-z_{t,e}$0 can preserve EQB’s global behavior while obtaining LEI’s local improvements without extensive tuning.

Conclusion

The paper presents a technically coherent separation of global and local MoE load balancing. EQB recovers exact global-batch BF16 quantiles with communication independent of token count, improving global balance and BPB over rank-averaged QB. LEI interprets local balancing as a straight-through hard-load objective and uses an identity-like surrogate to inject expert-specific residuals directly into router-score gradients. With bounded residuals, it improves local balance relative to GShard at comparable or slightly better downstream quality.

The strongest evidence is the reduction of Local MaxVio from 5.14 to 4.15 at 500B tokens alongside an increase in average accuracy from 48.25% to 48.56%. However, the accompanying increase in Global MaxVio, coefficient sensitivity, single-seed evaluation, and absence of measured dispatch throughput qualify the broader systems interpretation. The paper’s principal contribution is therefore not a universally dominant balancing objective, but a precise algorithmic decomposition: exact quantile control for global routing statistics and direct, stabilized gradient correction for local routing errors.