Balancing Experts at Two Scales: Exact Quantile Balancing and Load-Error Injection

We recently released Kolibri 1, a sparse Mixture-of-Experts (MoE) language model with 78B total parameters and roughly 3.5B active parameters per token, only 4.5% of the total parameters. Training a model this sparse stably at scale depends heavily on an effective expert load-balancing scheme.

In each of Kolibri’s MoE layers, a learned router selects 6 of 384 routed experts to process each token. During training, the router determines how much training data each expert receives and how evenly work is distributed across devices. If it repeatedly favors a few experts, others receive too few tokens to learn well, leaving some of the model’s capacity underused. Moreover, during training and inference, experts may run on different devices with expert parallelism. If too many tokens go to one device, the others have to wait for it to finish.

In this post, we describe the load-balancing recipe used in Kolibri pre-training and demonstrate its behavior on a small proxy model. The recipe addresses both problems with two mechanisms: Exact Quantile Balancing (EQB) for global balance across the full batch and Load-Error Injection (LEI) for local balance within each microbatch. EQB extends Quantile Balancing (QB): it sets per-expert routing biases from exact quantiles of the router logits over the global batch. LEI injects each expert’s load error in the local microbatch directly into the router’s gradients. Kolibri pre-training uses EQB together with bounded LEI. For the full derivations and more details, check out our paper.

• Background: MoE routing, how we measure imbalance, and existing load-balancing methods.
• Quantile Balancing: why setting biases from logit quantiles avoids competition between experts under top-$K$.
• Exact Quantile Balancing: computing QB’s global quantiles exactly across ranks.
• Load-Error Injection: gradient correction for local balance and its relation to the GShard auxiliary loss.
• Experiments: ablations on a 7.5B-parameter proxy model trained for 100B tokens.
• Takeaways: the main results and trade-offs.

Background

Consider a general MoE layer with $E$ routed experts, where each token is routed to $K$ of them. For each token $t$, a learned linear router maps its input $x_t$ to one logit per routed expert. For a batch of $T$ tokens, these form the logit matrix $Z\in\mathbb R^{T\times E}$, where $Z_{t,e}$ is the logit of token $t$ for expert $e$. To balance expert loads, the router adds expert biases $\beta\in\mathbb R^E$ to the logits before selecting the top-$K$ experts. The output weights are computed separately by applying sigmoid to the original, unbiased logits:

Here $\mathcal I_t$ is the set of selected expert indices for token $t$, and $S_{t,e}$, the sigmoid score, is the weight given to expert $e$’s output. The routed experts $F_e$ then contribute a weighted sum of their outputs:

In our setting, the selected weights are not renormalized, and the biases affect only the selection, not the weighting.

The task of load balancing is to spread token assignments evenly across experts. We target balance at two scales:

• Global balance: Across all tokens in a full optimizer step, so every expert receives sufficient gradient signal.
• Local balance: Within each microbatch, to avoid bottlenecks in expert-parallel execution.

We quantify imbalance with maximal violation (MaxVio), which measures the largest relative overload among experts. Let $\mathcal B$ be a set of tokens, which can be the global batch $\mathcal B_{\mathrm g}$ or a local microbatch $\mathcal B_{\mathrm l}$. Let $c_e(\mathcal B)$ count the token assignments to expert $e$ within $\mathcal B$. An even split gives each expert a target of $K|\mathcal B|/E$ assignments. Dividing each count by this target gives the relative load $f_e(\mathcal B)$, and subtracting 1 gives the relative load error $\rho_e(\mathcal B)$:

MaxVio is the largest relative load error across experts:

The busiest expert receives $1+\operatorname{MaxVio}$ times its even share: MaxVio $=0$ is perfect balance, and MaxVio $=0.5$ means 1.5× the even share.

Figure 1 shows that higher local imbalance is associated with lower training throughput in a single pre-training run.

