MoE Journey: 6. Optimal Assignment Promotes Balance
As we know, load balancing is a fundamental and critical piece of the MoE architecture, directly affecting model efficiency and performance. This series has already covered two mainstream approaches to load balancing: the classic Aux Loss scheme introduced in MoE Journey: 2. It's Not Scarcity but Inequality That Matters, and the Loss-Free scheme proposed by DeepSeek, discussed in MoE Journey: 3. A Different Way to Allocate. Each has its strengths, and each has its limitations.
This article explores a third approach: optimal assignment, which treats load balancing as a linear programming problem under equality constraints. In its final form it still belongs to the Loss-Free family, but it rests on a completely different principle, and provides a more accurate, hyperparameter-free update rule.
Recap of the Methods
Of the two existing methods, Aux Loss is the more straightforward: its core idea is "penalize wherever things are unstable," applying a regularization term to punish load imbalance. However, Aux Loss has two problems: first, the penalty coefficient is hard to tune — too large and it interferes with the optimization of the main loss, too small and the balancing effect is weak; second, behind Aux Loss lies the STE (Straight-Through Estimator), which means its gradient is suboptimal and may introduce unknown effects beyond load balancing. more
To address this, DeepSeek proposed the second scheme, Loss-Free, which introduces an extra bias term to assist with ranking, as shown below:
\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\qquad\to\qquad \boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho} + \boldsymbol{b}} \rho_i \boldsymbol{e}_i\end{equation}
Note that $\boldsymbol{b}$ is used only to adjust the ranking of Experts; what actually gets multiplied onto the Expert output is still $\rho_i$, so it does not directly participate in the model's computation and does not introduce interfering gradients. However, the fact that $\boldsymbol{b}$ has no gradient also means we need to design a separate update rule for it, and the idea is quite intuitive: the larger $b_i$ is, the more likely the $i$-th Expert is to be selected. So it first tallies the current load distribution $\boldsymbol{F}$; if $F_i$ exceeds the target value $1/n$ then $b_i$ is decreased, otherwise $b_i$ is increased, i.e.
\begin{equation}\newcommand{sign}{\mathop{\text{sign}}}\boldsymbol{b} \leftarrow \boldsymbol{b} - \gamma\sign(\boldsymbol{F} - 1/n)\end{equation}
Overall, Loss-Free is less "invasive" to the model and is indeed simpler and more elegant, but it's not perfect either. Although it removes the penalty coefficient of Aux Loss, it still has an $\gamma$ parameter to tune, which plays the role of a learning rate for $\boldsymbol{b}$. The paper recommends $\gamma=10^{-3}$, but this is actually tightly coupled with $\boldsymbol{\rho}$ using a Sigmoid activation — switch to a different activation function and $\gamma$ needs to be retuned.
Moreover, even when using a Sigmoid activation, the distribution of $\boldsymbol{\rho}$ in some layers can be rather "pathological," in which case the model becomes quite sensitive to $\gamma$, and a fixed $\gamma$ makes it hard to achieve load balance. This situation isn't rare — for instance, if MoE is used in the first few layers of a model, balancing tends to be difficult there, which is why the "first_k_dense" trick exists. Similarly, when the model is large or the total number of Experts $n$ is large, individual layers can also become hard to balance.
Linear Programming
Let's state the problem to be solved more precisely: suppose there are $m$ Tokens, and the Router's scores for the $i$-th Token over $n$ Experts are $\boldsymbol{s}_i = (s_{i,1},s_{i,2},\cdots,s_{i,n})$, giving a total of $mn$ scores. These scores may be positive or negative, and need not lie within any predetermined range. We want to devise an assignment scheme based on these scores, deciding which Experts each Token should activate.
We stipulate that each Token selects $k$ Experts, so a basic scheme would simply pick the top $k$ Experts by score. But this scheme can suffer from load imbalance — some Experts might be activated noticeably more or less often than others — so we further stipulate that each Expert is activated exactly $mk/n$ times. Under these two constraints, we seek the assignment scheme with the highest total score, formalized as
\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_j x_{i,j} = k,\quad \sum_i x_{i,j} = \frac{mk}{n}\label{eq:target}\end{equation}
where $x_{i,j}=1$ indicates that the $i$-th Token selects the $j$-th Expert, and $x_{i,j}=0$ indicates it does not. Note also that $mk/n$ must be an integer for the two equality constraints to hold exactly; let's assume this for now. Here, each Expert is activated exactly $mk/n$ times, which represents the most ideal, perfectly uniform state. In practice one could of course relax this somewhat, but for the theoretical derivation we adopt the strictest constraint.
In the expression above, $x_{i,j}$ can only take the value $0$ or $1$, making this an integer programming problem. Discrete optimization is generally hard, so we consider its relaxed version:
\begin{equation}\max_{x_{i,j}\in[0,1]} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_j x_{i,j} = k,\quad \sum_i x_{i,j} = \frac{mk}{n}\label{eq:relax}\end{equation}
Now $x_{i,j}$ can take any real value in $[0,1]$, while the other conditions remain unchanged. Note that $x_{i,j}$ is linear in both the objective function and the equality constraints, so this is a linear programming problem over a bounded region.
Minimax
Constrained optimization is also generally not easy, so let's consider the unconstrained $\max\text{-}\min$ form (a.k.a. the "Lagrange multiplier method"):
\begin{equation}\max_{x_{i,j}\in[0,1]}\min_{\alpha_i,\beta_j} \sum_{i,j} x_{i,j}s_{i,j} - \sum_i \alpha_i\left(\sum_j x_{i,j} - k\right) - \sum_j \beta_j\left(\sum_i x_{i,j} - \frac{mk}{n}\right)\label{eq:relax-max-min}\end{equation}
If $\sum_j x_{i,j} = k$ and $\sum_i x_{i,j} = mk/n$, then the expression above is equivalent to $\eqref{eq:relax}$, and the maximum is some finite value; but if either of these fails to hold, then the $\min$ step can be driven to negative infinity, making the maximum negative infinity. Since a finite value is clearly larger than negative infinity, only the former case is possible, meaning it is indeed equivalent to $\eqref{eq:relax}$.
The objective $\eqref{eq:relax-max-min}$ is linear in both $x_{i,j},\alpha_i,\beta_j$, and linear functions are both convex and concave, and $[0,1]$ is a convex set, so the conditions of the Minimax theorem are satisfied, and we can swap the order of $\max$ and $\min$, giving
\begin{equation}\min_{\alpha_i,\beta_j} \max_{x_{i,j}\in[0,1]} \sum_{i,j} x_{i,j}(s_{i,j} - \alpha_i - \beta_j) + k\sum_i \alpha_i + \frac{mk}{n}\sum_j \beta_j\label{eq:relax-min-max}\end{equation}
In the expression above we've already isolated the term involving $x_{i,j},\alpha_i,\beta_j$. Looking closely, we notice that the $\max$ step can actually be solved directly: when $s_{i,j} - \alpha_i - \beta_j > 0$, we simply take $x_{i,j}=1$ to maximize the objective; when it is negative, we take $x_{i,j}=0$ to maximize the objective; and when it equals zero, $x_{i,j}$ can take any value with the same result. That is,
\begin{equation}\left\{\begin{aligned}&\,x_{i,j}^* = 1, &\, s_{i,j} - \alpha_i - \beta_j > 0 \\ &\,x_{i,j}^* = 0, &\, s_{i,j} - \alpha_i - \beta_j < 0 \\ &\,x_{i,j}^* \in [0,1], &\, s_{i,j} - \alpha_i - \beta_j = 0 \end{aligned}\right.\end{equation}
Here $s_{i,j} - \alpha_i - \beta_j = 0$ is the special case; assuming its occurrence probability is negligible, $x_{i,j}^*$ is then either 0 or 1. Of course, even when $s_{i,j} - \alpha_i - \beta_j = 0$ occurs, thanks to the arbitrariness of $x_{i,j}^*$, we can still let $x_{i,j}^*$ take the value 0 or 1 to satisfy the constraints. This shows that although the form $\eqref{eq:relax}$ is a relaxation of the original problem $\eqref{eq:target}$, its optimal solution is also the optimal solution of the original problem — the two are fully equivalent.
Divide and Conquer
Substituting the $x_{i,j}^*$ found above into $\eqref{eq:relax-min-max}$ gives $x_{i,j}^*(s_{i,j} - \alpha_i - \beta_j) = \max(0, s_{i,j} - \alpha_i - \beta_j)$, so the optimization objective $\eqref{eq:relax-min-max}$ simplifies to
\begin{equation}\min_{\alpha_i,\beta_j} \sum_{i,j} \max(0, s_{i,j} - \alpha_i - \beta_j) + k\sum_i \alpha_i + \frac{mk}{n}\sum_j \beta_j\end{equation}
We'll solve this using an alternating-minimization approach: fix $\beta_j$ to solve for $\alpha_i$, then fix $\alpha_i$ to solve for $\beta_j$, and alternate between the two. Since $\alpha_i,\beta_j$ has an obvious symmetry, these two steps are essentially the same problem. Let's first look at fixing $\beta_j$ and solving for $\alpha_i$; the problem then becomes equivalent to
\begin{equation}\min_{\alpha_i} \sum_{i,j} \max(0, s_{i,j} - \alpha_i - \beta_j) + k\sum_i \alpha_i\end{equation}
We can also observe that each term $\alpha_i$ is summed independently, so we can decompose this into $m$ independent subproblems — meaning we can temporarily drop the subscript $i$ and simplify the problem further to
\begin{equation}\min_{\alpha} k\alpha + \sum_j \max(0, s_j - \beta_j - \alpha)\end{equation}
Sort all values of $s_j - \beta_j$ in descending order as $s_{\sigma_1} - \beta_{\sigma_1} \geq s_{\sigma_2} - \beta_{\sigma_2} \geq \cdots \geq s_{\sigma_n} - \beta_{\sigma_n}$, so that the $j$-th largest element is $s_{\sigma_j} - \beta_{\sigma_j}$. Now suppose we already know $s_{\sigma_l} - \beta_{\sigma_l} \geq \alpha \geq s_{\sigma_{l+1}} - \beta_{\sigma_{l+1}}$; then the objective function equals
\begin{equation}k\alpha + \sum_{j=1}^l (s_{\sigma_j} - \beta_{\sigma_j} - \alpha) = \left\{\begin{aligned} &\,\sum_{j=1}^k (s_{\sigma_j} - \beta_{\sigma_j}) + \sum_{j=k+1}^l \underbrace{(s_{\sigma_j} - \beta_{\sigma_j} - \alpha)}_{\geq 0},&\, l \geq k \\ &\,\sum_{j=1}^k (s_{\sigma_j} - \beta_{\sigma_j}) - \sum_{j=l+1}^k \underbrace{(s_{\sigma_j} - \beta_{\sigma_j} - \alpha)}_{\leq 0},&\, l \leq k \\ \end{aligned}\right.\end{equation}
This shows that both increasing $l > k$ and increasing $l < k$ would enlarge the objective, so its minimum can only be attained at $l=k$, at which point the result does not depend on the specific value of $\alpha$ — that is, $\alpha^*$ can be any value between the $k$-th and $k+1$-th largest elements of $s_j - \beta_j$. By convention, we take $\alpha^*$ to be the $k+1$-th largest element.
Alternating Iteration
Restoring the subscript $i$, we get: for any given $i$, $\alpha_i^*$ is the $k+1$-th element when all values of $s_{i,j} - \beta_j$ are sorted in descending order. Similarly, fixing $\alpha_i$ and solving for $\beta_i$ yields: for any given $j$, $\beta_j^*$ is the $mk/n+1$-th element when all values of $s_{i,j} - \alpha_i$ are sorted in descending order. We alternate between these two steps as our final solving algorithm.
Suppose we have found sufficiently accurate $\boldsymbol{\alpha}^*$ and $\boldsymbol{\beta}^*$. Then, based on the earlier analysis, $x_{i,j}^*$ automatically satisfies the "either 0 or 1" property along with the constraints, and $x_{i,j}^*=1$ corresponds to $s_{i,j} - \alpha_i^* - \beta_j^* > 0$. Combining this with the constraint $\sum_j x_{i,j}^* = k$, we can conclude that for every Token $i$, the Experts it selects must be the Top-$k$ according to $\boldsymbol{s}_i - \boldsymbol{\beta}^*$.
This tells us that inference only needs $\boldsymbol{\beta}^*$ — $\boldsymbol{\alpha}^*$ is merely an intermediate variable of the solving process, and once training is finished it can be ignored entirely. This is crucial, because the size of $\boldsymbol{\beta}$ is fixed at $n$, whereas the size of $\boldsymbol{\alpha}$ is $m$, with $m$ being the global batch size, which changes dynamically — so it wouldn't be a sensible format for inference otherwise. The solving procedure that exploits this property is shown in the diagram below, where $\mathop{\text{desc_sort}}$ denotes descending sort.
$$\begin{array}{|l|} \hline \text{Quantile Balancing (QB): problem}\eqref{eq:target}\text{alternating solving algorithm} \\[4pt] \hline \text{input: score matrix}\boldsymbol{s}\in\mathbb{R}^{m\times n} \\ \text{output: assignment}\boldsymbol{x}\in\{0,1\}^{m\times n} \\[4pt] \hline \begin{array}{ll} 1: & \text{Initialize }\boldsymbol{\beta} = \boldsymbol{0}_{1\times n} \\ 2: & \textbf{For }t=1,2,\cdots,T\textbf{ do } \\ 3: & \qquad \boldsymbol{\alpha} \leftarrow \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\beta}, \text{axis=1})_{[:, k:k+1]} \\ 4: & \qquad \boldsymbol{\beta} \leftarrow \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\alpha}, \text{axis=0})_{[mk/n:mk/n+1]} \\ 5: & \text{Output } x_{i,j}=1 \text{ if } j\in\mathop{\text{argtop}}_k \boldsymbol{s}_i - \boldsymbol{\beta} \text{ else } 0 \end{array} \\ \hline \end{array}$$
There's one more improvement we can make: using the concept of "quantiles," we can unify "the $k+1$-th largest of $n$ numbers" and "the $mk/n+1$-th largest of $m$ numbers" — both are actually the "$1-k/n$-th quantile" along their respective dimensions. Numerical frameworks like NumPy, JAX, and PyTorch all implement a "quantile" function, which lets us avoid a full sort of the data and saves some complexity.
This is exactly why we call this algorithm "Quantile Balancing (QB)."
Watch Out for the Trap
But it's not time to celebrate yet — there's a rather subtle trap here, and falling into it can produce results that are fundamentally wrong, so extra care is needed.
As mentioned above, QB inference only needs $\boldsymbol{\beta}$, and it suffices to store only $\boldsymbol{\beta}$ during training, so from the Loss-Free perspective, QB can also be viewed as providing a new bias-update method, which one might just as well call "Quantile Bias." Whether it's the original SignSGD rule or the Quantile approach of this article, both depend on the scores of the entire batch of Tokens, so the order of operations must not be mixed up: the old $\boldsymbol{\beta}$ must be used to select the Experts for the current batch, and only afterward should the $\boldsymbol{\beta}$ value be updated. This is the only way to guarantee there is no information leakage.
Some readers may wonder: how could a bias vector that doesn't directly participate in the forward computation leak anything? True, intuitively there's very little information it could leak, but the risk is nonetheless real. Perhaps when training small models we could experiment with a "leaky" version, but for large models we shouldn't take this risk, precisely because they're large and powerful enough to amplify any subtle bug. So maintaining train/inference consistency and eliminating any risk of information leakage is a basic requirement when training large models.
Adapting this to real training scenarios, QB can be tweaked further: instead of zero-initializing and iterating $T$ times, we can start from the previous step's bias and take just one iteration step each time. This avoids overfitting to the current batch and also reduces computational cost. Selecting the Top-$k$ based on $\boldsymbol{s} - \boldsymbol{\beta}$ is something MoE already has to do anyway, and now it's simply changed to selecting the Top-$(k+1)$, which adds almost no extra cost. So the only extra step is using $\boldsymbol{s} - \boldsymbol{\alpha}$ to find the "$1-k/n$-th quantile" of $\text{axis=0}$.
$$\begin{array}{|l|} \hline \text{Quantile Balancing (QB) practical form} \\[4pt] \hline \text{input: score matrix}\boldsymbol{s}\in\mathbb{R}^{m\times n}\text{, previous step}\boldsymbol{\beta}\in\mathbb{R}^n \\ \text{output: assignment}\boldsymbol{x}\in\{0,1\}^{m\times n}\text{, new}\boldsymbol{\beta}\in\mathbb{R}^n \\[4pt] \hline \begin{array}{ll} 1: & x_{i,j}=1 \text{ if } j\in\mathop{\text{argtop}}_k \boldsymbol{s}_i - \boldsymbol{\beta} \text{ else } 0 \\ 2: & \boldsymbol{\alpha} \leftarrow \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\beta}, \text{axis=1})_{[:, k:k+1]} \\ 3: & \boldsymbol{\beta} \leftarrow \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\alpha}, \text{axis=0})_{[mk/n:mk/n+1]} \\ 4: & \text{Output } \boldsymbol{x},\boldsymbol{\beta} \end{array} \\ \hline \end{array}$$
Even with just a single iteration, though, this new step is still fairly expensive — it requires finding the $mk/n+1$-th largest element among $m$ elements, where $m$ equals "global number of samples × sequence length," typically in the millions at minimum. Given the various parallelism strategies and gradient accumulation involved, an exact implementation is usually unacceptably costly. A compromise is to split the samples into the largest micro-batches we can tolerate, compute a $\boldsymbol{\beta}$ separately for each micro-batch according to the formula, and then average them to get the final result.
Once these issues are resolved, QB is basically all upside. First, it has no learning-rate-like hyperparameter to tune; second, it balances load remarkably fast, and it's especially good at handling extreme cases — for instance, when using it to train an all-MoE model, even the first MoE layer becomes very well balanced. That said, for layers that the original SignSGD rule can already balance well, QB usually doesn't offer much additional advantage.
MaxVio comparison for the first MoE layer
Demo Code
Here's a piece of demo code for interested readers to try out on their own:
import numpy as np
def quantile_bias(s, k, T=5):
"""交替quantile求最优bias
原理：https://kexue.fm/archives/11619
"""
m, n = s.shape
beta = np.zeros((1, n))
for _ in range(T):
alpha = np.quantile(s - beta, 1 - k / n, axis=1, keepdims=True)
# alpha = alpha.clip(0, np.inf) # BIP会多出这一步
beta = np.quantile(s - alpha, 1 - k / n, axis=0, keepdims=True)
# beta = beta.clip(0, np.inf) # BIP会多出这一步
return beta
def max_min_avg_vio(s, k):
"""计算max_vio、min_vio和avg_vio
其中 max_vio ≥ 0, avg_vio ≥ 0, -1 ≤ min_vio ≤ 0，三者都是越接近于0表示越均衡
"""
m, n = s.shape
topk = np.argsort(-s, axis=1)[:, :k]
f = np.bincount(topk.reshape(-1), minlength=n)
f = f / f.sum() * n - 1
return f.max(), f.min(), np.abs(f).mean()
m, n, k = 100000, 256, 8
s = np.random.rand(m, n) + np.random.rand(n) # 模拟一个不均匀的打分
b = quantile_bias(s, k, 5)
max_min_avg_vio(s, k) # 直接取top-k的max_vio、min_vio和avg_vio
max_min_avg_vio(s - b, k) # 减去偏置后top-k的max_vio、min_vio和avg_vio
Related Work
The idea of viewing MoE load balancing through the lens of optimal assignment first appears in the paper BASE Layers: Simplifying Training of Large, Sparse Models, but the first paper to give a complete general solution is Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models (abbreviated BIP), and QB is in fact an improvement upon BIP.
Compared with QB's objective $\eqref{eq:target}$, BIP replaces the equality constraint with an inequality constraint:
\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_j x_{i,j} \leq k,\quad \sum_i x_{i,j} \leq \frac{mk}{n}\label{eq:target-leq}\end{equation}
Repeating the same sequence of derivations, at the step where we form the $\max\text{-}\min$, an extra constraint $\alpha_i\geq 0, \beta_j\geq 0$ appears, i.e., we must consider
\begin{equation}\max_{x_{i,j}\in[0,1]}\min_{\alpha_i\geq 0,\beta_j\geq 0} \sum_{i,j} x_{i,j}s_{i,j} - \sum_i \alpha_i\left(\sum_j x_{i,j} - k\right) - \sum_j \beta_j\left(\sum_i x_{i,j} - \frac{mk}{n}\right)\end{equation}
in order for it to be equivalent to the relaxed version of problem $\eqref{eq:target-leq}$. Then, $\max$ and $\min$ can still be swapped, and repeating a similar derivation, the update rule for $\boldsymbol{\alpha},\boldsymbol{\beta}$ ends up with an extra clipping operation on $\max(0,\cdot)$ to keep it non-negative:
\begin{equation}\begin{aligned} \boldsymbol{\alpha} \leftarrow &\, \max(0, \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\beta}, \text{axis=1})_{[:, k:k+1]}) \\ \boldsymbol{\beta} \leftarrow &\, \max(0, \mathop{\text{desc_sort}}(\boldsymbol{s} - \boldsymbol{\alpha}, \text{axis=0})_{[mk/n:mk/n+1]}) \end{aligned}\end{equation}
However, in the author's own tests, this clipping operation noticeably slows down load balancing, and it often produces cases where you can "rob the rich" but not "help the poor" (extremely frequent Experts get suppressed, but extremely rare Experts can never be rescued), which suggests this clipping operation is actually counterproductive. For this reason, QB directly changes the original problem to use an equality constraint, which removes the non-negativity requirement on $\boldsymbol{\alpha},\boldsymbol{\beta}$ and simplifies the solution as much as possible.
Moreover, BIP appears to make the same mistake discussed in the previous section: its procedure updates $\boldsymbol{\beta}$ before selecting the Top-$k$, which violates both the causality principle (leaking future information) and the train/inference consistency principle (leaking across samples). Even so, BIP's thorough analysis of load balancing from the optimal-assignment perspective remains highly inspiring.
After some searching, the author found that following BIP, there have been further explorations along this direction, which overlap with this article to some extent but are not fully consistent with it. We won't discuss them in detail here:
Maximum Score Routing For Mixture-of-Experts
Selective Sinkhorn Routing for Improved Sparse Mixture of Experts
MicroMoE: Fine-Grained Load Balancing for Mixture-of-Experts with Token Scheduling
A Theoretical Framework for Auxiliary-Loss-Free Load Balancing of Sparse Mixture-of-Experts in Large-Scale AI Models
Gradient Descent
Finally, let's introduce one more solution scheme, situated somewhere between Loss-Free and QB. As we know, in QB's iterative scheme, computing $\boldsymbol{\alpha}$ is relatively cheap; what's genuinely expensive is computing $\boldsymbol{\beta}$, which requires some form of sorting across all Tokens. Suppose $\boldsymbol{\alpha}$ is given; then the optimization objective for $\boldsymbol{\beta}$ is
\begin{equation}\min_{\beta_j} \underbrace{\sum_{i,j} \max(0, s_{i,j} - \alpha_i - \beta_j) + \frac{mk}{n}\sum_j \beta_j}_{\text{denote}\ell}\end{equation}
Besides using the quantile to find its optimal solution, is there a cheaper way to obtain an approximate solution? As it turns out, yes! It's clear that the objective function $\ell$ is differentiable, so we can certainly consider gradient descent, whose gradient is
\begin{equation}\frac{\partial\ell}{\partial\beta_j} = \frac{mk}{n} - \sum_{i=1}^m \chi(s_{i,j} - \alpha_i - \beta_j > 0)\end{equation}
where $\chi$ is the indicator function, $\chi(\text{True})=1,\chi(\text{False})=0$. Clearly this gradient is also relatively cheap to compute, and once we have the gradient, we can perform gradient descent. To align with Loss-Free, we consider SignSGD:
\begin{equation}\beta_j \leftarrow \beta_j - \gamma\sign\left(\frac{\partial\ell}{\partial\beta_j}\right)\end{equation}
Using this to replace the Quantile-based update of $\boldsymbol{\beta}$ in QB gives a cost roughly comparable to Loss-Free, and empirically its effectiveness also falls somewhere between the two (assuming Loss-Free also uses a Sigmoid activation and the same value of $\gamma$). In fact, we can further prove that if the largest and second-largest elements in every row of the score matrix $\boldsymbol{s} - \boldsymbol{\beta}$ are always distinct ($k$ and $k+1$ never tie), this scheme is strictly equivalent to Loss-Free.
Summary
In this article, we examined the MoE load balancing problem from the perspective of optimal assignment, and arrived at a new Aux-Loss-free load balancing algorithm, Quantile Balancing. It is more stable and accurate than the existing Loss-Free scheme, applies to Router scores of any range, and requires no additional hyperparameters to tune.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.