PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Yang Zhou, Guanhua Chen, Ming Liu

arXiv:2604.01622v2cs.LGcs.CL

TL;DR

Diffusion language-model MoE systems inherit token-choice routing despite their parallel denoising structure. This paper studies expert-choice routing with timestep-dependent capacity and finds better balancing, throughput, convergence, and downstream accuracy than token-choice routing.

  • Problem

    Existing diffusion language-model MoE systems inherit token-choice routing, while systematic evidence on routing paradigms suited to parallel denoising remains limited.

  • Method

    The paper compares expert- and token-choice routing under controlled training, then schedules expert capacity by denoising timestep under matched FLOPs and tests router-only retrofits.

  • Results

    Expert-choice routing delivers better load balance, throughput, and convergence than token-choice routing, while low-mask-ratio scheduling and router replacement improve training or downstream accuracy.

  • Takeaways & Limitations

    Diffusion-language-model computation can be treated as an adaptive policy, with expert capacity allocated across denoising timesteps rather than fixed architecturally.

  • Takeaways & Limitations

    Full-sequence expert-choice routing does not apply across causally masked blocks in block-wise or semi-autoregressive decoding, requiring a hybrid design left for future work.

Abstract

from arXiv · show

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.

1 Introduction

The introduction argues that expert-choice routing better matches diffusion language models’ parallel denoising structure than token-choice routing, while enabling controllable computation allocation across timesteps. It presents evidence that EC improves load balance, throughput, convergence, and downstream finetuning, especially when capacity favors low-mask-ratio steps.

  • Motivation: Token-choice routing inherits autoregressive assumptions that overlook DLMs’ non-causal attention and simultaneous generation structure.TC lets each token independently select preferred experts, a design suited to causal generation where a global token view is unavailable.
  • Motivation: Token-choice routing creates expert load imbalance, requiring an auxiliary balancing loss while providing no hard guarantee of balanced utilization.Independent token decisions can overload some experts and leave others idle; the auxiliary objective adds compute and may interfere with language-model training.
  • Contributions: Expert-choice routing makes expert capacity externally controllable, enabling timestep-dependent computation across denoising steps with different masking ratios.This control allows capacity to vary during the explicit denoising loop rather than emerging from independent token choices.
  • Contributions: EC routing achieves better load balance, higher throughput, and faster convergence than TC routing for DLM MoE training.The paper frames EC as strictly superior to TC for DLM MoE training on these criteria.
  • Contributions: Order-of-magnitude higher learning efficiency in low-mask-ratio contexts explains why allocating more expert capacity to those timesteps yields the largest marginal return under matched FLOPs.The paper introduces timestep-dependent capacity scheduling and evaluates multiple allocation strategies.
  • Contributions: A simple router replacement lets pretrained TC DLMs achieve faster convergence and improved finetuning accuracy across diverse downstream tasks.The retrofit changes only the router while transferring existing pretrained models to EC routing.

2 Background

Masked diffusion language models generate sequences through iterative, bidirectional denoising in which each timestep presents a distinct masking-ratio-dependent task. In MoE layers, token-choice routing leaves expert loads uncontrolled, whereas expert-choice routing guarantees balanced capacity while allowing computation to scale externally.

  • Diffusion Language Models: DLMs iteratively denoise the entire sequence, masking a fraction of positions according to γ(t) across T decoding steps.Each step predicts all masked positions and unmasks a subset according to the masking schedule.
  • Diffusion Language Models: Each denoising step has a distinct masking ratio and qualitatively different task, while every forward pass processes the full L-token sequence.Bidirectional attention removes causal constraints during prediction.
  • Mixture-of-Experts Routing: Token-choice routing lets each token select top-k experts, but per-expert load remains uncontrolled, with overflow dropped beyond capacity.An auxiliary load-balancing loss encourages uniform utilization but is fragile in practice.
  • Mixture-of-Experts Routing: Expert-choice routing lets each expert select top-c tokens, guaranteeing exactly c assignments per expert without dropped tokens or an auxiliary loss.Setting c = kN/E matches TC top-k total token–expert pairs, while varying c scales computation proportionally.

3 Expert-Choice as the Proper Routing Paradigm

