Learning Unmasking Policies for Diffusion Language Models

• URL: http://arxiv.org/abs/2512.09106v2
• Date: Fri, 12 Dec 2025 12:48:32 GMT
• Title: Learning Unmasking Policies for Diffusion Language Models
• Authors: Metod Jazbec, Theo X. Olausson, Louis Béthune, Pierre Ablin, Michael Kirchhof, João Monteiro, Victor Turrisi, Jason Ramapuram, Marco Cuturi,
• Abstract summary: Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks.<n>One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary.<n>In this work, we propose to train sampling procedures using reinforcement learning.
• Score: 33.44995119635116
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.

Related papers

• Learn from Your Mistakes: Self-Correcting Masked Diffusion Models [31.536464269884103]
  Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models.<n>We propose a framework that trains a model to perform both unmasking and correction.<n>We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence.
  arXiv  Detail & Related papers  (2026-02-12T05:17:31Z)
• Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
  This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
  arXiv  Detail & Related papers  (2026-02-02T09:21:45Z)
• Causal Autoregressive Diffusion Language Model [70.7353007255797]
  CARD reformulates the diffusion process within a strictly causal attention mask, enabling dense, per-token supervision in a single forward pass.<n>Our results demonstrate that CARD achieves ARM-level data efficiency while unlocking the latency benefits of parallel generation.
  arXiv  Detail & Related papers  (2026-01-29T17:38:29Z)
• dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning [36.12942468805232]
  Masked diffusion language models offer potential for parallel token generation.<n>Open-source MDLMs decode fewer than 5 tokens per model forward pass.<n>dUltra learns unmasking strategies for efficient parallel decoding.
  arXiv  Detail & Related papers  (2025-12-24T23:31:48Z)
• MARS-Sep: Multimodal-Aligned Reinforced Sound Separation [72.85468563236005]
  MARS-Sep is a reinforcement learning framework for sound separation.<n>It learns a factorized Beta mask policy that is optimized by a clipped trust-region surrogate.<n>Experiments on multiple benchmarks demonstrate consistent gains in Text-, Audio-, and Image-Queried separation.
  arXiv  Detail & Related papers  (2025-10-12T09:05:28Z)
• Reward-Weighted Sampling: Enhancing Non-Autoregressive Characteristics in Masked Diffusion LLMs [44.55861996331439]
  Masked diffusion models (MDMs) offer a promising non-autoregressive alternative for large language modeling.<n>Standard decoding methods for MDMs select tokens independently based on individual token confidences at each diffusion step.<n>We propose Reward-Weighted Sampling (RWS) to provide a principled global signal during the iterative diffusion process.
  arXiv  Detail & Related papers  (2025-08-31T05:48:30Z)
• Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models [13.575063025878208]
  Masked diffusion language models promise fast, non-autoregressive text generation.<n>Existing samplers, which pick tokens to unmask based on model confidence, ignore interactions when unmasking multiple positions in parallel.<n>We propose the Dilated Unmasking Scheduler (DUS), an inference-only, planner-model-free method that partitions sequence positions into non-adjacent dilated groups and unmasked them in parallel.
  arXiv  Detail & Related papers  (2025-06-23T18:49:23Z)
• 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)
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking [17.511240770486452]
  Masked diffusion models (MDMs) have shown competitive performance compared to autoregressive models (ARMs) for language modeling.<n>We introduce EB-Sampler, a drop-in replacement for existing samplers, utilizing an Entropy Bounded unmasking procedure.<n> EB-Sampler accelerates sampling from current state of the art MDMs by roughly 2-3x on standard coding and math reasoning benchmarks without loss in performance.
  arXiv  Detail & Related papers  (2025-05-30T17:52:55Z)
• Generalized Interpolating Discrete Diffusion [65.74168524007484]
  Masked diffusion is a popular choice due to its simplicity and effectiveness.<n>We generalize a new family of general interpolating discrete diffusion (GIDD) which offers greater flexibility in the design of the noising processes.<n>Exploiting GIDD's flexibility, we explore a hybrid approach combining masking and uniform noise, leading to improved sample quality.
  arXiv  Detail & Related papers  (2025-03-06T14:30:55Z)
• Continual-MAE: Adaptive Distribution Masked Autoencoders for Continual Test-Time Adaptation [49.827306773992376]
  Continual Test-Time Adaptation (CTTA) is proposed to migrate a source pre-trained model to continually changing target distributions. Our proposed method attains state-of-the-art performance in both classification and segmentation CTTA tasks.
  arXiv  Detail & Related papers  (2023-12-19T15:34:52Z)
• Improving Self-supervised Pre-training via a Fully-Explored Masked Language Model [57.77981008219654]
  Masked Language Model (MLM) framework has been widely adopted for self-supervised language pre-training. We propose a fully-explored masking strategy, where a text sequence is divided into a certain number of non-overlapping segments.
  arXiv  Detail & Related papers  (2020-10-12T21:28:14Z)

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.