Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Shuibai Zhang; Caspian Zhuang; Chihan Cui; Zhihan Yang; Fred Zhangzhi Peng; Yanxin Zhang; Haoyue Bai; Zack Jia

cs.LGApr 2, 2026

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Authors: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia,

Organizations: University of Wisconsin-Madison · Scitix · Cornell University · Duke University ·

Abstract

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.

Figures & tables

Figure 1: Training loss vs. wall-clock time. EC reaches loss 3.75 in 10.6h,

2.0\times

faster in wall-clock time than TC (20.7h).

Panel of Figure 2 — Figure 2: Left: TC (top-1) vs. EC (capacity $c\!=\!2$ ) routing on a $6\!\times\!3$ gating score matrix. Both methods assign the same total of 6 token–expert pairs, but TC produces imbalanced per-expert loads (1/4/1) while EC guarantees uniform loads (2/2/2) by construction. Right: GPU memory snapshot during inference of LLaDA-2.0-mini (16B) with expert parallelism across 8 H100 GPUs. TC exhibits high variance (std 3.6 GB) with one GPU using 70.3 GB while others use ${\sim}$ 58–64 GB. EC maintains perfectly uniform memory (std 0.0 GB).

Figure 3: Linear-reverse scheduling: as mask ratio

r

decreases during denoising, per-expert capacity increases, concentrating compute on the most consequential predictions.

Scheduler	$s(r)$	Compute bias	PPL $\downarrow$
Linear-Rev.	$1-r$	Low mask ratio	36.5
Static ( $k$ =20)	–	–	37.1
Cosine-Rev.	$\tfrac{1}{2}(1+\cos\pi r)$	Low mask ratio	37.2
Gaussian	$\tilde{g}(r)$	Intermediate	37.3
Linear	$r$	High mask ratio	37.5
Gaussian-Rev.	$1-\tilde{g}(r)$	Extremes	37.6

Table 1: Scheduling strategies and final validation perplexity (30B tokens, OpenWebText). Expert capacity is

k(r)=k_{\min}+(k_{\max}-k_{\min})\cdot s(r)

with

k_{\min}{=}8

k_{\max}{=}32

; static baseline uses

k{=}20

\tilde{g}(r)

: normalized Gaussian centered at

r{=}0.5

with

\sigma{=}0.22

(Appendix H ). All schedulers match in expected FLOPs.

Figure 4: Scheduler comparison on OpenWebText (30B tokens, matched average FLOPs). Left: reverse schedulers; Right: forward schedulers. Reverse schedulers allocate more experts to low-mask-ratio steps and consistently outperform their forward counterparts.

Figure 5: 8B-A1B pretraining comparison: dynamic EC (linear-reverse,

k

=2–14) vs. static EC (

k

=8) on Nemotron-CC. Left: validation perplexity; Center: MMLU 5-shot accuracy; Right: ARC-Challenge 25-shot accuracy. Dynamic EC outperforms static EC at every checkpoint under matched average FLOPs. The linear-reverse schedule carries over from the small-scale bake-off in § 4.3 and was not re-evaluated at this scale; both curves are single runs.

Figure 6: Mechanistic analysis of learning efficiency (8B-A1B, Nemotron-CC). Left: Convergence rate $\eta_{r}$ (Eq. 2 ) by mask-ratio bin and training stage for the static EC baseline; higher values indicate faster learning. Low-mask-ratio bins learn over $10\times$ faster. Right: Ratio $\eta_{r}^{\text{dyn}}/\eta_{r}^{\text{static}}$ ; values ${>}1$ indicate dynamic EC learns faster. The advantage concentrates in low-mask-ratio bins.

Figure 7: Retrofitting a pretrained TC DLM (LLaDA-MoE) to EC and dynamic EC across four benchmarks. The

t{=}0

point of each EC curve is the post-swap, pre-finetuning baseline: the router swap alone costs a few points (HumanEval

0.42\to 0.38

, HumanEval+

0.33\to 0.30

Pass@1), which finetuning then recovers and exceeds. EC converges notably faster than TC in wall clock time while reaching comparable final accuracy, and dynamic EC achieves slightly higher final accuracy than both TC and static EC. See Appendix I for results plotted against training samples.

