Load Balancing Is All You Need

I have been on Zhihu for a few years, but this is the first time I am actually promoting my own paper.

This is my second first-author paper. The first one was about Speculative Decoding; I modified the algorithm for choosing among multiple drafts, but the work was honestly too niche. In real applications, nobody is really going to run multi-draft on top of batched inference and pay several times more parallel compute. Once that disappears, the algorithm basically loses its point, so there was not much reason to talk about it.

In 2025 I tried to train a Generative Reward Model with RL. The idea was simple: if the evaluation generated by the model agrees with humans, give a meta-reward of 1; otherwise give 0. Then, because I was too lazy and also got sick for a month, by the time I recovered DeepSeek had released a paper doing almost exactly the same thing, and of course my own version could not compare in quality. So I tried a few rescue attempts and forcibly injected some novelty, but the new RL algorithm kept collapsing during training; to this day I still do not know whether it was VeRL’s problem or mine, so I had to give up.

After failing this many times, I had no choice but to settle down. I realized that although I had many high-level ideas, my machine learning foundations were never solid, and I had not built good programming habits either. Fortunately, I interned at a company over the summer and had access to GPUs. I did not work very seriously on my boss’s tasks, but I did use the chance to really patch up my foundations. After returning to school, I also changed my old loose habits and slowly learned architecture, torch, communication, and optimization, and finally produced a paper that I personally think is not bad.

I had wanted to work on MoE for a long time. The first time the idea occurred to me was in 2024, when I saw Mixture of Depths (MoD), a Google paper about training a router that could skip both the MLP and Attention in a transformer, enabling dynamic compute allocation. Only later did I realize that although MoD sells dynamic compute allocation as its main point, that part is entirely due to Expert Choice routing. Its only difference is that it also tries to skip Attention, but that part does not look quite right to me, because skipping memory writes, namely Key and Value computation, and skipping memory reads, namely Query computation, are two different things and may not be implementable with the same router.

Let me first roughly explain the MoE architecture and routing design.

Simply put, MoE replaces the original fixed MLP in the model with many experts. Every time a token comes in, the router matches this token against all experts and decides which experts should process it. That is routing.

The mainstream methods on the market are Token Choice routing, where each token picks the experts with the highest matching scores. But if all tokens like the same few popular experts, the unpopular experts have nothing to do. To stop everyone from piling onto the same experts, people have to force all kinds of restrictions and penalties onto the model, pushing it toward load balancing. This makes the whole thing complicated.

The second approach is called Expert Choice routing. The idea is reversed: each expert actively chooses the highest-scoring tokens from the current large batch of tokens. Once allocated this way, every expert’s workload is kept perfectly even. So it naturally achieves perfect load balancing. At the same time, since any number of experts can choose the same token, the model can dynamically allocate compute according to token properties.

With these two advantages, you should understand why I am so obsessed with Expert Choice routing.

But in order to choose tokens, each expert has to wait until the whole batch of tokens is available before making a decision. If we are doing text generation, where tokens come out one by one, the model cannot obtain the routing scores of “future” tokens in advance for comparison, so Expert Choice cannot be used directly for large language models. At the same time, the model can also use routing as a channel for future information leakage, leaking the routing scores of future tokens to current tokens and cheating, which causes training to collapse. This is not what I want to see.

But can EC really not be made causal?

Mixture of Experts

For the following discussion, let us first define the notation and algorithms for Mixture of Experts.

In a Transformer, an MoE layer usually replaces the standard dense feed-forward network (FFN) with a router and a set of experts. Here, G denotes granularity, meaning the number of experts selected by each token. E denotes the expansion rate, meaning the ratio between the total MLP parameters and the activated parameters. Suppose we have a batch containing B tokens, and the feature representation of each token is x_i. The router first computes a score matrix S.

Based on these scores, the routing rule produces a binary assignment matrix A, where A_ij = 1 means the i-th token activates the j-th expert, and otherwise it is 0. Each selected expert computes its own output and multiplies it by a gating weight. So for token i, the output of the whole MoE layer is the weighted sum of all activated expert outputs: [formula omitted in source draft].

The routing rule that determines the assignment matrix is crucial. It controls both how compute is allocated and how the load is balanced among experts. We can formalize MoE routing as a constrained optimization problem: under a limited compute budget, find an assignment matrix that maximizes the total routing score. The higher the score, the better the match between token and expert; and through the gate, that expert contributes more to the final output.

