Mask-Aware Policy Gradients for Diffusion Language Models

AI summary

Overview

Research area: Reinforcement learning post-training for Masked Diffusion Language Models (MDLMs), applied to mathematical reasoning and code generation.

Technical level: Advanced (requires familiarity with policy gradient methods, diffusion language models, and likelihood estimation).

Scope: This paper derives a policy gradient estimator for masked diffusion language models that optimizes both token predictions and the order in which masked positions are unmasked, and evaluates it on four reasoning and coding benchmarks.

What This Paper Is About

Reinforcement learning works well for autoregressive language models because their likelihood is tractable, but masked diffusion language models generate text by iteratively unmasking positions from a fully masked sequence, and their sequence log-likelihood requires marginalizing over all possible generation trajectories, which is intractable. Existing reinforcement learning methods for MDLMs approximate the likelihood by modeling only which tokens get predicted, ignoring the second decision the model makes at every step: which positions to keep unmasked and which to return to [MASK]. The goal of this paper is to include that position-selection decision in the policy gradient so that the unmasking order itself becomes something the model is trained to get right.

Key Contributions

1. A two-stage action MDP formulation of MDLM generation. The authors formalize each denoising step as first predicting tokens at all masked positions and then selecting which of those positions to unmask, and show that the resulting policy gradient decomposes into a token term and a masking (position) term.

2. A probabilistic replacement for greedy top-K remasking. Greedy top-K selection is non-differentiable, so the authors sample the unmasking set from a Plackett–Luce distribution over the model's own token log-likelihoods, with a temperature parameter that recovers greedy selection as it goes to zero. This yields a differentiable distribution over position subsets.

3. Position log-probabilities computed for free. The masking term is derived as a softmax over unmasking scores computed from the model's own logits, requiring no additional parameters, no architectural changes, and no forward passes beyond those already needed for the token likelihood.

4. A theoretical demonstration that token-only gradients are insufficient. In Appendix G the authors give a finite-horizon MDLM counterexample with a linear policy in which the token-only gradient is identically zero along a direction that strictly increases the true objective, while the full position-aware gradient is not.

5. State-of-the-art empirical results across two mathematical reasoning and two code generation benchmarks, plus training-efficiency and block-size analyses.

Main Findings

• Best result in every configuration tested. Using LLaDA-8B-Instruct as the base model, the method achieves 81.0, 85.9, and 87.1 on GSM8K at generation lengths 128, 256, and 512; 37.4, 42.2, and 44.2 on MATH500; 33.2, 43.1, and 44.0 on HumanEval; and 47.1, 52.8, and 53.4 on MBPP.

• Gains of up to 4.0% on MATH500, 2.9% on HumanEval, and 2.5% on GSM8K and MBPP over prior methods. At MATH500 the improvement is +4.0 at generation length 128 and +2.4 at length 512.

• The masking term is what drives the gains. The comparison against StepMerge is the key ablation: both methods share the same trajectory-based likelihood framework and the same K subsampling, differing only in whether the masking log-probability term is present. Removing it costs 2–4 points consistently across all configurations.

• Improvements over evidence-bound methods are consistent. Against GDPO and SPG w/ Mixture, the method shows gains particularly on MATH500 and the code generation benchmarks. The paper notes that SPG does not report code generation results, so the authors evaluated SPG's official code themselves on HumanEval and MBPP, and that SPG's originally reported GSM8K number at length 256 was 86.1 rather than the 83.9 reproduced in the table.

• Trained faster despite lower per-step throughput. On GSM8K with 8 H100 GPUs at generation length 128, the method runs 290 steps/h versus SPG's 360 steps/h, but converges in 6,000 steps versus 6,500, reaching a final accuracy of 81.0 versus 78.5. It matches SPG's final accuracy in roughly 15 hours versus 18 hours, and reaches 74% accuracy 1.9× faster. It also reaches 79% and 80% accuracy, which SPG never reaches. Across three random seeds final accuracy varies by at most 0.2 points.

• The advantage grows with inference block size. At generation length 256 with confidence-based unmasking, comparing block sizes 32, 64, and full sequence: against SPG the margin grows from +2.0 to +4.3 on GSM8K and from +2.2 to +3.5 on MATH500; against StepMerge it grows from +2.6 to +4.4 on GSM8K and from +3.1 to +3.7 on MATH500. The authors attribute this to position selection mattering more as the diffusion window expands.

• Generalization beyond the main setting is reported in appendices. The method is reported to transfer to a second base model, Dream-7B (Appendix C), and to two planning tasks, Sudoku and Countdown (Appendix B).

