Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2604.02340v2
• Date: Sat, 11 Apr 2026 10:21:32 GMT
• Title: Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models
• Abstract summary: We exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps.<n>We show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs.<n>Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.
• Score: 6.4453302264198165
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Recent advances in masked diffusion language models (MDLMs) narrow the quality gap to autoregressive LMs, but their sampling remains expensive because generation requires many full-sequence denoising passes with a large Transformer and, unlike autoregressive decoding, cannot benefit from KV caching. In this work, we exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps. Across models trained on OpenWebText and LM1B, we show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs with only modest degradation in generative perplexity under both unconditional and prefix-conditional generation, while preserving sample diversity. We support these findings with a step-importance analysis based on loss and KL divergence between small and large models across timesteps, as well as an exhaustive search over coarse step segments, both of which identify the middle of the diffusion trajectory as most sensitive consistently across datasets. Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.

Related papers

• ASSERT: Adaptive Stochastic Sampling for Robust Diffusion Models on Analog Compute-in-Memory Hardware [13.123394042691642]
  Diffusion models achieve strong image generation quality but incur high iterative denoising costs.<n>We show that the early, high-noise denoising stage is substantially more vulnerable than the final refinement stage.<n>We propose ASSERT, a training sampler that uses higher deterministicity early and smoothly denoising.
  arXiv  Detail & Related papers  (2026-09-01T09:13:10Z)
• Continuous Diffusion Scales Competitively with Discrete Diffusion for Language [48.290726603587444]
  We construct RePlaid as a likelihood-based continuous diffusion language model (DLM)<n>We show that RePlaid exhibits a compute gap of only $20times$ compared to autoregressive models.<n>We also offer theoretical insights to understand the advantage of likelihood-based training.
  arXiv  Detail & Related papers  (2026-05-18T15:15:24Z)
• Simple Denoising Diffusion Language Models [32.08084107137018]
  We propose a simplified denoising-based loss for Uniform-state Diffusion Models (USDMs)<n>We introduce a simple modification to our denoising loss with contrastive-inspired negative gradients, which is practical and yield additional improvements in generation quality.
  arXiv  Detail & Related papers  (2025-10-27T02:05:26Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Accelerating Diffusion LLMs via Adaptive Parallel Decoding [60.407727995313074]
  We introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel.<n>APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.
  arXiv  Detail & Related papers  (2025-05-31T06:10:10Z)
• FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion [22.207275433870937]
  Diffusion language models offer parallel token generation and inherent bidirectionality.<n>State-of-the-art diffusion models (e.g., Dream 7B, LLaDA 8B) suffer from slow inference.<n>We introduce Guided Diffusion, a training-free method that uses a lightweight pretrained autoregressive model to supervise token unmasking.
  arXiv  Detail & Related papers  (2025-05-27T17:39:39Z)
• Variational Autoencoding Discrete Diffusion with Enhanced Dimensional Correlations Modeling [48.96034602889216]
  Variencoding Discrete Diffusion (VADD) is a novel framework that enhances discrete diffusion with latent variable modeling.<n>By introducing an auxiliary recognition model, VADD enables stable training via variational lower bounds and amortized inference over the training set.<n> Empirical results on 2D toy data, pixel-level image generation, and text generation demonstrate that VADD consistently outperforms MDM baselines.
  arXiv  Detail & Related papers  (2025-05-23T01:45:47Z)
• One-Step Diffusion Model for Image Motion-Deblurring [85.76149042561507]
  We propose a one-step diffusion model for deblurring (OSDD), a novel framework that reduces the denoising process to a single step.<n>To tackle fidelity loss in diffusion models, we introduce an enhanced variational autoencoder (eVAE), which improves structural restoration.<n>Our method achieves strong performance on both full and no-reference metrics.
  arXiv  Detail & Related papers  (2025-03-09T09:39:57Z)
• Conditional GAN for Enhancing Diffusion Models in Efficient and Authentic Global Gesture Generation from Audios [10.57695963534794]
  Methods based on VAEs are accompanied by issues of local jitter and global instability. We introduce a conditional GAN to capture audio control signals and implicitly match the multimodal denoising distribution between the diffusion and denoising steps.
  arXiv  Detail & Related papers  (2024-10-27T07:25:11Z)
• One More Step: A Versatile Plug-and-Play Module for Rectifying Diffusion Schedule Flaws and Enhancing Low-Frequency Controls [77.42510898755037]
  One More Step (OMS) is a compact network that incorporates an additional simple yet effective step during inference. OMS elevates image fidelity and harmonizes the dichotomy between training and inference, while preserving original model parameters. Once trained, various pre-trained diffusion models with the same latent domain can share the same OMS module.
  arXiv  Detail & Related papers  (2023-11-27T12:02:42Z)

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.