Mask-Aware Policy Gradients for Diffusion Language Models
Provides a principled RL method for MDLMs, improving reasoning and coding performance for language model practitioners.
Reinforcement learning for Masked Diffusion Language Models is improved by formalizing generation as a two-stage action MDP and optimizing both token and masking policy gradients, achieving 87.1% on GSM8K and 53.4% on MBPP.
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.