Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing
Authors: Hanchi Sun, Yixin Liu, Yonghui Wu, Lichao Sun
Organizations: Computer Science and Engineering, Lehigh University, Bethlehem, PA, USA · MD-HOBI-BIOMED INFORMATICS, University of Florida, Gainesville, FL, USA
Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6× fewer tokens.
Figures & tables
Figure 1: evaluation loss for Dense, TC, and ET. Compared to TC, ET achieves a 0.067 final loss gap (TC vs ET), or equivalently reaching same performance level with 1.6x few tokens.
Figure 2: Illustration of TC, EC, and ET routing mechanisms and their routing pools. Left: TC routes each token independently to its top- G experts, causing load imbalance. Middle: EC has each expert select its top- k tokens from the batch, requiring access to all tokens including future ones (non-causal). Right: ET routes each token independently by comparing its score against the population’s top- (1/E) quantile estimated by an EMA-tracked threshold ci , enabling fully causal routing over the population.
Method
Batch
CE loss ( ↓ )
CORE ( ↑ )
dense
—
3.002
15.743
TC
—
2.893
17.983
TC aux
64k
2.892
15.894
TC loss-free
512k
2.898
18.031
EC
2k
2.910
17.91
EC
8k
2.845
18.83
Table 1: Main results comparing Expert Choice (EC), Token Choice (TC), and Expert Threshold (ET) routing. Batch : token routing pool size. EC uses global selection batch, TC uses per-step batch, ET reports effective EMA pool size N/(1−β) . TC variants: no load-balancing, auxiliary loss ( α=0.001 ), or loss-free ( u=0.005 ). We report validation cross-entropy (CE) loss ( ↓ ) and CORE Eval score ( ↑ ).
Method
Batch
CE loss ( ↓ )
CORE ( ↑ )
dense
—
2.751
20.43
TC aux
32k
2.687
22.31
EC
256k
2.621
24.98
ET ( β=0.999 +warmup)
256k → 500M
2.620
25.14
Table 2: d20 results.
Figure 3: Cutoff stability vs expert usage tradeoff. Top Signed cutoff deviation relative to the EMA for EC at 512k batch size. ET stays at zero because routing uses the cutoff EMA directly. Bottom Expert usage for EC at 512k and ET. ET varies around the capacity target while EC remains constant.
Figure 4: Expert specialization analysis. (a) Token-level activation intensity on a GSM8K passage, colored by total fanout (sum of experts activated across layers). The model assigns more computation to structurally important tokens (punctuation, sentence boundaries, numerical results) than to common content words. (b) Expert token ratio heatmaps for HumanEval (code) and GSM8K (math). Top: EC (batch size 2k). Bottom: ET. ET achieves sharper patterns in expert activation, suggesting more domain-aware routing and specialization.
Figure 5: Activation dynamics for EC (2k) and ET. Both methods allocate more computation to early positions, with EC (2k) showing a sharper spike. Faint dashed curves show per-layer means and solid red curves show the global mean. When binned by loss, EC (2k) fanout rises monotonically while ET peaks early before declining.
Figure 6: EC performance across routing batch sizes. Training CE loss decreases and CORE Eval score increases with larger batches.
Figure 7: Train loss EMA and eval loss for EC at different batch sizes and ET. Solid lines show train loss EMA and dashed lines show eval loss. EC (2k) shows a large train-eval discrepancy, while EC (512k) and ET remain closely aligned.
Figure 8: Within-family checkpoint-pair routing consistency on a fixed validation stream, measured by weighted Jaccard. ET is consistently more stable than EC 2k and stays close to EC 64k. TC is competitive on nearby checkpoints but degrades more on longer ranges.
Micro Batch/Seq
Batch
Population
Aux loss
Global LBL
–
–
–
LossFree
Seq EC
Batch EC
ET (ours)
Table 3: Taxonomy of load balancing methods by scope (Aux loss ( Lepikhin et al., 2021 ) ; Global LBL ( Qiu et al., 2025 ) ; LossFree ( Wang et al., 2024 ) ; Seq EC ( Zhou et al., 2022 ) ; Batch EC ( Ludziejewski et al., 2024 ) ).
ET
Similar To
Connection
Cutoff-EMA ci
LossFree bias bi
Per-expert scalar; no aux loss
1−β
LossFree μ
Update rate
Table 4: Conceptual connections between ET and recent work (LossFree ( Wang et al., 2024 ) ; GShard ( Fedus et al., 2022 ) ).
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Illustration of the binary-search encoding strategy. The cutoff value c partitions the interval [0,1] into regions corresponding to different expert selection patterns, enabling N bits of future information to be encoded in a single real-valued threshold.
Feature
GPT-2
d12
d20
Tokenization
GPT-2
RustBPE
Vocab Size
50,257
65,536
Activation
GELU
ReLU 2 ( So et al., 2021 )
Normalization
LayerNorm
RMSNorm ( Zhang and Sennrich, 2019 )
FFN Dimension
4×dmodel
4×dmodel
Linear Layer Bias
Yes
No
Appendix
Table 5: Model architecture and size configurations. Architecture features are shared between d12 and d20 (nanochat-style). For MoE variants with G=1 , E=16 : 16 routed experts + 1 shared = 17 total experts. Total params include all expert parameters; active params include only the shared expert plus on average one routed expert per token.
Hyperparameter
Value
Total tokens
10B / 11.2B
Batch size (tokens)
524,288 (0.5M)
Sequence length
2048
Muon Warmup steps
No warmup
AdamW Warmup steps
250
Learning rate schedule
Linear decay
Appendix
Table 6: Training hyperparameters.
Parameter
Init.
LR
Opt.
WE (embed.)
N(0,1)
0.2λ
AdamW
Wlm (head)
0
0.004λ
AdamW
WQKV (attn)
Asp.-ratio
0.02λ
Muon
WO (attn proj)
0
0.02λ
Muon
Wrouter
N(0,d−21)
0.02λ
Muon
W↑(e) (exp. up)
Asp.-ratio
0.02λ
Muon
Appendix
Table 7: Parameter initialization and optimizer configuration. Aspect-ratio scaled init. uses std=din−1/2⋅min(1,dout/din) . LRs include μ P ( Yang et al., 2022 ) scaling λ=(dmodel/768)−1/2 .
Figure 10: Effect of TopK warmup on ET training dynamics (first 8k steps). Before 4k steps, ET no warmup exhibits unstable threshold routing: (a) the cutoff-EMA lags behind the actual cutoff, (b) raw expert usage is low, and (c) starvation rate is high as the capacity lower bound is frequently triggered. ET no warmup relies only on the capacity factor during this stage, which is suboptimal because it does not provide a stable threshold estimate or balanced expert learning signal. With warmup, the cutoff-EMA trajectory stabilizes (d), and router outputs show lower variance in both logits (e) and gates (f). Note: ec_shared_bsz512k does not log raw usage and underflow metrics, so panels (b) and (c) show only the two ET runs.
Figure 11: Layer-1 cutoff-EMA (expert 0) under EC/ET vs DeepSeek loss-free load balancing.
Method
Shared
CE
CORE
EC (bsz 512k)
Yes
2.843
19.94
EC (bsz 512k)
No
2.862
16.307
ET no warmup ( β=0.999 )
Yes
2.844
16.867
ET no warmup ( β=0.999 )
No
2.862
18.515
Appendix
Table 8: Ablation on the shared expert mechanism. In both ET no warmup and EC, shared expert improves loss by roughly 0.02.
Figure 12: Comparison of evaluation loss with and without normalization. The configuration without normalization (blue) consistently achieves lower loss than the fanout-normalized variant (orange).
Figure 13: Capacity constraint behavior during ET training (from step 4k onward, after warmup). (a) Raw expert usage before capacity capping. (b) Saturation rate: fraction of selected tokens dropped due to capacity limits. (c) Starvation rate: fraction of unused expert capacity. Both saturation and starvation rates remain low, confirming minimal train-inference mismatch.
Figure 14: Within-family checkpoint-pair routing consistency on a fixed validation stream, measured by joint JSD. Lower values indicate more stable routing. The same broad story remains. ET separates clearly from EC 2k and stays close to EC 64k.
Figure 15: Layerwise mean fanout versus loss bin for the two main runs. EC 2k shows a stronger positive dependence in several layers, while ET remains more mixed across depth.
Figure 16: Standalone EC 8k activation dynamics in the same form as the main figure. Faint dashed gray curves show the per-layer means and the solid red curve shows the global mean across layers. Compared with EC 2k, EC 8k is noticeably flatter across the loss range.
Figure 17: Inverse layerwise activation diagnostics for the remaining three runs. Each panel plots mean loss against fanout by layer, highlighting how the loss–fanout relation remains setup dependent across ET no warmup, EC 64k, and EC 512k.
Figure 18: Router logit histograms from the warmup ET run. The bulk of the distribution is roughly bell shaped, with a heavier right tail. This asymmetric tail is consistent with activated tokens receiving reinforcing gradient signals that can further increase their logits. We view this figure as a qualitative appendix diagnostic rather than a core result.
Figure 19: Per-layer expert fanout on GSM8K and HumanEval passages. Each cell shows the number of experts activated for a given token at a given layer. Numerical and code-specific tokens receive substantially higher fanout than function words.
Figure 20: Full expert routing on GSM8K and HumanEval passages. Each panel shows binary expert activation (black = activated) across all layers and experts for every token. Routing patterns reveal domain-specific structure.
Figure 21: Token activation intensity on a HumanEval passage. Each token is colored by total fanout (sum of experts activated across all layers). Code-specific tokens (variable names, operators, keywords) receive higher activation than boilerplate text.
Figure 22: Expert activation heatmaps across routing configurations. Each row corresponds to a routing variant (EC with batch sizes 2k, 8k, 64k, 512k, and ET). Columns show HumanEval (code) and GSM8K (math) domains. Color intensity indicates expert token ratio. Specialization sharpens with larger EC batch sizes, and ET achieves comparable patterns without batch size dependence.
Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.
Shuibai Zhang, Caspian Zhuang, Chihan Cui +8
University of Wisconsin-Madison · Scitix · Cornell University +4
Mixture-of-Experts (MoE) language models route each token to a small subset of experts, but whether the routes selected by a trained top-k router are good ones is rarely evaluated directly. Holding the model fixed, we compare each standard route against sampled equal-compute alternatives for the same token and score each by the next-token probability it assigns to the realized token in a verified reasoning trajectory. The result is sharply token-conditional: the standard router is well-aligned with route utility on confident tokens but uninformative on the fragile tokens that drive hard reasoning, where lower-loss equal-compute routes consistently exist inside the frozen model but are not selected. The same pattern holds across Qwen3-30B-A3B, GPT-OSS-20B, DeepSeek-V2-Lite, and OLMoE-1B-7B, and follows structurally from how standard top-k training evaluates routing decisions: the language modeling loss scores only the executed route, and load balancing depends only on aggregate routing statistics. A minimal router-only update to the final-layer router, leaving every expert and every other router frozen, is sufficient to shift pass@K on AIME 2024+2025 and HMMT 2025 for both Qwen3-30B-A3B and GPT-OSS-20B, suggesting that at least part of the failure reflects router-reachable misallocation rather than expert capacity alone.
Youngsik Yoon, Siwei Wang, Wei Chen +1
Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea
Standard Mixture-of-Experts (MoE) models rely on centralized routing mechanisms that introduce rigid inductive biases. We propose Routing-Free MoE which eliminates any hard-coded centralized designs including external routers, Softmax, Top-K and load balancing, instead encapsulating all activation functionalities within individual experts and directly optimized through continuous gradient flow, enabling each expert to determine its activation entirely on its own. We introduce a unified adaptive load-balancing framework to simultaneously optimize both expert-balancing and token-balancing objectives through a configurable interpolation, allowing flexible and customizable resource allocation. Extensive experiments show that Routing-Free MoE can consistently outperform baselines with better scalability and robustness. We analyze its behavior in detail and offer insights that may facilitate future MoE design ad optimization.
Yilun Liu, Jinru Han, Sikuan Yan +2
Ludwig Maximilian University of Munich · Munich Center for Machine Learning · University of California, Los Angeles