ArXiv

Mask-Aware Policy Gradients for Diffusion Language Models

2026-07-16 · 16:57 UTC ·Haran Raajesh, Kulin Shah, Adam Klivans... ·1 min read
Authors
Haran Raajesh, Kulin Shah, Adam Klivans...
Categories
cs.CL, cs.AI, cs.LG
arXiv
https://arxiv.org/abs/2607.15200v1
PDF
https://arxiv.org/pdf/2607.15200v1

Brief

Mask-aware policy gradients formalize MDLM generation as a two-stage MDP—choosing tokens and choosing which positions to remask—and derive a policy-gradient decomposition into token and masking terms. Optimizing both terms produces state-of-the-art performance on math and coding benchmarks (87.1% GSM8K, 53.4% MBPP). Only the abstract was provided here; full text was not included.

Why it matters

Proposes mask-aware policy gradients for Masked Diffusion Language Models (MDLMs), formalizing generation as a two-stage action MDP (token placement then which positions to remask) and showing the policy gradient decomposes into a token term and a masking term.

Key details

  • Optimizing both token and masking terms yields state-of-the-art results on reasoning and coding benchmarks: 87.1% on GSM8K and 53.4% on MBPP; paper by Haran Raajesh, Kulin Shah, Adam Klivans, and Philipp Krähenbühl was accepted at COLM 2026.
  • Identifies a shortcoming of prior MDLM RL methods that approximate the log-likelihood by modeling only token predictions and ignore the order of position unmasking; the new method explicitly models both decisions.
Source evidence

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.

Comment: Accepted at COLM 2026