---
title: "Mask-Aware Policy Gradients for Diffusion Language Models"
source: "arXiv AI/CL/LG（论文）"
source_tier: "t2"
url: "https://arxiv.org/abs/2607.15200v1"
permalink: "https://feed.astralor.com/items/itm_06291c4b697f29c22c46b224"
published: "2026-07-16T16:57:34.000Z"
score: 83
channel: "ai"
tags: ["GSM8K","MBPP","Reinforcement Learning","Masked Diffusion Language Model","Policy Gradient"]
---

# Mask-Aware Policy Gradients for Diffusion Language Models

## AI commentary

> This paper is valuable because it precisely identifies the long-overlooked 'unmasking order' problem in RL training of diffusion language models and provides an elegant theoretical solution. For teams exploring post-training or reasoning enhancement of diffusion models, this work is more fundamentally important than mere sampling tuning. It demonstrates that 'what to choose' and 'when to choose' must be optimized together, which will likely inspire other diffusion-based generation tasks.

## At a glance

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 only model token predictions, ignoring the unmasking order during generation. This paper observes that MDLM generation involves two decisions: what tokens to place and which positions to remask. The authors formalize it as a two-stage action MDP, showing the policy gradient naturally decomposes into a token term and a masking term. Joint optimization achieves state-of-the-art results on GSM8K (87.1%) and MBPP (53.4%).

## Key points

- Existing methods ignore the unmasking order and only approximate log-likelihood via token predictions.
- Formalizes MDLM generation as a two-stage action MDP, showing policy gradient decomposes into a token term and a masking term.
- Joint optimization achieves state-of-the-art on GSM8K (87.1%) and MBPP (53.4%) for mathematical reasoning and code generation.

## Full text

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.

---

Verify against the original before citing; curated by Astral Feed: https://feed.astralor.com/items/itm_06291c4b697f29c22c46b224