Understanding, extending and improving quantile balancing

We release code and configurations for this post at this link.

Introduction

When training MoEs, we want some amount of load balancing across experts, i.e., routed experts should activate at similar rates. There are at least two reasons for this - a training-related reason and a systems reason.

The training-related reason is, as discussed in e.g. OLMoE, Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, and Expert Choice Routing, that not having load balancing tends to concentrate tokens on a small set of experts, and the starved experts remain under-trained, potentially leading to a feedback loop where better optimized experts get chosen more often.

The systems reason is that we want the model to be easy to serve (and train) - in an expert-parallel setup, where experts are sharded across GPUs/nodes, if there is a single replica of each expert, then the expert having to serve $1 + x$ times the mean number of tokens will roughly take $1 + x$ times more time than the average expert (as a first approximation). Having such stragglers is not optimal and from a systems perspective, having tokens spread out evenly across experts is desirable. One caveat here is that now there are already libraries for expert parallelism that dynamically replicate experts to have perfect physical load balancing across ranks (not logical load balancing across experts), such as MoonEP, and this straggler mechanism used to be more of an issue with DeepEP and such. Expert replication is still kind of messy, so some may still prefer load balancing for systems reasons too.

So we'll talk about how load balancing is done, what our desiderata might be, and two improvements on top of quantile balancing - one for improved balance/steerability, and one for flexibility (quantile balancing in itself does not have a way to control the amount of balancing).

Choosing the balancing granularity

Background: balancing granularity and distribution shift

Now that we understand that we need to balance token allocation across experts, let's think about the granularity we want the balance to hold on.

There are multiple possible granularities. From a serving perspective, one would want perfect load balance across the inference token distribution, and the ideal case is having such balance in every time interval, across all user queries.

In training, load balancing is done either within the microbatch (or sequence-wise) or the whole global batch.

Note that a learned router would have at least some kind of specialization (not necessarily interpretable) as far as token/context distribution goes (otherwise we could allocate experts uniformly at random without quality loss). Intuitively, a router could learn to route to certain experts for certain kinds of documents/prompts (dictated by activation geometry). So under a distribution shift, we would see an increase in imbalance in a globally load balanced MoE. Typically, the finer the granularity, the lower one would expect the variance in the load to be, and a sequence or microbatch balanced MoE would be more robust to such distribution shifts - more specifically, let's say for each "subgroup" $s$, we have a normalized load vector $p_s$ satisfying $\|p_s - p^*\| \le \varepsilon$, then for any reweighting of these subgroups, the final load imbalance is $\|\sum_{s} w_s p_s - p^*\| \le \sum_s w_s \|p_s - p^*\| \le \varepsilon$. But this also trades off expert specialization and quality, so there may be an optimal non-zero extent of such balancing that can help the model in general.

As an aside, strong expert specialization can sometimes imply the existence of experts that are always active or are important in some other sense, so if a small set of experts covers the majority of experts used for a certain type of task, then things such as pruning or even loading these experts into faster memory (in e.g. heterogeneous memory setups such as local inference of large models with CPU offloading) can be considered. However, it does not seem likely that optimizing for this kind of human-interpretable expert specialization is optimal as far as loss/benchmark performance is concerned. One example is REAP, an expert-pruning method that seems to preserve performance on certain domains after pruning experts.

What should expert loads look like?

For now, we ignore the systems concerns (naively delegating all of them to something like MoonEP), as well as sequence-wise and microbatch load balancing, and think about how things should be, in an intuitive manner.

Let's fix a certain data distribution (say, the training distribution). If there is expert specialization, the load distributions over experts that we see in loss-wise optimal setups are not necessarily uniform. There might also be more than one of these optimum distributions, and load distribution is just a low dimensional observation of something much more complex. It would generally depend on the full training dynamics, through things such as the job of the rest of the architecture (e.g. whether or not we have shared experts or other similar mechanisms), what kinds of feedback loops exist/what instabilities we face during training (and which optima are actually reachable through training), how expert specialization happens (and in general, how the router is learned and its geometry).

One way to think about the optimal load distribution intuitively is by coarsely splitting experts into high and low frequency experts. For high-frequency experts, the routing for them can be approximated by them being always on (e.g., shared experts in a single layer or alternating dense/MoE FFN layers if we consider the whole network), and the rest of the experts are routed experts. Other aspects can also be important for routing - one can imagine a bias on top of a linear router or even a nonlinear router to improve routing quality as well as balance (and the ZAYA1-8B technical report shows that increasing the expressivity of routers empirically seems to help both quality as well as balance).

One might want to "factorize" out as much of this kind of behavior as possible in different ways (such as the shared experts in DeepSeekMoE or the alternating dense/MoE FFN layers in the MAI-Thinking-1 report), and keep the routed part of the MoE fairly balanced, because we want to both avoid routing collapse (because it leads to reduced expressivity) as well as not lose out on expressivity due to reduced expert specialization (as far as more obvious experts that can be shared/factored out in some other way are concerned). The former is the much bigger concern, because training stability in general governs many of the decisions made for large scale training runs.