Load-balancing mechanisms generally take one of two forms. Auxiliary losses, such as the GShard loss, add a balancing term to the training objective, so the router learns balance through gradients. Bias adjustment instead sets per-expert routing biases outside gradient descent. It includes DeepSeek’s auxiliary-loss-free controller, which moves each bias by a fixed step in the direction that reduces its expert’s load error, and QB, proposed by Jianlin Su and used in MARIN and Kimi K3, which sets the biases directly from quantiles of the router outputs. Our recipe uses one method from each family: EQB builds on QB, and LEI is a gradient-based correction closely related to the GShard loss.

Quantile Balancing

Under top-$K$, experts compete for each token’s $K$ slots: raising one expert’s bias can displace assignments from other experts. DeepSeek’s auxiliary-loss-free controller updates each expert’s bias from that expert’s own load error, with an update rate controlling the size of the adjustment. These updates do not explicitly account for the displaced assignments, so correcting one expert’s load can change the load errors of others. A small update rate limits these effects, but loads approach balance only gradually, over many steps.

Following Jianlin Su and Kimi K3 (Appendix C), we interpret QB in the context of threshold routing. Top-$K$ can be rewritten as a threshold: shift each token’s logits by a per-token bias so that its $K$th and $(K{+}1)$th largest lie on either side of zero, then select every logit above zero. With these token biases held fixed, each expert’s load can be adjusted independently. QB needs no update rate because it sets each expert’s bias directly from the logit quantile corresponding to its target load.

Top-$K$ routing on the bias-adjusted logits $Z_{t,:}+\beta^\top$ gives the assignment mask $M\in\{0,1\}^{T\times E}$:

For a vector $v\in\mathbb R^n$ and $0\le k<n$, let the quota cutoff $Q_k(v)$ be the $(k+1)$th largest entry of $v$, so that exactly $k$ entries lie strictly above it. Subtracting $Q_k(v)$ from $v$ therefore leaves exactly $k$ positive entries. For simplicity, we assume no ties at the cutoff.

We now add token biases $\alpha\in\mathbb R^T$ and define the logits with both biases as

Threshold routing selects every pair $(t,e)$ with $A_{t,e}>0$. Setting $\alpha_t=-Q_K(Z_{t,:}+\beta^\top)$ moves each token’s cutoff to zero, so threshold routing on $A$ selects exactly the same pairs as top-$K$ routing on $Z_{t,:}+\beta^\top$:

With the token biases $\alpha$ held fixed, changing one expert’s bias moves only that expert’s logits relative to the zero threshold, without directly displacing assignments from other experts. Figure 2 illustrates this difference between top-$K$ and fixed-threshold routing: raise expert A’s bias under each rule and compare the loads of the other experts.

Under top-$K$, every token selects exactly $K$ experts by construction. With threshold routing, we express this requirement as a quota on each row of $M$ (the total experts per token) and expert load balance as a quota on each column (the total tokens per expert). The targets are $K$ experts per token and $r=TK/E$ tokens per expert. We assume $1\le K<E$ and that $r$ is an integer.

Load balancing then amounts to choosing $\alpha$ and $\beta$ so that both quotas hold. QB solves this approximately by alternating between rows and columns, holding the router logits fixed. The token update shifts each row so that exactly $K$ of its logits are positive. The expert update then shifts each column so that exactly $r$ are positive. Both updates change only the biases, not the model parameters:

Figure 3 illustrates the expert update: adding $\beta_e=-q_e$ moves each expert’s cutoff $q_e=Q_r(Z_{:,e}+\alpha)$ to zero, so every expert selects the same fraction of tokens even though their logit distributions differ. The token update applies the same idea across experts for each token.

Across QB iterations, experts still interact indirectly through the token-bias updates, which depend on all expert logits. Each update meets its own quota exactly but can disturb the other. After the expert update, for example, a token may no longer have exactly $K$ adjusted logits above zero. Repeating the two updates brings both quotas closer, but is not guaranteed to satisfy them exactly. In practice, we perform a single round per optimizer step, as described below. Figure 4 walks through one round on a toy example with $K=1$: it aligns the token cutoffs, then the expert cutoffs, and returns to top-1 routing.

