Mask-Aware Policy Gradients for Diffusion Language Models
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
The authors studied how to teach a special type of language model, called a Masked Diffusion Language Model (MDLM), to make better decisions when filling in missing words. They found that the model not only decides which words to place but also which positions to hide again as it generates text. By treating these two choices separately in a reinforcement learning framework, the authors improved the model's reasoning abilities on math and coding tests. Their approach achieved high accuracy on challenging benchmarks.
Authors
Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl
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.