Jianlin Su's blog series on MoE has yet another view of MoE as some kind of a "low-rank" approximation to the full dense MLP. In that sense, one can think of some frequent experts corresponding to certain "important" or "common" features (or something analogous to SVD modes), and the less frequent ones (or the ones to be balanced) as the more specific features that don't get activated all the time (but when they do, they end up being in the top few experts that matter as opposed to always having a tiny contribution).

Note that when we have shared experts, we can think of the model as a dense backbone plus some extra experts that need to be routed to - these are the extra parameters that potentially behave like a content-addressable memory (where the addressing is done by the router) and help reduce loss by storing "knowledge" without necessarily being active all of the time. The factorization might be important because if we do not do so, and have perfect load balancing, then if there is some shared work that is necessary to be done for every token, every chosen subset of experts has to do at least part of the work that this shared expert would have done, reducing the available capacity for expert specialization. Tuning the number of shared experts can help.

It is still unclear whether this 2-way factorization is optimal or not - it might very well be possible that the optimal way would have more tiers to such routing while still controlling how the router works (we can have multiple routers, each balanced separately, with different expert subsets with different sparsities, a special case being shared experts corresponding to the trivial router plus the rest of the experts with another router, or a load balancing setup that has different load targets/target ranges for different experts, all balanced at the same time, or even learned-routing-free methods such as hash routing for some subsets).

A brief history of MoE load balancing

Understanding and extending quantile balancing

There are two ways to understand quantile balancing. The first is an ad hoc way and is a bit easier to think about in an elementary manner, and the second one is a more theoretically grounded way and also helps us generalize quantile balancing (which we will do in this section).

To gain some intuition, let's first look at expert-choice routing (what we had been talking about so far was token-choice routing). In this type of routing, experts choose tokens based on their allocated capacity. In this setting, we can have perfect load balancing, but at the cost of causal leakage - a routing decision made for an earlier token based on a later token can act as a side channel for information, and for large models it can let the model use this to predict the next token much better, while not being compatible with (and worsening the performance at) inference. It also causes information leakage across different sequences in a batch. Note that in the case of token-choice routing, dropping tokens itself (when experts have hard capacity constraints) can lead to causal leakage too, depending on implementation (in general, strictly causal decisions on dropping tokens can lead to a strong positional bias in which tokens are dropped or lead to wastage of capacity if some kind of budgeting is done, so load balancing is important here too).

Still, we would like something that has as good of a balance as possible. Assuming that the token distribution doesn't change much from batch to batch, it might be fine if we accept a delay of a single step when optimizing for perfect load balancing for a top-$K$ token-choice MoE.

Let's say we have the scores $S_{ij}$ for an expert $j$ for a token $i$ (these can be completely unrelated to the gate values, but in the standard quantile balancing setup, they are taken to be the same). Let's interpret the biases as "costs" (and use $S_{ij} - b_j$ this time, so that demand for an expert increases its cost).

We now want to update the biases from $b_j$ to $b'_j$ so that after the update, we have reasonable balance.

Now, consider a single token. To be in the top $K$ experts (out of $N$) for a token, a given biased score needs to beat the $K+1$-th largest value of all the expert biased scores $S_{ij} - b_j$ (taken over $j$). This gives us a threshold $a_i$ for each token as a quantile of the biased-score distribution of experts for that token. Now, suppose $b_j$ is fixed. In the spirit of alternating optimization algorithms such as EM, let's try to "solve" for $b'$ assuming $b$ is constant (that is, the thresholds are frozen). For a given expert $j$, we want $MK/N$ tokens (out of a total of $M$) to have their cutoff be satisfied by this expert. As earlier, this happens if $S_{ij} - b'_j > a_i$, or equivalently $b'_j < S_{ij} - a_i$, for $MK/N$ tokens out of a total of $M$ tokens. Thus, the new bias is again a quantile but this time of a distribution of biased scores over tokens for that expert.

In most cases, intuitively, this iteration, if carried out on the old batch, should get sufficiently close to the bias that solves the problem optimally. However, this needs us to compute statistics over a batch repeatedly and has non-trivial communication overhead, so the quantile balancing method does only one step (instead of iterating to convergence) and tries to amortize across steps in a similar vein to minibatch stochastic optimization methods.

Let's now look at it from a slightly different perspective. The following approach, for the special case of perfect balancing, is mostly covered in the Kimi K3 report, but we will generalize it below.

Matching tokens to experts is a bipartite matching problem but with more general degree constraints instead of the usual unit degree constraints. To see why, let's draw a bipartite graph, with tokens on the left and experts on the right. Each expert has a specific number of tokens it is assigned, which will be its degree in the matching, and each token has $K$ experts that it will be assigned to (i.e., their degrees must be $K$). The Gale-Ryser theorem gives us a condition for when a balanced matching exists and we'd need to verify it to check feasibility. It states that non-increasing non-negative integer sequences $(r_1,\ldots,r_M)$ and $(d_1,\ldots,d_N)$ are the two degree sequences of a simple bipartite graph if and only if their sums agree and, for every $m \in \{1,\ldots,M\}$,

$$\sum_{i = 1}^{m} r_i \leq \sum_{j = 1}^{N} \min(d_j, m)$$

When we require perfect balance (i.e., all experts have the same degree and all tokens have the same degree), the condition in the theorem is trivially true, and a balanced matching is guaranteed. Even when we have arbitrary degree constraints on experts while all tokens must have the same degree $K$, it is fairly mechanical to verify that the Gale-Ryser condition is satisfied as long as the expert degrees sum to $KM$ and the degrees are bounded by their trivial bounds.

