Expert Threshold Routing: A Causal Reinterpretation of Expert Choice for Autoregressive MoE Language Models

Executive Summary

TL;DR

This paper targets a very specific but important flaw in sparse MoE routing for autoregressive language models: the best load-balanced routing rule, Expert Choice (EC), is not causal. The authors propose Expert Threshold (ET) routing, where each expert maintains an EMA-tracked score cutoff approximating a global quantile, and tokens are routed independently by thresholding against that cutoff. This preserves dynamic computation allocation, achieves load balance in expectation without auxiliary losses, and keeps training and inference routing rules identical. Empirically, ET beats Token Choice (TC) by 0.067 CE loss at 2.4B parameters, which the paper interprets as needing 1.6x fewer tokens to reach the same performance.

Background Positioning

This is not merely another MoE benchmark tweak. Its real contribution is to identify that EC’s benefit comes from approximating a population-level routing threshold, then to replace the non-causal batch top-k operation with a causal online estimator of that threshold. In that sense, the paper is a strong methodological repair of EC for autoregressive LLMs, with some genuine SOTA-style gains over TC in the authors’ training setup.

Problem & Motivation

Sparse MoE works because it decouples parameter count from active FLOPs. But routing is the entire game: if the router collapses onto a few experts, the rest of the model becomes dead weight.

Why Token Choice is not enough

The dominant MoE recipe in LLMs is Token Choice:

• each token picks exactly top-G experts;
• this enforces a fixed per-token compute budget;
• load balancing is then patched with auxiliary losses or controller-like heuristics.

The paper’s critique is sharp: this mixes two objectives that are not naturally aligned.

• Sparsity constraint: each token must use exactly G experts.
• Load balancing constraint: each expert should see roughly the same number of tokens.

These two constraints make the ideal routing problem combinatorial. In practice, TC keeps the sparsity constraint exact and only approximates the balancing constraint, which is why load imbalance remains a persistent systems and optimization issue.

Why Expert Choice is attractive but flawed

Expert Choice flips the problem:

• each expert picks its top-k tokens;
• load balancing becomes exact by construction;
• per-token compute becomes dynamic, since a token may be selected by 0, 1, or many experts.

This is conceptually elegant. But for autoregressive generation it is fundamentally problematic, because whether token t gets selected depends on the scores of future tokens in the same batch. That creates:

• information leakage during training,
• non-causal dependence during inference,
• and batch-composition dependence, which is operationally awkward and sometimes undesirable.

So the central question becomes:

Can we keep EC’s dynamic allocation and balancing behavior, but eliminate its dependence on future tokens?

That is the core motivation for ET.

Methodology - The Core

From batch top-k to population quantiles

The key move is to stop insisting on exact per-batch load balance. The authors argue that for avoiding routing collapse, what we really need is:

• not exact equality in each batch,
• but balanced activation in expectation over the data distribution.

This reframes the routing problem.

Under EC, expert i selects the top k = N/E tokens in a batch. Equivalently, it selects tokens whose scores exceed that batch’s kth-largest score. That kth-largest score is just a batch-specific quantile threshold.

ET asks: why not estimate the corresponding global quantile threshold over the population and use that directly?

Then routing becomes:

• compute router logit r_{t,i},
• compare it with expert threshold c_i,
• route if r_{t,i} > c_i.

Formally, the routing rule is:

z_{t,i} = 1{ r_{t,i} > c_i }

where c_i is an EMA estimate of the expert’s top-1/E population cutoff.

This looks simple, but it changes the causal structure completely:

• EC: routing decision depends on other tokens in the current batch.
• ET: routing decision depends only on the current token score plus a slowly updated global statistic.

That is the whole trick.

Intuition: ET as the infinite-batch limit of EC

This is the paper’s deepest conceptual insight.

If EC is run on larger and larger batches, each single token has less influence on the batch cutoff. In the limit, the cutoff becomes effectively independent of any particular token and converges to a population statistic. ET simply approximates this limit by keeping an EMA of the cutoff.

So ET is not an arbitrary new heuristic. It is a causalized approximation to large-batch EC.

This perspective also explains the empirical findings:

• small-batch EC is noisy;
• large-batch EC improves;
• ET matches large-batch EC without requiring large-batch coordination at inference.

Formal view: what constraint is being relaxed?

The paper’s optimization framing is useful.

