Learning Unmasking Policies for Diffusion Language Models

AI summary

Overview

Research area: Machine learning — discrete/masked diffusion language models (dLLMs), inference-time sampling, and reinforcement learning for policy learning.

Technical level: Advanced. The paper assumes familiarity with masked diffusion, Markov decision processes, and policy-gradient RL (GRPO).

Scope: The paper proposes training a small transformer policy with reinforcement learning to decide which tokens a diffusion language model should unmask at each sampling step, replacing hand-tuned confidence heuristics.

What This Paper Is About

Diffusion language models generate text by starting from a fully masked sequence and unmasking tokens over a series of steps, and the rule that decides which tokens to unmask strongly affects both answer quality and speed. Existing approaches are handcrafted heuristics, such as unmasking every token whose confidence exceeds a threshold (Fast-dLLM), but these require manual tuning, degrade with larger block sizes, and often depend on semi-autoregressive (block) generation. The authors recast unmasking as a sequential decision-making problem and train a lightweight policy to solve it directly with reinforcement learning.

Key Contributions

1. An MDP formalization of dLLM sampling. The authors define masked diffusion sampling as a Markov decision process where the pretrained dLLM is the environment, the action is a binary vector of unmasking decisions over the sequence, and reward is given only at the final generation step.

2. A lightweight confidence-based policy. The policy is a single-layer transformer that maps token confidences, a mask indicator vector, and the timestep to per-position unmasking logits, sampled through a Bernoulli likelihood available in closed form. Its size is less than 0.01% of the pretrained dLLMs used in the experiments.

3. An RL training pipeline based on GRPO. Group relative policy optimization is used with greedy dLLM decoding (τ = 0) so that variation within a group of rollouts comes only from unmasking actions, with a multiplicative reward that combines task correctness with a computational penalty controlled by α.

4. Experiments across regimes and transfer settings. Learned policies match Fast-dLLM-style heuristics in semi-autoregressive generation and outperform them in the full-diffusion setting, with additional studies of model transfer, domain transfer, sequence-length transfer, and non-greedy decoding.

Main Findings

• Matching heuristics in semi-autoregressive generation. With block length BL = 32 on LLaDA-8B-Instruct, learned policies exceed the random baseline and high-confidence unmasking, and match Fast-dLLM on GSM8K and MATH-500. The authors suggest Fast-dLLM may be near-optimal in this regime.

• Outperforming heuristics in full diffusion. With BL = L = 256, all methods degrade relative to semi-AR, but the learned policies show the smallest decline and the best overall performance, reaching roughly 50% accuracy at about 12 NFEs on GSM8K, compared to at most 30% for the heuristic methods regardless of semi-AR use.

• Different unmasking behavior than Fast-dLLM. The policy unmasks adjacent tokens much less frequently and distributes compute more uniformly across blocks, whereas Fast-dLLM allocates more compute to earlier blocks. The fast (α = 10) policy instead spends most compute in the final block when generating numerical answers.

• A test-time efficiency knob. Scaling the Bernoulli probabilities as u ∼ Ber(min{1, β·s}) for β > 0 (following Chen et al., 2025a) gives a smoother accuracy-efficiency traversal than changing α at training time. On MATH-500 at around 25 NFEs, the policy reaches 20% accuracy versus Fast-dLLM's 10%.

• α controls speed but is hard to tune. Policies trained with higher α are faster but less accurate, and α = 10.0 showed greater training instability. Varying α traverses the Pareto frontier less smoothly than varying Fast-dLLM's threshold λ; α = 3.0 behaves nearly identically to α = 1.0, and extending the grid in Appendix B.5 (α ∈ {10.0, 9.0, … 1.0, 0.3, 0.0}) found α ≥ 4.0 converges to either an α = 3.0-like or α = 10.0-like policy with nothing in between.

• Expert steering closes much of the full-diffusion gap. Steering exploration with semi-AR samples from Fast-dLLM lets RL nearly close the gap to the best semi-AR accuracy at mid-to-high NFEs (about 80% on GSM8K and about 35% on MATH), while mostly retaining low-NFE performance — but it introduces significant training instability and reduces controllability through α.

• Reversal of generation order matters. In full diffusion (L = 256), Fast-dLLM unmasks right-to-left, which the authors attribute to LLaDA's corrupted confidence on padding tokens; the expert-steering policy learns to unmask earlier positions first on average, which likely explains its superior performance.

• Partial transfer across models, weaker transfer across domains. Policies trained on LLaDA nearly match Fast-dLLM on Dream and perform similarly to policies trained on Dream directly, with α = 10 as the exception. Transfer from math to code (HumanEval, MBPP) is weak, especially on HumanEval; training a new policy on KodCode-RL-10K narrows the HumanEval gap and improves low-NFE performance on MBPP.

