Token Balancing

Mixture of Experts (MoEs) are very popular. They've been popular for a long time, but I'm currently catching up, so there's a lot of new things here for me.

I think we generally all understand why imbalanced experts are bad; both from a learning perspective (experts don't improve if they don't get enough gradients) and a performance perspective (unbalanced routing means many GPUs idle).

Recently I was reading about SonicMoEs, which is focused on improving arithmetic intensity and activation memory usage as the MoE becomes more finegrained. Throughout that the paper talks a lot about the expert routing, but doesn't focus on the routing problem - there's an S: (TK,) which is just score of the $t$th token with the $k$th expert (assuming we've alread filtered the list down to topk) and the corresponding indices (indices).

It turns out we can compute those scores in a few different ways, but the routing is pretty much all the same - the auxiliary-loss-free load balancing introduced in this paper by the DeepSeek folks (i.e. a form of "smart routing"). We use the token activation itself to route it to the right expert:

# assume you have `hidden` which is (B x T, H) - batch x tokens, hidden_size
# score_function can be F.softplus(x).sqrt(), or F.sigmoid
s = score_function(hidden @ router_weight) # (B x T, E)

# compute assignment
topk_idx = (s + expert_bias).topk(k, dim=-1).indices # (B x T, k)

# compute weights
weights = s.gather(-1, topk_idx) # bias excluded, (B x T, k) 
weights = weights / weights.sum(-1, keepdim=True)

What do major open source LLMs use for their score function? DeepSeek V4's is F.softplus(x).sqrt(). Inkling (and DSV3) uses x.sigmoid() (with shared experts separated out, which is common, but then reweighted together with the chosen experts).

Stabilizing training

One trick here is how to determine expert_bias, which is just an offset per expert (i.e. (E,)-shaped), as we're training. Let's simplify down to just maintaining balance on a batch-by-batch basis, rather than sequence-level balancing, which seems quite a bit harder.

The DeepSeek V3 method (which is also used in V4) is to actually use a hash-based routing metod for the first 3 MoE layers, which skips "smart routing" - each token is essentially quasirandomly assigned to a different expert, which naturally balances the routing.

The DeepSeek V3 method (which is also used in V4) is to have expert bias scheme above, and update it with flat steps: if an expert is getting more than the correct average number of tokens in one batch, in the next batch its expert bias is reduced by $\eta = 0.001$ (similarly increased if it's not getting enough).

This is apparently pretty difficult to tune, and it can suffer from slow convergence (here, slow convergence to equal token counts across experts). It's also common to change the routing scheme as we get further down into layers; with DeepSeek V3 the first few layers are dense FFNs rather than MoEs. In DeepSeek V4, they added hash-based routing (on token ID) for the first 3 MoE layers, i.e. routing that is not token-value-dependent to force the tokens to balance across experts.

The new thing lately in the past year has been quantile routing, introduced by Jianlin Su's blog post and later explored by the Marin team. This is pretty interesting - it basically says that if you look at a given expert's allocation in a particular batch, and look at the topk cutoff routing score for each token (i.e. looking at s + expert_bias, and picking the $k/E$th quantile), and say "hey, what would have been a good expert score to get me to the right number of tokens?" (i.e. $Tk/E$), each expert can kind of decide what to set their expert bias to be in the next step. This is parameter-free! And hopefully, the routing stabilizes, and we get to a fixed point, where the expert bias isn't updating every batch (or we can update it after accumulating $N$ batches of statistics, etc.).

There are some really nice plots in the Marin blog post.