# Mask-Aware Policy Gradients for Diffusion Language Models

arXiv: https://arxiv.org/abs/2607.15200 · HTML: https://mlsanity.com/paper/2607.15200/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 48 of its 59 references that are in the map.
- First posted on arXiv 16 Jul 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 30 of 48)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Group Sequence Policy Optimization](https://mlsanity.com/paper/2507.18071/) (2025)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models](https://mlsanity.com/paper/2210.08933/) (2022)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Scaling Diffusion Language Models via Adaptation from Autoregressive Models](https://mlsanity.com/paper/2410.17891/) (2024)
- [MMaDA: Multimodal Large Diffusion Language Models](https://mlsanity.com/paper/2505.15809/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Autoregressive Diffusion Models](https://mlsanity.com/paper/2110.02037/) (2021)
- [d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2504.12216/) (2025)
- [dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Caching](https://mlsanity.com/paper/2506.06295/) (2025)
- [dKV-Cache: The Cache for Diffusion Language Models](https://mlsanity.com/paper/2505.15781/) (2025)

## Built on by (most cited first, 3 of 3)

- [CanvasAnneal: Curriculum Reinforcement Learning for Diffusion Language Models](https://mlsanity.com/paper/2609.13060/) (2026)
- [Fork-dLLM: Avoiding the Flexibility Trap in Diffusion Language Models](https://mlsanity.com/paper/2609.39859/) (2026)
- [RL Post-Training Diffusion Language Models Along Their Own Generation Paths](https://mlsanity.com/paper/2610.17266/) (2026)

## Abstract (arXiv)

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.