Controlled pretraining shows that expert-choice routing is consistently preferable to token-choice routing in diffusion language-model MoEs. EC removes structural load imbalance by construction, producing higher throughput and faster wall-clock convergence than TC variants.

  • Controlled comparison: Controlled pretraining varied only the routing mechanism while holding architecture, data, and hyperparameters constant.All models were trained from scratch.
  • Convergence and throughput: 10.6h: EC reached loss 3.75, approximately 2.0× faster than dropless TC and capacity-bounded TC (cap=1.25) at ∼20h.The comparison uses training loss as a function of wall-clock time.
  • Load balance: 70.3 GB: one TC GPU used substantially more memory than others at ∼58–64 GB, creating a structural straggler penalty.Per-GPU memory standard deviation under TC was 3.6 GB.
  • Load balance: 0.0 GB: EC maintained uniform per-GPU memory usage, eliminating the structural imbalance that auxiliary losses cannot fully remove under TC.Under TC, tokens choose experts independently; EC removes the imbalance by construction.

4 Timestep-Adaptive Expert Capacity

Because expert capacity is explicit in EC routing, it can vary with denoising timestep; under matched FLOPs, allocating more capacity to low-mask-ratio steps consistently improves performance and transfers to larger-scale downstream evaluations. Low-mask-ratio bins learn substantially faster, explaining why concentrating computation there yields the greatest benefit.

  • Capacity scheduling: EC routing enables timestep-dependent expert capacity because capacity is an explicit hyperparameter rather than an emergent quantity.The capacity schedule is defined as k(r) = clamp(kmin + (kmax − kmin) · s(r), kmin, kmax).
  • Scheduler comparison: Under matched FLOPs, linear-reverse achieves the lowest validation perplexity, while low-mask-ratio-favoring schedulers outperform high-mask-ratio- and intermediate-favoring schedulers.Dynamic variants use kmin=8, kmax=32 with E[k]=20, matching the static k=20 baseline.
  • Scale transfer: Dynamic linear-reverse EC consistently outperforms static EC on validation perplexity, MMLU 5-shot accuracy, and ARC-Challenge 25-shot accuracy at 8B-A1B scale.The comparison uses Nemotron-CC with dynamic kmin=2, kmax=14 and static k=8, matched at E[k]=8; both curves are single runs.
  • Mechanistic explanation: Low-mask-ratio bins converge up to 7× faster than high-mask-ratio bins, and dynamic EC further improves convergence where the baseline learning efficiency is highest.Dynamic EC has convergence-rate ratios above 1 in low-mask-ratio bins and below 1 in high-mask-ratio bins.
  • Limitations: The convergence-rate analysis explains why low-mask-ratio-favoring schedules win, but it does not explain the residual gap between linear-reverse and cosine-reverse.The two reverse schedules show nearly identical per-bin convergence rates despite different capacity allocations near r=0 and r=1.

5 Retrofitting Pretrained TC DLMs

Pretrained TC DLMs can be retrofitted to EC through a router-only conversion followed by standard task-specific finetuning, preserving or improving accuracy while accelerating convergence and decoding. Dynamic EC adds no parameters and achieves the highest average accuracy across the evaluated benchmarks.

  • Conversion: The retrofit replaces token-choice routing with expert-choice routing while preserving the pretrained router projection and all other model components.TC uses per-row top-k selection, whereas EC uses per-column top-c selection; no architectural changes beyond the router are required.
  • Finetuning: Standard task-specific finetuning is required after conversion to adapt expert FFNs to EC’s changed token distribution.Dynamic EC only makes capacity c a function of masking ratio and has the same trainable parameter set as static EC.
  • Results: 54.9% vs. 53.6% for static EC and 52.6% for TC: dynamic EC achieves the highest average accuracy across four benchmarks.EC converges notably faster than TC while reaching comparable final accuracy, and dynamic EC further improves final performance.
  • Finetuning: The router swap alone causes HumanEval 0.42 →0.38 and HumanEval+ 0.33 →0.30 Pass@1, but finetuning recovers and exceeds the lost accuracy.The initial drop occurs because experts were optimized under TC’s load distribution and must adapt to EC’s per-step token mixture.
  • Results: 1.3–1.5× faster decoding: EC and dynamic EC match or exceed TC’s peak accuracy while requiring less evaluation time.The inference speedup results from deterministic load balancing eliminating TC routing’s straggler effects.

