MoE Tour: 8. Forced Sequence-Level Balancing

MoE环游记：8、强制序列级均衡 苏剑林 · 2026-05-22 · Machine Learning statistics loss function RNN MoE

So far, the "MoE Tour" series has covered seven posts, five of which revolve around MoE routing and load balancing. In terms of routing form, these can be split into static computation and dynamic computation; in terms of how load balancing is achieved, they split into Aux-Loss and Loss-Free approaches, with Loss-Free further dividing into DeepSeek's SignSGD scheme and the QB scheme I proposed.

One detail worth noting is that the Loss-Free schemes we've discussed so far all achieve balance in the sense of the global batch, whereas in practice we often add a sequence-level Aux Loss to avoid extreme imbalance within a single sequence. If we believe that every Aux-Loss scheme has a corresponding Loss-Free counterpart, how should we go about implementing sequence-level Loss-Free balancing? That's the topic explored in this post.

Recap

Let's first recap our previous discussions of load balancing. In MoE Tour: 2. Fairness Matters More Than Scarcity, we first discussed the MoE load-balancing problem, and the solution we proposed was Aux Loss — adding an extra loss function to encourage balance. The drawback of Aux Loss is that it introduces an extra weighting coefficient to tune; the upside is that it can flexibly control the granularity of balancing, achieving balance at the sequence, group, or global level (though of course, the more global, the more communication overhead). more

Starting from MoE Tour: 3. A Different Way to Allocate, we discussed Loss-Free load-balancing strategies, which introduce an extra bias term to control the degree of balance, updated via a hand-designed SignSGD-style rule. This approach is more flexible than Aux Loss and more infra-friendly. In MoE Tour: 4. Harder Tasks Deserve More Input, we generalized it to MoE with a dynamic number of activated experts.

Then, in MoE Tour: 6. Optimal Allocation for Balance, we viewed the load-balancing problem from the perspective of optimal allocation, and by alternately solving a dual objective, we obtained a new load-balancing strategy called Quantile Balancing (QB). Purely in terms of the result, this can be seen as a new and more accurate way of solving for the bias term in the Loss-Free scheme. Then, in MoE Tour: 7. A Minimalist Solution for Dynamic Activation, we likewise generalized QB to dynamic MoE.

There's a subtlety here: the bias term introduced by Loss-Free is shared globally, so all current Loss-Free schemes, without exception, can only achieve global load balancing. As mentioned above, if needed, Aux Loss can also be used to promote sequence-level balance — so does Loss-Free have a corresponding scheme for sequence-level balance?

Local Centering

Coincidentally, @ChangJonathanC explored something similar in the article Causal Routing Bias for Aux-Loss-Free MoE Training. The author proposes two approaches; let's briefly go over them first, which will also help us compare with our own approach later. Let $\boldsymbol{s}\in\mathbb{R}^{l\times n}$ denote the Router output for a sequence of length $l$, and $\boldsymbol{s}_i\in\mathbb{R}^n$ denote the Router output for the $i$-th token. The first approach proposed in that article is to subtract from $\boldsymbol{s}$ its exponential moving average (EMA):

\begin{equation}\hat{\boldsymbol{s}}_i = \boldsymbol{s}_i - \bar{\boldsymbol{s}}_i,\qquad \bar{\boldsymbol{s}}_i = \gamma\bar{\boldsymbol{s}}_{i-1} + (1 - \gamma)\boldsymbol{s}_i\end{equation}

We then use $\hat{\boldsymbol{s}}$ to decide which experts to activate (e.g., choosing the top-$k$, possibly combined with techniques such as QB). The idea behind this is quite simple: intuitively, if using $\boldsymbol{s}$ for the decision causes Expert $j$ to be activated more often, that means its score $\boldsymbol{s}_{:,j}$ tends to be larger on average, so to encourage balance we should penalize it accordingly — and the penalty amount is exactly the local average estimated by the EMA.

In this way, in the new $\hat{\boldsymbol{s}}$, each expert's score locally has approximately zero mean, preventing any single expert from becoming too "prominent," thereby achieving sequence-level balance. However, this is merely an empirical trick, with no theoretical guarantee that it will actually improve balance. In my own simple tests, in extremely imbalanced scenarios such as the first MoE layer, it didn't provide substantial help. Perhaps the author was also aware of this limitation, which is why a second approach was proposed.

Test-Time Training

The second approach is actually fairly intuitive too. In MoE Tour: 3. A Different Way to Allocate, didn't we introduce a bias term to control global balance, updated via SignSGD? Given that, we can borrow the idea of TTT — "test-time training" — and activate experts token-by-token along the sequence dimension, updating this bias term based on the distribution of experts already activated so far.

