Mask-Aware Policy Gradients for Diffusion Language Models

Source: arXiv:2607.15200 · Published 2026-07-16 · By Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

TL;DR

This paper addresses the challenge of applying reinforcement learning (RL) to Masked Diffusion Language Models (MDLMs), which unlike autoregressive models have an intractable log-likelihood estimation due to their generation procedure that involves both predicting tokens and choosing which masked positions to unmask (remask). Existing RL approaches approximate the log-likelihood using only token predictions, ignoring the choice of unmasking order. The authors model MDLM generation as a two-stage Markov Decision Process (MDP) that explicitly accounts for both token prediction and position selection at each step. They propose a novel policy gradient estimator that decomposes into a token gradient and a masking (position) gradient by replacing deterministic top-K position selection with a probabilistic Plackett–Luce distribution over unmasking positions. This enables end-to-end optimization of both aspects without additional model parameters or forward passes.

Empirically, this mask-aware policy gradient achieves state-of-the-art results on four benchmarks: mathematical reasoning datasets GSM8K and MATH500, and code generation datasets HumanEval and MBPP, improving accuracy by 2–4% absolute over strong baselines including recent ELBO-based methods and trajectory-based methods that ignore position gradients. The approach converges faster in wall-clock time despite slightly lower per-step throughput and generalizes to different base models and planning tasks. The findings demonstrate that jointly optimizing the position selection in MDLM generation unveils significant performance gains that previous token-only RL methods miss.

Key findings

• Mask-aware policy gradients that jointly optimize token predictions and unmasking positions improve GSM8K accuracy by +2.5% absolute over StepMerge baseline at 128 tokens.
• On MATH500, the method achieves a +4.0% accuracy gain (from 33.1% to 37.4%) at generation length 128 vs. StepMerge.
• For code generation, the approach improves HumanEval by +2.2% and MBPP by +2.5% accuracy at 128 tokens compared to baselines.
• The position (masking) gradient provides necessary training signal; token-only RL gradients can miss directions that improve expected return (shown by a minimal counterexample).
• Replacing greedy top-K position selection with a probabilistic Plackett–Luce distribution over masked positions yields a differentiable masking term in the policy gradient.
• Joint optimization converges faster and to higher final accuracy than SPG (a strong ELBO baseline), reaching SPG’s final accuracy 3 hours earlier and higher peak accuracy (81.0% vs 78.5%) on GSM8K.
• Performance gains increase with larger inference block size, indicating the importance of mask-aware optimization grows as more masked positions are considered each step.
• The masking gradient adds no computational overhead beyond token likelihood computation since position logits are derived from the model’s own output logits.

Threat model

n/a; this paper is focused on improving training methods for masked diffusion language models rather than security threat models or adversarial scenarios.

Methodology — deep read

1. Threat model & assumptions: The paper assumes an RL framework where the language model aims to maximize expected reward on downstream tasks (mathematical reasoning and code generation). The adversary is conceptualized as the policy gradient optimizer seeking to improve the generation policy πθ. The model assumes access only to sampled trajectories of the MDLM generation, without direct tractable likelihoods, and must estimate gradients of the expected reward given the intractable marginalization over possible unmasking sequences.

2. Data: Experiments use public benchmarks GSM8K (8,000 examples), MATH500, HumanEval (164 hand-authored code problems), and MBPP (5,000+ Python code generation examples). The base pretrained model is LLaDA-8B-Instruct (8B parameter masked diffusion LM). Code training uses KodCode-Light-RL-10K dataset. Evaluation is zero-shot, with standard train-test splits from prior work. No mention of additional data preprocessing beyond standard tokenization and masked diffusion inputs.

3. Architecture / algorithm: Masked Diffusion Language Models iteratively unmask masked token positions over T denoising steps. At each step, the MDLM predicts tokens at all masked positions, then selects a subset of these positions to remain unmasked (remasking the rest). The key novelty is modeling generation as a two-stage MDP where at each step the policy outputs (a) token predictions for masked positions, (b) a distribution over subsets of positions to unmask. The masking selection is modeled via a Plackett–Luce probabilistic distribution derived from the token log-likelihoods, making the position selection differentiable. The full trajectory likelihood is the product of token prediction probabilities and the masking probabilities over remasking subsets. The policy gradient thus decomposes exactly into a sum of gradients from tokens and from the masking distribution.

4. Training regime: RL fine-tuning uses GSPO, a recent policy gradient optimizer. Model weights are adapted with Low-Rank Adaptation (LoRA) of rank 128 and scaling 64 on attention and MLP layers, quantized to 4-bit (NF4), training in bfloat16 on 8 Nvidia H100 GPUs. Batch sizes and epochs are not precisely stated but training is monitored for convergence and stability across 3 random seeds. The probabilistic masking temperature τ is tuned (appendix details).

5. Evaluation protocol: Metrics are task-specific accuracies (% solved problems on GSM8K, MATH500, HumanEval, MBPP). Baselines include ELBO-based approaches (d1, wd1, GDPO, SPG) and trajectory-based StepMerge. Ablations show impact of including masking gradient and changes in inference block size. Wall-clock time comparisons are plotted for training efficiency on GSM8K. Statistical tests are not explicitly reported, but results are averaged across multiple runs with low variance.