Routing	Pass@1	Time	Pass@1	Time	Acc	Time	Acc	Time	Acc	Time
	HumanEval		HumanEval+		GSM8K		MedQA		Avg
TC	53.9	1369	46.1	1369	74.8	672	35.6	1885	52.6	1324
EC	55.5	1008	48.4	1008	73.8	459	36.7	1371	53.6	962
Dynamic EC	58.6	1056	51.6	1056	73.8	461	35.7	1378	54.9	988

Table 2: Peak accuracy (%) and evaluation decode time (s) for TC-to-EC retrofitting. Peak values are the best scores during SFT (within the time window of Figure 7 ). EC and dynamic EC achieve comparable or higher accuracy while decoding

1.3

–

1.5\times

faster.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Per-layer token drop ratio at step 300K for static and dynamic EC (8B-A1B, Nemotron-CC). Middle layers drop

<

1.1% of tokens under static EC. Dynamic EC has higher drop ratios due to reduced capacity at high-mask-ratio steps.

Panel of Figure 9 — Figure 9: TC routing ablations. (a) Removing or replacing the auxiliary load-balancing loss does not close the gap with EC ( $1.7\times$ faster). (b) Increasing the capacity factor slows TC further due to padding overhead ( $2.0\times$ faster for EC vs. TC cap=1.5).

Figure 10: Per-step training loss. The auxiliary load-balancing loss improves per-step convergence (TC dropless LB reaches loss 3.75 in fewer steps), but this gain is offset by reduced throughput (Table 3 ), resulting in slower wall-clock convergence (Figure 9(a) ).

Routing	Throughput	Relative to EC
EC	52.1	1.00 $\times$
TC (dropless, no LB)	44.4	0.85 $\times$
TC (dropless, loss-free bias)	38.6	0.74 $\times$
TC (cap=1.0, LB)	35.4	0.68 $\times$
TC (cap=1.25, LB)	27.0	0.52 $\times$
TC (cap=1.5, LB)	25.9	0.50 $\times$

Appendix

Table 3: Average training throughput (TFLOP/s/GPU) over the first 35h. Higher is better. EC achieves

1.5

–

2.1\times

higher throughput than all TC variants.

Figure 11: Convergence rate comparison between linear-reverse and cosine-reverse (both from H100 runs, evaluated up to 128K steps). Left/Center: per-bin $\eta_{r}$ for each scheduler. Right: ratio $\eta_{r}^{\text{Lin-Rev}}/\eta_{r}^{\text{Cos-Rev}}$ ; values near 1.0 indicate no meaningful difference. The two schedulers have nearly identical per-bin learning dynamics.

Stage (steps)	$\eta_{r}$ , $[0,0.25)$	$\eta_{r}$ , $[0.75,1.0)$	Low/high
16K–32K	62.2	9.8	6.35 $\times$
32K–64K	22.5	2.9	7.76 $\times$
64K–128K	7.6	1.1	6.91 $\times$
128K–256K	2.5	0.3	8.33 $\times$

Appendix

Table 4: Convergence rate

\eta_{r}

for the lowest and highest mask-ratio bins across training stages (8B-A1B, static EC, Nemotron-CC), in the units plotted in Figure 6 . The low/high ratio is stable and slightly increasing, so low-mask-ratio bins do not saturate relative to high-mask-ratio ones.

Scheduler	$\mathbb{E}[s(r)]$	$\mathbb{E}[k(r)]$	$\Delta$ vs. static
Static ( $k=20$ )	0.5000	20.00	0.000
Linear	0.5000	20.00	0.000
Linear-reverse	0.5000	20.00	0.000
Cosine	0.5000	20.00	0.000
Cosine-reverse	0.5000	20.00	0.000
Gaussian	0.5010	20.02	$+$ 0.025

Appendix

Table 5: Expected top-

k

for each scheduler under

r\sim\mathrm{Uniform}(0,1)