Taking a $\boldsymbol{s}\in\mathbb{R}^{l\times n}$, Top-$k$ MoE as an example: in the original Loss-Free scheme, we introduce a global bias vector $\boldsymbol{\beta}\in\mathbb{R}^n$, and change the activation mechanism to selecting Top-$k$ per $\boldsymbol{s}_i - \boldsymbol{\beta}$. Here, the idea is instead to assign a $\boldsymbol{\beta}_i\in\mathbb{R}^n$ to every $\boldsymbol{s}_i$, and change the activation mechanism to Top-$k$ per $\boldsymbol{s}_i - \boldsymbol{\beta}_i$. The update rule for $\boldsymbol{\beta}_i$ can be a sign-gradient ascent:

\begin{equation}\boldsymbol{\beta}_i = \boldsymbol{\beta}_{i-1} + \eta\mathop{\text{sign}}(\boldsymbol{f}_{\leq i-1} - 1/n)\end{equation}

Here $\boldsymbol{f}_{\leq i-1}$ is the distribution of already-activated experts up to the $i-1$-th token. In the author's final version, some adjustments were made: the sign function was dropped, and $\boldsymbol{f}_{\leq i-1}$ was changed to $\boldsymbol{f}_{i-1}$ (the activation distribution of the $i-1$-th token), in order to better achieve local balance:

\begin{equation}\boldsymbol{\beta}_i = \boldsymbol{\beta}_{i-1} + \eta(\boldsymbol{f}_{i-1} - 1/n)\end{equation}

This approach preserves the causal structure of the sequence and can also guarantee near-perfect sequence-level balance. However, since it needs to perform a Top-$k$ selection at every step and update the bias term based on the result $\boldsymbol{f}$, this essentially introduces a nonlinear RNN. Nonlinear RNNs cannot be parallelized, so it's foreseeable that this will become a bottleneck for long sequences. Of course, one could consider updating by chunks (mini-batch TTT) to improve efficiency, but that feels a bit less elegant.

The Optimal Solution

Next, the approach proposed in this post is called "Moving Quantile Balancing (MQB)," which, as the name suggests, is an evolution built on top of "Quantile Balancing (QB)." So let's first briefly recap QB.

In the order I originally wrote them, MoE Tour: 6. Optimal Allocation for Balance came before MoE Tour: 7. A Minimalist Solution for Dynamic Activation, but in fact the latter is much simpler both conceptually and methodologically, so let's start from there. Again, given $\boldsymbol{s}\in\mathbb{R}^{l\times n}$, we introduce a bias term $\boldsymbol{\beta}\in\mathbb{R}^n$. The static version of MoE activates the top $k$ experts by $\boldsymbol{s}_i - \boldsymbol{\beta}$ for each token, while dynamic activation means activating every expert with $\boldsymbol{s}_i - \boldsymbol{\beta} > 0$, with the count varying.

To simultaneously ensure load balancing and control the average number of activations to be $k$, we need to choose an appropriate $\boldsymbol{\beta}$. Remarkably, in this setting, the optimal solution for $\boldsymbol{\beta}$ can be expressed exactly:

\begin{equation}\boldsymbol{\beta} = \mathop{\text{desc_sort}}(\boldsymbol{s}, \text{axis=0})_{[lk/n:lk/n+1]} = \mathop{\text{quantile}}(\boldsymbol{s}, 1-k/n, \text{axis=0})\end{equation}

That is, if we sort $\boldsymbol{s}$ along the sequence dimension (currently $\text{axis=0}$) in descending order and take the $lk/n$-th largest element, that is the optimal solution for $\boldsymbol{\beta}$ — corresponding to the "$1-k/n$-quantile."

