Mask-Aware Policy Gradients for Diffusion Language Models

TL;DR - A method to apply reinforcement learning to Masked Diffusion Language Models (MDLMs) by properly accounting for the token-placement and position-unmasking decisions during generation, improving reasoning and coding performance.

• Reframes MDLM generation as a two-stage action MDP, where each step decides both what tokens to place at masked positions and which positions to remask.
• Shows the policy gradient decomposes into a token term and a masking term; prior approaches ignored the unmasking order, only modeling token predictions.
• Optimizing both terms jointly yields state-of-the-art results: 87.1% on GSM8K (math) and 53.4% on MBPP (coding).
• Addresses a core challenge in RL for diffusion LMs: the intractable log-likelihood estimation.