Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

Hanchi Sun; Yixin Liu; Yonghui Wu; Lichao Sun

cs.AIMar 12, 2026

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

Abstract

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 $\times$ 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.

Panel of Figure 2 — 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 $c_{i}$ , enabling fully causal routing over the population.

Method	Batch	CE loss ( $\downarrow$ )	CORE ( $\uparrow$ )
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-\beta)

. TC variants: no load-balancing, auxiliary loss (

\alpha{=}0.001

), or loss-free (

u{=}0.005

). We report validation cross-entropy (CE) loss (

\downarrow

) and CORE Eval score (

\uparrow

Method	Batch	CE loss ( $\downarrow$ )	CORE ( $\uparrow$ )
dense	—	2.751	20.43
TC aux	32k	2.687	22.31
EC	256k	2.621	24.98
ET ( $\beta{=}0.999$ +warmup)	256k $\rightarrow$ 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.

Panel of Figure 4 — 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.

Panel of Figure 5 — 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 $c_{i}$	LossFree bias $b_{i}$	Per-expert scalar; no aux loss
$1-\beta$	LossFree $\mu$	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\times d_{\text{model}}$	$4\times d_{\text{model}}$
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.
$W_{E}$ (embed.)	$\mathcal{N}(0,1)$	$0.2\lambda$	AdamW
$W_{\text{lm}}$ (head)	$\mathbf{0}$	$0.004\lambda$	AdamW
$W_{QKV}$ (attn)	Asp.-ratio	$0.02\lambda$	Muon
$W_{O}$ (attn proj)	$\mathbf{0}$	$0.02\lambda$	Muon
$W_{\text{router}}$	$\mathcal{N}(0,d^{-\frac{1}{2}})$	$0.02\lambda$	Muon
$W^{(e)}_{\uparrow}$ (exp. up)	Asp.-ratio	$0.02\lambda$	Muon

Appendix

Table 7: Parameter initialization and optimizer configuration. Aspect-ratio scaled init. uses

\text{std}=d_{\text{in}}^{-1/2}\cdot\min(1,\sqrt{d_{\text{out}}/d_{\text{in}}})

. LRs include

\mu

P ( Yang et al., 2022 ) scaling

\lambda=(d_{\text{model}}/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 ( $\beta{=}0.999$ )	Yes	2.844	16.867
ET no warmup ( $\beta{=}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.

Panel of Figure 15 — 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.

Panel of Figure 17 — 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.

Panel of Figure 18 — 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.

Panel of Figure 19 — 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.

Panel of Figure 20 — 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.

Expert Routing

Expert routing in mixture-of-experts models selects which subnetworks process each input, balancing accuracy, computation, and expert workload.

See all 118 papers

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
Apr 2, 2026Shuibai Zhang, Caspian Zhuang, Chihan Cui +8Expert Routing Mixture-of-Experts Language Models
University of Wisconsin-Madison · Scitix · Cornell University +4
When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models
May 8, 2026Youngsik Yoon, Siwei Wang, Wei Chen +1Mixture-of-Experts Language Models LLM Routing
Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea
Routing-Free Mixture-of-Experts
Apr 1, 2026Yilun Liu, Jinru Han, Sikuan Yan +2Mixture-of-Experts Models Expert Load Balancing
Ludwig Maximilian University of Munich · Munich Center for Machine Learning · University of California, Los Angeles

Apr 2, 2026J/K move · Enter open · S save

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang+7

University of Wisconsin-Madison · Scitix · Cornell University+4

Expert Routing Mixture-of-Experts Language Models Diffusion Language Models Adaptive Computation

Open paper AskPDF arXiv

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.

Abstract

Expert Routing

Explore similar work

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Routing-Free Mixture-of-Experts