MoE Journey: 4. More Effort Where It's Hardest

MoE环游记：4、难处应当多投入苏剑林 · 2025-03-28 · Machine Learning optimization gradient MoE dynamics

In the previous two posts we've been discussing load balancing. In MoE Journey: 3. A Different Approach to Allocation, when introducing the Loss-Free scheme, I left a bit of suspense: the Bias term it introduces has a redundant degree of freedom, and this degree of freedom can be used for another interesting purpose. This post is about exactly that.

As we know, MoE computes only with the top-$k$ best-matching Experts for each Token, which saves computation while still increasing the parameter count. However, if we think about it carefully, there's an obvious way this strategy could be improved: intuitively, not every Token is equally difficult, so a more sensible scheme would allocate more compute to harder Tokens and less to easier ones — this might maximize performance under the same limited budget.

And the extra degree of freedom in the Bias mentioned above turns out to be exactly what we need to achieve this goal simply.more

Design idea

First, let's recall that 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}

Load imbalance is a common problem in MoE training, and researchers proposed the Aux Loss to address it — we covered this in MoE Journey: 2. It's Not Scarcity but Inequality We Should Worry About. Additionally, in MoE Journey: 3. A Different Approach to Allocation we introduced DeepSeek's Loss-Free scheme, which turns MoE into

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

and then achieves load balancing by adjusting the newly introduced Bias term $\boldsymbol{b}$. In order to let each Token select a dynamic number of Experts, my approach is to slightly modify the Loss-Free formulation:

\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argwhere}} \boldsymbol{\rho} + \boldsymbol{b} > 0} \rho_i \boldsymbol{e}_i\end{equation}

That is, any Expert satisfying $\rho_i + b_i > 0$ is selected. This way, the number of Experts chosen for each Token is naturally dynamic, and there's no need for sorting — in some sense the whole thing becomes even simpler.

Optimization objective

The optimization objective of $\boldsymbol{b}$ has two parts: first, as in Loss-Free, we want load balance; second, we want to control the average number of Experts selected per Token to be $k$, which we can call budget control — otherwise the model could just $b_i = \infty$ select all the Experts, which is not what we want.

For load balancing we still follow the Loss-Free training approach. Let's define the notation $\boldsymbol{f} = [f_1, f_2, \cdots, f_n]$