• TC: exact per-token sparsity + exact per-expert balance is combinatorial, so practice keeps top-G and approximates balance.
• EC: drops fixed per-token sparsity, keeps exact per-batch expert balance.
• ET: drops exact per-batch expert balance too, and only requires balance in expectation.

That sounds like a relaxation, but it is the right relaxation for autoregressive LM, because it removes the dependence on future tokens.

The tradeoff is clean:

• EC gives exact per-batch utilization but unstable thresholds.
• ET gives stable thresholds but slightly fluctuating per-batch utilization.

The authors argue, plausibly, that for deployment and causal consistency, the latter is preferable.

Algorithmic details

For each expert i:

1. Route token t if r_{t,i} > c_i.
2. During training, update c_i using an EMA of the batch kth-largest router logit.
3. During inference, keep using the same thresholding rule.

This symmetry matters a lot. There is no special inference-time surrogate router. No extra predictor. No batch-level workaround.

Warmup is not optional

There is one practical complication: early in training the router score distribution is unstable, so the cutoff EMA is unreliable. If ET is used immediately, many tokens fail to cross threshold and experts starve.

The fix is straightforward:

• run standard EC for the first 4k steps,
• let the thresholds accumulate meaningful statistics,
• then switch to ET.

This is one of those details that is easy to dismiss but probably central in practice. It also reveals an important fact: ET depends on a reasonably calibrated score distribution. Before that emerges, thresholding is too brittle.

Experiments & Results

Main quantitative results

The paper evaluates on GPT-like autoregressive LMs trained on FineWeb-Edu, with two scales:

• d12: 575M total, 195M active params
• d20: 2.4B total, 561M active params

Configuration highlights:

• 16 routed experts
• G=1
• one always-on shared expert
• sigmoid gates instead of softmax gates

d20 results

The most important table is the 2.4B-scale result:

• Dense: 3.002 CE, 15.743 CORE
• TC: 2.893 CE, 17.983 CORE
• TC aux: 2.892 CE, 15.894 CORE
• TC loss-free: 2.898 CE, 18.031 CORE
• EC 512k: 2.843 CE, 19.94 CORE
• ET: 2.844 CE, 19.876 CORE

The headline is clear:

• ET improves over TC by 0.049 to 0.054 CE depending on TC variant, and the paper emphasizes a 0.067 gap in its final loss trajectory comparison.
• ET nearly matches the best large-batch EC result while being causal.

d12 results

At smaller scale:

• Dense: 2.751 CE, 20.43 CORE
• TC aux: 2.687 CE, 22.31 CORE
• EC 256k: 2.621 CE, 24.98 CORE
• ET: 2.620 CE, 25.14 CORE

Here ET slightly beats EC on CORE and matches it on CE.

Why the results are believable

The result pattern aligns well with the paper’s thesis.

1. EC improves with batch size

This is exactly what should happen if the real object is the population quantile and small-batch top-k is only a noisy estimator of it.

The reported EC trend on d20:

• 2k batch: 2.910 CE
• 8k batch: 2.845 CE
• 64k batch: 2.841 CE
• 512k batch: 2.843 CE

The performance saturates around large batches, which fits the “infinite-batch limit” interpretation nicely.

2. ET matches large-batch EC

This is the strongest evidence that the method is not just another balancing heuristic. If ET were only a controller trick, matching large-batch EC so closely would be less likely. Instead, the data suggests ET is approximating the same routing regime.

3. Small-batch EC has a train-eval mismatch

This is a very important operational point. If EC is trained with small routing batches and evaluated causally with thresholds, the train and inference routing distributions diverge substantially. ET eliminates this because its training rule is already causal.

Cutoff stability vs usage stability

One of the paper’s better analysis sections explains the core systems tradeoff.

• EC: fixed expert usage, varying cutoff.
• ET: fixed cutoff, varying expert usage.

This is subtle but crucial.

In EC, the top-k rule guarantees each expert gets exactly the target number of tokens in a batch, but the threshold itself moves around depending on batch composition. In ET, the threshold is stable because it is EMA-based, but expert usage fluctuates around the target.

For language modeling deployment, stable routing semantics are often more valuable than exact per-batch utilization, especially when inference is token-by-token.

Dynamic computation allocation

A major advantage of ET over TC is that it is not locked into fixed per-token sparsity. Some tokens can trigger:

• no routed experts,
• one routed expert,
• or multiple routed experts.

