PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Unmasking Policies for Diffusion Language Models

Metod Jazbec, Theo X. Olausson, Louis Béthune, Pierre Ablin, Michael Kirchhof, João Monteiro, Victor Turrisi, Jason Ramapuram, Marco Cuturi

arXiv:2512.09106v4cs.LG

TL;DR

Heuristic unmasking strategies for diffusion language models are difficult to tune and can degrade outside semi-autoregressive generation. This paper trains reinforcement-learning unmasking policies, which match or exceed heuristics across evaluated settings and can transfer across models.

  • Problem

    Handcrafted unmasking heuristics are difficult to tune and sensitive to sampling configuration, especially outside semi-autoregressive generation.

  • Method

    The paper formalizes dLLM sampling as a Markov decision process and trains a lightweight transformer-based unmasking policy with reinforcement learning.

  • Results

    Learned policies match or exceed recently proposed sampling heuristics across the evaluated diffusion language model settings.

  • Takeaways & Limitations

    Reinforcement learning can automate the discovery of scalable and robust unmasking strategies for diffusion language models.

  • Takeaways & Limitations

    During early training, the speed penalty can favor faster samples even when they are incorrect.

Abstract

from arXiv · show

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 critical design aspect of dLLMs is the sampling procedure that selects which tokens to unmask at each diffusion step. Indeed, recent work has found that heuristic strategies such as confidence thresholding improve both sample 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 block 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 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 (block) generation, while outperforming them in the full-diffusion setting.

1 Introduction

Diffusion language models offer parallel generation and competitive performance, but their sampling heuristics are difficult to tune and work best with semi-autoregressive generation. This work formulates unmasking as a Markov decision process and learns lightweight transformer-based policies with reinforcement learning, matching heuristic samplers such as Fast-dLLM.

  • Background: Masked diffusion language models have demonstrated impressive performance, with LLaDA and Dream matching similarly sized autoregressive language models.
  • Motivation: Diffusion language models can generate multiple tokens in parallel from a fully masked sequence, offering potential inference-speed advantages over left-to-right autoregressive generation.Masked diffusion models iteratively unmask a fixed number of tokens at randomly sampled positions during each sampling step.
  • Motivation: Although increasingly sophisticated sampling heuristics improve the state of the art, they are difficult to tune and appear most effective with small blocks unmasked sequentially.This semi-autoregressive setup imposes sequential unmasking over small token blocks.
  • Contribution: The paper proposes reinforcement-learned unmasking policies that formalize diffusion sampling as a Markov decision process and use a lightweight transformer-based policy.The policy is trained through an RL pipeline based on group relative policy optimization.
  • Empirical findings: Learned policies match heuristic samplers such as Fast-dLLM in experiments while addressing challenges faced by heuristics outside semi-autoregressive generation.

2 Background

Masked diffusion models generate sequences by iteratively predicting token distributions and selecting still-masked positions to unmask. Sampling strategy is therefore central to quality and efficiency, while confidence-based heuristics can be sensitive to configuration choices.

  • Training: MDM training uses a BERT-style masked predictor trained to reverse a forward process that independently masks tokens according to the diffusion timestep.
  • Generation: At each sampling step, an MDM predicts distributions for all positions and selects a subset of still-masked tokens to unmask.Generation begins from an all-masked sequence and produces partially denoised sequences iteratively.
  • Generation: Different unmasking strategies affect both generation quality and efficiency, motivating the development of improved sampling techniques.
  • Confidence-based sampling: Confidence-based heuristics select tokens using predicted uncertainty, including fixed-count high-confidence unmasking and threshold-based variable-count unmasking.High-confidence unmasking chooses K highest-confidence tokens, whereas Fast-dLLM compares confidences with a threshold λ.
  • Limitations: Handcrafted confidence samplers require design choices such as confidence measures or thresholds and can be highly sensitive to the sampling configuration.The paper specifically highlights sensitivity in settings using semi-autoregressive block generation.

3 Learning Unmasking Policies

