PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2509.10396v1cs.LG

TL;DR

Sparse verifiable rewards make it difficult for policies to discover correct solutions, wasting online-RL samples. IGPO injects partial reasoning hints during sampling, and its training recipe improves mathematical benchmarks, reaching state-of-the-art results among full-attention masked dLLMs.

  • Problem

    On challenging tasks, policies often fail to discover correct solutions, leaving binary rewards with little learning signal and wasting online-RL samples.

    Source: §1 · p. 2

  • Method

    IGPO injects partial reasoning hints when a policy is unlikely to solve a problem, then verifies completions and uses successful ones for policy optimization.

    Source: §1 · p. 2, §3.1 · p. 6

  • Results

    The full recipe improves all three benchmarks, including +4.9% on GSM8K relative to LLaDA-Instruct, and achieves state-of-the-art among full-attention dLLMs.

    Source: §1 · p. 3

  • Takeaways & Limitations

    IGPO can restore learning signals for group-based policy optimization when standard sampling yields uniform outcomes.

    Source: §3.1 · p. 4, §6 · p. 11

Abstract

from arXiv · show

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 three mathematical benchmarks--GSM8K, Math500, and AMC--achieving new state-of-the-art results for full-attention masked dLLMs.

1 Introduction

Sparse binary rewards can leave difficult RL problems with few useful learning signals and wasted samples; IGPO uses dLLM inpainting to guide exploration, alongside concise-trace SFT and entropy filtering.

  • On challenging tasks, policies often fail to find correct solutions, so binary rewards provide little learning signal and waste online-RL samples.Source: §1 · p. 2
  • IGPO injects partial reasoning when a policy struggles, verifies completions, and optimizes successful samples while keeping non-injected tokens on-policy.Source: §1 · p. 2
  • IGPO reduces all-wrong groups by ∼60% while preserving output diversity, restoring non-degenerate gradients in group-based optimization.Source: §1 · p. 2
  • Length-Aligned SFT uses synthetically rewritten concise reasoning traces to align SFT lengths with RL sampling and evaluation, strengthening downstream RL initialization.Source: §1 · p. 3
  • The full recipe improves over LLaDA-Instruct by +4.9% on GSM8K, +8.4% on Math500, and +9.9% on AMC, achieving state-of-the-art among full-attention dLLMs.Source: §1 · p. 3
  • Ablations find partial inpainting outperforms full ground-truth inpainting by staying closer to the policy distribution, while entropy-based filtering stabilizes training dynamics.Source: §1 · p. 3

2 Preliminaries

Section 2 describes masked diffusion as iterative masking and unmasking, then presents GRPO’s group-based objective and the challenge of estimating dLLM probabilities for policy optimization.

2.1 Masked Diffusion Large Language Models

  • Masked dLLMs progressively corrupt tokens, train to recover masked tokens with a weighted NLL objective, and generate by iteratively unmasking continuation tokens.The masking probability decreases with timestep, reaching complete masking at t = 1.Source: §2.1 · p. 3

2.2 Policy Optimization for Masked Diffusion Large Language Models

  • GRPO samples a response group and shares each response’s sequence-level advantage across its tokens, using clipping and reverse-KL regularization.Group statistics avoid separate value-function training.Source: §2.2 · pp. 3–4
  • Unlike autoregressive models, dLLMs lack sequential probability factorization, making exact tokenwise probabilities costly because they require marginalizing over denoising trajectories.Generation composes the unmasking predictor across multiple denoising steps.Source: §2.2 · p. 4
  • DiffuGRPO uses mean-field estimators for single-pass token- and sequence-level terms, trading bias against exact diffusion-policy probabilities for practical GRPO optimization.The method is adopted for token-level importance ratios.Source: §2.2 · p. 4

3 Methods

The methods use inpainting to repair all-wrong GRPO groups with verified responses, and rewrite verbose reasoning traces into concise SFT data aligned with generation lengths.

