MoE Journeys: 2. It's Not Scarcity but Inequality That Should Worry Us

MoE环游记：2、不患寡而患不均 苏剑林 · 2025-02-21 · Machine Learning loss function gradient sparse MoE

In the previous post, MoE Journeys: 1. Starting from Geometric Meaning], we introduced a geometric interpretation of MoE, aiming to derive and understand MoE by starting from the best approximation of a dense model. At the end of that post, we also noted that writing down MoE's computational formula is only the beginning — training an actually effective MoE model requires filling in a lot of details, such as the load balancing problem we'll discuss in this post.

Load balancing — as the saying goes, "it's not scarcity but inequality that should worry us" — in plain terms means making sure every Expert is doing work, and as much as possible an equal share of the work, so as to avoid wasting compute on some Experts. Load balancing is both a requirement for making full use of training compute, and a requirement for realizing as much as possible the potential of MoE's large parameter count.

Requirements Analysis

As we know, the basic form of MoE is

\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\end{equation}more

For a conventional MoE, $\boldsymbol{\rho}$ is a probability distribution (Router), $\boldsymbol{e}_i=\boldsymbol{v}_i$, and $\boldsymbol{v}_i$ is the output of a small FFN (Expert); whereas for the geometric MoE we derived in the previous post, $\boldsymbol{\rho}$ has no normalization requirement — it predicts the norm of the Expert, while $\boldsymbol{e}_i=\boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$ predicts the direction of the Expert.

Regardless of which formulation of MoE we use, the actual behavior is roughly the same — they just differ in how we understand them. But note that although MoE's formula gives the impression that "for every token encountered, we go find the corresponding Expert to compute it," in actual training it's really the other way around: first we allocate the appropriate compute to each Expert, and then tokens are routed to their assigned Experts for parallel computation. This is exactly why the component responsible for scoring, $\boldsymbol{\rho}$, is called the Router.

Given this, if Experts are allocated unevenly, we may end up in a situation like this: some Experts (Dead Experts) sit nearly idle almost all the time, wasting compute; while some Experts have too many tokens to handle and simply can't keep up, forcing them to drop tokens (i.e., abandon processing some tokens). In theory, the appearance of Dead Experts means MoE has failed to reach its intended parameter count — that is, we've spent the GPU memory for a large parameter count but ended up training something with the effectiveness of a small parameter count model.

So, whether from the perspective of training or of performance, we want to ensure load balancing across Experts.

Auxiliary Loss

The conventional approach to promoting load balancing is to add a related loss function, which we usually call the "Aux Loss (Auxiliary Loss)." The Aux Loss in mainstream use today can be traced back to the 2020 paper GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding].

Before introducing the Aux Loss, we need to first bring in some new notation. First, as we've already mentioned, for a general MoE, $\boldsymbol{\rho}$ need not be a probability distribution. We denote the normalized version of $\boldsymbol{\rho}$ as $\boldsymbol{p}=[p_1,p_2,\cdots,p_n]$, and its Top-$k$ version as $\boldsymbol{f}=[f_1,f_2,\cdots,f_n]$, where

\begin{equation}p_i = \frac{\rho_i}{\sum_{i=1}^n \rho_i},\qquad f_i = \left\{\begin{aligned}1/k, \quad i\in \mathop{\text{argtop}}\nolimits_k \boldsymbol{\rho} \\ 0, \quad i\not\in \mathop{\text{argtop}}\nolimits_k \boldsymbol{\rho}\end{aligned}\right.\end{equation}

Next we define $\boldsymbol{P}=\mathbb{E}[\boldsymbol{p}],\boldsymbol{F}=\mathbb{E}[\boldsymbol{f}]$, where $\mathbb{E}$ denotes averaging over all tokens of all samples. It's not hard to see that $\boldsymbol{F}$ is precisely the Expert's current load distribution, while $\boldsymbol{P}$ is essentially a smooth approximation of $\boldsymbol{F}$.

With this notation in hand, we can write the Aux Loss as:

\begin{equation}\mathcal{L}_{\text{aux}} = \boldsymbol{F}\cdot \boldsymbol{P} = \sum_{i=1}^n F_i P_i\label{eq:aux-loss}\end{equation}

The literature generally defines the Aux Loss with an extra factor of $n$, i.e., their Aux Loss equals $n \mathcal{L}_{\text{aux}}$ as defined here. Additionally, some large-scale MoEs may compute the Aux Loss on a per-device basis, to achieve intra-device balance and reduce inter-device communication — these are implementation choices left to individual practitioners. However, more recent experiments have shown that forcing local balance in this way is quite likely to hurt the model's final performance.

Straight-Through Estimation

I wonder if anyone has noticed something odd: whether it's the original source, subsequent papers, or popular-science articles — at least in everything I've read — citations of the Aux Loss never come with a proof. It seems everyone just takes it for granted that the Aux Loss above obviously promotes balance. But is it really that obvious?

Personally, I couldn't see why, so below I'll give a derivation of formula $\eqref{eq:aux-loss}$, from which we can also design other forms of Aux Loss ourselves. First, define the uniform distribution $\boldsymbol{Q}=(1/n,1/n,\cdots,1/n)$. As mentioned, $\boldsymbol{F}$ is the current load distribution, so load balancing is equivalent to $\boldsymbol{F}=\boldsymbol{Q}$. This gives us the following fairly intuitive Aux Loss:

\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\label{eq:aux-loss-2}\end{equation}

The problem is that $\boldsymbol{F}$ is obtained from $\mathop{\text{argtop}}_k$ via an argmax-like (discrete) operation, which means the expression above is not directly a differentiable objective we can use. How do we solve this? The answer is the STE (Straight-Through Estimator)] trick, which designs separate functions for the forward and backward passes. Specifically, $\boldsymbol{F}$ is not differentiable, while $\boldsymbol{P}$, as its smooth approximation, is differentiable. So during backpropagation we simply replace $\boldsymbol{F}$ with $\boldsymbol{P}$, i.e.,