6. Reproducibility: Code is publicly available at https://github.com/Haran71/mask-aware-policy-gradients. The datasets and base models are public or previously published. Detailed appendices provide hyperparameter settings, temperature tuning, and exact approximations used (StepMerge). Models trained with quantization and LoRA fine-tuning to reduce compute are described clearly.

Example end-to-end: Generating a sequence from prompt c proceeds by starting with a fully masked sequence z0, then at each diffusion step t, the model predicts token logits ˆzt for masked positions, uses a Plackett–Luce distribution parameterized by token logits to sample which positions Ut to unmask, remasks the other positions, producing zt. The full trajectory ˆz over token predictions and position subsets is scored by a reward R(c, zT). Policy gradients are computed jointly over the log-likelihoods of token choices and the probability of sampled positions under the Plackett–Luce model, enabling gradient updates that improve both token prediction accuracy and unmasking order selection for better downstream task performance.

Technical innovations

• Formalization of MDLM generation as a two-stage-action Markov Decision Process with explicit token prediction and position (mask) selection steps.
• Derivation of a decomposed policy gradient estimator separating token and masking terms, with the masking term computed from the model’s own token logits without additional parameters.
• Replacing deterministic greedy top-K position selection with a differentiable probabilistic unmasking via a Plackett–Luce distribution sampled without replacement.
• Joint optimization of both token prediction and position selection gradients yields improved expected reward and state-of-the-art results on reasoning and code benchmarks.

Datasets

• GSM8K — ~8,000 mathematical reasoning problems — public
• MATH500 — 500 math problems — public
• HumanEval — 164 hand-authored Python coding problems — public
• MBPP — 5,000+ Python code generation problems — public
• KodCode-Light-RL-10K — 10,000 code samples for RL training — public

Baselines vs proposed

• StepMerge: GSM8K accuracy (128 tokens) = 78.0% vs Proposed: 81.0%
• StepMerge: MATH500 accuracy (128) = 33.1% vs Proposed: 37.4%
• StepMerge: HumanEval accuracy (128) = 30.5% vs Proposed: 33.2%
• StepMerge: MBPP accuracy (128) = 44.9% vs Proposed: 47.1%
• SPG: GSM8K accuracy (128) = 78.5% vs Proposed: 81.0%
• GDPO: MATH500 accuracy (128) = 33.2% vs Proposed: 37.4%
• SPG: MBPP accuracy (128) = 44.3% vs Proposed: 47.1%

Limitations

• The approach relies on approximate trajectory likelihoods using StepMerge and K-subsampling, which may introduce bias unquantified beyond theoretical bounds.
• Evaluation is zero-shot; performance under fine-tuning or few-shot prompting is not explored.
• No explicit adversarial or robust testing of the policy gradient is reported; vulnerability to reward hacking or policy collapse is unstudied.
• The computational overhead per optimization step is higher due to repeated forward passes, despite faster convergence overall.
• Only certain generation lengths and block sizes are evaluated; scalability to very long sequences or other modalities is not shown.
• The masking distribution temperature τ and other hyperparameters require tuning and the method’s sensitivity to these is only partially explored.

Open questions / follow-ons

• How does the mask-aware policy gradient method perform across a wider variety of diffusion LM architectures and sizes, including very large-scale models?
• Can this approach be extended to multimodal diffusion models beyond text, such as image or audio generation?
• How sensitive is the method to the temperature parameter τ and other hyperparameters governing the position sampling distribution?
• What is the impact of the approach in the presence of distributional shift or noisy reward signals, and how robust is training stability?

Why it matters for bot defense

From a bot-defense and CAPTCHA perspective, this paper’s insights about leveraging the richer trajectory structure of Masked Diffusion Language Models have several implications. It highlights that sequence generation policies can benefit significantly by optimizing not only output tokens but also the latent decision processes like position selection during decoding. This suggests that advanced language models underlying modern CAPTCHA solvers or adversarial bots might employ similarly complex generation strategies that are partially observable. Reinforcement learning techniques that do not account for such intertwined latent controls may underestimate model capabilities or fail to influence their behavior effectively. Practitioners developing bot defenses could explore these internal generation dynamics as potential features for bot detection or design harder CAPTCHAs resistant to diffusion-based generation methods. On the other hand, the existence of policy gradients over position choices also indicates that attack methods might improve by jointly optimizing token outputs and structural decoding decisions, raising the bar for defense strategies relying solely on token-level analysis.

Cite

@article{arxiv2607_15200,
  title={ Mask-Aware Policy Gradients for Diffusion Language Models },
  author={ Haran Raajesh and Kulin Shah and Adam Klivans and Philipp Krähenbühl },
  journal={arXiv preprint arXiv:2607.15200},
  year={ 2026 },
  url={https://arxiv.org/abs/2607.15200}
}

Read the full paper

• arXiv:2607.15200 (abstract)
• arXiv:2607.15200 (PDF)