\begin{equation}f_i = \left\{\begin{aligned}1, \quad \rho_i + b_i > 0 \\ 0, \quad \rho_i + b_i \leq 0\end{aligned}\right.\end{equation}

and let $\tilde{\boldsymbol{F}}=\mathbb{E}[\boldsymbol{f}]$, so that $\boldsymbol{F} = \tilde{\boldsymbol{F}}/|\tilde{\boldsymbol{F}}|$ is the current Expert distribution, where $|\tilde{\boldsymbol{F}}|$ is the sum of the components of $\tilde{\boldsymbol{F}}$. Loss-Free's proposed update rule is:

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

where $\boldsymbol{Q}=(1/n, 1/n, \cdots, 1/n)$ is the target uniform distribution. As we've mentioned several times, $\boldsymbol{b}$ has a redundant degree of freedom, reflected in the fact that adding the same constant to all components of $\boldsymbol{b}$ leaves the ranking result unchanged. This means we can rewrite the update rule $\eqref{eq:aux-loss-free}$ as

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

Here, a bar over a vector denotes the mean of all its components — a scalar — and subtracting a scalar from a vector means subtracting that scalar from every component. This way the resulting $\boldsymbol{b}$ necessarily satisfies $\overline{\boldsymbol{b}}=0$, without affecting the load-balancing effect. So we can $\overline{\boldsymbol{b}}$ reserve this degree of freedom for budget control.

How should we understand this? Clearly, if we add the same positive number to all of $b_i$, then the probability of satisfying $\rho_i + b_i > 0$ will increase, and hence the total budget will also increase. So the approach is straightforward: first compute the current average budget, which turns out to be exactly $|\tilde{\boldsymbol{F}}|$; if it's larger than $k$, then decrease $\boldsymbol{b}$ a bit, and vice versa. Folding this into equation $\eqref{eq:aux-loss-free-2}$ gives

If we only want to ensure that the budget doesn't exceed $k$, without requiring it to equal $k$ exactly, we can instead leave things unchanged when $|\tilde{\boldsymbol{F}}| < k$:

Attempting to simplify

Looking closely at equation $\eqref{eq:aux-loss-free-3}$, we notice it's doing two things: first, pushing $\boldsymbol{F}=\tilde{\boldsymbol{F}}/|\tilde{\boldsymbol{F}}|$ toward $\boldsymbol{Q}$, and second, pushing $|\tilde{\boldsymbol{F}}|$ toward $k$. These two seem like they could be merged into one: pushing $\tilde{\boldsymbol{F}}$ toward $\tilde{\boldsymbol{Q}}=k\boldsymbol{Q}=(k/n,k/n,\cdots,k/n)$. So equation $\eqref{eq:aux-loss-free-3}$ can be simplified to

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

I ran experiments with both equation $\eqref{eq:aux-loss-free-3}$ and equation $\eqref{eq:aux-loss-free-5}$, and found their performance to be roughly comparable. However, the two metrics — load balance and budget control — fluctuate much more in the early stages of training under equation $\eqref{eq:aux-loss-free-5}$. So readers who value stability may prefer equation $\eqref{eq:aux-loss-free-3}$ or $\eqref{eq:aux-loss-free-4}$, while readers who value simplicity can go with equation $\eqref{eq:aux-loss-free-5}$.

Given that $\mathop{\text{sign}}$ only keeps the sign of $\tilde{F}_i - \tilde{Q}_i$ while discarding its magnitude, I also tried replacing $\mathop{\text{sign}}$ with RMS Norm:

\begin{equation}\boldsymbol{b}\leftarrow \boldsymbol{b} - \gamma (\tilde{\boldsymbol{F}} - \tilde{\boldsymbol{Q}})/\Vert\tilde{\boldsymbol{F}} - \tilde{\boldsymbol{Q}}\Vert_{RMS}\end{equation}

where the $\Vert\cdot\Vert_{RMS}$ of a vector refers to the square root of the sum of squares of its components. Clearly the RMS of $\mathop{\text{sign}}$ is 1, and after RMS Norm the RMS is also 1, so the two updates are of the same order of magnitude and can share the same $\gamma$. Since RMS Norm preserves the relative magnitude of $\tilde{F}_i - \tilde{Q}_i$, smaller errors lead to smaller updates, so it fluctuates somewhat less than $\mathop{\text{sign}}$ — though not by much.

Of course, using RMS Norm to replace $\mathop{\text{sign}}$ for extra stability is a general trick — equations $\eqref{eq:aux-loss-free}$, $\eqref{eq:aux-loss-free-2}$, $\eqref{eq:aux-loss-free-3}$, or $\eqref{eq:aux-loss-free-4}$ can all be modified this way; it's really a matter of taste, and in any case it only helps a little.

Initialization scheme

Having settled the update rule for $\boldsymbol{b}$, let's now consider the initialization of $\boldsymbol{b}$ — an interesting but not particularly critical question.

With the conventional approach — initializing $\boldsymbol{b}$ to all zeros and using a Sigmoid activation for $\boldsymbol{\rho}$ — the initial stage would select all $n$ Experts, clearly exceeding the budget of $\leq k$, which would lead to a great many Token Drops. That said, if we're not too fussy about it, this isn't really a serious problem, because while most other model parameters get a Warmup, $\boldsymbol{b}$ typically doesn't, so within the first few Warmup steps the model will naturally fix this issue on its own.

If we do care about this, we can control the initial budget by adjusting the initialization of $\boldsymbol{b}$. Suppose the Router's input is a $d$-dimensional vector with zero mean and unit variance (approximately true thanks to RMSNorm), and the Router's weights are initialized with variance $\sigma^2$; then the Router's logits are approximately zero-mean with variance $\sigma^2 d$. Given this, we can use a normal approximation together with bisection to estimate an initial $\boldsymbol{b}$:

import numpy as np

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def b_init(n, k, d, sigma, eps=0.1):
    b1, b2 = -1, 0
    std = sigma * d**0.5
    logits = np.random.randn(10000, n) * std
    scores = sigmoid(logits)
    while True:
        b = (b1 + b2) * 0.5
        c = ((scores + b) > 0).sum(1).mean()
        if -eps < c - k < eps:
            return b
        elif c > k:
            b2 = b
        else:
            b1 = b

b_init(32, 4, 1024, 6e-3)

The code above assumes a Sigmoid activation, so the search interval is $[-1, 0]$; adjust it accordingly for other activation functions. That said, the recommendation here is the same as in MoE Journey: 3. A Different Approach to Allocation: adding $\boldsymbol{b}$'s $\boldsymbol{\rho}$ can uniformly use a Sigmoid activation, and it's only when multiplying by the Expert's $\boldsymbol{\rho}$ that other activation functions are worth considering.

Before this post, some prior work had already attempted MoE designs with a dynamic number of selected Experts. Below I list a few pieces of work I found, along with some brief personal commentary from my own aesthetic standpoint.

A relatively straightforward approach is taken by AdaMoE and MoE++, which mix in some low-cost Experts among the regular ones — such as null Experts, copy Experts, or constant Experts — while also encouraging load balance. This way, when a Token selects one of these simple Experts, it's effectively selecting fewer of the standard Experts, indirectly achieving a dynamic count. The benefit of this approach is that it can reuse the existing infrastructure of Top-$k$ MoE, but it lacks some flexibility.

Another straightforward idea is to change the Top-$k$ selection into a Top-$p$ selection, as in Harder Tasks Need More Experts: Dynamic Routing in MoE Models. This transformation seems natural at first glance, but in practice it has quite a few issues — for instance, it's hard to precisely control the average budget, because when $\boldsymbol{\rho}$ is close to a uniform distribution, the proportion selected by Top-$p$ becomes very large. Because of this, the original paper adds an extra entropy loss to push $\boldsymbol{\rho}$ away from uniformity. Overall, I feel the problems this introduces outweigh the benefits.

A rather distinctive approach is Ada-K Routing, which adds a new module to predict the number of Experts to activate, and trains it via reinforcement learning. This is sound in principle, but introducing reinforcement learning inevitably adds training complexity. DA-MoE, on the other hand, uses Attention scores to identify important Tokens and assigns them more Experts, but this feels somewhat ad hoc, since "MoE" in principle isn't restricted to FFN layers — once applied to Attention, wouldn't there be no Attention scores left to use?

The approach most similar in form to the one in this post is probably ReMoE, which likewise selects Experts based on a zero threshold, but achieves load balance and budget control via an Aux Loss, additionally mixing in hand-crafted gradients to control the Aux Loss weight — overall it feels a bit more patched-together. This post, by contrast, continues the Loss-Free line of thinking, using the extra degree of freedom in $\boldsymbol{b}$ to regulate this threshold, thereby achieving a dynamic number of Experts with minimal modification.

Summary

This post proposed an MoE design with a dynamically chosen number of Experts. The main idea is to slightly modify the Loss-Free MoE formulation and adjust the update rule for the Bias term, using its extra degree of freedom to simultaneously achieve load balancing and budget control.