In fact, we can make this a bit more general. Let's say that token $i$ can be assigned to between $l_i$ and $u_i$ experts (this specific generalization is not useful for usual top-$K$ MoEs, but we keep this regardless for generality), and expert $j$ can be assigned to between $l'_j$ and $u'_j$ tokens (this specific generalization is important if we want to have slack in our expert loads or arbitrary exact loads and so on), and assume that these thresholds are such that there exists a solution to this assignment problem. (Also, we assume that these bounds are integers and are bounded by the natural bounds one would have for degrees). The case where $l_i = u_i = K$ and $l'_j = u'_j$ corresponds to the situation above. The analysis for the special case of perfect balancing (plus $l'_j = u'_j = q$) is a bit easier, but we will deliberately use the more general case to be able to generalize quantile balancing further.

Now we need to choose one among these balanced matchings. To do so, we need to define the quality of a matching. We define it as

$$\sum_{ij} S_{ij} X_{ij}$$

where $S_{ij}$ is the score given by token $i$ to expert $j$, and $X_{ij}$ is $1$ if token $i$ is routed to expert $j$ and $0$ otherwise. Note that $S_{ij}$ is not necessarily related to the gating score in this formulation, but it is taken to be equal to the gating score in quantile balancing, intuitively because a better expert should have a higher score. Later, we will see how taking $S_{ij}$ to be the post-sigmoid gate and the pre-sigmoid gate logit makes things different, so it is worth thinking about $S_{ij}$ as a value we use for defining an objective function, separately from its utility as a gating function (which is also helpful for gradient flow through the router).

Note that $X_{ij}$ is a discrete binary variable here. Let's for the moment relax its domain to $[0, 1]$. For the moment, for perfect balancing, assume that $q = MK/N$ is an integer.

We want to show that all the vertices of the polyhedron corresponding to the equivalent linear programming problem are integral, so that for each linear "quality" function, the optimum value of the function is achieved at an integral point (as opposed to only at points where $X_{ij}$ is fractional for some $i, j$). This follows because the optimum of any linear function is achieved at an extreme point of the polyhedron.

There are two ways to approach this.

So, we have shown that it does not matter if we relax our integer programming problem (corresponding to assigning experts to tokens) to a linear programming problem. This means that we can use standard ideas from optimization.

Let's write the original linear programming problem for a top-$K$ MoE again: $\max \sum_{ij} S_{ij} X_{ij}$ subject to $0 \le X_{ij} \le 1$, $\sum_k X_{ik} = K$, $l'_j \le \sum_k X_{kj} \le u'_j$. Here we enforced each token to use exactly $K$ experts to reflect that we use top-$K$ MoEs as well as to simplify the analysis.

We now think about the Lagrangian. To do this, introduce a variable $\alpha_i$ for each constraint of the form $\sum_k X_{ik} = K$, $\lambda_j \ge 0$ for $\sum_k X_{kj} \le u'_j$, and $\mu_j \ge 0$ for $\sum_k X_{kj} \ge l'_j$.

Consider the sum

$$\mathcal{L}(X, \alpha, \lambda, \mu) = \sum_{ij} S_{ij} X_{ij} + \sum_i \alpha_i (K - \sum_k X_{ik}) + \sum_{j} \lambda_j (u'_j - \sum_{k} X_{kj}) + \sum_{j} \mu_j (\sum_{k} X_{kj} - l'_j)$$

Define $g(\alpha, \lambda, \mu) = \max_{X \in [0, 1]^{M \times N}} \mathcal{L}(X, \alpha, \lambda, \mu)$. Note that under the constraints on $\alpha, \lambda$ and $\mu$, we have, for any feasible $X$, that $\mathcal{L}(X, \alpha, \lambda, \mu) \ge \sum_{ij} S_{ij} X_{ij}$, so $g(\alpha, \lambda, \mu)$ is an upper bound on the value of any feasible matching. So we can minimize this over the multipliers $\alpha, \lambda$ and $\mu$ to get the best possible upper bound on the value of any feasible matching. Strong duality says that, since this is a feasible linear program, this bound is actually tight. That is,

$$\max_{\text{feasible } X} \sum_{ij} S_{ij} X_{ij} = \min_{\lambda \ge 0, \mu \ge 0, \alpha} \max_{X \in [0,1]} \mathcal{L}(X, \alpha, \lambda, \mu)$$

Now, let's rearrange and perform the inner maximization to get the maximum value as:

$$\min_{\lambda \ge 0, \mu \ge 0, \alpha} \max_{X \in [0, 1]} \sum_{ij} (S_{ij} - \alpha_i + \mu_j - \lambda_j) X_{ij} + K \sum_{i} \alpha_i + \sum_{j} (\lambda_j u'_j - \mu_j l'_j)$$

which equals (by choosing $X_{ij}$ appropriately)

$$\min_{\lambda \ge 0, \mu \ge 0, \alpha} \sum_{ij} (S_{ij} - \alpha_i + \mu_j - \lambda_j)_{+} + K \sum_{i} \alpha_i + \sum_{j} (\lambda_j u'_j - \mu_j l'_j)$$