The paper formulates dLLM sampling as a reverse-time MDP and learns a lightweight confidence-based unmasking policy with reinforcement learning. The policy uses Bernoulli actions and a reward balancing correctness with sampling efficiency.

  • 3.1 Markov decision process: The sampler models partially masked generation as a reverse-time MDP whose actions select token positions to unmask, with rewards based on correctness and efficiency.States contain the prompt and current generation; actions are binary unmasking vectors, and the reward is provided at the final generation step.
  • 3.2 Sampling policy: A lightweight transformer maps token confidences, mask indicators, and the timestep to per-position unmasking logits, adding less than 0.01% of the pretrained dLLM size.The policy operates on readily available token confidences rather than hidden states, minimizing computational overhead during sampling.
  • 3.2 Sampling policy: The method samples independent Bernoulli unmasking actions with a closed-form likelihood, favoring this formulation over comparable dynamic Plackett–Luce sampling for simplicity.A test-time temperature controls policy sharpness and can trade off accuracy against efficiency.
  • 3.3 Reinforcement-learning training: The policy is trained with GRPO using a multiplicative reward that combines task correctness with a step-count penalty controlled by α.An additive penalty caused reward hacking because incorrect but faster samples could receive positive advantages early in training.

4 Experiments

Across generation settings, learned RL sampling policies match strong heuristics in semi-autoregressive generation and outperform them in full diffusion, while expert steering further closes the performance gap. Experiments also show transfer limitations, controllability trade-offs, and benefits from confidence-based policy inputs and stochastic decoding.

  • Short-block generation: With BL = 32, learned policies outperform random and high-confidence sampling while matching Fast-dLLM, suggesting near-optimal semi-autoregressive performance in the mid-to-high NFE range.The comparison is reported for Figures 4a and 4c.
  • Test-time control: Varying β provides a smooth test-time control over the accuracy-efficiency frontier without retraining, with policies matching Fast-dLLM at mid-to-high NFEs and outperforming it at lower NFEs.Higher β increases the average number of tokens unmasked per step, while lower β decreases it.
  • Long-block generation: In full-diffusion generation with BL = 256, RL policies experience the smallest performance decline and achieve the best overall performance, reaching ∼50% GSM8K accuracy at ∼12 NFEs.Heuristic methods remain at ≤30% accuracy regardless of semi-autoregressive use.
  • Expert steering: Expert steering improves exploration, enabling RL to discover policies that almost close the performance gap to the best setting.The method uses semi-autoregressive samples generated by Fast-dLLM during training.
  • Transfer: RL policies transfer across generation lengths and models through confidence-only inputs, but transfer degrades on coding tasks and high-α policies can overfit LLaDA-specific confidence patterns.Coding-specific training narrows the transfer gap, while sequence-length transfer remains effective without retraining.

5 Related work

Prior work develops heuristic samplers and reinforcement-learning methods for diffusion language models, alongside orthogonal efficiency techniques and adaptive-computation policies. This work positions itself as learning end-to-end, adaptive sampling policies for dLLMs via reinforcement learning.

  • Heuristic samplers for dLLMs: Heuristic dLLM samplers improve decoding without training, using spatial or temporal information, alternative confidence measures, and explicit token-dependency modeling.Related heuristics also address remasking and dynamically adjusting generation length.
  • Reinforcement learning post-training for dLLMs: Prior reinforcement-learning work post-trains dLLMs, but some methods unmask fixed token counts while this work dynamically adapts unmasking through a Bernoulli formulation.Seed Diffusion identifies computation-aware reinforcement learning as important for efficient closed-source coding dLLMs.
  • Reinforcement learning for adaptive computation: This work connects dLLM sampling to adaptive computation and is presented as the first to learn adaptive sampling policies with reinforcement learning for dLLMs.Related approaches learn input-dependent computation through stochastic gating, dynamic block skipping, and reinforcement-learning early-exit decisions.
  • Reinforcement learning for adaptive computation: Concurrent DiFFPO predicts adaptive thresholds used like Fast-dLLM’s fixed thresholds, whereas this work learns its sampling policy end-to-end.Other efficiency efforts include KV caching, speculative decoding, separate decoder modules, and diffusion forcing.

