mlsanityData to 7 Oct 2026

Mask-Aware Policy Gradients for Diffusion Language Models

arXiv 2607.15200 ↗First posted 16 Jul 2026cs.CL cs.AI cs.LGcited by 2

On the map: LLM reasoning → Diffusion reasoning and test-time scaling

also in 1 more
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.

Key facts. Cited by 2 papers in the map (data through 7 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.

Builds on

48 of 59 references.

Closest ancestors · 5

Same direction, newest first

  1. 2025Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization9 Oct
  2. 2025DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning2 Oct
  3. 2025wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models7 Jul
  4. 2025Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models15 May
  5. 2025d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning16 Apr

Other references · 43

  1. 2022Training language models to follow instructions with human feedbackcited by 15,712
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  4. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 8,851
  5. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  6. 2017Deep reinforcement learning from human preferencescited by 3,846
  7. 2023Let's Verify Step by Stepcited by 3,598
  8. 2021Program Synthesis with Large Language Modelscited by 3,509
  9. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  10. 2015Trust Region Policy Optimizationcited by 2,503
  11. 2025Understanding R1-Zero-Like Training: A Critical Perspectivecited by 1,276
  12. 2025Kimi k1.5: Scaling Reinforcement Learning with LLMscited by 1,193
31 more
  1. 2022MaskGIT: Masked Generative Image Transformercited by 1,082
  2. 2022Diffusion-LM Improves Controllable Text Generationcited by 1,006
  3. 2025Large Language Diffusion Modelscited by 872
  4. 2024Simple and Effective Masked Diffusion Language Modelscited by 800
  5. 2025Group Sequence Policy Optimizationcited by 788
  6. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distributioncited by 712
  7. 2024Simplified and Generalized Masked Diffusion for Discrete Datacited by 606
  8. 2022DiffuSeq: Sequence to Sequence Text Generation with Diffusion Modelscited by 439
  9. 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 404
  10. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 385
  11. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Datacited by 371
  12. 2024Scaling Diffusion Language Models via Adaptation from Autoregressive Modelscited by 268
  13. 2025MMaDA: Multimodal Large Diffusion Language Modelscited by 241
  14. 2021Autoregressive Diffusion Modelscited by 208
  15. 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusionscited by 208
  16. 2025dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Cachingcited by 168
  17. 2025dKV-Cache: The Cache for Diffusion Language Modelscited by 140
  18. 2025KodCode: A Diverse, Challenging, and Verifiable Synthetic Dataset for Codingcited by 115
  19. 2025LaViDa: A Large Diffusion Language Model for Multimodal Understandingcited by 81
  20. 2025Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Modelscited by 80
  21. 2025Path Planning for Masked Diffusion Model Samplingcited by 79
  22. 2024Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Designcited by 72
  23. 2025Unified Multimodal Discrete Diffusioncited by 46
  24. 2025FUDOKI: Discrete Flow-based Unified Understanding and Generation via Kinetic-Optimal Velocitiescited by 46
  25. 2025SPG: Sandwiched Policy Gradient for Masked Diffusion Language Modelscited by 38
  26. 2025Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspectivecited by 30
  27. 2025Learning Unmasking Policies for Diffusion Language Modelscited by 29
  28. 2025Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Stepcited by 16
  29. 2025Soft Policy Optimization: Online Off-Policy RL for Sequence Modelscited by 12
  30. 2025dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learningcited by 10
  31. 2026Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Modelscited by 9

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026CanvasAnneal: Curriculum Reinforcement Learning for Diffusion Language Modelscited by 0
  2. 2026Fork-dLLM: Avoiding the Flexibility Trap in Diffusion Language Modelscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models
  2. 2026Diffusion-State Policy Optimization for Masked Diffusion Language Models
  3. 2026Register Tokens for Bounded-State Reasoning in Diffusion Language Models
  4. 2025d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation
  5. 2025On the Reasoning Abilities of Masked Diffusion Language Models
  6. 2026UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching
  7. 2026Parallelism and Generation Order in Masked Diffusion Language Models: Limits Today, Potential Tomorrow
  8. 2025Activation Steering for Masked Diffusion Language Models