Let's define $\beta_j = \lambda_j - \mu_j$ as the signed "price" (we explain this terminology later) of the expert $j$ (for now, this is just how $\lambda$ and $\mu$ appear in the non-linear part of the objective) - then the summand on the rightmost end is minimized when $\lambda_j = (\beta_j)_{+}$ and $\mu_j = (-\beta_j)_{+}$. Substituting and defining $h_j(\beta) = u'_j \beta$ for $\beta \ge 0$ and $l'_j \beta$ otherwise, we get

$$\min_{\alpha, \beta} \sum_{ij} (S_{ij} - \alpha_i - \beta_j)_{+} + K \sum_{i} \alpha_i + \sum_{j} h_j(\beta_j)$$

As an aside, let's see what $h_j$ is really about. In the case where $l'_j = u'_j$, $h_j$ becomes a linear function. Otherwise, it has an interesting interpretation - if we instead parameterize the slack interval as being centered $q_j = \frac{l'_j + u'_j}{2}$ and within $\delta_j = \frac{u'_j - l'_j}{2}$ of it, then $h_j(\beta)$ is $q_j \beta + \delta_j |\beta|$. This is equivalent to adding an L1 penalty on $\beta$ on top of the original problem, leading to some bias towards $\beta$ sparsity.

Let's first try to eliminate $\alpha$. For fixed $\beta$, it is a convex function in each $\alpha_i$ separately, with the (global and local) optimum being when the slope (subgradient) is $0$, i.e., $S_{ij} - \alpha_i - \beta_j$ is non-negative (define slope at $0$ arbitrarily) at $K$ values of $j$, i.e., $\alpha_i$ is anywhere between the $K$-th and $K+1$-th highest values in $S_{i \cdot} - \beta$ (so it corresponds to setting the $X$ corresponding to the top $K$ to $1$ and the rest to $0$, modulo ties). So the minimization objective becomes $\sum_{i} \mathrm{TopKSum}_K(S_{i \cdot} - \beta) + \sum_j h_j(\beta_j)$.

Before deriving quantile balancing here, we pause to see what happens if we throw gradient descent at this problem, as noted in the Kimi K3 paper for exact balance without slack. In the case of exact balance without slack (i.e., $h_j$ is linear and equals $q_j \beta$), away from ties in routing, the gradient becomes exactly the negative of the imbalance, and the gradient descent step increases $\beta$ if the expert is overloaded, and reduces $\beta$ if it is starved. SignSGD instead of gradient descent gives DeepSeek's aux-loss-free load balancing, and e.g., SMEBU from Arcee Trinity Large is a soft version of it, somewhat reminiscent of an Adam-atan2 + LaProp style version of the algorithm.

However, if we had some slack, the gradient step would look like this: if $\beta > 0$ (i.e., the cost of choosing an expert is positive), then we look at the upper bound - if it is violated, then we increase $\beta$ by a value proportional to the violation, and if it is not, we reduce it by the same amount (note that we are being deliberately conservative in updating $\beta$ so that only the bare minimum for keeping the load within bounds is done). Something analogous happens for the lower bound. However, if $\beta = 0$, then we update $\beta$ only if there is any violation, and update it in the correct direction; if there is no violation then $\beta$ remains $0$. Doing a SignSGD step here like in the aux-loss-free load balancing family of algorithms would work too. Note how we get an analogue for a with-slack version of load balancing almost immediately.

One may ask, why not solve the optimization problem optimally in the first place, instead of using gradient descent? After all, gradient descent or SignSGD have the issue that the step sizes for $\beta$ are not sensitive to the score density near the margins (and an update size that does not take this into account can lead to either a drastic change or a negligible change in load) - standard Newton-ish methods could be worth trying too, though second-order estimation could be an issue (and besides, this is a pretty discrete problem, so it would need some kind of continuous approximation).

The naive way is to try and solve $\min_\beta \sum_i \mathrm{TopKSum}_K(S_{i \cdot} - \beta) + \sum_j h_j(\beta_j)$ by noting that this is a convex optimization problem. However, this does not have a closed form, and typical convex optimization algorithms are iterative in nature. In particular, if they require computing the objective function or a subgradient many times, it could get expensive quite quickly. Also, if the batch is small enough, fitting to the current batch perfectly is likely to be suboptimal for the next batch (and would necessitate some kind of proximal penalty for things to be stable across steps). So we would like to do a single iteration (or a few iterations tops) of any such algorithm, taking inspiration from stochastic minibatch optimization algorithms. Note that approaches such as bundle/ADMM/primal-dual algorithms/LP solvers/Sinkhorn-style approaches can be quite expensive (computationally or memory/communications-wise), so we don't pursue them here.

Since this form of the problem doesn't look as tractable, let's go back to

$$\min_{\alpha, \beta} \sum_{ij} (S_{ij} - \alpha_i - \beta_j)_{+} + K \sum_{i} \alpha_i + \sum_{j} h_j(\beta_j)$$