6 Related Work

Prior work has advanced diffusion language models from foundational discrete-text diffusion methods to large-scale systems and adaptations of pretrained autoregressive models. Related research also addresses sparse MoE load balancing and timestep-dependent computation, motivating this paper’s focus on adaptive denoising computation.

  • Diffusion language models: Diffusion language models have progressed from early discrete-text foundations to large-scale systems that rival autoregressive LLMs.Recent work also adapts pretrained autoregressive models into diffusion language models, while commercial deployments include Mercury.
  • Mixture-of-experts: Sparse MoE architectures scale model capacity with sublinear compute cost, but token-choice routing can produce uneven expert loads.Prior mitigations include auxiliary losses, optimal assignment, bias correction, and ReLU-based approaches.
  • Timestep-dependent computation: The finding that low-mask-ratio steps benefit most from additional computation aligns with continuous-diffusion analyses under the correspondence between low noise and low mask ratio.Prior work reports that smaller-t denoising tasks are harder to learn, while the high-noise regime converges easily and is oversampled.

7 Conclusion

Expert-choice routing is consistently preferable to token-choice routing for diffusion language model mixture-of-experts models, improving throughput through deterministic load balancing and enabling timestep-dependent expert capacity scheduling. Low-mask-ratio denoising steps have an order-of-magnitude higher learning efficiency, explaining why the linear-reverse scheduler performs best under matched FLOPs.

  • 7 Conclusion: Expert-choice routing consistently outperforms token-choice routing for DLM MoE models through deterministic load balancing and higher throughput.It also enables timestep-dependent expert capacity scheduling.
  • 7 Conclusion: Low-mask-ratio denoising steps have an order-of-magnitude higher learning efficiency than other steps.This mechanistic finding explains why concentrating capacity according to denoising timestep is beneficial.
  • 7 Conclusion: The linear-reverse scheduler outperforms all alternatives under matched FLOPs.Its advantage follows from the higher learning efficiency of low-mask-ratio denoising steps.

Ethics Statement … D Token Coverage in Expert-Choice Routing

The paper reports responsible-use considerations, summarizes appendix coverage, identifies limitations and future extensions, and analyzes token coverage and information preservation in expert-choice routing.

  • Ethics Statement: The experiments use public datasets, involve no human subjects or private data, and introduce no risks beyond those inherent to large language models.The authors encourage responsible deployment when applying the methods to production systems.
  • Contents of the Appendix: The appendix covers limitations, future directions, language-model use, token coverage, experimental details, routing ablations, convergence analysis, scheduler definitions, retrofitting results, and adaptive computation.These topics are distributed across Sections A–J.
  • A Limitations: Capacity schedules are hand-designed, and the optimal schedule may depend on model scale, dataset, and task.The paper proposes learned capacity predictors or reinforcement-learning policies as extensions.
  • B Future Directions: Future work includes learned capacity, continued pretraining before finetuning, EC-specific inference mechanisms, and block-wise diffusion extensions.Proposed mechanisms include timestep-aware speculative decoding, per-prompt adaptive capacity, and adaptive parallel decoding composition.
  • D Token Coverage in Expert-Choice Routing: Static EC with k=8 drops fewer than 1.1% of tokens in middle layers and has a 2.7% mean drop ratio across all layers.These measurements are from step 300K using a 5-point average.
  • D Token Coverage in Expert-Choice Routing: Unrouted tokens are still processed by shared expert FFNs at every layer, so token drops remove them from the routed pathway rather than computation altogether.This architectural design addresses the concern that expert-choice routing may cause information loss.
  • D Token Coverage in Expert-Choice Routing: Under perfectly correlated drops, the probability of a token being unrouted at every layer is bounded by the per-layer marginal rate: ∼2.7% for static EC and ∼8.0% for dynamic EC.The bound does not assume independent drop events because router decisions are correlated through the residual stream.

E Experimental Setup … F.2 Capacity Factor