\begin{equation}\mathcal{L}_{\text{aux}} = \frac{1}{2}\Vert \boldsymbol{P} + \text{sg}[\boldsymbol{F}-\boldsymbol{P}] - \boldsymbol{Q}\Vert^2 = \frac{1}{2}\sum_{i=1}^n (P_i + \text{sg}[F_i - P_i] - 1/n)^2\label{eq:aux-loss-3}\end{equation}

where $\text{sg}[]$ is the stop-gradient operator, which keeps the forward output unchanged but forces the gradient to zero. After this modification, $\mathcal{L}_{\text{aux}}$ becomes a genuinely usable Aux Loss. Let's try computing its gradient:

\begin{equation}\begin{aligned} \nabla_{\boldsymbol{\theta}}\mathcal{L}_{\text{aux}} =&\, \frac{1}{2}\nabla_{\boldsymbol{\theta}}\sum_{i=1}^n (P_i + \text{sg}[F_i - P_i] - 1/n)^2 \\ =&\, \sum_{i=1}^n (P_i + \text{sg}[F_i - P_i] - 1/n) \nabla_{\boldsymbol{\theta}}(P_i + \text{sg}[F_i - P_i] - 1/n)\\ =&\, \sum_{i=1}^n (F_i - 1/n) \nabla_{\boldsymbol{\theta}}P_i = \nabla_{\boldsymbol{\theta}}\sum_{i=1}^n (F_i - 1/n) P_i\\ =&\, \nabla_{\boldsymbol{\theta}}\left(\sum_{i=1}^n F_i P_i\right) \end{aligned}\end{equation}

Here $\boldsymbol{\theta}$ denotes the model parameters. The final result shows that the gradient of formula $\eqref{eq:aux-loss-3}$ equals the gradient of formula $\eqref{eq:aux-loss}$, which means using formula $\eqref{eq:aux-loss}$ as the Aux Loss is gradient-equivalent to formula $\eqref{eq:aux-loss-3}$ — and this is exactly how the Aux Loss of formula $\eqref{eq:aux-loss}$ arises.

However, formula $\eqref{eq:aux-loss}$ only has meaning as an equivalent gradient — it doesn't have meaning as a Loss, and isn't really a genuine loss function. For instance, when $\boldsymbol{F} = \boldsymbol{P}$, we can compute that formula $\eqref{eq:aux-loss}$ equals $1/n$, but in fact we can construct some $\boldsymbol{F}$ not equal to $\boldsymbol{P}$ that makes it smaller than $1/n$. So formula $\eqref{eq:aux-loss}$ doesn't behave like a normal loss where smaller is always better, nor is its minimum attained at $\boldsymbol{F} = \boldsymbol{P}$.

The General Form

The derivation above actually gives us a general recipe for constructing an Aux Loss: first build a loss based on $\boldsymbol{F}$ that meets our requirements, and then, at implementation time, replace $\boldsymbol{F}$ with $\boldsymbol{P} + \text{sg}[\boldsymbol{F}-\boldsymbol{P}]$. For example, since we know that maximizing entropy can also push a distribution toward balance, we can likewise construct an Aux Loss using the negative entropy:

\begin{equation}\mathcal{L}_{\text{aux}} = \sum_{i=1}^n (P_i + \text{sg}[F_i - P_i])\log(P_i + \text{sg}[F_i - P_i])\end{equation}

The expression above can be directly used in code. Of course, if we want to simplify it, we can similarly compute its gradient, giving the result

\begin{equation}\nabla_{\boldsymbol{\theta}}\mathcal{L}_{\text{aux}} = \nabla_{\boldsymbol{\theta}}\sum_{i=1}^n(P_i + \text{sg}[F_i - P_i]) \log(P_i + \text{sg}[F_i - P_i]) = \nabla_{\boldsymbol{\theta}}\sum_{i=1}^n P_i \log F_i\end{equation}

In both of these gradient-simplification steps, we made use of the following identity

\begin{equation}\sum_{i=1}^n \nabla_{\boldsymbol{\theta}}P_i = \nabla_{\boldsymbol{\theta}}\sum_{i=1}^n P_i = \nabla_{\boldsymbol{\theta}}1 = \boldsymbol{0}\end{equation}

which relies on the facts that $\boldsymbol{P}$ is a probability distribution and that the target distribution $\boldsymbol{Q}$ is uniform. But if we don't insist on the simplified equivalent form, and instead directly use an Aux Loss of the form $\boldsymbol{F}\to \boldsymbol{P} + \text{sg}[\boldsymbol{F}-\boldsymbol{P}]$, we're not bound by these two constraints.

For example, the fact that $\boldsymbol{P}$ serves as a smooth approximation of $\boldsymbol{F}$ only relies on the property that "when $P_i$ is large, $F_i$ is usually also large." So using an unnormalized $\mathbb{E}[\boldsymbol{\rho}]$ as $\boldsymbol{P}$ is usually fine too — and this can actually matter in certain special situations (for example, when $\boldsymbol{\rho}$ has both positive and negative values), since in that case it cannot be normalized into a probability distribution. Similarly, take the target $\Vert\boldsymbol{F} - \boldsymbol{Q}\Vert^2$: clearly, this can push $\boldsymbol{F}$ toward any target distribution $\boldsymbol{Q}$ we want, not necessarily a uniform one.

Summary

This post introduced the load balancing problem in MoE, and presented a general recipe for constructing an Aux Loss. Besides the Aux Loss, there are other approaches to promoting load balancing — we'll discuss those next time.