6 Conclusion … B.6 Impact of policy temperature τπ

The paper concludes that reinforcement-learned unmasking policies can match or exceed sampling heuristics, while highlighting per-α training costs and test-time temperature control as important deployment considerations. The appendices provide the sampling algorithm and extensive robustness, transfer, qualitative, and implementation analyses.

  • 6 Conclusion: Learned unmasking policies match or exceed recent sampling heuristics and support automated discovery of scalable, robust sampling mechanisms.The experiments provide quantitative and qualitative insight into training and deploying these policies across reasoning and code domains.
  • Appendix: The appendix documents the policy-sampling algorithm and additional experimental, qualitative, trajectory, methodological, background, implementation, architectural, and tabular analyses.Appendix B covers replicated figures, temperature effects, transfer, decoding, and policy-design ablations, while later appendices address sampling alternatives, expert steering, diffusion background, implementation, architecture, and tabulated results.
  • B.1 Figure 1 replicated for {LLaDA-8B-Instruct, Dream-7b-Instruct} × {GSM8K, MATH-500}: Across LLaDA and Dream on GSM8K and MATH, confidence heuristics perform well with semi-AR generation but degrade significantly in full diffusion.The replication compares BL = 32 semi-AR and BL = 256 full-diffusion settings.
  • B.2 Figure 4 replicated using additional baselines: With additional baselines, the learned policies mainly match the best baselines under semi-AR generation and outperform them in full diffusion.The comparison uses BL = 32 for semi-AR and BL = 256 for full diffusion.
  • B.3 Figure 4 replicated using wall-clock time as the efficiency measure: Measuring efficiency by wall-clock time instead of NFEs produces minimal or nonexistent changes, indicating negligible policy overhead from the 300K-parameter unmasking model versus the 8B-parameter base dLLM.Experiments ran on A100 GPUs.
  • B.4 Figure 4 replicated for Dream-7b-Instruct: Dream experiments evaluate policies across semi-AR and full-diffusion regimes, using α ∈ {10, 3, 1, 0.3, 0} with temperature settings tailored to block size.They use τπ = 0.5 for BL = 32 and τπ = 1 for BL = 256.
  • B.5 Controllability of learning unmasking policies with RL (via α): Varying α reveals non-monotonic NFE behavior and training-seed variation, so policies with the same α can differ in accuracy and speed.The denser α grid shows convergence near either the α = 3.0 or α = 10.0 policy for α ≥ 4.0.
  • B.6 Impact of policy temperature τπ: Policy temperature can create compute–performance trade-offs; τπ = 0.5 is optimal for semi-AR generation, whereas τπ = 1 performs best in full diffusion.For α = 3, 0.3, and 0, the study varies τπ ∈ {1.5, 1.0, 0.5}; α = 0 with BL = 32 shows an effective trade-off.

B.7 Model transfer results · B.8 Sequence-length transferability

The learned unmasking policies transfer from LLaDA to Dream with performance comparable to Dream-specific training. Policies trained at sequence length 256 also transfer to length 512 with nearly unchanged performance, whereas heuristics degrade further.

  • B.7 Model transfer results: The model-transfer evaluation compares Random, High Confidence, Fast-dLLM, and Ours.
  • B.7 Model transfer results: Transferred LLaDA-trained policies achieve results comparable to policies trained specifically on Dream.
  • B.7 Model transfer results: The transfer experiment evaluates LLaDA-trained policies on Dream with τπ = 0.5.
  • B.7 Model transfer results: The α = 10 policy is plotted separately to avoid misleading interpolation to α = 3.
  • B.8 Sequence-length transferability: Learned policies trained with BL = L = 256 retain almost identical performance when evaluated at twice the sequence length, BL = L = 512.
  • B.8 Sequence-length transferability: At the doubled sequence length, heuristic methods degrade further relative to L = 256, while the learned policies do not.

B.9 Non-greedy decoding results … C.2 BL32-fast