3.1 IGPO: Inpainting Guided Policy Optimization

  • IGPO triggers only for all-wrong groups and varies hint ratios and chunk sizes, excluding final-answer tokens from hints to discourage reasoning bypass.Selected ground-truth chunks remain fixed while the model denoises the other positions.Source: §3.1 · p. 5
  • When a group has zero advantages, IGPO replaces up to a fraction of its incorrect responses with correctness-verified inpainted responses, restoring reward variance.The replacement count is bounded by min of the verified correct inpainted responses and ⌊λG⌋.Source: §3.1 · pp. 4–6
  • For injected hint tokens, IGPO updates only the highest-entropy positions, focusing learning where the model is uncertain amid ground-truth and policy distribution mismatch.The filter selects the top τ percentile of hint positions by entropy.Source: §3.1 · pp. 6–7

3.2 Length-Aligned SFT via Concise Reasoning Trace Rewriting

  • Full-attention masked dLLMs face a length mismatch: RL generation is limited to 256 tokens, evaluation can use 1024, and verbose SFT traces can exceed 10k.Repeated reflections in these traces are also unsuitable for limited-context generation.Source: §3.2 · p. 7
  • Length-Aligned SFT rewrites verbose traces into concise, structured reasoning that preserves essential logic and empirically outperforms training on the original verbose traces.The rewritten traces remove redundant reflections and condense elaborations into mathematically rigorous statements.Source: §3.2 · p. 7

4 Experiments

Experiments test whether length-aligned SFT followed by IGPO improves mathematical reasoning and reduces all-wrong groups, while examining partial hints, entropy filtering, and rewritten traces.

4.1 Complete Training Recipe

  • The two-stage recipe first applies length-aligned SFT to rewritten OpenR1-Math traces, then uses MetaMathQA reasoning traces for IGPO inpainting.The SFT stage uses the dataset’s default split of 94,000 math problems.Source: §4.1 · p. 7

4.2 Experimental Setup

  • Experiments train from LLaDA-8B-Instruct on 12,794 deduplicated MetaMathQA questions and evaluate on GSM8K, MATH500, and AMC.Online RL generation uses temperature 1.2.Source: §4.2 · p. 8

4.3 Main Results

  • The full recipe improves GSM8K by 4.9%, MATH500 by 8.4%, and AMC by 9.9% relative to LLaDA-Instruct, establishing a new state of the art among full-attention masked dLLMs.Table 1 reports AMC at 24.4% (avg@16); the GSM8K and MATH500 gains are reported for pass@1.Source: Table 1 · p. 8, §4.3 · p. 9
  • IGPO shows more stable training than standard GRPO from both base and SFT initializations and reduces the all-wrong group ratio by around 60%.Figure 3 averages results over three seeds and three math benchmarks.Source: Figure 3 · p. 8, §4.3 · p. 9

4.4 Analysis and Ablation Studies

  • Partial hint injection outperforms full injection and no-hint sampling, with self-generated reasoning providing a learning signal between current policy behavior and target traces.The comparison uses partial injection η ∼ U[0.2, 0.6] versus full injection η = 1.0.Source: Figure 4 · p. 9, §4.4 · p. 9
  • On GSM8K, learning from the top 20% highest-entropy hint positions (τ = 0.2) performs best, while τ = 0.8 or 1.0 produces less stable training.The experiment uses temperature 0.1 and generation length 256.Source: §4.4 · p. 9, Figure 5 · p. 10
  • Rewritten traces capped at 1024 tokens yield stronger SFT checkpoints than original traces truncated at 4096 tokens, and stronger rewritten-trace initialization leads to higher final RL performance.Across both initializations, IGPO maintains pass@5 while standard RL shows declining diversity.Source: §4.4 · p. 10, Figure 6 · p. 10

5 Related Work

Related work reviews the development of diffusion language models and methods for reinforcement learning with them. It positions this work’s inpainting approach alongside existing optimization and exploration strategies.

5.1 Diffusion Language Models

  • Masked diffusion LLMs include full-attention models such as LLaDA and hybrid block diffusion, while recent work also targets faster inference through caching and parallel decoding.The paper focuses on full-attention masked dLLMs.Source: §5.1 · p. 11