In fact, this is nothing more than another way of expressing Expert Choice. As we know, the problem with Expert Choice is that it breaks causality (it's non-causal), which is also evident from the fact that it computes $\boldsymbol{\beta}$ along the sequence dimension. If we're in an encoder scenario, such a non-causal scheme is naturally not a problem, but for a unidirectional autoregressive model, it's not acceptable.

Sliding Quantiles

The "test-time training" approach from @ChangJonathanC discussed earlier gave me some inspiration. We can compute the bias term token-by-token along the sequence dimension based on QB. Since QB can express the optimal solution for the bias term directly based on $\boldsymbol{s}$, without needing to know the actual expert activation outcomes, this opens up some possibilities for parallelization.

My initial idea was $\boldsymbol{\beta}_i = \mathop{\text{quantile}}(\boldsymbol{s}_{[:i]}, 1-k/n, \text{axis=0})$ — that is, computing the bias term over all token scores up to and including the current position, so that each position gets its own bias term, and the optimality of QB guarantees balance up to the current position. Since this requires computing quantiles progressively along the sequence dimension, let's call this operation "Cumulative Quantile," and the corresponding scheme "Cumulative Quantile Balancing (CQB)."

CQB is feasible in principle, but since quantiles cannot be updated incrementally, CQB has to read the full data and recompute the quantile at every step, so the per-step computational cost grows linearly, giving a total complexity that's quadratic. To alleviate this, and to better reflect local balance, my idea is to set a window $w$, so that each token only computes the optimal bias term within a window not exceeding $w$, i.e.

\begin{equation}\boldsymbol{\beta}_i = \mathop{\text{quantile}}(\boldsymbol{s}_{[i-w:i]}, 1-k/n, \text{axis=0})\end{equation}

This changes the operation from "Cumulative Quantile" to "Moving Quantile," and the corresponding scheme becomes the prototype of "Moving Quantile Balancing (MQB)." It's quite similar to SWA (Sliding Window Attention) — the complexity is linear, and in principle it can be parallelized. The reason it's still just a "prototype" is that quantiles still cannot be updated incrementally, so the per-step complexity remains relatively high, and further improvement is still needed.

Bucketed Estimation

The various obstacles described in the previous section are, fundamentally, all limited by the fact that "quantiles are nonlinear and cannot be updated incrementally." Is there any way we could linearize this — even approximately? Fortunately, it turns out we can!

Illustration of histogram-based approximate quantile estimation

A key fact is: as long as we know the distribution of a variable, we can compute quantiles via cumulative probability — and a distribution can be updated incrementally! So the key here is to convert the quantile estimation problem into a distribution estimation problem, which we do via bucket counting, also known as "histogram approximation." First, assume the Router scores $\boldsymbol{s}$ all lie within $[0, 1]$, which can be ensured by adding an activation function such as Sigmoid; then, divide $[0, 1]$ evenly into $b$ buckets, discretizing each value of $\boldsymbol{s}$ into an integer, and then converting it into the corresponding one-hot vector.

This gives us a 0/1 array of shape $l\times n\times b$. If we average this along the sequence dimension $l$, we get the score distribution for each expert (represented as a $b$-dimensional vector), from which we can compute each expert's quantile. However, to preserve causality, we can't simply average over the entire sequence — we can only do something like a "cumulative average." To avoid the trouble caused by the strict boundary of "Moving Quantile," and to make the whole process "smoother," here we choose to use an EMA instead.

That is, after bucketing $\boldsymbol{s}$ into one-hot form, we compute an EMA along the sequence dimension, so that each token gets a local distribution, from which we can compute the corresponding $1-k/n$-quantile — this is the key variable in MQB, $\boldsymbol{\beta}_i$. Since EMA is a linear, parallelizable operation, we've achieved sequence-level balancing in a relatively efficient, Loss-Free manner. This is the final version of MQB.

$$\begin{array}{|l|} \hline \text{Moving Quantile Balancing (MQB)} \\[4pt] \hline \text{input: score matrix}\boldsymbol{s}\in [0,1]^{l\times n} \\ \text{output: corrected score matrix}\hat{\boldsymbol{s}}\in \mathbb{R}^{l\times n} \\[4pt] \hline \begin{array}{ll} 1: & \boldsymbol{h}_{i,j} = \mathop{\text{one_hot}}(\lfloor s_{i,j} \times b\rfloor)\in\{0,1\}^b \\ 2: & \bar{\boldsymbol{h}}_{i,j} = \gamma\bar{\boldsymbol{h}}_{i-1,j} + (1-\gamma)\boldsymbol{h}_{i,j} \\ 3: & m_{i,j}^* = \min\{m \,|\, \sum_{t=0}^m \bar{h}_{i,j,t} \geq 1 - \frac{k}{n}\} \\ 4: & \beta_{i,j} =(m_{i,j}^* + 1/2) / b \\ 5: & \hat{\boldsymbol{s}}_i = \boldsymbol{s}_i - \boldsymbol{\beta}_i, \end{array} \\ \hline \end{array}$$

The General Case

So far, our discussion has been based on the dynamic-activation version of MoE, i.e., activating every expert with $\boldsymbol{s}_i - \boldsymbol{\beta}_i > 0$, where the count is variable. From MoE Tour: 6. Optimal Allocation for Balance, we know that for a Top-$k$ version of MoE, its $\boldsymbol{\beta}$ doesn't have a simple closed-form solution and requires alternating iteration.

So how should Top-$k$ MoE implement sequence-level load balancing in a Loss-Free way? In fact, the fact that MQB can achieve sequence-level balance under dynamic activation shows that it has already "flattened" the local abnormal spikes in the router. At this point, taking the Top-$k$ of $\boldsymbol{s}_i - \boldsymbol{\beta}_i$ may still be imbalanced, but this imbalance can only be global in nature at most — so we just need to add one additional step of global balancing via QB on top of it, to restore balance.

That is to say, for the Top-$k$ version of MoE, our forced sequence-level balancing scheme is MQB+QB: on top of the $\boldsymbol{s}_i - \boldsymbol{\beta}_i$-processed sequence, we apply QB once more (SignSGD would also work), achieving sequence-level load balancing. But note that perfect sequence-level load balancing often noticeably hurts performance; in most cases global balancing is already sufficient, and sequence-level balancing is only meant to prevent extreme cases — so it shouldn't be applied too aggressively.

To reduce the impact on the main task, whereas the Aux Loss scheme can lower its loss coefficient, MQB can instead consider taking a weighted average between the current and preceding scores:

\begin{equation}\lambda(\boldsymbol{s}_i - \boldsymbol{\beta}_i) + (1-\lambda)\boldsymbol{s}_i = \boldsymbol{s}_i - \lambda \boldsymbol{\beta}_i,\qquad \lambda\in[0, 1]\end{equation}

That is, by introducing $\lambda < 1$ to weaken the degree of sequence balancing, and then layering QB on top to ensure global balance (since $\lambda < 1$ alone cannot guarantee balance, in this case both the dynamic version and the Top-$k$ version need the additional QB step).

Other Details

After bucketing, the core operation of MQB also reduces to an EMA — so, after all this back and forth, we end up back at EMA again, in a form similar to the first scheme proposed by @ChangJonathanC, the difference being what the EMA is applied to. The scheme of @ChangJonathanC applies EMA directly to the raw score of size $l\times n$, whereas MQB first expands it to size $l\times n\times b$ before applying EMA.

This change is somewhat like how linear attention expands capacity via an outer product to remember more information, and the optimality of QB also guarantees that it achieves sequence-level balance, rather than remaining a mere empirical trick. Another similarity to linear attention: EMA is itself fundamentally an RNN, so at inference time we need to additionally store a state vector of size $n\times b$. In practice we've found that $b=100$ already gives good results, so a state of this size should be quite acceptable.

It's also worth pointing out that, whether QB or MQB, they are only used for the expert activation decision; the final gating value multiplied onto the experts is still computed from the original router scores, before subtracting the bias term. This is a feature common to all Loss-Free schemes — ensuring that the balancing intervention term does not directly alter the form of the MoE itself, so as to minimize its impact on the main model.

Finally, regarding the bucketing scheme: my current experiments use a Sigmoid activation followed by uniform discretization at $[0,1]$, and there's likely still room for experimentation on finer-grained bucketing strategies — I'll leave that to interested readers.

Experimental Results

For MQB, I ran some simple experiments, using a model with roughly 3B total parameters and a 128-choose-4 MoE configuration. MQB's moving-average coefficient was 0.99, the default number of buckets was 100, and MaxVio was used as the load-balancing metric. Since the first MoE layer is the most imbalanced, all results below show only the first layer's behavior.

First, the full-strength MQB with $\lambda=1$:

Comparison of MQB(λ=1) versus QB and SignSGD

We can see that $\lambda=1$ achieves near-perfect load balancing, at the cost of a Loss increase of 0.06 — a very noticeable drop, indicating that overly aggressive sequence-level balancing hurts performance. But if we tune $\lambda$ down to around 0.3, the Loss barely drops at all, while still improving load balancing, as shown below:

Load balancing under different values of λ

Further Thoughts

From the experiments above we can further confirm that perfect sequence-level balancing ($\lambda=1$) noticeably hurts performance, whereas taking $\lambda=0.3$ roughly preserves performance while still improving the balance at each layer.

Overall, this post only answers the question of "how to implement sequence-level load balancing in a Loss-Free way." Whether sequence-level load balancing is actually needed, why it might be needed, and to what degree it's needed — these are questions we can't give a definitive answer to. One tentative understanding is that extreme sequence-level imbalance might harm the quality of the sequence itself (collapsing into effectively a smaller model), so it's desirable to encourage some degree of sequence-level balance.

Besides MQB, there are actually other schemes that can also achieve sequence-level balance in a Loss-Free way — for example, the Hash Routing approach we introduced in How Did DeepSeek V4's tid2eid Come About?, which can also be seen as a simple scheme for forcing sequence-level balance. But Hash Routing redefines the routing mechanism entirely, making it a rather unconventional approach.

The value of an approach like MQB lies in the fact that it remains more compatible with conventional MoE formulations, and it converts "tuning the Aux Loss coefficient" into "tuning an interpretable local bias term," giving us more intuitive control over the strength of the intervention, while its parallelizability also ensures efficiency. Since it only affects the router's decision without injecting gradients, its impact on the main task's performance is kept to a minimum as well.

Summary

This post has focused on "how to implement sequence-level load balancing in a Loss-Free way." Starting from the original Quantile Balancing (QB), we gradually derived a method called Moving Quantile Balancing (MQB) that successfully achieves this goal. However, whether sequence-level balancing is truly necessary, and to what degree it should be applied, remains an open question.