• Sequence-length and non-greedy generalization. Policies trained at L = 256 retain similar performance at L = 512 while baselines degrade further. With τ > 0, policy sampling yields higher pass@k than Fast-dLLM, with the gap widening in k (+0.98% for k = 1 versus +2.56% for k = 32 on average across datasets), and also dominates under self-consistency or an external reward model selection.

• Multiplicative reward avoids reward hacking. The proposed multiplicative reward, r · (1 − (T − t)/T)^α, is compared to an additive alternative; the additive version collapses to a very fast but often incorrect policy that unmasks everything at once, with sampling steps collapsing to the minimum for all inputs.

Methodology in Plain English

The authors treat the diffusion language model as a fixed environment and learn a separate, tiny decision-maker on top of it.

At each sampling step, the dLLM produces a probability distribution over the vocabulary at every still-masked position. From these distributions the authors extract a confidence value per position. Their policy — a single-layer transformer with adaptive layer normalization for conditioning — reads the confidence vector, a binary vector marking which positions are still masked, and the timestep, and outputs an unmasking score per position.

Those scores are turned into independent Bernoulli probabilities, so each masked position is independently chosen for unmasking. This design has a convenient property: the probability of any particular unmasking action can be written in closed form, so no approximation is needed to compute the RL objective.

Training uses GRPO. For each prompt, the model generates a group of trajectories; the dLLM temperature is fixed at 0 (greedy) so that differences within the group come only from the unmasking choices. Each trajectory receives a reward only when generation finishes; that reward is the correctness of the answer multiplied by a factor that penalizes the fraction of the horizon used, raising the penalty's sharpness with the hyperparameter α. Advantage is the reward minus the group mean, without standard-deviation normalization, and the reward signal is propagated back to all earlier timesteps in the trajectory. The likelihood ratio is clipped as in PPO-style objectives, and the KL term is removed because the policy is trained from scratch.

The policies were trained on LLaDA-8B-Instruct across five values of α (10, 3, 1, 0.3, 0), semi-autoregressively at block length 32, on one epoch of a mixture sampled proportionally from the GSM8K and MATH training sets — roughly 15,000 training samples — then evaluated on the test sets of GSM8K and MATH.

Why This Matters

Impact on research. The work reframes dLLM decoding from hyperparameter tuning into a learnable control problem, and shows that a policy with under 0.01% of the base model's parameters can be trained and transferred. It also connects policy sampling to the reported diversity collapse of deterministic confidence-based samplers, since stochastic Bernoulli unmasking produces higher pass@k rates and better self-consistency and reward-model selection.

Real-world applications (drawn from the domains the paper evaluates, plus the general setting):

• Mathematical reasoning assistants that must answer quickly under a fixed compute budget, as tested on GSM8K and MATH / MATH-500.
• Code generation and completion, as tested on HumanEval and MBPP, where the authors train a dedicated policy on KodCode-RL-10K.
• Any deployment of open dLLMs such as LLaDA-8B-Instruct or Dream-7B-Instruct where token throughput per model call determines serving cost.
• Batch or long-context generation where sequence length varies (the paper demonstrates transfer from L = 256 to L = 512), removing the need to retrain a sampler per configuration.

Industry relevance. The method targets inference cost rather than model capability, keeps the base dLLM frozen, and adds negligible compute overhead, which makes it attractive for serving systems. The released code is at https://github.com/apple/ml-rl-dllm.

Future Directions

• Stabilizing expert steering. Expert steering closes much of the full-diffusion accuracy gap but introduces significant training instability and causes multiple values of α to collapse to near-identical policies; the authors leave stabilizing it for future work.
• Better controllability of the accuracy-efficiency trade-off. Training-time α traverses the Pareto frontier unevenly, and the paper's fix is a test-time scaling parameter β; a more principled training-time control remains open.
• Closing the gap between BL = L and BL = 32 policies. Policies trained in full diffusion still underperform semi-AR-trained policies at mid-to-high NFEs, which the authors read as evidence that the policies remain far from optimal and that exploration is insufficient.
• Broader transfer. Policies trained on mathematical data do not fully transfer to code, especially on HumanEval, motivating more diverse training mixtures and further study of where model, domain, and sequence-length transfer break down.

Target Audience

Researchers and engineers working on diffusion language models, inference-time efficiency, or RL for language model post-training will benefit most. It is also relevant to practitioners deploying open dLLMs who want faster generation without changing the underlying model, and to readers interested in whether handcrafted decoding heuristics can be replaced by learned control policies.

Authors’ 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 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.