We can try and solve this problem using blockwise coordinate descent (i.e., alternating between optimizing over $\alpha$ and $\beta$ while keeping the other fixed), and it has a simple and fast closed form iteration. For the special case of perfect balance, we see that this recovers our old balancing - optimizing over $\alpha$ while keeping $\beta$ fixed gives the top-K sum as earlier, and the $\alpha$ is the threshold we need to cross to get into the top $K$ of the biased scores $S_{ij} - \beta_j$ (consider what happens when we pair up each copy of $\alpha_i$ with each active expert). The same argument when we keep $\alpha$ fixed gives us the equivalence for the $\beta$ update in our old analysis and this setup.

Let's reason about it again in the most general case, to generalize quantile balancing. When optimizing for $\alpha$, nothing changes because $h_j$ is separate, and it leads to the same definition that leads to them being the top-$K$ cutoffs for a given token. When optimizing for $\beta$, things change a bit.

Note that $\beta = 0$ is a privileged value in the general case (unlike the standard quantile balancing setup, where we can shift all the $\beta_j$ by the same constant and it would not affect the result) - this is because $\beta = \lambda - \mu$, and at most one of the $\lambda$ and $\mu$ is non-zero as we saw earlier - positive $\lambda$ corresponds to a positive $\beta$, which is used to penalize an expert that would be over-selected if no penalty was applied, and a negative $\beta$ encourages an otherwise under-selected expert to be selected. If both are $0$, it means the expert is already fine and needs no help. In some way, since $\alpha$ is the threshold vector for each token, and it is decided by the value of $S - \beta$, $\beta$ can be thought of as a price that one has to pay for an expert (while gaining $S$ from being allotted that expert) and this $\beta$ is decided by the demand/supply situation due to the forces wanting to have expert balance.

Keeping $\alpha$ fixed, the problem for a given expert becomes, defining $R_{ij} = S_{ij} - \alpha_i$ as the margins for token $i$ for expert $j$, $\min_{\beta_j} \sum_i (R_{ij} - \beta_j)_{+} + h_j(\beta_j)$. This is a convex problem, and its local minimizers are also its global minimizers, and have zero in the subgradient. In what follows, we will ignore ties to make our life easier. Let $N_j(x)$ be the number of tokens whose current margins $R_{ij}$ for this expert exceed $x$ - this is a non-increasing function of $x$. Then for $x < 0$, the derivative of the objective is $l'_j - N_j(x)$, and for $x > 0$, the derivative is $u'_j - N_j(x)$.

Now we derive the update - if we have $l'_j \le N_j(0) \le u'_j$, we can just set $\beta_j = 0$. If $l'_j > N_j(0)$, since $N_j$ is non-increasing, the minimizer satisfies $\beta_j < 0$ - in which case we find $x$ such that $N_j(x) = l'_j$. Similarly, if $u'_j < N_j(0)$, we find $x$ such that $N_j(x) = u'_j$. This is the same as finding the (usually non-unique) quantiles corresponding to either $l'_j$ or $u'_j$, and is a natural generalization of the standard quantile balancing update from Kimi K3 (modulo how the exact inverse-quantile value is chosen - this is non-unique due to gaps between token margins).

An important detail to be kept in mind is that since this is a non-smooth problem that is also not "separable" into the blocks of $\alpha$ and $\beta$ that we perform block-coordinate descent on, this algorithm does not necessarily converge, either for standard quantile balancing or our version. There may be non-benign oscillations that don't converge to the optimum (there do not seem to be any convergence guarantees from theory), but also, the convention we use can also get stuck at a non-optimal point (for the latter, consider $M = N = 3$, $K = q = 1$, $S = [[-3, 1, 0], [1, 3, -3], [-1.9, -2, 2]]$ - starting at $\beta = (0, 0, 0)$ gives $\alpha = (0, 1, -1.9)$ and $\beta = (0, 1, 0)$ as a fixed point, with the dual objective $5$, whereas the optimal assignment is $1 \to 2$, $2 \to 1$, $3 \to 3$ whose score is $4$, so the solution found was suboptimal).

One way to reduce noisy $\beta$ assignments could be to check at the old $\beta_j$ first, to avoid jitter from our quantile/0 assignments (or even overfitting to the new batch), but we don't pursue this. Another way is adding a penalty such as $\gamma |\beta_j - \beta_j_{\mathrm{old}}|^p$ for $p = 1$ or $2$ or something similar - this still keeps the problem convex as well as solvable in the same step (potentially needing a binary search), but again has the issue of $\gamma$ not being equivalent across experts - ideally such a penalty has a geometry that respects the relevant score density, which we will not pursue further due to its complexity.

