A Tour of MoE: 3. A Different Way to Allocate

MoE环游记：3、换个思路来分配 苏剑林 · 2025-03-05 · Machine Learning optimal loss function gradient MoE

This post continues our exploration of MoE load balancing. In the previous post, MoE Wanderings: 2, Better Balanced Than Not, we mainly discussed the idea of using an Aux Loss to promote load balancing. Aux Loss is simple and intuitive, but it has an obvious drawback — the weight is hard to tune: too low and it fails to promote balance, too high and it easily hurts the LM Loss. So the community has long been looking for alternatives.

What we'll share in this post is a scheme called "Loss-Free," proposed by DeepSeek in Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. Compared to DeepSeek's many dazzling open-source works, this paper may seem less flashy, but in the author's view, its potential academic impact could far exceed that of the others, because the method it proposes is not only simple and effective but also highly general — a true classic.

The Gist of the Method

Faced with load imbalance, Aux Loss's response is to use an extra loss to guide the Router toward giving balanced scores, whereas Loss-Free's idea is to switch to a new allocation approach — that is, without changing the Router's existing scores, it changes how $\mathop{\text{argtop}}_k \boldsymbol{\rho}$ is used for allocation. more

Actually, there have been earlier efforts in this direction too. For instance, in 2021 Facebook proposed the BASE Layer, treating expert allocation as a linear assignment problem: with load balancing as the constraint, one seeks the allocation that maximizes the Router's total score under that constraint, which can be solved with something like the Hungarian algorithm. However, this scheme requires knowing the scores for all tokens, so for an autoregressive LLM, it only applies to training — inference still has to rely on $\mathop{\text{argtop}}_k \boldsymbol{\rho}$, creating an inconsistency between training and inference. Moreover, due to limitations of current solving algorithms, it only applies to the $k=1$ scenario.

By contrast, Loss-Free's approach is very simple yet effective. It notices a fact: we can always introduce a bias term $\boldsymbol{b}$ such that the allocation given by $\mathop{\text{argtop}}_k \boldsymbol{\rho} + \boldsymbol{b}$ is balanced. So it changes the form of MoE to

\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}

Here $\boldsymbol{b}$ is a vector independent of the input, determined during training; once training is finished it stays fixed, so it can also be used at inference time — in other words, training and inference share the same form. Note that what gets multiplied by $\boldsymbol{e}_i$ is still $\rho_i$, not $\rho_i + b_i$ — that is, $\boldsymbol{b}$ only participates in the allocation process, not in the MoE forward computation, so we have no particular requirement on the sign of $\boldsymbol{b}$ or $\boldsymbol{\rho} + \boldsymbol{b}$.

Hand-Crafting the Gradient

How do we train $\boldsymbol{b}$? We know that the optimization direction for $\boldsymbol{b}$ should naturally promote load balance. To that end, following the notation from the previous post, let's first define $\boldsymbol{f}=[f_1,f_2,\cdots,f_n]$:

