Technical Summary: Mask-Aware Policy Gradients for Diffusion Language Models

Problem Statement

Reinforcement learning (RL) has successfully enhanced reasoning capabilities in autoregressive Large Language Models (LLMs) by optimizing policy gradients based on tractable log-likelihoods. However, extending these methods to Masked Diffusion Language Models (MDLMs) remains challenging. Unlike autoregressive models, which generate tokens sequentially from left to right, MDLMs generate text by iteratively unmasking tokens from a fully masked sequence.

The core difficulty lies in the intractability of the log-likelihood estimation for MDLMs. Existing approaches approximate this likelihood using Evidence Lower Bound (ELBO) variants or by modeling only the token prediction component of the generation trajectory. These methods ignore a critical aspect of MDLM generation: the order in which positions are unmasked. In MDLMs, the generation trajectory involves two distinct decisions at each step: (1) predicting which token to place at masked positions, and (2) deciding which of those positions to keep unmasked (and which to remask). Prior methods that ignore the second decision fail to utilize the full signal available in the generation process for policy optimization.

Methodology

The authors formalize MDLM generation as a two-stage action Markov Decision Process (MDP). At each denoising step $t$, the model performs:

1. Token Prediction: Predicting tokens for all currently masked positions.
2. Position Selection (Remasking): Deciding which subset of predicted positions to reveal and which to return to the [MASK] state.

To enable policy gradient optimization over this process, the authors introduce Probabilistic Remasking. Standard MDLM inference often uses a deterministic, greedy top-$K$ strategy to select positions to unmask based on confidence scores. This greedy selection is non-differentiable, preventing gradient-based optimization of the unmasking order. The proposed method replaces this with a probabilistic variant:

• Instead of deterministically selecting the top-$K$ positions, the model samples a subset of positions $U_t$ from a Plackett–Luce distribution.
• The probabilities for this distribution are derived directly from the model's own token log-likelihoods (logits), scaled by a temperature parameter $\tau$.
• This formulation creates a differentiable distribution over position subsets, allowing the log-probability of the unmasking decision to be included in the trajectory likelihood.

Policy Gradient Decomposition

By modeling the full trajectory likelihood $\pi_\theta(\hat{z}|c)$ as a product of token prediction probabilities and unmasking probabilities, the policy gradient naturally decomposes into two terms:
$$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Token Gradient}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Masking Gradient}} \right) \right]$$

• Token Term: Reinforces the correctness of the predicted tokens.
• Masking Term: Reinforces the selection of the unmasking order.
• Efficiency: The masking term is computed from the model's existing logits used for token prediction, requiring no additional parameters, architectural changes, or extra forward passes.

The authors implement this using the Group Sequence Policy Optimization (GSPO) algorithm, clipping the importance ratios for both token and position terms separately. To handle the computational cost of evaluating full trajectories, they employ the StepMerge approximation, grouping denoising steps into segments to reduce the number of forward passes while preserving trajectory structure.

Key Contributions

1. Formalization of MDLM as a Two-Stage MDP: The paper provides a theoretical framework showing that MDLM generation involves both token selection and position selection, and that the policy gradient decomposes accordingly.
2. Probabilistic Remasking: The derivation of a differentiable unmasking distribution (Plackett–Luce) based on the model's own logits, enabling the optimization of the unmasking order without architectural modifications.
3. Theoretical Insight: Demonstration that ignoring the position component of the policy gradient can result in missing directions that improve the expected return, even when token probabilities remain unchanged.
4. State-of-the-Art Performance: Empirical validation showing that jointly optimizing both terms yields consistent improvements over existing trajectory-based and ELBO-based methods.

Experimental Results

The method was evaluated on mathematical reasoning (GSM8K, MATH500) and code generation (HumanEval, MBPP) benchmarks, using LLaDA-8B-Instruct as the base model.

• Performance Gains: The proposed method achieved new state-of-the-art results across all benchmarks.
  • GSM8K: 87.1% accuracy (at generation length 128), a +2.5% improvement over the previous best (StepMerge).
  • MATH500: 53.4% accuracy, a +4.0% improvement.
  • HumanEval: 47.1% accuracy, a +2.2% improvement.
  • MBPP: 53.4% accuracy, a +2.5% improvement.
• Ablation Studies:
  • Block Size: The performance gap between the proposed method and baselines widened as the inference block size increased (from 32 to full sequence), suggesting that the benefit of optimizing position selection grows with the complexity of the unmasking decision space.
  • Efficiency: The method converged faster in wall-clock time compared to the strong ELBO-based baseline (SPG), reaching SPG's final accuracy in ~15 hours versus ~18 hours, and achieving higher final accuracy (81.0% vs 78.5% on GSM8K).
  • Generalization: The approach generalized to a different base model (Dream-7B) and planning tasks (Sudoku, Countdown), consistently outperforming baselines.
• Comparison to DCoLT: When compared to LLaDOU (which uses a separate head for position selection), the proposed method achieved comparable or better accuracy with significantly lower training costs (~160 GPU hours vs ~800), as it requires no additional modules.

Significance and Claims

The paper claims that the flexibility of MDLMs comes with a richer trajectory structure than autoregressive models, specifically the decision of when to unmask positions. By treating this decision as part of the policy and optimizing it via a differentiable probabilistic framework, the authors demonstrate that joint optimization of token predictions and position selection is crucial for maximizing the performance of diffusion language models in RL settings.

The authors emphasize that their approach is practically efficient, as it leverages existing model logits without requiring new parameters or architectural changes, making it compatible with standard MDLM training infrastructures like StepMerge. The results suggest that the "masking gradient" provides a meaningful training signal that prior methods, which focused solely on token likelihoods, failed to capture.