Another aspect to think about is that hard constraints might not necessarily be good, and that the effective $h_j$ might be better off being smoother, and it is possible that in the current formulation for infinitely large batches, expert load gets stuck at the boundaries for many experts. One way of thinking about this could be the following: as we already noted, $h_j$ is a sum of a linear term and an L1 penalty - this L1 penalty controls how much load balance matters for us (and the shape of $h_j$ is a plausible way of controlling what load profiles matter for us). To control this "profile importance", we might want some kind of principled way of thinking about $h_j$'s shape - it should presumably be convex (to keep the problem convex). As it turns out, there indeed is a way of doing so, and that way derives $h_j$ as a convex conjugate of a penalty being subtracted. Instead of hard constraints, we add a penalty term on the per-expert load in the primal objective that we want to maximize, say $-\sum_j \phi_j(d_j(X))$, where $d_j(X) = \sum_i X_{ij}$ is the load on expert $j$, and we also keep the primal top-$K$ constraint (i.e., $X \in [0, 1]^{M \times N}$ becomes $X \in [0, 1]^{M \times N}$ with $\sum_j X_{ij} = K$, call this constrained set $\mathcal{X}$) instead of dualizing it. Note that integrality of the optimum doesn't carry over for general $\phi$, but we will still think about this nevertheless. If $\phi_j^*$ is the convex conjugate of $\phi_j$, then $-\phi_j(d) = \inf_\beta (\phi_j^* (\beta) - \beta d)$. Plugging this in, we get $\sup_{X \in \mathcal{X}} \inf_{\beta} \langle S, X \rangle - \sum_j \beta_j d_j(X) + \sum_j \phi_j^*(\beta_j)$. If this satisfies the finite dimensional Fenchel-Rockafellar conditions, then we can swap the $\sup$ and $\inf$, to get $\inf_\beta \sup_{X \in \mathcal{X}} \sum_{ij} (S_{ij} - \beta_j) X_{ij} + \sum_j \phi_j^*(\beta_j)$, which shows that $h_j$ corresponds to the role of a convex conjugate $\phi_j^*$ of the negative of the penalty, and indeed, the convex conjugate of $h_j$ is the indicator function of the interval ($\infty$ outside the interval, $0$ inside the interval). Of course, if $h_j$ is not piecewise linear, one would in general need a binary search (or similar) to find the zero subgradient points.

It is also possible that noise softens the boundaries more than the hard constraints tell us, because we have to satisfy the "oracle" balance plus the noise, instead of just the oracle balance. Solving the problem exactly is also likely not ideal, because as we mentioned earlier, the optimal $β$-s can be very different between batches (and also be overfitted) unless we have some kind of regularization/damping towards the previous $\beta$.

We tried two special cases of this generalization and compared it to the usual perfect-balance QB setup. In the slack setup, we use $[l'_j, u'_j] = [\mathrm{round}((1 - l)q), \mathrm{round}((1 + u)q)]$. The other setup has no slack but a different exact load assignment for each expert, proportional to $1/(c + j)$ for the $j$-th expert, for a specified value of $c$. The results are discussed later.

Improving quantile balancing: design and results

Let's first list down some considerations one may have for a method of solving the load balancing problem, and think about them in relation to quantile balancing.

For an online/stochastic optimization algorithm, there are some things we care about, such as robustness to noise, convergence speed, proximity to the previous iterate/staying within a trust region, and flexibility/generality in a meaningful way. We also care about the implementation being reasonably fast (either exactly or while being a reasonable approximation).

Let's come back to quantile balancing. Specific to quantile balancing, there are a few things: the objective function (linear sum of scores), the constraints (exact balancing), and the actual optimization algorithm (histogram-based block coordinate descent like stochastic algorithm amortized over the run).

Choosing the assignment score: sigmoid vs. logit

As far as the objective function goes, it is unclear what the scores $S_{ij}$ for a given token-expert assignment should be. Standard quantile balancing (as done in Kimi K3 and the Jianlin Su blog posts) uses $S_{ij}$ as the gate value. A motivation for the score being monotonic in the value of the gate is that the gate affects the expert's magnitude of contribution, so it is in some sense an importance score. However, this does not tell us that it must be the gate value itself - why not a higher power of the gate (such as its square), or the gate's pre-sigmoid logit? One way to interpret the objective function is how it scores things - for a sigmoid gate, it may not differentiate much between a gate value of $0.99$ and $0.999$, but the pre-sigmoid logit score would (which could be either beneficial or not). In the end, the ideal case is that we have good (or at least steerable) load balance while also having better loss.

