Mask-Aware Policy Gradients for Diffusion Language Models

2026-07-16

AI Summary

Existing reinforcement learning approaches struggle to directly optimize masked diffusion language models due to the intractability of their log-likelihood and the neglect of how token unmasking order affects generation quality. This work formalizes the generation process of such models as a two-stage Markov decision process that explicitly separates token selection from mask position decisions. Building on this formulation, the authors introduce a mask-aware policy gradient decomposition method that enables joint optimization of content generation and unmasking sequence. The proposed approach achieves state-of-the-art performance, attaining 87.1% accuracy on GSM8K and 53.4% on MBPP, thereby setting new records on both benchmarks.

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.

Problem

• Masked Diffusion Language Models
• Reinforcement Learning
• Log-likelihood Estimation
• Token Prediction
• Unmasking Order

Innovation

• Mask-Aware Policy Gradient
• Masked Diffusion Language Models
• Two-Stage Action MDP
• Reinforcement Learning for LLMs
• Log-Likelihood Estimation