Inpainting-Guided Policy Optimization for Diffusion Large Language Models
Siyan Zhao, Mengchen Liu, Jing Huang, Miao Liu, Chenyu Wang, Bo Liu, Yuandong Tian, Guan Pang, Sean Bell, Aditya Grover, Feiyu Chen
Abstract
Masked diffusion large language models (dLLMs) are emerging as promising alternatives to autoregressive LLMs, offering competitive performance while supporting unique generation capabilities such as inpainting. We explore how inpainting can inform RL algorithm design for dLLMs. Aligning LLMs with reinforcement learning faces an exploration challenge: sparse reward signals and sample waste when models fail to discover correct solutions. While this inefficiency affects LLMs broadly, dLLMs offer a distinctive opportunity—their inpainting ability can guide exploration. We introduce IGPO (Inpainting Guided Policy Optimization), an RL framework that strategically inserts partial ground-truth reasoning traces during online sampling. Unlike providing full solutions, inpainting steers exploration toward promising trajectory spaces while preserving self-generated reasoning, bridging supervised fine-tuning and reinforcement learning. We apply IGPO to group-based optimization methods such as GRPO, where exploration failures cause zero advantages and gradients. IGPO restores meaningful gradients while improving sample efficiency. We also propose supervised fine-tuning on synthetically rewritten concise traces that better align with dLLM generation patterns. With additional techniques including entropy-based filtering, our training recipe yields substantial gains across four mathematical benchmarks—GSM8K, Math500, AMC and Minerva—achieving new state-of-the-art results for full-attention masked dLLMs.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext fc9123a9-f24d-4115-b16b-c19d50a9170eCited by top-tier papers5
- Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language ModelsSiyan Zhao, Zhihui Xie, Mengchen Liu, Jing Huang et al.ICML 2026 · 245 citations
- Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order AlgorithmsYinuo Ren, Haoxuan Chen, Yuchen Zhu, Wei Guo et al.NeurIPS 2025 · 51 citations
- Lavida-R1: Advancing Reasoning for Unified Multimodal Diffusion Language ModelsShufan Li, Yuchen Zhu, Kangning Liu, Zhe Lin et al.ICML 2026 · 4 citations
- Unifying Deductive and Abductive Reasoning in Knowledge Graphs with Masked Diffusion ModelYisen Gao, Jiaxin Bai, Yi Huang, Xingcheng Fu et al.WWW 2026 · 3 citations
- Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from -ParityJianhao Huang, Baharan MirzasoleimanICML 2026 · 2 citations
Builds on27
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
- Diffusion-LM Improves Controllable Text GenerationXiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang et al.NeurIPS 2022 · 1,546 citations
- Large Language Diffusion ModelsShen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang et al.NeurIPS 2025 · 949 citations
- Simple and Effective Masked Diffusion Language ModelsSubham S. Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan et al.NeurIPS 2024 · 929 citations
Related papers
- d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement LearningSiyan Zhao, Devaansh Gupta, Qinqing Zheng, Aditya GroverNeurIPS 2025 · 191 citations
- AG-GRPO: Answer-Guided GRPO for Masked Diffusion Language ModelsJuhyeong Kim, Gyunyeop Kim, Sangwoo KangACL 2026
- Improving Reasoning for Diffusion Language Models via Group Diffusion Policy OptimizationKevin Rojas, Jiahe Lin, Kashif Rasul, Anderson Schneider et al.ICLR 2026 · 34 citations
- Reinforcement Learning Meets Masked Generative Models: Mask-GRPO for Text-to-Image GenerationYifu Luo, Xinhao Hu, Keyu Fan, Haoyuan Sun et al.NeurIPS 2025 · 12 citations
- The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language ModelsZanlin Ni, Shenzhi Wang, Yang Yue, Tianyu Yu et al.ICML 2026 · 4 citations