Note that different score functions will generally lead to different optimal assignments. In the case where the optimal choice of score function is not clear from a loss or training dynamics perspective, we may want to turn to other desiderata, as long as it does not adversely affect the loss and other things we care about. So in some sense, we have some freedom of choosing $S_{ij}$ - for our experiments we choose the pre-sigmoid logit as one choice. On a skim, it might seem like a natural choice since it is more likely to behave like a typical pre-activation of a neural network (which in literature is typically assumed to be Gaussian-like or just constant std with both positive and negative values and doesn't necessarily saturate) than the sigmoid. But there are other more important reasons behind this choice, which are: compatibility with implementation, robustness to noise, and equivariance under linear transforms, and we discuss these below.

Consider the implementation for quantile balancing - it requires us to compute quantiles of a quantity over all tokens, for each expert separately (the other quantile computation is fairly simple to do locally). However, doing this exactly has a non-trivial communication overhead, especially in large scale training. The Kimi K3 paper's solution to this is to use a histogram - since the scores are all sigmoids (and hence bounded), one can just compress things into a $\beta_{\min} - \Theta(1)$ to $\beta_{\max} + \Theta(1)$ range. The paper claims that the error in $\beta$ is governed by the bin sizing. However, this is not the full story - consider a router with a low logit temperature (or low entropy in some sense, for a single token). Scores then concentrate around 0 or 1. Similarly (but a bit less crucially), when the router has a high logit temperature, the bounds on sigmoid are not tight (the actual part of the support being used is small), and many bins are wasted. If you recall the discussion on Newton-ish methods from the previous section, for good load balancing, it is not only these bin sizes that matter, and we need to also care about how the biased scores are packed within bins, for good quantile estimation as far as the resulting load error is concerned (e.g., in an extreme case, a bin can have the entire distribution and not provide good enough resolution even after interpolation, which assumes a uniform distribution inside the bin). So for a loose analogy, this is like using Newton's method or some other second order algorithm for an ill-conditioned problem under low precision.

For a given expert $j$, we want to compute a quantile of $S_{ij} - \alpha_i$ over all tokens $i$. When $S_{ij}$ is a pre-sigmoid logit, one issue we face while trying to apply compression is that it is not bounded naturally like sigmoid is. So we need to compress the quantiles from $\mathbb{R}$ to some bounded interval (say $[-1, 1]$), while also preserving the resolution at the tails more than e.g. sigmoid, for extreme sparsity. One way of doing so is to use the map $x \mapsto \frac{x}{1 + |x|}$ (we name this "compact"), and indeed we use this for most of our experiments. However, there is a better way. For a sensible optimization algorithm, one would expect that $\beta_j$ and the distributions of $S_{ij} - \alpha_i$ do not change too much between consecutive batches. In that case, the relevant quantile of $S_{ij} - \alpha_i - \beta_{\mathrm{old}, j}$ would roughly be zero. So if we compress these values instead of $S_{ij} - \alpha_i$, with an $\mathbb{R} \to [-1, 1]$ compression function which has a good resolution in $[-\delta, \delta]$ where $\delta$ is the typical change in $\beta$ between consecutive steps, we have already solved this problem to some extent (it is likely possible to improve it by estimating $\delta$, but we just use the same compression function $x / (1 + |x|)$ after this shift; we call this method "relative-compact"). We notice the most stark benefits of this in settings where the desired expert loads are highly non-uniform.

So the above deals with histogram resolution as a potential issue, which can also be verified using numerical simulations. But what happens if instead of using a histogram, we just used the entire distribution, with the caveat that we still have minibatch noise (so perhaps it does not matter as much at large pretraining batch sizes)? Does sigmoid still fare poorly? The answer seems to be yes, at least a bit worse, if one checks numerical simulations for instance. A somewhat nice theoretical aside is that when we transform the logits by a positive scaling and a translation, $S_{ij}$ and $\alpha_i$ both translate by the same amount, so the things we need to take the quantile of scale by the same amount. Under the exact setting (no histogram), this changes the "conditioning" everywhere uniformly, unlike in the sigmoid case, where things can start to concentrate in different places.

As an aside, note that, as mentioned earlier, we don't necessarily have to limit ourselves to global load balancing, and that it might be good to have some "local" load balancing such as across sequences or microbatches. Moving Quantile Balancing (MQB) by Jianlin Su is one way to do so for sequence-level load balancing. It maintains a histogram (using an EMA) that is causal and its strength is controllable by a multiplier $\le 1$. We don't focus on it here, however.

When I looked at the public implementations of quantile balancing after running the experiments below, the situation was a bit varied. The initial Megatron implementation uses QB with raw router logits, as we propose, but it also does not use a histogram (and instead uses a crude approximation by averaging per-microbatch quantiles). It cites Marin which implemented quantile balancing, and one of their agent-generated reports says that using the sigmoid gate instead of pre-sigmoid logits was catastrophically worse (in their setup), and compression also made things much worse for them. Separately, NVIDIA's Kimi K3 support tracker links to a global-batch histogram implementation PR which tries to match K3 by applying QB to raw sigmoid scores and using histograms too.

Now let's look at how sigmoid vs logit fare, at four different sparsity scales. The hypothesis here is that logit would be better than sigmoid as far as balance is concerned, with potential benefits in loss, but not guaranteed.

Experimental setup

Do logit scores improve balance and loss?

We compare the sigmoid and logit versions of QB below. The loss difference is logit minus sigmoid (so negative loss difference means logit is better). "Perfect compact" refers to perfect balance using $x / (1 + |x|)$ for logit and the standard setup for sigmoid, "perfect relative" refers to centering using the old bias before using $x / (1 + |x|)$ for logit and the standard setup for sigmoid, "slack" refers to relaxed load bounds, "inverse" refers to the setting with the $\propto 1/(c + j)$ expert loads, and $E_r$ is the number of routed experts.

For exact uniform matching (perfect balance), we notice that logit has lower U-MAD/B-MAD throughout compared to sigmoid, showing better balance than sigmoid in general. When slack/inverse settings are comparable, it has lower B-MAD too. U-MAD is proportional to a total-variation distance to the uniform distribution over experts, and MaxVio is more of an $L_{\infty}$ version of that - we see much larger relative gains in MaxVio in our experiments than in U-MAD.

For loss, logit seems either comparable or better across all settings with perfect balance, but the story is a bit mixed across slack settings. So the primary benefit of logit seems to be better/more steerable balance (less violation of desired bounds) compared to sigmoid.

Relaxing exact load constraints

Now that we've considered the objective function, let's come to the constraints.

The constraints in quantile balancing are quite exact, and there is no way to control how strong the load balancing is. As discussed in the previous section, perfect balance is not really needed if things such as MoonEP can reduce the serving/training penalty of not having perfect balance. There are at least two orthogonal-looking ways of deviating from uniformity - either we give experts some slack, or assign exact loads that are not uniform. The former might be more unpredictable, but for the latter, under perfect balance, it is possible to statically allocate experts and have load balancing across different ranks, if the exact load assignment is chosen appropriately (there is a wide range of choices that lets this work).

One intuition is that a model learns "greedily" at each step, which may lead to feedback loops leading to expert starvation. Load balancing forces it to have a good balance, because in the longer term it may be better to share capacity across experts. One of the predictions we can make is that for extremely sparse MoEs, a lower bound on the expert load is much more important than upper bounding. Controlling the slack we give to these bounds (or controlling the deviation from uniformity) can be a good way to balance the model's greedy choices with the model's long term health/training performance (and also load-balance-wise and loss-wise optimality). It could even be possible to schedule this rather than have it be static, but this would be quite complicated so we don't pursue this idea further.

So we try the following two things: first, instead of load being exactly $q$, we assign bounds $(1 - l) q$ and $(1 + u) q$ where $l, u \ge 0$ and $l \le 1$ (to see if we benefit from slack), and second, we assign non-uniform loads proportional to $1/(c + j)$ for expert rank $j$, with $c = k E_r$ for a tuned multiplier $k$ on top of the total number of routed experts $E_r$ and appropriate rounding (to see if we benefit from a certain non-uniform load distribution). The goal is to check whether this translates to loss improvements.

Does load slack improve loss?

We see that slack tends to help compared to enforcing perfect balance via quantile balancing, as far as loss goes.

The selected values of $(l, u)$ (and relative performance of other runs) show that it seems generally more important to keep the lower bounds somewhat tight while the upper bounds might be relaxed, especially for very sparse MoEs. The idea is that for very sparse MoEs, it is much more important to avoid underloaded experts (perhaps more learning signal is the bottleneck for these experts) than to tightly constrain the upper bounds on the load (which also gives evidence for the idea that a single shared expert might not be the optimal way to model high frequency experts).

Do inverse-rank load targets help?

We don't see a clear general benefit (and it doesn't harm loss either), so it does not seem that enforcing a distribution improves (or worsens) the loss as compared to perfect balancing. So it is possible that having slack matters more than an exact distribution, but I would also expect this setup to be at least slightly more sensitive to hyperparameters than the slack setup. Figuring out/learning the optimal exact distribution and/or its schedule (or in general, the exact lower and upper bounds), or further tuning, is out of scope, but could be worth trying.