\begin{equation}f_i = \left\{\begin{aligned}1/k, \quad i\in \mathop{\text{argtop}}\nolimits_k \boldsymbol{\rho}+\boldsymbol{b} \\ 0, \quad i\not\in \mathop{\text{argtop}}\nolimits_k \boldsymbol{\rho}+\boldsymbol{b}\end{aligned}\right.\end{equation}

as well as $\boldsymbol{F}=\mathbb{E}[\boldsymbol{f}]$, where $\boldsymbol{F}$ is naturally the current load distribution of the experts under bias $\boldsymbol{b}$. Next we define the uniform distribution as $\boldsymbol{Q}=(1/n,1/n,\cdots,1/n)$, so load balancing amounts to minimizing

\begin{equation}\mathcal{L}_{\text{aux}} = \frac{1}{2}\Vert\boldsymbol{F} - \boldsymbol{Q}\Vert^2 = \frac{1}{2}\sum_{i=1}^n (F_i - 1/n)^2\end{equation}

This objective is non-differentiable, but based on our experience from the previous post, we know that a Straight-Through Estimator (STE) can solve this problem. The key to STE is to find a differentiable quantity that moves in the same direction as $\boldsymbol{F}$, to serve as a smooth approximation of $\boldsymbol{F}$. Here our only optimization parameter is $\boldsymbol{b}$, and it happens to have exactly the property we want (increasing $b_i$ raises the probability that $i$ is selected, which in turn makes $F_i$ larger), so the answer suggests itself:

\begin{equation}\mathcal{L}_{\text{aux}} = \frac{1}{2}\Vert\boldsymbol{b} + \text{sg}[\boldsymbol{F}-\boldsymbol{b}] - \boldsymbol{Q}\Vert^2 = \frac{1}{2}\sum_{i=1}^n (b_i + \text{sg}[F_i - b_i] - 1/n)^2\end{equation}

Its gradient is

\begin{equation}\nabla_{\boldsymbol{b}}\mathcal{L}_{\text{aux}} = \frac{1}{2}\nabla_{\boldsymbol{b}}\Vert\boldsymbol{b} + \text{sg}[\boldsymbol{F}-\boldsymbol{b}] - \boldsymbol{Q}\Vert^2 = \boldsymbol{F} - \boldsymbol{Q}\end{equation}

So using gradient descent (SGD) to update $\boldsymbol{b}$ gives

\begin{equation}\boldsymbol{b}\leftarrow \boldsymbol{b} - \gamma (\boldsymbol{F} - \boldsymbol{Q})\end{equation}

Here $\gamma$ is the learning rate for $\boldsymbol{b}$. However, the update rule that Loss-Free ultimately settled on is slightly different — it uses signed gradient descent (SignSGD):

\begin{equation}\boldsymbol{b}\leftarrow \boldsymbol{b} - \gamma \mathop{\text{sign}}(\boldsymbol{F} - \boldsymbol{Q})\label{eq:aux-loss-free}\end{equation}

This result is actually easy to understand: if $F_i$ is larger than $1/n$, then decrease $b_i$ a bit; otherwise increase $b_i$ a bit.

An Improved Version

Besides SignSGD with an added $\mathop{\text{sign}}$, the author has found that applying RMS Norm directly to $\boldsymbol{F} - \boldsymbol{Q}$ (i.e., Normalized SGD) often achieves a better balancing effect for the same $\gamma$:

\begin{equation}\boldsymbol{b}\leftarrow \boldsymbol{b} - \gamma\frac{\boldsymbol{F} - \boldsymbol{Q}}{\text{RMS}(\boldsymbol{F} - \boldsymbol{Q})}\end{equation}

Here $\text{RMS}$ stands for "Root Mean Square," defined as

\begin{equation}\text{RMS}(\boldsymbol{F} - \boldsymbol{Q}) = \sqrt{\frac{1}{n}\sum_{i=1}^n (F_i - Q_i)^2}\end{equation}

It's not hard to see that both $\mathop{\text{sign}}(\boldsymbol{F} - \boldsymbol{Q})$ (with $\mathop{\text{sign}}$ added) and $\frac{\boldsymbol{F} - \boldsymbol{Q}}{\text{RMS}(\boldsymbol{F} - \boldsymbol{Q})}$ (with RMS Norm added) have $\text{RMS}$ equal to 1, so the two are roughly on the same scale, which means we can use the same $\gamma$ for both.

Put simply, the problem with $\mathop{\text{sign}}$ is that it applies the same update magnitude regardless of how close $F_i$ is to the target $Q_i$, which means a value of $F_i$ that has already gotten close to $Q_i$ is prone to being pushed away from the balance it has already achieved, causing oscillation. RMS Norm, on the other hand, preserves the relative magnitudes among $F_i-Q_i$, making the update magnitude more adaptive; in theory this is more conducive to promoting balance, and empirically it also tends to perform better.

Cut from the Same Cloth

When introducing Loss-Free, the original paper doesn't go through the Aux Loss derivation above; instead it directly gives the update rule in equation $\eqref{eq:aux-loss-free}$, giving the impression that the gradient $\mathop{\text{sign}}(\boldsymbol{F} - \boldsymbol{Q})$ for $\boldsymbol{b}$ was "hand-crafted" — which is also where the name "Loss-Free" comes from.

However, from the derivation given in this post, we can see that the update rule $\eqref{eq:aux-loss-free}$ can also be obtained entirely from the Aux Loss perspective — the two are cut from the same cloth. It seems the most direct benefit of Loss-Free is that you don't need to tune the Aux Loss weight anymore, but it actually still has a learning-rate parameter $\gamma$ that needs tuning. Even though the original paper has already searched for a default value of $\gamma=0.001$, it's undeniable that this hyperparameter still exists.

In the author's view, the essential innovation of Loss-Free is not the absence of an Aux Loss, but rather the separation of the optimization parameters for the Aux Loss and the LM Loss, thereby achieving both load balance and model capability without compromise. The most crucial step is noticing the fact that "a single bias term is enough to achieve load balance," and then letting the Aux Loss optimize only the newly introduced bias $\boldsymbol{b}$, while the LM Loss optimizes the remaining parameters — minimizing the negative impact of the Aux Loss on the LM Loss.

By contrast, the conventional Aux Loss scheme requires all parameters to promote load balance, while the LM Loss also optimizes all parameters — so the two optimization directions may not be fully compatible, making it harder to find an optimal balance point. Therefore, by exploiting the fact that "a single bias term is enough to achieve load balance" to separate the optimization parameters of the two losses, Loss-Free provides an elegant solution to the load balancing problem.

Related Details

Although Loss-Free is already simple and clear enough, there are still a few details worth paying attention to in practice.

First, for each batch of data, we should update the model parameters according to the LM Loss first, and only then update $\boldsymbol{b}$ according to equation $\eqref{eq:aux-loss-free}$. This is because the update of $\boldsymbol{b}$ depends on statistics $\boldsymbol{F}$ gathered from all tokens; if we update $\boldsymbol{b}$ before updating the rest of the model's parameters, there is in principle a risk of leaking future information. Although intuitively a single vector $\boldsymbol{b}$ can't leak much information, the risk nonetheless exists, so it should be avoided as much as possible.

Second, we just said that the original paper has already tuned $\gamma=0.001$, but this result may be tied to the paper's choice of Sigmoid as the Router's $\boldsymbol{\rho}$ activation function. The reason isn't hard to see: after applying Sigmoid, each $\rho_i$ is relatively independent and lies within $(0,1)$, so $\gamma=0.001$ corresponds to an update magnitude of roughly one-thousandth per step. If you switch to Softmax, ReLU, or some other activation function, you might need to re-tune $\gamma$.

To address this, the author's suggested approach is to decouple the activation functions used for the Gate and the Bias, i.e.,

\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho} + \boldsymbol{b}} \rho_i \boldsymbol{e}_i\qquad\to\qquad \boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}^{(\sigma)} + \boldsymbol{b}} \rho_i^{(h)} \boldsymbol{e}_i\end{equation}