Biases in Training and Inference

Each optimizer step performs one token-bias update followed by one expert-bias update. The resulting $\beta$ is applied in the next step’s forward pass, which still selects the top-$K$ experts from $Z_{t,:}+\beta^\top$. Inference uses the same rule with $\beta$ fixed. Top-$K$ depends only on rankings within each token, so it ignores $\alpha_t$. Selecting the top-$K$ experts on the next batch therefore acts as one more token-bias update: it restores the token quota exactly but no longer guarantees exact expert balance. In addition, $\beta$ is computed from the current batch but applied to the next one, whose logits differ because the data and the router parameters have changed.

Expert loads on the next batch are therefore balanced only approximately: closely when one round already brings loads near their targets and the logit distributions change little from step to step.

Because top-$K$ ignores a common offset in $\beta$, we center $\beta$ by subtracting its mean, which keeps this offset from drifting to large values and improves numerical stability.

Exact Quantile Balancing

For each optimizer step, the global batch $\mathcal B_{\mathrm g}$, with $T=|\mathcal B_{\mathrm g}|$ tokens, consists of all tokens processed across data-parallel ranks and gradient-accumulation microbatches. Each rank holds the router logits $Z_{t,:}$ of its local tokens for all experts, together with the shared expert biases $\beta$. The token bias $\alpha_t$ can therefore be computed locally from $Z_{t,:}+\beta^\top$.

Updating an expert bias $\beta_e$, however, requires a quantile of the token-adjusted logits $Z_{:,e}+\alpha$ across all tokens in the global batch. These values are distributed across ranks and microbatches, so their statistics must be combined. Gathering all token-adjusted logits would communicate $T\times E$ values per layer.

MARIN averages quantiles computed separately on each rank. This is cheaper, but it does not generally recover the global quantile and makes the result depend on the token partition.

Kimi K3 accumulates histograms over microbatches and combines them with one all-reduce per layer at the end of each optimizer step. This aggregates the global distribution with communication independent of token count, but the quantile estimate is limited by the bin width.

EQB computes the global bfloat16 (BF16) quota cutoff for token-adjusted logits with two-pass radix selection. The first pass identifies a range containing the cutoff; the second resolves its exact value within that range. Both passes combine histogram counts across ranks without gathering the full logits.

For each expert, the quota cutoff $Q_r(Z_{:,e}+\alpha)$ is the $(r+1)$th largest token-adjusted logit in the global batch, or equivalently the $p$th smallest, with $p=T-r$.

We encode each BF16 token-adjusted logit as an order-preserving 16-bit key, so comparing keys is equivalent to comparing logits. Each key has a high and a low byte, and we find the target key one byte at a time:

1. Locate the high byte. Each rank accumulates a full 65,536-bin histogram of its keys, per expert and across its local tokens and microbatches. Summing it into 256 high-byte bins and all-reducing gives the global counts, whose cumulative sums identify the bin containing the $p$th smallest key and the number $c$ of keys in lower bins.
2. Resolve the low byte. Each rank takes the 256 low-byte counts of the selected high-byte bin from its local histogram. A second all-reduce gives the global counts within that bin, and their cumulative sums locate the target, the $(p-c)$th smallest key in that bin, yielding the low byte.

The two bytes identify the complete key. Decoding it gives the exact BF16 cutoff, whose negative becomes the expert bias $\beta_e$.

Figure 5 illustrates the same procedure with two-digit decimal numbers: tens correspond to the high byte, and ones to the low byte. The cutoff is the 12th-smallest of 24 values, leaving 12 values above it. The first pass places it in the 50–59 bin; seven values precede that bin, so the second pass finds its fifth-smallest value, 55.

The result is exactly the cutoff that sorting all BF16 token-adjusted logits of the global batch would give. EQB thus removes the distributed quantile approximation; quantization and the limitations of a single QB update remain.

The total reduction payload per layer is

