Research map / Reinforcement learning / RL from feedback

Diffusion-State Policy Optimization for Masked Diffusion Language Models

NeurIPS 2026 · Reinforcement learning › RL from feedback · approach: RL fine-tuning of generative models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DiSPO, a plug-in credit-assignment objective that branches from intermediate masked states of masked diffusion LMs using cached logits and adds a same-state policy-gradient term to terminal-feedback RL (diffu-GRPO, SPG), improving LLaDA-8B-Instruct on math and planning tasks at matched rollout compute.

Problem

Terminal-only rewards on final completions give coarse credit assignment to the intermediate mask-filling decisions made during masked diffusion LM generation.

Approach

At selected intermediate masked states, resample the currently masked positions from rollout-cached logits, score the branched completions with the terminal reward, and apply a policy-gradient update only to the newly filled tokens, reusing the same rollouts as terminal-feedback optimization.

Key results

• On LLaDA-8B-Instruct, DiSPO with diffu-GRPO raises Sudoku average accuracy from 14.1 to 31.4 and Countdown average from 33.9 to 53.9 under matched compute.
• DiSPO with SPG raises Sudoku average from 26.2 to 43.5 and GSM8K average from 79.9 to 80.8.
• Under 32-token block rollouts, 17.8% of uniformly sampled intermediate states already yield different rewards across same-state branches.

Surprising findings

• Step-wise loss alone (L_step only) outperforms the terminal base loss alone on Sudoku (20.3 vs 14.1 average).
• Current implementation is about 0.4x baseline wall-clock throughput despite no extra rollouts or optimizer steps.

Limitations

• Wall-clock overhead from reward and surrogate evaluations (~0.4x throughput in current implementation).
• Gains on GSM8K and MATH500 are small compared with planning tasks.
• Evaluated only on LLaDA-8B-based models and two base optimizers.

Experimental setup

Datasets: GSM8K MATH500 Sudoku (4x4 training set) Countdown (Countdown-Tasks-3to4) s1k

Benchmarks: GSM8K MATH500 Sudoku Countdown

Metrics: exact-match accuracy

Baselines: diffu-GRPO SPG SAPO LLaDA-8B-Instruct SFT & diffu-GRPO

Builds on: diffu-GRPO SPG GRPO policy gradient theorem

Closest papers by meaning

• Unifying Masked Diffusion Models with Various Generation Orders and Beyond (ICML 2026 · similarity 0.90)
• On Powerful Ways to Generate: Autoregression, Diffusion, and Beyond (ICLR 2026 · similarity 0.90)
• Theoretical Benefit and Limitation of Diffusion Language Model (NeurIPS 2025 · similarity 0.89)
• From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.89)
• DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models (NeurIPS 2026 · similarity 0.89)
• On the Reasoning Abilities of Masked Diffusion Language Models (ICLR 2026 · similarity 0.88)