The experiments compare otherwise matched expert-choice and token-choice models, then ablate auxiliary loss and capacity factor to isolate TC’s convergence and throughput bottlenecks. Results attribute EC’s advantage primarily to TC’s inherent load imbalance, which worsens as capacity increases despite reducing token dropping.

  • E Experimental Setup: EC and TC are compared using identical architectures, data, learning rate, and hyperparameters, differing only in the routing mechanism.The models use 16 Transformer layers, hidden size 2048, 64 fine-grained experts, and 2 shared experts, trained on Nemotron-CC at 2×10−4.
  • E Experimental Setup: Scheduler experiments use OpenWebText with dynamic kmin=8 and kmax=32 variants against a static constant-k=20 baseline.These models use 16 layers, hidden size 512, 512 fine-grained experts, sequence length 513, global batch size 256, and 30B training tokens.
  • E Experimental Setup: The 8B-A1B pretraining setup compares a static constant-k=8 baseline with a linear-reverse dynamic variant using kmin=2 and kmax=14.The setup uses 16 layers, hidden size 2048, 64 fine-grained experts, sequence length 2049, global batch size 288, and Nemotron-CC.
  • F TC Routing Ablations: TC routing ablations measure wall-clock time as cumulative per-iteration training time, excluding evaluation and checkpointing.The ablations test whether auxiliary loss or capacity factor explains EC’s convergence advantage.
  • F.1 Auxiliary Loss: 1.7× faster convergence remains for EC than the slowest dropless TC variant after removing or replacing the auxiliary load-balancing loss.This rules out the auxiliary loss as TC’s primary weakness and identifies fundamental token-choice load imbalance as the dominant bottleneck.
  • F.1 Auxiliary Loss: 36% fewer iterations let TC (dropless, LB) reach loss 3.75 in 50.6k versus 78.6k steps for TC (dropless, no LB), but throughput falls from 44.4 to 24.9 TFLOP/s/GPU.The auxiliary loss improves per-step optimization while worsening system-level load imbalance, producing a 1.78× slowdown and slower wall-clock convergence.
  • F.2 Capacity Factor: Increasing capacity factor from 1.0 to 1.5 slows wall-clock convergence because higher expert capacity amplifies straggler effects despite reducing token dropping.The maximum load per expert becomes capacity factor × N/E, forcing all GPUs to wait for the most overloaded expert; EC is 2.0× faster than TC at cap=1.5.

F.3 Throughput Analysis · G Convergence Rate Computation · G.1 Linear-Reverse vs. Cosine-Reverse

EC delivers structurally higher training throughput by enforcing deterministic load balancing, while convergence-rate analysis finds nearly identical per-bin learning dynamics for linear-reverse and cosine-reverse. Thus, per-bin convergence rate does not explain their reported PPL gap.

  • F.3 Throughput Analysis: Throughput is reported in TFLOP/s/GPU using forward-pass FLOPs per training step and measured wall-clock time per step.Because all models share the same architecture, forward-pass FLOPs are identical across runs, so throughput differences reflect hardware-utilization efficiency.
  • F.3 Throughput Analysis: EC achieves 52.1 TFLOP/s/GPU, outperforming every TC configuration by 1.2×–2.1%.Average throughput is measured over the first 35 hours of training.
  • F.3 Throughput Analysis: Deterministic load balancing gives all GPUs the same token count per step, eliminating idle time and driving EC’s throughput advantage.The advantage is structural rather than dependent on auxiliary objectives or capacity bounds.
  • F.3 Throughput Analysis: Load imbalance, rather than auxiliary loss, is the dominant throughput bottleneck across TC variants.Removing auxiliary loss or replacing it with loss-free bias does not close the gap with EC; the auxiliary loss has negligible computational cost.
  • F.3 Throughput Analysis: Capacity-bounded TC throughput decreases monotonically with capacity factor, from 35.4 to 27.0 to 25.9.The passage attributes throughput differences to varying degrees of load imbalance across configurations.
  • G Convergence Rate Computation: Per-token cross-entropy loss is recorded by masking-ratio bin, and ηr is estimated by least-squares fitting ln Lr within geometrically spaced training stages.The stages span 16K–32K, 32K–64K, 64K–128K, and 128K–256K steps.
  • G Convergence Rate Computation: ηr measures fractional loss decrease per step, enabling comparison across bins with different loss scales while describing optimization speed rather than task importance.Different irreducible loss floors mean that “learns faster” refers to fractional optimization speed.
  • G.1 Linear-Reverse vs. Cosine-Reverse: The 0.97–1.04 convergence-rate ratio range shows nearly identical per-bin learning dynamics for linear-reverse and cosine-reverse, so ηr cannot explain their PPL gap of 36.5 vs. 37.2.Both schedulers use the same data and are evaluated up to 128K steps.