With 256 experts, this is 0.5 MiB per layer, or 24 MiB across our 48 MoE layers per optimizer step, independent of token count. Both passes aggregate global counts, so the answer is invariant to how the same logits are partitioned.

EQB and Kimi’s histogram approach trade off quantile accuracy, communication, and local memory. EQB reduces 512 counts per expert in total: 256 in each of two sequential all-reduces. For the 1,000-bin example in Kimi’s Appendix D, this is about half the reduction payload at the same expert count and counter precision, while recovering an exact BF16 quantile.

This smaller payload comes at a memory cost. The full-resolution histogram takes $65{,}536\times E\times4$ bytes per layer, 64 MiB with 256 experts, independent of batch size, whereas Kimi’s 1,000-bin histogram needs under 1 MiB.

Load-Error Injection

Since the expert biases in EQB are shared among all microbatches, an expert can be overloaded in one microbatch and underloaded in another; those errors can cancel in the global count. Figure 6 shows a minimal example: two microbatches overload different experts, and their combined load is perfectly uniform.

A shared bias cannot give an expert opposite corrections in different microbatches, so local balance must come from the router itself, trained on each microbatch’s load errors. The standard way to provide such a signal is a GShard-style auxiliary loss. The original GShard loss uses softmax probabilities as a differentiable surrogate for discrete expert loads. Sigmoid router scores, however, are not normalized across experts, so sigmoid routers such as DeepSeek-V3 add a separate normalization for this loss. LEI skips the loss and adds each expert’s load error directly to the gradients of its sigmoid scores.

Injecting the Local Load Error

Let $\rho\in\mathbb R^E$ be the load errors from Background, computed on the local microbatch $\mathcal B_{\mathrm l}$. LEI adds them to the incoming gradient $G_{t,e}$ of the training loss with respect to the scores $S_{t,e}$, for every token in the microbatch:

where $\lambda>0$ sets the correction strength. Gradient descent then lowers the scores of overloaded experts ($\rho_e>0$) and raises those of underloaded ones ($\rho_e<0$). A lower score then makes the expert less likely to be selected. Each microbatch uses its own load errors, and all of its tokens receive the same correction for a given expert.

Bounding the correction. Unbounded, LEI and GShard can destabilize training: attention log-sum-exp (LSE) grows sharply (Figure 7). We therefore rescale the whole vector $\rho$ by one shared factor, which keeps its direction and zero sum, so that its largest entry becomes $\tanh(\lVert\rho\rVert_\infty)$ instead of $\lVert\rho\rVert_\infty$. Small load errors pass almost unchanged, and no entry ever exceeds 1. Kolibri pre-training uses $\lambda=10^{-5}$ (technical report, Section 2.1.5.2).

GShard and LEI as Straight-Through Estimators

GShard and LEI can both be read as straight-through estimators (STEs) of the same objective, the squared load error $\frac12\sum_e\rho_e^2$. The forward pass measures the true hard loads, and the backward pass routes each load error to the scores through a differentiable surrogate for those loads. The two methods differ only in this surrogate. GShard uses scores normalized across experts, so changing one score changes every normalized score for that token, and each expert’s load error also reaches the scores of other experts. LEI uses the raw sigmoid scores, so each score receives only its own expert’s load error, which is exactly the update above. Since sigmoid scores are computed independently for each expert, this unnormalized surrogate is arguably the more natural choice for sigmoid routers. In our experiments, this direct path gives better local balance than either GShard variant.

Experiments

We train a 7.5B-parameter MoE with 0.5B active parameters per token: two dense layers followed by 48 MoE layers, each with 256 routed experts, 6 active experts, and one shared expert. We compare six configurations at 100B training tokens, using 64 data-parallel ranks, a global batch of 576 sequences, and a sequence length of 16,384.

We measure imbalance with MaxVio, as defined in Background. Global MaxVio counts assignments across the whole optimizer step; Local MaxVio measures imbalance within the microbatches dispatched during that step. We average Local MaxVio across rank-local microbatches and ranks. For both Global and Local MaxVio, we report the worst MoE layer at each step, averaged over the final 1,000 steps.