Token Choice

The optimization objective of standard Token Choice (TC) routing can be written as:

[formula omitted in source draft]

Here, the “sparsity constraint” guarantees that each token must select exactly G experts, while the “load balancing constraint” guarantees that each expert must process exactly its assigned number of tokens.

The problem is that exactly solving this optimization problem with two hard constraints requires a combinatorial algorithm such as Hungarian matching, with complexity as high as [formula omitted in source draft]. This is obviously impossible in practical training. So most Token Choice methods, such as Switch Transformer, compromise. They strictly satisfy the sparsity constraint by directly letting each token take the Top-G experts, and then introduce an auxiliary loss or some loss-free load balancing strategy to approximate the load balancing constraint as much as possible.

Expert Choice

Although the load balancing constraint is crucial for preventing routing collapse, where all tokens flood into a tiny number of experts, the “sparsity constraint” does not really bring any substantive benefit. So Expert Choice (EC) routing simply upends the setup: it completely removes the sparsity constraint and only enforces load balancing inside the batch. At this point, the original optimization problem simplifies to:

[formula omitted in source draft]

This relaxed optimization problem has an extremely simple closed-form solution: each expert directly selects the Top-K tokens from the current batch, where K is that expert’s capacity.

This design brings two huge benefits.

First, perfect load balancing: by construction, each expert is forced to process exactly K tokens.

Second, dynamic compute allocation: some tokens may not be chosen by any expert, so their computation is skipped; some tokens may be selected by multiple experts at the same time. This realizes adaptive compute allocation based on token importance.

However, in order to satisfy load balancing within each batch, EC routing introduces a fatal causality problem. Look closely: since experts choose tokens, whether a token is selected, namely the value of A_ij, depends not only on its own score but also on the scores of all other tokens in the current batch. Naturally, this includes “future” tokens that do not even exist yet during autoregressive generation. Although some later work, for example extending it to batch-level top-k, can alleviate this problem to some extent, as long as the routing decision still depends on the composition of the current batch, the causality problem cannot be completely removed.

Expert Threshold

Regularization is a common technique in machine learning for preventing collapse. But what regularization do we actually need? Which constraints are useless?

Back to the discussion above. First, the sparsity constraint in Token Choice is meaningless. We do not need to hard-code how much compute each token must use. It does not help load balancing, and it restricts the model’s flexibility.

The load balancing constraint in Expert Choice is meaningful. Without it, routing collapses, meaning all tokens rush toward only a few experts. But is it too strict to require every batch to be perfectly load balanced? After all, different batches have different data distributions. Forcing every batch to be load balanced may instead make the decision boundary unstable.

Therefore, we propose Expert Threshold (ET) routing. We neither require each token to activate a fixed number of experts, nor require each batch to be strictly load balanced. We only require that, in expectation, every expert processes the same proportion of tokens.

In other words, compared with EC, we no longer select each expert’s Top-K tokens within each batch. Instead, we select each expert’s top fraction of tokens over the entire data distribution. This means we no longer need to depend on the current batch composition, which solves the causality problem.

The concrete method is that for each batch’s cutoff, namely the K-th largest routing score for an expert, we compute an exponential moving average (EMA). This gives us each expert’s threshold, which estimates that expert’s global top-fraction cutoff over the full data distribution. During both training and inference, we only need to use this threshold to decide whether a token should be activated. If its score is greater than the threshold, it belongs to that top fraction and is activated; otherwise, it is not activated.

In this way, we solve the causality problem while still guaranteeing load balancing. We also no longer depend on the composition of the current batch, which is precisely what removes the causality problem.

The Connection Between ET and EC

Conceptually, ET can be understood as doing Expert Choice on an infinitely large batch.

EC lets each expert choose the Top-K tokens from the current batch of B tokens. When the batch is small, swapping in just one token may change the position of the cutoff. But when the batch size tends toward infinity, the cutoff converges to a fixed quantile of the routing score distribution. Then each token’s routing decision only depends on its own score and no longer depends on who else is in the batch.

What ET does is directly approximate this limit. We do not wait for the batch to become infinitely large; instead, we use EMA to estimate that global quantile threshold, and then use it for routing.