where $\boldsymbol{\rho}^{(\sigma)} = \sigma(\boldsymbol{x}\boldsymbol{W}^{(R)}), \boldsymbol{\rho}^{(h)} = h(\boldsymbol{x}\boldsymbol{W}^{(R)})$, $\sigma(\cdot)$ is the Sigmoid function, and $h(\cdot)$ is any monotonic function with non-negative range. In plain terms, what gets added to $\boldsymbol{b}$ is the Sigmoid-activated score, so we can reuse $\gamma=0.001$; as for the factor multiplying the Expert Gate, we can use a different activation function, as long as its monotonicity matches that of Sigmoid.

Also, since the update rule $\eqref{eq:aux-loss-free}$ includes a $\text{sign}$ function, it's possible to end up training a $b_i$ whose absolute value exceeds 1, and its overall magnitude might keep growing — this is all normal and has no effect on model performance. In fact, $\boldsymbol{b}$ has one redundant degree of freedom, because adding the same constant to all $b_i$ doesn't change the result of $\mathop{\text{argtop}}_k \boldsymbol{\rho} + \boldsymbol{b}$. We can put this extra degree of freedom to use for other fun purposes (more on that next time).

Further Thoughts

Beyond MoE load balancing, the idea behind Loss-Free can be applied to many similar problems — for example, codebook collapse in VQ-VAE can be addressed with the same idea, and it turns out to be more natural and general than the previously introduced "rotation trick" or "linear transformation trick." Indeed, the opening remark of this post — that "Loss-Free's potential academic impact could far exceed that of other works" — is precisely based on this generality of Loss-Free.

Setting aside the specific application context, mathematically speaking, Loss-Free's contribution can be understood as providing a way to solve assignment problems via gradient descent. A classic linear assignment problem can be expressed as:

\begin{equation}\min_f \sum_{i=1}^n c_{i, f(i)}\end{equation}

where $c_{i,j}$ is a given cost function and $f$ is a bijection from $\{1,2,\cdots,n\}$ to itself. In the context of this post, isn't $c_{i,j}$ exactly the scores for $n$ tokens against $n$ experts, and isn't the sought-after $f$ exactly a load-balanced allocation scheme? The usual approach to solving such problems is to search for as good a solution as possible within the space that satisfies the constraints, whereas Loss-Free flips this around: it first constructs a solution that is optimal but not necessarily satisfying the constraints:

\begin{equation}f(i) = \mathop{\text{argmin}}_j c_{i,j}\end{equation}

This solution is certainly optimal in terms of score, but it doesn't necessarily satisfy the bijection condition — and here, failing to satisfy the bijection is equivalent to load imbalance. So we introduce a bias:

\begin{equation}f(i) = \mathop{\text{argmin}}_j c_{i,j} + b_j\end{equation}

$b_j$ is initialized to zero, and then updated according to equation $\eqref{eq:aux-loss-free}$. In plain terms, the update rule says: whichever $j$ appears most often, reduce the corresponding $b_j$; otherwise increase it — until a bijection emerges.

Summary

This post introduced the Loss-Free method for MoE load balancing, proposed by DeepSeek, whose core idea is to achieve load balance by introducing a simple bias term. We further examined its connection to Aux Loss, as well as its potential for application to similar mathematical problems.