Methodology in Plain English

The researchers start from the observation that a diffusion language model makes two choices per denoising step, not one. First it fills in guessed tokens at every masked slot. Then it decides which of those guesses to keep and which to throw back to the masked state. Only the first choice had been included in prior reinforcement learning objectives.

To make the second choice trainable, they replace the usual deterministic "keep the K most confident positions" rule with a sampling rule. Positions are drawn without replacement, with probability proportional to the exponentiated log-likelihood of the token predicted there, divided by a temperature. This is a Plackett-Luce distribution, and its log-probability is differentiable and computable from the same forward passes already used for the token term.

With that in place, the trajectory likelihood becomes a product of three factors per step (token predictions, the unmasking selection, and the deterministic remasking operator). The remasking operator has no parameters, so it drops out of the gradient. What remains is a reward-weighted sum of a token gradient and an unmasking gradient, which can be plugged into any standard policy gradient algorithm. The authors use GSPO with separately clipped importance ratios for the two terms.

Experimentally they fine-tune LLaDA-8B-Instruct with LoRA (rank 128, scaling factor 64, dropout 0.05) on all attention and MLP projections, keeping base weights frozen in 4-bit NF4 quantization, using Flash Attention 2 and bfloat16 throughout. Rollouts use a sequence length of 256 tokens, 128 diffusion steps, and block size 32, with a token sampling temperature of 0.9 and a remasking temperature of 0.5. For both likelihood terms they use the StepMerge approximation with N=32 segments and K=12 subsampled boundaries. Evaluation is zero-shot on GSM8K and MATH500 for math (following the train-test splits and rewards of d1, wd1, and SPG) and HumanEval and MBPP for code (training on KodCode-Light-RL-10K, following GDPO).

Why This Matters

Impact on research. The paper argues that the standard way of adapting reinforcement learning to diffusion language models was solving an incomplete version of the problem: the trajectory log-likelihood had been modeled as a sequence of token predictions, when it is genuinely a sequence of token predictions and position selections. Showing that the extra term is free to compute, and giving a theoretical counterexample for why omitting it can zero out useful gradient directions, reframes how likelihood estimators for MDLMs should be constructed. It also connects to a separate line of work on learning unmasking orders, but argues the order term should arise from the policy gradient itself rather than from a separately trained planner head.

Real-world applications:

• Mathematical reasoning assistants built on diffusion language models, where the reported gains on GSM8K and MATH500 would translate into fewer wrong answers on step-by-step arithmetic and competition problems.
• Code generation and code completion tools, given the reported improvements on HumanEval and MBPP.
• Constrained planning tasks such as Sudoku and Countdown, which the authors report in the appendix.
• Any pipeline that runs parallel decoding on masked diffusion models, since the benefit of the method reportedly widens as the inference block size grows.

Industry relevance. The efficiency results matter for deployment: reaching a target accuracy in 15 hours instead of 18 on the same 8-GPU hardware lowers the cost of post-training runs, and the method requires no new parameters or architecture changes, so it can be dropped into existing MDLM training code. The finding that the advantage grows with block size is relevant to inference-time cost, because larger blocks mean fewer sequential decoding steps.

Future Directions

• Extending the analysis beyond the linear-policy counterexample. The appendix proof is a minimal finite-horizon construction; whether the same failure mode of token-only gradients occurs in realistic MDLM settings is not established.

• Understanding when position selection matters most. The paper observes that the gap widens as block size grows and attributes this to more masked positions being available to choose from, but it does not isolate which properties of a task or a decoding schedule make mask-aware optimization more or less valuable.

• Choosing the sampling temperature and unmasking-set size. The authors set a remasking temperature of 0.5 and report ablations on the sampling temperatures and on the size of the unmasking set |U_t| in Appendix E, leaving open how these should be set for new base models or tasks.

• Combining with separately learned unmasking policies. The related-work section contrasts this method with approaches that train a dedicated planner or an explicit unmasking policy, and notes the closest relative (LLaDOU/DCoLT) trains a separate module for position logits while this work reuses the base model's own logits. Whether the two approaches are complementary is not reported.

Target Audience

This paper is aimed at machine learning researchers working on reinforcement learning for language models, and more specifically at those working on discrete diffusion and masked diffusion language models who need tractable policy gradient estimators. It will also be useful to engineers post-training or deploying diffusion LLMs who care about training cost and inference block size, and to theorists interested in the structure of likelihood estimation when the generation order is itself a learned policy. Readers should be comfortable with policy gradient derivations, the ELBO objective for masked diffusion, and the mechanics of remasking-based generation.

Authors’ abstract

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.