We have already seen that load balance is a fundamental and critical component of the MoE architecture, directly affecting the model’s efficiency and performance. Two articles in this series have introduced two mainstream approaches to achieving load balance: the classic Aux Loss scheme introduced in MoE Tour: 2, Worry Not About Scarcity but About Inequality, and the Loss-Free scheme proposed by DeepSeek in MoE Tour: 3, A Different Way of Assignment. Each has its strengths and its limitations.
This article explores a third approach: optimal assignment, which treats load balance as a linear programming problem with equality constraints. In its final form, it is still Loss-Free, but it is based on a completely different principle and provides a more accurate, hyperparameter-free update rule.
Method Review
Of the two existing methods, the idea behind Aux Loss is relatively simple: “penalize wherever there is imbalance,” imposing a penalty on load imbalance via a regularization term. However, Aux Loss has two problems: first, the penalty coefficient is hard to tune—too large interferes with the optimization of the main loss, too small yields poor balance; second, behind Aux Loss is STE (Straight-Through Estimator), which means its gradients are suboptimal and may bring unknown effects beyond load balancing.
To address this, DeepSeek proposed a second scheme, Loss-Free, which introduces an extra bias term to assist with sorting, 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 only used to adjust the sorting of the Experts; what gets multiplied onto the Experts is still \(\rho_i\), so it does not directly participate in the model’s computation and produces no interfering gradients. However, the fact that \(\boldsymbol{b}\) has no gradient means we need to design a custom update rule for it, and the idea is very intuitive: the larger \(b_i\) is, the more likely the \(i\)-th Expert is to be selected. So we first compute the current load distribution \(\boldsymbol{F}\); if \(F_i\) exceeds the expected value \(1/n\), we shrink \(b_i\), otherwise we increase \(b_i\): \[\begin{equation} \boldsymbol{b} \leftarrow \boldsymbol{b} - \gamma\,\mathop{\text{sign}}(\boldsymbol{F} - 1/n) \end{equation}\] Overall, Loss-Free is less “intrusive” to the model and can indeed be considered simpler and more elegant, but it is not perfect. It lacks the penalty coefficient of Aux Loss, but it still has a \(\gamma\) parameter to tune, which acts as the learning rate of \(\boldsymbol{b}\); the paper recommends \(\gamma=10^{-3}\), which is in fact tightly coupled with the use of Sigmoid activation for \(\boldsymbol{\rho}\)—once the activation function is changed, \(\gamma\) needs to be re-tuned.
In addition, even if we use Sigmoid activation, the distribution of \(\boldsymbol{\rho}\) in some layers can be rather “deformed,” in which case the model becomes quite sensitive to \(\gamma\), and a constant \(\gamma\) struggles to achieve load balance. This situation is not rare: for example, if the first few layers of a model use MoE, they often fail to balance easily, which is why the “first_k_dense” operation exists; also, when the model is relatively large or the total number of Experts \(n\) is relatively large, individual layers can also be hard to balance.
Linear Programming
Let us describe the problem to be solved more precisely: suppose there are \(m\) tokens, and the Router score of the \(i\)-th token over the \(n\) Experts is \(\boldsymbol{s}_i = (s_{i,1},s_{i,2},\cdots,s_{i,n})\), so there are \(mn\) scores in total. These scores may be positive or negative, and are not necessarily within some prescribed range. We wish to devise an assignment plan based on these scores, deciding which Experts each token should activate.
We stipulate that each token can only select \(k\) Experts, so a basic scheme is to pick the \(k\) Experts with the highest scores. But this scheme may suffer from load imbalance, i.e., some Experts may be activated noticeably too often or too rarely. So we further stipulate that each Expert can only be activated \(mk/n\) times. Under these two constraints, we seek the assignment 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\) means the \(i\)-th token has selected the \(j\)-th Expert, and \(x_{i,j}=0\) means it has not. Also note that \(mk/n\) must be an integer for the two equality constraints to hold strictly; we assume this for now. Here each Expert is activated exactly \(mk/n\) times, which is the ideal state of absolute uniformity; in practice this can of course be relaxed, but in the theoretical derivation we adopt the strictest constraint.
In the above, \(x_{i,j}\) can only take the value \(0\) or \(1\), making this an integer programming problem; discrete optimization is usually 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]\), and the other conditions remain unchanged. Note that \(x_{i,j}\) is linear in both the objective function and the constraint equations, so this is a linear programming problem on a bounded region.
Max-Min
Constrained optimization is usually not easy either, so we consider the constraint-free \(\max\text{-}\min\) form (i.e., the “method of Lagrange multipliers”): \[\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 above is equivalent to Eq. \(\eqref{eq:relax}\), and the maximum is a finite value; but if either is violated, then the \(\min\) step can reach negative infinity, and the maximum would be negative infinity. Clearly negative infinity is not as large as a finite value, so only the former can occur, i.e., it is equivalent to Eq. \(\eqref{eq:relax}\).
The objective \(\eqref{eq:relax-max-min}\) is linear in \(x_{i,j},\alpha_i,\beta_j\); a linear function is both convex and concave, and \([0,1]\) is a convex set, so it satisfies the conditions of the Minimax theorem, and we can swap the order of \(\max\) and \(\min\); i.e., it equals \[\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 above we have separated out the terms involving \(x_{i,j},\alpha_i,\beta_j\). Careful observation shows that the \(\max\) step can actually be solved directly: when \(s_{i,j} - \alpha_i - \beta_j > 0\), we take \(x_{i,j}=1\) to maximize the objective; when it is less than 0, we take \(x_{i,j}=0\) to maximize the objective; when it equals 0, \(x_{i,j}\) can take any value without changing the result. That is, \[\begin{equation} \left\{\begin{aligned} &\,x_{i,j}^* = 1, \quad s_{i,j} - \alpha_i - \beta_j > 0 \\ &\,x_{i,j}^* = 0, \quad s_{i,j} - \alpha_i - \beta_j < 0 \\ &\,x_{i,j}^* \in [0,1], \quad s_{i,j} - \alpha_i - \beta_j = 0 \end{aligned}\right. \end{equation}\] where \(s_{i,j} - \alpha_i - \beta_j = 0\) is a rather special case; assuming its probability of occurrence is negligible, \(x_{i,j}^*\) is then either 0 or 1. Of course, even if \(s_{i,j} - \alpha_i - \beta_j = 0\), due to the arbitrariness of \(x_{i,j}^*\), we can also make \(x_{i,j}^*\) take 0 or 1 according to the constraints. From this we see that although the form \(\eqref{eq:relax}\) is a relaxation of the original problem \(\eqref{eq:target}\), its optimal solution is also an optimal solution of the original problem; the two are completely equivalent.
Divide and Conquer
Substituting the above \(x_{i,j}^*\) into Eq. \(\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 will solve it using the idea of alternating minimization: first fix \(\beta_j\) and solve for \(\alpha_i\), then fix \(\alpha_i\) and solve for \(\beta_j\), alternating between the two. Since \(\alpha_i,\beta_j\) have an obvious symmetry, these two steps are really the same problem. Let us first look at fixing \(\beta_j\) and solving for \(\alpha_i\); the problem is then 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 further observe that each \(\alpha_i\) term accumulates independently, so we can decompose it into \(m\) independent sub-problems, which means 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}\] Arrange all \(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}\), i.e., the \(j\)-th largest element is \(s_{\sigma_j} - \beta_{\sigma_j}\). Then suppose we already know that \(s_{\sigma_l} - \beta_{\sigma_l} \geq \alpha \geq s_{\sigma_{l+1}} - \beta_{\sigma_{l+1}}\); 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}, \quad 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}, \quad l \leq k \end{aligned}\right. \end{equation}\] This shows that whether \(l > k\) or \(l < k\) enlarges the objective function, so the minimum can only be attained at \(l=k\), in which case the result does not depend on the specific value of \(\alpha\); i.e., \(\alpha^*\) can be any number 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 obtain: for any given \(i\), \(\alpha_i^*\) is the \((k+1)\)-th element when all \(s_{i,j} - \beta_j\) are sorted in descending order. Similarly, the result of fixing \(\alpha_i\) and solving for \(\beta_j\) is: for any given \(j\), \(\beta_j^*\) is the \((mk/n+1)\)-th element when all \(s_{i,j} - \beta_j\) are sorted in descending order. We alternate these two steps as the final solution algorithm.
Suppose we have found sufficiently accurate \(\boldsymbol{\alpha}^*\) and \(\boldsymbol{\beta}^*\); then by the earlier analysis, \(x_{i,j}^*\) automatically satisfies the 0-or-1 property and the constraints, and \(x_{i,j}^*=1\) corresponds to \(s_{i,j} - \alpha_i^* - \beta_j^* > 0\). Combined with the constraint \(\sum_j x_{i,j}^* = k\), we can conclude that for each token \(i\), the Experts it selects must be the Top-\(k\) of \(\boldsymbol{s}_i - \boldsymbol{\beta}^*\).
This tells us that the inference stage only needs \(\boldsymbol{\beta}^*\); \(\boldsymbol{\alpha}^*\) is merely an intermediate variable of the solution process and can be ignored after training. This point is crucial, because the size of \(\boldsymbol{\beta}\) is the fixed \(n\), while the size of \(\boldsymbol{\alpha}\) is \(m\), where \(m\) is the global batch size, which changes dynamically—so it is not a sensible inference format. The solution procedure exploiting this is shown below, where \(\mathop{\text{des\_sort}}\) denotes descending-order sorting.
\[\begin{array}{|l|} \hline \text{Quantile Balancing (QB): alternating solution algorithm for problem \eqref{eq:target}} \\[4pt] \hline \text{Input: score matrix } \boldsymbol{s}\in\mathbb{R}^{m\times n} \\ \text{Output: assignment plan } \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{des\_sort}}(\boldsymbol{s} - \boldsymbol{\beta}, \text{axis=1})_{[:, k:k+1]} \\ 4: & \qquad \boldsymbol{\beta} \leftarrow \mathop{\text{des\_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}\]
One further improvement is that we can use the concept of “quantiles” to unify “the \((k+1)\)-th largest among \(n\) numbers, and the \((mk/n+1)\)-th largest among \(m\) numbers”; they are in fact both the “\(1-k/n\) quantile” of their respective dimensions. Numerical frameworks such as Numpy, Jax, and Torch all have “quantile” function implementations; using them avoids fully sorting the data and saves some complexity.
For this very reason, we call this algorithm “Quantile Balancing (QB)”.
Beware the Pitfall
But it is not yet time to celebrate; there is a rather inconspicuous pitfall here, and falling into it may yield fundamentally wrong results, so be especially careful.
As just said, QB inference only uses \(\boldsymbol{\beta}\), and storing only \(\boldsymbol{\beta}\) during training suffices. So from the Loss-Free perspective, QB provides a new way of updating the Bias, and it could well be called “Quantile Bias”. Both the original SignSGD and the Quantile of this article rely on the scores of all tokens, so the order must not be mistaken: you must use the old \(\boldsymbol{\beta}\) to select the Experts for the current batch of data, and only then update the value of \(\boldsymbol{\beta}\); only in this way can you guarantee that there is no information leakage.
Some readers may wonder: what can a Bias vector that does not directly participate in forward computation possibly leak? Admittedly, intuitively the leakable information seems very limited, but the risk does exist. Perhaps for training small models we could try the leaking version, but training large models should not take this risk, because a large model is capable enough to amplify any subtle bug. Therefore, keeping training and inference consistent and eliminating any risk of information leakage is a basic tenet of training large models.
Combined with real training scenarios, QB can be further adjusted: instead of initializing at zero and iterating \(T\) times, we can start from the previous step’s Bias and iterate only once per step, which avoids overfitting to the current batch while reducing the computational cost. Selecting Top-\(k\) based on \(\boldsymbol{s} - \boldsymbol{\beta}\) is something vanilla MoE already does; now it merely changes to selecting Top-\((k+1)\), adding almost no cost. So the extra step is finding the “\(1-k/n\) quantile” along \(\text{axis=0}\) using \(\boldsymbol{s} - \boldsymbol{\alpha}\).
However, even with just one iteration, 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,” usually on the order of millions at minimum. Under various parallelism strategies and gradient accumulation, an exact implementation is usually unacceptable. A compromise is to split the samples into the largest mini-batches we can accept, compute a \(\boldsymbol{\beta}\) separately for each mini-batch according to the formula, and then average them as the final result.
Once these issues are resolved, QB is almost all advantages: first, it has no hyperparameters like a learning rate to tune; second, it balances very quickly and is especially good at handling extreme cases—for example, using it to train an all-MoE model, even the first MoE layer becomes extremely balanced. However, for layers that the original SignSGD rule can already balance, QB usually offers no advantage.
Demo Code
Here is a piece of demo code; interested readers can experiment with it:
import numpy as np
def quantile_bias(s, k, T=5):
"""Alternating quantile to find the optimal bias
Principle: 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 has this extra step
beta = np.quantile(s - alpha, 1 - k / n, axis=0, keepdims=True)
# beta = beta.clip(0, np.inf) # BIP has this extra step
return beta
def max_min_avg_vio(s, k):
"""Compute max_vio, min_vio and avg_vio
where max_vio >= 0, avg_vio >= 0, -1 <= min_vio <= 0;
for all three, the closer to 0 the more balanced
"""
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) # simulate uneven scores
b = quantile_bias(s, k, 5)
max_min_avg_vio(s, k) # max_vio, min_vio and avg_vio of directly taking top-k
max_min_avg_vio(s - b, k) # max_vio, min_vio and avg_vio of top-k after subtracting the biasGradient Descent
Finally, we introduce another solution scheme that sits between Loss-Free and QB. As we know, in QB’s iteration scheme, computing \(\boldsymbol{\alpha}\) is relatively cheap; what is truly expensive is computing \(\boldsymbol{\beta}\), which requires some kind of sorting across all tokens. Given \(\boldsymbol{\alpha}\), 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 Quantile to find its optimal solution, is there a cheaper way to find an approximate solution? There is indeed! Clearly the objective \(\ell\) is differentiable, so we can perfectly well consider gradient descent; its 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\). Computing this gradient is also relatively cheap; with the gradient in hand, we can do gradient descent. To align with Loss-Free, we consider SignSGD: \[\begin{equation} \beta_j \leftarrow \beta_j - \gamma\,\mathop{\text{sign}}\left(\frac{\partial\ell}{\partial\beta_j}\right) \end{equation}\] Using it to replace the Quantile update step for \(\boldsymbol{\beta}\) in QB achieves a cost close to Loss-Free, and in tests its performance is also in between the two (with both it and Loss-Free using Sigmoid activation and the same \(\gamma\)).
Article Summary
In this article we explored MoE load balancing from the optimal-assignment perspective, deriving a new auxiliary-loss-free load balancing algorithm, Quantile Balancing. It is more stable and accurate than existing Loss-Free schemes, applies to Router Scores with arbitrary value ranges, and has no extra hyperparameters to tune.
When reprinting, please include the address of this article: https://kexue.fm/archives/11619
For more detailed reprinting terms, please refer to: Scientific Spaces FAQ
Comments