From another angle, ET and EC make different tradeoffs between “stable thresholds” and “stable load.” EC lets the threshold change with the batch to guarantee perfectly balanced load in every batch, at the cost of routing decisions fluctuating with batch composition. ET does the reverse: it fixes the threshold to keep routing decisions stable, at the cost of a little fluctuation in expert utilization within each batch.

Interestingly, precisely because of this equivalence, ET threshold routing can be used directly as a causal inference scheme for models trained with EC. As long as we record each expert’s cutoff EMA as the threshold after EC training, we can switch seamlessly to ET at inference time without any retraining.

Actually, my initial idea came from this direction: find a way to make EC’s batch size as large as possible. At the time I was reading the DiffMoE paper on a plane, and they emphasized the importance of batch size. Many related experiments also show that the larger the batch size, the better the model performance. So I simply pushed it all the way to infinity, and that solved the causality problem.

Experimental Results

The experiments were done on Nanochat, Karpathy’s open-source GPT-2 training code, at two scales. The small d12 model has 575M parameters, with 195M active. The large d20 model has 2.4B parameters, with 561M active. Each MoE layer used 16 routed experts plus 1 shared expert, with expansion rate 16, meaning that on average each token activates only 1 routed expert plus the shared expert. The training data came from FineWeb-Edu 100B. The models were trained for 10B and 11.2B tokens respectively.

I compared three broad categories of routing methods. All models had exactly the same architecture and parameter count. Only the routing rule differed.

Main conclusions:

ET consistently outperforms Token Choice at both scales. On d12, CE loss is lower by 0.05 and the CORE benchmark is higher by 1.89. On d20, the gap is larger: CE loss is lower by 0.067 and CORE is higher by 2.83.

When the batch is large enough, 512k tokens, EC’s training loss is basically tied with ET. This validates the theory above. Explicit large-batch selection and EMA threshold estimation reach the same destination by different paths.

At the d20 scale, ET slightly exceeds the best EC: CE loss 2.620 vs. 2.621, CORE 25.14 vs. 24.98.

One detail worth noting is that although EC training is not causal, during inference we uniformly use ET’s cutoff EMA as the threshold for causal inference. So all EC and ET CORE scores above are evaluated under causal conditions. There is no information leakage.

There is another interesting phenomenon. EC performs very poorly with a small batch, 2k tokens, and is even worse than Token Choice without any load balancing. But as the batch size increases, its performance steadily improves. By 512k, it is about the same as ET. This curve exactly confirms what we said earlier: ET is the infinite-batch limit of EC.

When using ET for inference, sequence EC, with a batch size of 2k tokens, has severe train-inference mismatch.

For more experimental results, please read my paper.

I also have lots of goodies here, namely interpretability plots.

Conclusion

The core observation of this paper is actually very simple. The reason Expert Choice routing is not causal is that it performs top-k selection inside each batch, so the routing decision depends on other tokens in the same batch. But if the batch becomes infinitely large, this dependency disappears. Each token’s routing decision depends only on itself and a global threshold. ET uses EMA to approximate this limit.

At the end of the day, we do not need every batch to be strictly load balanced. It is enough to be balanced in expectation over the whole data distribution. Once this one constraint is relaxed, the causality problem solves itself naturally, and performance does not drop. It even improves.

I hope this work makes people re-examine Expert Choice routing. It has long been considered incompatible with autoregressive language models, but actually it was only one step away.

Postscript

After submitting the paper at the end of January, I wanted to polish it a bit more, add some visualization and interpretability experiments, and adjust the story. So the preprint kept getting delayed until now. I thought that since EC was proposed in 2022 and nobody had really followed up for more than three years, probably nobody would race me on this. Unexpectedly, in the last two weeks, Su Jianlin of kexue.fm came up with almost the same idea. Of course the perspective is different, but the algorithmic core is the same. For people who are good at math, the same algorithm can be told through five or six different stories, but they are actually equivalent. I will cite Su Jianlin again here for reference and comparison.

Appendix

How much information does an Expert Choice model leak, exactly?

DeepSeek previously gave a trivial conclusion: at most, it leaks the information contained in the routing combination of choosing k from N.

Here I give a constructive communication method showing that, assuming infinite precision, this upper bound is achievable. In other words, even if the batch size is pushed to infinity, there can still be leakage.

Fortunately, the real world uses finite precision, and models are not that clever.