🤖 machine learning

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

This paper proposes Expert-Choice (EC) routing as a superior alternative to Token-Choice (TC) routing for Diffusion Language Models, demonstrating that EC's deterministic load balancing, combined with adaptive timestep-dependent expert capacity, significantly improves throughput, convergence speed, and downstream accuracy while enabling the retrofitting of existing models.

Original authors: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

Published 2026-04-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-end restaurant kitchen. Your goal is to cook a complex, multi-course meal (generating a long piece of text) for a group of guests.

In the world of AI, there are two main ways to run this kitchen: the Old Way (Autoregressive) and the New Way (Diffusion).

The Old Way: You cook one dish at a time, from appetizer to dessert. You can't start the main course until the appetizer is done. It's slow but steady.
The New Way (Diffusion): You put all the ingredients on the table at once, but they are all covered in a thick fog (masks). Your job is to clear the fog and reveal the dish step-by-step. In each step, you look at the whole table and fix a few blurry spots until the whole meal is clear. This is much faster because you can work on the whole table at once.

This paper is about how to organize the chefs (the "Experts") in this Diffusion kitchen to make it run even faster and better.

The Problem: The "Pick-Your-Own-Chef" Chaos

Currently, most Diffusion kitchens use a system called Token-Choice (TC).
Imagine every single ingredient (token) on the table gets to pick which chef it wants to work with.

