Inpainting-Guided Policy Optimization for Diffusion Large Language Models

ICLR 2026 Conference Submission21202 Authors

Keywords: Diffusion Large Language Models, Reinforcement Learning, Inpainting, Group Relative Policy Optimization

TL;DR: IGPO, an RL method for diffusion LLMs that uses inpainting to inject partial reasoning hints when stuck with all-wrong responses, achieving SoTA results on math benchmarks for masked diffusion LLMs

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.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 21202