The paper evaluates learned unmasking policies across decoding regimes and input ablations, finding competitive or improved accuracy-efficiency behavior while revealing distinct token-ordering and compute-allocation strategies.

  • B.9 Non-greedy decoding results: Figure 19 compares policy sampling with Fast-dLLM under non-greedy BL = 32 decoding, reporting better pass@k scaling for policy sampling.The passage attributes this scaling to additional stochasticity from Bernoulli sampling of unmasking decisions.
  • B.10 Additional policy input ablations: Figure 20 tests removing time, mask, or both inputs from a policy trained with α = 1 on LLaDA with BL = 32, observing reduced performance in nearly all cases.Zeroing time alone is the exception on GSM8K; the figure also reports changes in average NFEs.
  • • BL256 (: The BL256 analysis compares Fast-dLLM, policy sampling, and expert steering in the full-diffusion regime.The supplied passages identify this as full diffusion with BL = L = 256 but do not provide the corresponding outcome values.
  • • BL256 (: Unmasking order is defined by each token’s first unmasking timestep, with simultaneous unmasking represented as a partial order.The analysis reverses the time axis for presentation, using t = 1, 2, …, T.
  • C.1 BL32-slow: In BL32-slow, the learned policy nearly matches Fast-dLLM in both accuracy and efficiency while using a more dispersed within-block unmasking strategy.The policy simultaneously unmasks adjacent tokens for around 20% of tokens, compared with around 65% for Fast-dLLM.
  • C.1 BL32-slow: In BL32-slow, the policy distributes computation almost uniformly across blocks, using approximately 10 NFEs per block, whereas Fast-dLLM concentrates computation earlier.Fast-dLLM uses approximately 17 NFEs in the first block and decreases to approximately 4 in the final block.
  • C.2 BL32-fast: Policy sampling achieves approximately 38% accuracy versus approximately 18% for Fast-dLLM in the BL32-fast setting.Both methods are very fast, but policy sampling substantially improves accuracy by slowing down near the final block and generating the numerical answer later.

C.3 BL256 · D Generation Trajectories Visualizations · E Dynamic Plackett-Luce Sampling

In full-diffusion generation, the learned policy follows substantially different unmasking orders from Fast-dLLM and avoids confidence-driven padding-token errors, while trajectory visualizations illustrate these differences across settings. The appendix also specifies Dynamic Plackett-Luce Sampling, which supports variable-length unmasking through a STOP token and an efficient Gumbel-argsort implementation.

  • C.3 BL256: Fast-dLLM often unmasks tokens right-to-left, apparently because padding-token overfitting makes confidence-based methods prioritize padding early.The appendix attributes this behavior to LLaDA’s inclusion of padding tokens in the SFT loss, which contributes to poor confidence-based performance.
  • C.3 BL256: Policy sampling shows little spatial structure in unmasking order, suggesting it can disregard inflated padding-token confidence across semi-autoregressive and full-diffusion settings.This interpretation is presented as evidence that the policy learns robustness without explicit supervision.
  • C.3 BL256: In full diffusion, the strategies use comparable NFEs but differ significantly in accuracy and unmasking order.Fast-dLLM and policy sampling have Spearman correlations of 0.26 ± 0.21, while Fast-dLLM and expert-steered policy sampling have −0.24 ± 0.51.
  • D Generation Trajectories Visualizations: The trajectory visualizations compare token identities and unmasking times across semi-autoregressive and full-diffusion generation for policy, Fast-dLLM, and expert-steered policy sampling.Cells are color-coded by unmasking time, with blue indicating earlier and red indicating later generation, across L = 256 positions.
  • E Dynamic Plackett-Luce Sampling: Dynamic Plackett-Luce Sampling replaces fixed-length selection with a variable-length ordered unmasking sequence, allowing between 1 and L selected tokens.It interprets policy logits as unnormalized utilities and samples indices without replacement, with selection probability proportional to exponentiated utility.
  • E Dynamic Plackett-Luce Sampling: DPLS introduces a STOP token with fixed utility 0 and repeatedly samples tokens until STOP is selected, thereby determining the unmasking length dynamically.Previously selected indices are masked from subsequent sampling, and non-STOP selections are added to the ordered set.
  • E Dynamic Plackett-Luce Sampling: Although the variable-length procedure is not directly GPU-friendly, DPLS can be implemented efficiently as a Gumbel-argsort with post-STOP masking and likelihood marginalization.This preserves efficient likelihood computation by marginalizing over actions sampled after STOP.

F Expert Steering · G Extended Background

Expert Steering improves confidence-policy training by mixing deterministic expert samples into GRPO, addressing the scarcity of autoregressive-like rollouts. The background contrasts absorbing, uniform-state, and continuous-input diffusion, motivating the absorbing formulation used here.

  • F Expert Steering: Full-diffusion policies trained naively beat heuristic methods but underperform policies trained in the semi-AR setting.This comparison is reported with reference to Figure 4.
  • F Expert Steering: The failure of naive full-diffusion training is attributed to the vanishingly small chance of encountering autoregressive-like rollouts by chance.Semi-AR decoding lies within the function class representable by the learned policy, so the issue is exploration rather than representational capacity.
  • F Expert Steering: Expert Steering replaces the learned policy with a train-time mixture of that policy and deterministic expert policies.The method is designed to expose training to expert-guided behavior while retaining the learned policy for deployment.
  • F Expert Steering: The method augments each GRPO group with expert samples and uses the mixture policy for likelihood ratios to avoid instability.This prevents near-zero likelihoods under the learned policy for samples drawn from deterministic Dirac experts.
  • F Expert Steering: Experiments use one Fast-dLLM expert with λ = 0.9 and BL = 32 per group, biasing the policy toward the heuristic when it performs worse.If the learned policy is better, the expert sample receives negative advantage and becomes less likely.
  • G Extended Background: Masked diffusion has monotonic reverse dynamics from M to tokens, whereas uniform-state diffusion permits token↔token substitutions and broader exploration.The masked process has a point-mass stationary distribution at M; the uniform process has a uniform stationary distribution.
  • G Extended Background: The work adopts absorbing diffusion, whose objective reduces to weighted MDM and naturally supports the unmask-only sampling policies studied here.Continuous-time analyses describe a schedule-invariant weighted cross-entropy objective without changing the prediction target.
  • G Extended Background: Continuous-input diffusion operates on token embeddings with continuous-time and continuous-input trajectories, using ODE/SDE samplers but relaxing exact discreteness.It introduces an embedding decoding interface and can incorporate classifier-free guidance.

H Training and Policy Network Configuration

This section describes the training and policy-network configuration used in the main experiments.

  • H Training and Policy Network Configuration: Table 1 summarizes the training and policy configuration for the main experiments.

I Policy Architecture Diagram

The policy architecture diagram defines shorthand for confidence, masking, timesteps, top-k selection, and hidden dimensionality, while indicating associated tensor shapes.

  • I Policy Architecture Diagram: The diagram abbreviates confidence as c, mask as m, timestep as t, top-k as k, and hidden dimensionality as d_hidden.
  • I Policy Architecture Diagram: The architecture notation links top-k selection and hidden dimensionality to the policy’s implementation vocabulary.
  • I Policy Architecture Diagram: It labels the principal dimensions using L, k, and singleton dimensions.

J Tabular Reference for Main LLaDA Experiments

This appendix provides tabular reference results for the main LLaDA experiments, enabling easier comparison of exact values. The tables cover GSM8K and MATH at block lengths 32 and 256, with values corresponding directly to plotted points.

  • J Tabular Reference for Main LLaDA Experiments: The appendix tabulates the main LLaDA experiments so future work can compare exact numerical results more easily.The reported values correspond directly to the plotted points.
  • J Tabular Reference for Main LLaDA Experiments: Tables 2 and 3 report GSM8K and MATH results at BL = 32.Table 2 covers GSM8K, while Table 3 covers MATH.
  • J Tabular Reference for Main LLaDA Experiments: Tables 4 and 5 report GSM8K and MATH results at BL = 256.Table 4 covers GSM8K, while Table 5 covers MATH.
Loading 2512.09106v4…