The Chaos: Some ingredients are very popular (like "the" or "cat") and all rush to the same 2 or 3 superstar chefs. Those chefs get overwhelmed, sweating and dropping orders. Meanwhile, other chefs are standing around doing nothing, twiddling their thumbs.
The Fix (that doesn't work well): The kitchen manager tries to shout, "Hey, stop picking the same chef!" or adds a penalty for picking the popular ones. But it's messy. The kitchen gets clogged, the line moves slowly, and the whole team has to wait for the slowest, busiest chef to finish before they can move to the next step.

The Solution: The "Chef-Choice" System

The authors propose switching to Expert-Choice (EC) routing.
Instead of ingredients picking chefs, the chefs pick the ingredients.

How it works: Each chef says, "I can handle exactly 5 ingredients." They scan the table and grab the 5 ingredients that need them the most.
The Result: Every chef has exactly 5 ingredients. No one is overloaded; no one is idle. The kitchen runs like a perfectly synchronized dance.
The Benefit: Because everyone is working at the same speed, the whole kitchen moves 2x faster. There is no waiting for the "straggler" chef.

The Secret Sauce: "Time-of-Day" Scheduling

The authors found a second, even cooler trick. They realized that not all steps in the cooking process are equally hard.

Early Steps (High Fog): When the table is covered in thick fog, the chefs are just guessing. They don't have much context. Giving them extra help doesn't help much because they are working in the dark.
Late Steps (Low Fog): As the fog clears, the chefs can see the dish clearly. They are refining the details. This is where the magic happens. If you give your best chefs extra hands right now, the dish gets perfect much faster.

So, they introduced Timestep-Dependent Capacity.

The Strategy: As the fog clears (the "mask ratio" goes down), they give the chefs more ingredients to handle.
The Analogy: Imagine a construction crew. When the building is just a pile of bricks (early stage), you don't need 100 architects; you need a few laborers. But when the building is almost done and you are polishing the windows (late stage), you want your best, most detailed-oriented architects working on it.
The Result: By focusing their energy on the "polishing" phase, the model learns 10 times faster during those critical moments.

Can We Fix Old Kitchens?

The best part? You don't have to rebuild the whole restaurant to use this.
If you have an existing kitchen using the "Pick-Your-Own-Chef" system, you can simply swap out the manager's rulebook.

You keep all the same chefs and ingredients.
You just change the rule so the chefs pick the ingredients instead.
Result: The old kitchen immediately becomes faster, learns better, and produces higher-quality meals without any expensive renovations.

Summary

Stop letting ingredients pick chefs. Let chefs pick ingredients (Expert-Choice). This stops the traffic jams and makes the kitchen run 2x faster.
Don't treat every moment the same. Give your chefs the most work when the fog is clearing (low-mask steps), because that's when they learn the most.
It's a quick upgrade. You can take any existing AI model and swap this "manager" rule in to make it instantly smarter and faster.

In short: Stop the chaos, let the experts choose, and focus your energy where it matters most.

1. Problem Statement

Diffusion Language Models (DLMs) offer a promising alternative to autoregressive models by enabling parallel, non-autoregressive text generation. However, existing DLMs that utilize Mixture-of-Experts (MoE) architectures have inherited Token-Choice (TC) routing from autoregressive systems. This inheritance creates two critical issues:

Load Imbalance: In TC routing, tokens independently select their preferred experts. This leads to stochastic load imbalance where some experts are overloaded while others remain idle. Mitigating this requires auxiliary load-balancing losses, which consume compute and introduce competing gradient signals, yet fail to provide hard guarantees.
Rigid Computation: TC routing treats expert capacity as an emergent property of token selection, making it impossible to dynamically control how much computation is allocated to specific stages of the denoising process.
Inefficiency: The load imbalance causes "straggler" effects in parallel training (some GPUs wait for overloaded ones), significantly reducing throughput and slowing down convergence.

2. Methodology

The authors propose replacing Token-Choice (TC) routing with Expert-Choice (EC) routing and introducing Timestep-Dependent Expert Capacity.

A. Expert-Choice (EC) Routing

Instead of tokens choosing experts, experts choose a fixed number of tokens (capacity $c$ ) based on gating scores.

Deterministic Load Balancing: By design, every expert processes exactly $c$ tokens. This eliminates load imbalance without requiring auxiliary loss functions.
Global View Compatibility: DLMs process the entire sequence non-causally at each denoising step, making them naturally compatible with EC, which requires a global view of all tokens to perform expert-side selection.

B. Timestep-Dependent Capacity Scheduling

Since DLMs iterate through denoising steps with varying masking ratios ( $r$ ), the authors treat expert capacity ( $k$ ) as a dynamic variable $k(r)$ rather than a fixed constant.

Hypothesis: Different masking ratios present qualitatively different learning tasks. The authors investigate whether allocating more compute to specific masking ratios yields higher returns.
Scheduling Strategies: They test various functions $s(r)$ $s (r)$ mapping the masking ratio to capacity, including:
- Linear-Reverse: Allocates more capacity to low masking ratios ( $r \to 0$ ).
- Forward/Intermediate: Allocates more capacity to high masking ratios or intermediate steps.
Constraint: All dynamic schedules are calibrated to match the expected FLOPs of a static baseline to ensure fair comparison.

3. Key Contributions

Superiority of EC Routing: Demonstrated that EC routing is strictly superior to TC for DLMs, providing deterministic load balancing, higher throughput, and faster convergence.
Adaptive Computation Policy: Introduced the concept of scheduling expert capacity based on the denoising timestep.
Mechanistic Insight: Identified that low-mask-ratio steps (where most tokens are visible) have an order-of-magnitude higher learning efficiency. Consequently, concentrating compute on these steps yields the largest marginal return.
Retrofitting Capability: Showed that existing pretrained TC DLMs can be converted to EC by replacing only the router, achieving faster convergence and improved accuracy without retraining from scratch.

4. Experimental Results

Training Efficiency and Throughput

Speedup: EC routing achieved 2.0× faster wall-clock convergence compared to TC variants.
- EC reached a training loss of 3.75 in 10.6 hours.
- TC (dropless, load-balanced) took 20.7 hours.
Throughput: EC achieved 52.1 TFLOP/s/GPU, which is 1.5–2.1× higher than all TC configurations. This is attributed to the elimination of the straggler effect (all GPUs finish processing tokens simultaneously).
Memory Stability: EC maintained perfectly uniform GPU memory usage (std 0.0 GB), whereas TC showed high variance (std 3.6 GB) due to load imbalance.

Timestep-Adaptive Capacity

Best Strategy: The Linear-Reverse scheduler (allocating more capacity to low-mask-ratio steps) consistently achieved the lowest validation perplexity under matched FLOPs.
- Linear-Reverse PPL: 36.5
- Static Baseline PPL: 37.1
- High-Mask-Favoring PPL: 37.5
Mechanism: Analysis of convergence rates ( $\eta_r$ ) revealed that tokens in low-mask-ratio contexts learn ~7–10× faster than those in high-mask-ratio contexts. Allocating more experts to these high-efficiency steps maximizes the marginal return on compute.

Retrofitting Pretrained Models

The authors converted a pretrained TC model (LLaDA-MoE) to EC by swapping the router.
Results: The EC version converged significantly faster during fine-tuning and achieved 1.3–1.5× faster decoding at inference time.
Accuracy: Dynamic EC (Linear-Reverse) achieved the highest average accuracy across four benchmarks (HumanEval, GSM8K, MedQA, etc.), outperforming both static EC and the original TC model.

5. Significance and Conclusion

This paper fundamentally shifts the paradigm for MoE in Diffusion Language Models:

From Fixed to Adaptive: It establishes that computation in DLMs should be treated as an adaptive policy (varying by timestep) rather than a fixed architectural constant.
Structural Efficiency: It proves that Expert-Choice routing is the natural fit for non-autoregressive, parallel generation, solving the load imbalance problem structurally rather than via auxiliary losses.
Practical Impact: The ability to retrofit existing TC models with minimal changes (router replacement) offers a low-cost path to immediate performance gains in throughput, convergence speed, and inference latency for deployed DLMs.

The code is available at https://github.com/zhangshuibai/EC-DLM.