At 100B tokens, we evaluate on MMLU, ARC, HellaSwag, HumanEval, MBPP, and GSM8K using eval-framework. We report mean bits per byte (BPB) across the six benchmarks. PriorBPB, developed concurrently, reduces plain BPB’s answer-length confound and would be the better choice for future comparisons.

Max LSE reports the per-step maximum sliding-window attention log-sum-exp (LSE), averaged over the final 1,000 steps, using the metric plotted in Figure 7.

EQB improves both Global and Local MaxVio over rank-averaged QB, as well as BPB on all six benchmarks. Adding bounded LEI lowers Local MaxVio from 5.38 to 3.52 (busiest expert at 4.5× instead of 6.4× its even share), but raises mean BPB from 0.6570 to 0.6720. Bounded LEI gives better balance and mean BPB than either GShard variant, though EQB alone has the lowest mean BPB at 100B tokens.

Figure 7 shows the evolution of Global MaxVio, Local MaxVio, and attention LSE over training for all six runs.

In the unbounded GShard and LEI runs, attention LSE grows sharply. Bounding the corrections keeps LSE stable, with little change in Local MaxVio for LEI and a further reduction for GShard.

Takeaways

• Viewing QB as threshold routing explains why quantile-based biases avoid the competition that couples per-expert bias controllers under top-$K$.
• EQB computes QB’s global quantiles exactly for BF16 logits, with communication independent of token count. At 100B tokens, it improves balance over rank-averaged QB and gives the lowest mean BPB in our comparison.
• Global balance does not imply local balance. Bounded LEI lowers Local MaxVio from 5.38 to 3.52 while keeping attention LSE stable, at a cost of 0.015 mean BPB. Whether that trade is worthwhile depends on how much local imbalance limits throughput in a given expert-parallel setup (Figure 1).

Citation

For the full derivations and experiments, see our paper on arXiv. If you use EQB or LEI in your work, please cite:

@misc{neitemeier2026eqb,
  title         = {Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts},
  author        = {Pit Neitemeier and Jiaze Li and Alessio Serra and Philipp Scholl and Sohir Maskey},
  year          = {2026},
  eprint        = {2609.28053},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2609.28053}
}

Acknowledgements

We thank our co-authors on the paper, Alessio Serra, Philipp Scholl and Sohir Maskey, as well as the pre-training, infrastructure and evaluation teams whose work made the Kolibri training runs and the ablations in this post possible. We are also grateful to our reviewers for their careful reading and valuable feedback on this post.

References

1. Alabi, T., Blanchard, J. D., Gordon, B., Steinbach, R. (2012). Fast k-Selection Algorithms for Graphics Processing Units. ACM Journal of Experimental Algorithmics, 17(4).
2. Aleph Alpha (2026). Kolibri: A Sovereign European Model on the Pareto Frontier. Technical report.
3. Aleph Alpha Research. eval-framework. GitHub repository.
4. Burns, T. F., Teja S, P., Schneider, F., Reinhardt, M. (2026). PriorBPB: Reducing the Answer Length Confound of Bits-per-Byte While Improving Accuracy Prediction. NeurIPS 2026 Trust-AI-Eval Workshop.
5. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
6. Dial, L. (2026). Mixture of Experts Quantile Balancing: Validated at 32B-A5B (1e22 FLOPs) Scale. OpenAthena blog.
7. Kimi Team (2026). Kimi K3: Open Frontier Intelligence. arXiv:2607.24653.
8. Lepikhin, D. et al. (2021). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021.
9. Neitemeier, P., Li, J., Serra, A., Scholl, P., Maskey, S. (2026). Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts. arXiv:2609.28053.
10. Su, J. (2026). Travels in MoE: 6. Promoting Load Balance via Optimal Assignment. Blog post (in Chinese).
11. Wang, L., Gao, H., Zhao, C., Sun, X., Dai, D. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. arXiv:2408.15664.