5.2 Reinforcement Learning for Diffusion Language Models

  • RL for dLLMs must handle intractable likelihoods and estimation challenges; unlike Prefix-RFT’s left-to-right prefix guidance, inpainting can use bidirectional conditioning.The cited approaches include approximated policy ratios, preference optimization, weighted likelihood objectives, and stepwise alignment.Source: §5.2 · p. 11

6 Conclusion

The conclusion presents IGPO, length-aligned SFT, and entropy-based filtering as a combined recipe for mathematical reasoning with full-attention masked dLLMs.

6 Conclusion

  • IGPO uses inpainting hints to steer masked-diffusion RL, while length-aligned SFT and entropy filtering accompany reported state-of-the-art results across multiple math benchmarks.The conclusion describes IGPO as addressing zero-advantage cases by creating reward variance for policy-gradient updates.Source: §6 · p. 11

A Experiments Hyperparameters

The appendix section is titled “Experiments Hyperparameters.”

A Experiments Hyperparameters

  • The appendix includes a table labeled “Training Hyperparameters.”Source: Table 2 · p. 15

B Inpainting Generation Qualitative Examples

The example shows a question being inpainted with partial reasoning hints, followed by a completed solution that connects those hints with additional reasoning.

B Inpainting Generation Qualitative Examples

  • The inpainting example begins with a circle problem asking for the largest inscribed triangle with a diameter as one side.The prompt asks for the area given a circle of radius 6 cm.Source: B Inpainting Generation Qualitative… · p. 16
  • At denoising step 0, the model receives a partially masked input containing the question and an initial reasoning hint.The visible hint states that the triangle’s base is twice the circle’s radius, giving a 12 cm base.Source: B Inpainting Generation Qualitative… · p. 16
  • After denoising, the generated solution extends the hint by choosing a 6 cm height and calculating an area of 36 square centimeters.The example presents the answer in the requested boxed format.Source: B Inpainting Generation Qualitative… · p. 16

C Length-Aligned SFT: SFT trace revision length distribution comparison

The revised SFT traces are much shorter than the original OpenR1-Math-220K traces, addressing length mismatch with RL sampling and evaluation.

C Length-Aligned SFT: SFT trace revision length distribution comparison

  • Rewriting constrains every trace to under 1500 tokens, compared with original traces that could exceed 20,000 tokens.The original dataset contains 94k math problems.Source: C Length-Aligned SFT: SFT trace… · p. 17, Figure 7 · p. 17
  • Original traces extend beyond 10,000 tokens, while LLaDA’s maximum context length is 4096 tokens, so naive SFT would truncate many sequences.Even traces within the context limit have a length mismatch with 256-token RL sampling and 512-token evaluation.Source: C Length-Aligned SFT: SFT trace… · p. 17
  • The revised distribution removes the original extreme-length range and aligns SFT traces with the shorter RL-sampling and evaluation phases.The figure compares token-length distributions before and after revision.Source: Figure 7 · p. 17, C Length-Aligned SFT: SFT trace… · p. 17

D Prompt for SFT traces revision

The revision prompt requests concise, understandable solutions that retain reasoning details, and the examples contrast verbose original traces with rewritten ones.

D Prompt for SFT traces revision

  • The revision prompt asks for a more concise, understandable solution without changing details, while specifying a random target length of 100–1000 words.It also requires reasoning and a boxed numerical answer in designated tags.Source: D Prompt for SFT traces revision · p. 18
  • The original prime-listing trace uses extended self-checking and repeated reasoning before reaching its answer.It begins by recalling what a prime number is and checking numbers in sequence.Source: D Prompt for SFT traces revision · p. 19
  • Checking each number from 1 through 15 classifies 2, 3, 5, 7, 11, and 13 as prime; the shown response separately boxes 14 as its answer.The Sieve of Eratosthenes confirms the listed primes.Source: D Prompt for SFT traces revision · pp. 18, 21
Loading 2509.10396v1…