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

arXiv: https://arxiv.org/abs/2608.01784 · HTML: https://mlsanity.com/paper/2608.01784/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 0 of its 35 references that are in the map.
- First posted on arXiv 3 Aug 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)
- Mixture-of-experts / [Vision mixture-of-experts](https://mlsanity.com/moe/vision-mixture-of-experts/)

## Abstract (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.