That means the model can spend compute unevenly across tokens.

The paper’s qualitative analyses suggest:

• structurally important tokens,
• numerals,
• operators,
• boundaries,
• and harder tokens

often receive more fanout.

This is the right inductive bias. In language modeling, not every token deserves the same amount of computation. Fixed top-1 or top-2 routing is computationally neat, but it imposes a uniform budget where the task is highly non-uniform.

The interesting nuance is that ET’s dynamic allocation is somewhat smoother than small-batch EC, whose token fanout can become overly spiky, especially at the first token. That suggests ET is not only causal but also more stable as a computation allocator.

Expert specialization

The expert heatmaps are also revealing. ET develops specialization patterns comparable to large-batch EC and sharper than small-batch EC.

That matters because specialization is often the hidden reason MoE helps at all. If routing is noisy or unstable, experts cannot consistently carve out subdomains.

The paper shows ET experts specializing across domains like:

• HumanEval for code,
• GSM8K for math.

This is a useful result because it indicates ET’s thresholding does not wash out expert identity. On the contrary, stable thresholds may actually support stronger specialization.

Ablation insights

Warmup matters

Without EC warmup, ET suffers from:

• unstable cutoff estimates,
• low raw expert usage,
• high starvation,
• noisier logits and gates.

This is not a cosmetic ablation; it is a structural one. It tells us ET is only as good as its threshold estimate.

Shared expert helps

The shared expert improves CE by roughly 0.02 in some settings. This is consistent with the now-common intuition that always-on dense pathways provide a stable base representation, while sparse experts add specialized computation.

Capacity mismatch appears manageable

Since ET allows per-batch expert usage to fluctuate, there is a legitimate concern about hardware overflow. The paper uses a capacity factor of 0.5, and reports that saturation and starvation become infrequent after warmup. So the train-inference mismatch introduced by capacity clipping seems limited in their runs.

Critical Analysis & Conclusion

What is the real contribution?

The real contribution is not just “threshold routing works.”

It is this stronger idea:

Expert Choice can be reinterpreted as routing by a population-level quantile, and once you see that, the non-causal top-k operator becomes unnecessary.

That is a clean conceptual simplification. It turns a batch-coupled routing mechanism into a token-local one without abandoning the optimization principle that made EC appealing.

Why this works

The method works because it aligns three things that are usually in tension:

• causality: decision depends only on current token and past statistics,
• load balancing: each expert fires with roughly target frequency over the population,
• dynamic compute: token fanout is not fixed in advance.

TC usually sacrifices dynamic compute. EC sacrifices causality. ET is compelling because it preserves all three approximately enough to matter.

Limitations

The paper is strong, but a few limitations are worth stating directly.

1. It is still an approximation to exact balancing

ET only balances in expectation, not per batch. In large-scale distributed systems, even small fluctuations can matter for kernel efficiency, memory planning, and all-to-all communication.

2. Warmup adds procedural complexity

The need for 4k-step EC warmup means ET is not fully standalone from step 0. A more principled cold-start mechanism would strengthen the method.

3. Evidence is pretraining-centric

The experiments are convincing for pretraining CE and CORE, but there is still limited evidence on:

• downstream instruction tuning,
• RLHF or post-training,
• longer-context generation,
• latency-sensitive serving.

4. Comparison scope is selective

The strongest comparisons are against TC and EC variants. The broader dynamic routing space now includes Top-P routing, ReLU-based differentiable routing, null experts, sequence-level allocation, and multimodal adaptations. ET’s relative position in that broader frontier remains open.

Future Work

Several directions look natural:

• Better cold-start threshold estimation, possibly with analytically initialized cutoffs or bias-based bootstrapping.
• Hybrid ET + controller methods, where thresholding provides causal semantics and lightweight corrections manage hardware variance.
• Extension beyond FFN experts, e.g. thresholded routing for attention heads, depth skipping, or multimodal expert modules.
• Serving-oriented studies, especially whether ET reduces inference-time coordination overhead enough to translate into real throughput gains.

Final Takeaway

This paper is one of those works whose simplicity is deceptive. The final routing rule is almost trivial, but the conceptual move behind it is not: replace batchwise top-k with a population-level threshold estimate, and you recover most of EC’s advantages while restoring causality. For autoregressive MoE language models, that is an important step toward routing mechanisms that are not only performant, but also logically aligned with how generation actually happens.