REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models
2026-08-03 • Artificial Intelligence
Artificial IntelligenceComputation and Language
AI summaryⓘ
The authors study how Mixture-of-experts (MoE) models allocate computation in diffusion language models (DLMs), where all tokens are refined together but have different needs for expert help. They find that the usual method assigns experts uniformly, which doesn't match the tokens' varying refinement demands. To fix this, they propose REFLEX, a way to rearrange expert usage during inference without retraining, prioritizing tokens based on how much refinement they need. Their tests show REFLEX reduces computation by about 15% while maintaining or improving output quality compared to standard methods.
Mixture-of-experts (MoE)Diffusion language models (DLMs)Conditional computationToken refinementExpert allocationRoutingInference optimizationAutoregressive modelsREFLEX methodComputation-quality trade-off
Authors
Xiang Xia, Cheng Yan, Yiming Zhang, Jiazheng Liu, Hongyu Zhang, Wuyang Zhang
Abstract
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 (\textbf{RE}finement-aware \textbf{FLEX}ible 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.