with

k_{\min}=8

k_{\max}=32

Figure 12: SFT retrofitting results vs. number of training samples for TC, EC, and dynamic EC across four benchmarks. Results are consistent with the wall clock time view in Figure 7 .

Figure 13: TC routing ablation: comparing four TC variants (no_lb_dropless, lb_dropless, no_lb_cf1.0, no_lb_cf2.0) against EC and dynamic EC across four benchmarks. All TC variants show similar convergence behavior.

	HumanEval		HumanEval+		GSM8K		MedQA		Avg
Routing	Pass@1	Time	Pass@1	Time	Acc	Time	Acc	Time	Acc	Time
TC (no_lb, dropless)	53.9	1369	46.1	1369	74.8	672	35.5	1884	52.6	1324
TC (lb, dropless)	54.7	1220	50.0	1211	74.2	664	36.1	1852	53.8	1237
TC (no_lb, cf=1.0)	54.7	1274	47.7	1261	74.4	543	31.8	1645	52.2	1181
TC (no_lb, cf=2.0)	53.9	1288	47.7	1288	75.0	647	36.3	1830	53.2	1263
EC	55.5	1008	48.4	1008	73.8	459	36.7	1371	53.6	962

Appendix

Table 6: Peak accuracy (%) and evaluation decode time (s) for all TC variants, EC, and dynamic EC across four SFT benchmarks. EC and dynamic EC match or exceed all TC variants in accuracy while decoding faster.

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

REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models
Aug 3, 2026Xiang Xia, Cheng Yan, Yiming Zhang +3Expert Load Balancing Mixture-of-Experts Inference
University of Science and Technology of China, Hefei, China
Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing
Mar 12, 2026Hanchi Sun, Yixin Liu, Yonghui Wu +1Expert Routing Mixture-of-Experts Models
Computer Science and Engineering, Lehigh University, Bethlehem, PA, USA · MD-HOBI-BIOMED INFORMATICS, University of Florida, Gainesville, FL, USA
ALoDLM: Adaptively Looped Diffusion Language Models
Oct 3, 2026Liancheng Fang, Zhuowei Li, Youngeun Kim +10Adaptive Computation for Diffusion Models Diffusion Language Models
University of Illinois Chicago · Amazon AGI · Korea University

Aug 3, 2026J/K move · Enter open · S save

REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models

Xiang Xia, Cheng Yan, Yiming Zhang, Jiazheng Liu+2

University of Science and Technology of China, Hefei, China

Expert Load Balancing Mixture-of-Experts Inference Diffusion Language Model Inference Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Mixture-of-experts (MoE) models increase parameter capacity by activating only a small subset of experts for each token. This conditional-computation paradigm has enabled autoregressive language models to scale model capacity without a proportional increase in per-token computation. In diffusion language models (DLMs), however, each denoising forward jointly revisits all token positions despite their sharply different refinement demands, while the default fixed token-choice routing assigns them a uniform expert budget, creating a mismatch between expert computation and refinement demand. We argue that MoE inference in DLMs should therefore be viewed as refinement-aware compute allocation across heterogeneous token refinement states. We propose REFLEX (REfinement-aware FLEXible expert allocation), a training-free method that keeps the default router unchanged while reorganizing expert computation around the evolving refinement process. Specifically, REFLEX introduces a coarse-to-fine hierarchy for expert-budget allocation that aligns computation with block-relative refinement roles while using the Frontier-Progress Score to resolve active-block priorities. Across multiple widely used benchmarks on two representative MoE-based DLMs, LLaDA-MoE and LLaDA2.0-mini, REFLEX reduces allocated expert computation by 15% on average while preserving or even improving generation quality on most benchmarks relative to default routing. Compared with autoregressive-style variable-expert routing methods, REFLEX also yields a more consistent quality--computation trade-off, further supporting the importance of allocating expert computation according to the heterogeneous refinement demands exposed within each denoising forward.

Abstract

Expert Routing

Explore similar work

REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models

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

ALoDLM: Adaptively Looped Diffusion Language Models