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

• URL: http://arxiv.org/abs/2608.01784v1
• Date: Mon, 03 Aug 2026 06:59:06 GMT
• Title: REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models
• Abstract summary: Mixture-of-experts (MoE) models increase parameter capacity by activating only a small subset of experts for each token.<n>In diffusion language models (DLMs), however, each denoising forward jointly revisits all token positions despite their sharply different refinement demands.<n>We argue that MoE inference in DLMs should therefore be viewed as refinement-aware compute allocation across heterogeneous token refinement states.
• Score: 15.630466998208261
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts [48.294824672909876]
  Inference-time dynamic top-$k$ routing can reduce computation without retraining.<n>We show that reducing the number of activated experts consistently increases the RMS scale and variance of SMoE outputs.<n>We propose Layer-wise Distribution Alignment (LDA), a lightweight inference-time correction that uses layer-wise calibration statistics to align reduced-routing representations with the default configuration.
  arXiv  Detail & Related papers  (2026-09-08T06:23:07Z)
• SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs [10.607329413550223]
  We propose SoftMoE, which replaces discrete routing with a truncated soft top-$k$ LapSum relaxation.<n>We further parameterize the mean number of active experts per layer and impose a global budget constraint.<n>SoftMoE remains fully compatible with autoregressive modeling and performance comparable to or better than sparse MoE on language modeling and downstream tasks.
  arXiv  Detail & Related papers  (2026-06-16T14:05:41Z)
• Tying the Loop -- Tied Expert Layers in Mixture-of-Experts Language Models [50.34585122347149]
  We introduce Expert Tying, an architectural modification that shares expert parameters across consecutive transformer layers.<n>We evaluate this approach across common, state-of-the-art architectures, including OLMoE, Qwen3, and DeepSeek-style MoEs.
  arXiv  Detail & Related papers  (2026-06-15T15:08:09Z)
• Value-and-Structure Alignment for Routing-Consistent Quantization of Mixture-of-Experts Models [6.871645384293096]
  Mixture-of-Experts (MoE) models scale foundation models efficiently by activating only a subset of experts for each token.<n>MoE models are sensitive to routing instability: small quantization-induced perturbations can change the top-$k$ expert selection.<n>We propose Value-and-Structure Routing Alignment for Quantization (VSRAQ) to preserve expert-selection behavior under quantization.
  arXiv  Detail & Related papers  (2026-06-04T04:13:05Z)
• DAG-MoE: From Simple Mixture to Structural Aggregation in Mixture-of-Experts [56.175408382091796]
  We show that replacing the standard weighted-summation aggregation with structural aggregation expands the expert-combination space without altering the experts or router.<n>We propose DAG-MoE, a sparse MoE framework that employs a lightweight module to automatically learn the optimal aggregation structure among the selected experts.
  arXiv  Detail & Related papers  (2026-05-31T07:08:16Z)
• UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models [2.3099144596725565]
  We propose Universally Balanced Sparse Mixture-of-Experts (UB-SMoE) to address expert utilization imbalance and Top-K routing problems.<n>UB-SMoE achieves up to $45.0%$ computational reduction on low-resource clients while improving their performance by $8.7 times$ compared to existing heterogeneous LoRA-rank methods.
  arXiv  Detail & Related papers  (2026-05-15T23:06:59Z)
• Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts [27.40662720929157]
  Mixture-of-Experts (MoE) has become the dominant architecture for scaling large language models.<n>We propose expert upcycling, a method for progressively expanding MoE capacity by increasing the number of experts.<n>In our 7B-13B total parameter experiments, the upcycled model matches the fixed-size baseline on validation loss while saving 32% of GPU hours.
  arXiv  Detail & Related papers  (2026-04-21T05:53:33Z)
• Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models [11.628969213956502]
  We show that expert-choice routing is a better fit for Diffusion language models than token-choice routing.<n>We introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step.<n>Results establish EC routing as a superior paradigm for DLM MoE models.
  arXiv  Detail & Related papers  (2026-04-02T05:01:36Z)
• A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs [64.8510381475827]
  Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently.<n>SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized.<n>We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set.
  arXiv  Detail & Related papers  (2026-02-23T15:11:16Z)
• DSMoE: Matrix-Partitioned Experts with Dynamic Routing for Computation-Efficient Dense LLMs [86.76714527437383]
  This paper proposes DSMoE, a novel approach that achieves sparsification by partitioning pre-trained FFN layers into computational blocks.<n>We implement adaptive expert routing using sigmoid activation and straight-through estimators, enabling tokens to flexibly access different aspects of model knowledge.<n>Experiments on LLaMA models demonstrate that under equivalent computational constraints, DSMoE achieves superior performance compared to existing pruning and MoE approaches.
  arXiv  Detail & Related papers  (2025-02-18T02:37:26Z)
• FactorLLM: Factorizing Knowledge via Mixture of Experts for Large Language Models [50.331708897857574]
  We introduce FactorLLM, a novel approach that decomposes well-trained dense FFNs into sparse sub-networks without requiring any further modifications. FactorLLM achieves comparable performance to the source model securing up to 85% model performance while obtaining over a 30% increase in inference speed.
  arXiv  Detail & Related papers  (2024-08-15T16:45:16Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.