G.2 Convergence-Rate Trajectories Across Training … H.2 FLOPs Equivalence Across Schedulers

Across training, low-mask-ratio bins do not show relative saturation, while scheduler designs preserve nearly identical expected computation under uniform masking. Gaussian schedulers are normalized over the full capacity range, with only minimal deviation from the static baseline.

  • H.1 Gaussian Scheduler: The normalized Gaussian scheduler maps endpoint values to 0 and the midpoint to 1, spanning the full [kmin, kmax] capacity range.The experiments use σ = 0.22; Gaussian-reverse is defined as 1 − g̃(r).
  • G.2 Convergence-Rate Trajectories Across Training: The low/high convergence-rate ratio remains stable and slightly increases across four training stages, contradicting the saturation trend toward 1.This result indicates that low-mask-ratio bins do not saturate relative to high-mask-ratio bins in the 8B-A1B run.
  • H Scheduler Definitions and FLOPs Equivalence: Under uniform masking ratios, every scheduler has the same expected top-k as the static baseline.The experimental bounds are kmin = 8 and kmax = 32, with static k = 20.
  • H Scheduler Definitions and FLOPs Equivalence: MoE-layer FLOPs are proportional to k(r), so scheduler equivalence follows from matching the expected capacity E[k(r)].The expected capacity is computed as kmin + (kmax − kmin) · E[s(r)].
  • H.2 FLOPs Equivalence Across Schedulers: 20.00 is the expected top-k for cosine and cosine-reverse schedulers, matching the static baseline.Both schedulers have E[s(r)] = 0.5 under uniform masking ratios.
  • H.2 FLOPs Equivalence Across Schedulers: 20.02 and 19.98 are the expected top-k values for Gaussian and Gaussian-reverse schedulers with σ = 0.22.Their deviations from the static k = 20 baseline are ±0.02, or ±0.1% of expected computation.
  • H.2 FLOPs Equivalence Across Schedulers: Table 5 confirms that all schedulers match the static baseline in expected FLOPs under r ∼ Uniform(0, 1).This summarizes the FLOPs-equivalence analysis across the scheduler family.

I Additional SFT Results · J Additional Related Work

Additional SFT results show that EC converges faster and that dynamic EC reaches the highest final accuracy while EC variants decode faster than TC variants. Related work situates timestep-adaptive computation as distinct from depth-based adaptation and alongside parallel DLM decoding methods.

  • I Additional SFT Results: Across HumanEval, HumanEval-Plus, GSM8K, and MedQA, EC converges faster, while dynamic EC achieves the highest final accuracy.Figure 12 reports these trends against training samples and finds them consistent with wall-clock results.
  • I Additional SFT Results: All four TC configurations reach similar peak accuracy, indicating limited final-quality impact from TC load-balancing choices.The compared variants are no lb dropless, lb dropless, no lb cf1.0, and no lb cf2.0.
  • I Additional SFT Results: EC and dynamic EC match or exceed TC accuracy across four SFT benchmarks while requiring less evaluation decode time.Table 6 reports peak accuracy (%) and evaluation decode time (s).
  • J Additional Related Work: Prior adaptive-computation methods primarily vary computation along model depth through early exit, confidence halting, depth routing, or recursive transformers.The cited works include Schuster et al. (2022), Elhoushi et al. (2024), Raposo et al. (2024), Bae et al. (2025b), and others.
  • J Additional Related Work: This approach instead adapts computation along the timestep dimension, which is unique to iterative generative models.The paper characterizes timestep adaptation as orthogonal to depth-based computation adaptation.
  • J Additional Related Work: Diffusion-model research reports timestep-dependent learning dynamics, including findings that low-noise timesteps are harder to learn and methods addressing gradient conflicts.The cited approaches include process increment analysis, curriculum-based difficulty measurement, Min-SNR weighting, and noise schedule analysis.
  • J Additional Related Work: Parallel-decoding research accelerates DLM inference through KV caching, confidence-aware decoding, block-wise generation, and methods targeting parallel decode quality.Examples include Fast-dLLM, Fast-dLLM v2, Esoteric LMs, and dParallel.
Loading 2604.01622v2…