Alternatives to the quantile update

The remaining part is the actual algorithm being used to update expert prices. We've seen most of the relevant tradeoffs, so we will summarize those and talk about some in detail. We saw earlier how, under the same linear programming setup, aux-loss-free load balancing is a way to do gradient descent like algorithms on the dual and is quite cheap, but their step size doesn't handle the density of scores near the boundary, for which a Newton-like update could be plausibly thought of (though this problem is non-smooth, and needs approximation) - a diagonal Hessian-like approximation could be plausible. QB instead approximately solves each $\beta$ coordinate problem so it automatically adapts to such conditioning in some way, though it does not necessarily converge to an optimal solution. Optimizing using only a single batch's worth of data could also overfit on the batch and leave the estimate useless for the next update, so damping the update or checking whether the previous value gives a reasonable load or adding a proximal term that penalizes distance from previous parameter could help. These are relatively cheap and can be done using a histogram like standard QB. As we mentioned earlier, there seems to not be a good way of using stuff like bundle methods, ADMM, primal dual methods, LP/MCF solvers, Sinkhorn-style methods and so on, so currently there does not seem to be an obviously theoretically superior solver once training cost is taken into account, among the general purpose convex optimization solvers. Figuring this out could be interesting future work.

Acknowledgements

I would like to acknowledge various people in the Eleuther AI discord for discussing mixture of experts recently, GPT 5.6 Sol for writing up the outline section, creating images and writing most of the code.

References

• Conditional Computation in Neural Networks for Faster Models
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• OLMoE: Open Mixture-of-Experts Language Models
• Mixture-of-Experts with Expert Choice Routing
• Hash Layers for Large Sparse Models
• DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
• Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models
• MoE Travelogue 6: Optimal Assignment Promotes Balance (Quantile Balancing)
• MoE Travelogue 8: Enforcing Sequence-Level Balance (Moving Quantile Balancing)
• Kimi K3: Open Frontier Intelligence
• ZAYA1-8B Technical Report
• MAI-Thinking-1: Building a Hill-Climbing Machine
• Arcee Trinity Large Technical Report (SMEBU)
• REAP the Experts: Why Pruning Prevails for One-Shot MoE Compression
• MoonEP: A Perfectly Balanced Expert Parallelism Library via Dynamic Redundant Experts
• Megatron-LM baseline Quantile Balancing implementation, global-batch histogram implementation, and Kimi K3 support tracker
• Marin MoE experiment report

@online{quantile-balancing,
  author    = {nor},
  title     = {Understanding, extending and improving quantile balancing},
  year      = {2026},
  month     = {09},
  day       = {02},
  url       = {https://nor-blog.pages.dev/posts/2026-09-02-quantile-balancing/},
}