# Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models

arXiv: https://arxiv.org/abs/2606.08501 · HTML: https://mlsanity.com/paper/2606.08501/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 54 of its 65 references that are in the map.
- First posted on arXiv 7 Jun 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)

## Builds on (most cited first, 30 of 54)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [s1: Simple test-time scaling](https://mlsanity.com/paper/2501.19393/) (2025)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Models](https://mlsanity.com/paper/2505.19223/) (2025)
- [Scaling Diffusion Language Models via Adaptation from Autoregressive Models](https://mlsanity.com/paper/2410.17891/) (2024)
- [MMaDA: Multimodal Large Diffusion Language Models](https://mlsanity.com/paper/2505.15809/) (2025)
- [Scaling up Masked Diffusion Models on Text](https://mlsanity.com/paper/2410.18514/) (2024)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2504.12216/) (2025)
- [Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning](https://mlsanity.com/paper/2410.14157/) (2024)
- [SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation](https://mlsanity.com/paper/2510.06303/) (2025)
- [Right Question is Already Half the Answer: Fully Unsupervised LLM Reasoning Incentivization](https://mlsanity.com/paper/2504.05812/) (2025)
- [KodCode: A Diverse, Challenging, and Verifiable Synthetic Dataset for Coding](https://mlsanity.com/paper/2503.02951/) (2025)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [SEED-GRPO: Semantic Entropy Enhanced GRPO for Uncertainty-Aware Policy Optimization](https://mlsanity.com/paper/2505.12346/) (2025)
- [Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models](https://mlsanity.com/paper/2509.06949/) (2025)
- [Lumina-DiMOO: An Omni Diffusion Large Language Model for Multi-Modal Generation and Understanding](https://mlsanity.com/paper/2510.06308/) (2025)
- [Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design](https://mlsanity.com/paper/2410.13643/) (2024)
- [Boosting the Generalization and Reasoning of Vision Language Models with Curriculum Reinforcement Learning](https://mlsanity.com/paper/2503.07065/) (2025)

## Built on by (most cited first, 4 of 4)

- [Towards Unbiased On-Policy Distillation for Block Diffusion Language Models](https://mlsanity.com/paper/2610.05373/) (2026)
- [MaskSense: Confronting the Visual Exploration Trap in Masked Image Generation](https://mlsanity.com/paper/2610.17009/) (2026)
- [Match the Canvas, Not Just the Answer: Predictability-Aware SFT for Diffusion Language Models](https://mlsanity.com/paper/2610.18001/) (2026)
- [Backpropagating Verifier Rewards Through Continuous Flow Language Models](https://mlsanity.com/paper/2610.19080/) (2026)

## Abstract (arXiv)

Reinforcement learning (RL) holds immense promise for enhancing the reasoning capabilities of diffusion large language models (dLLMs). However, progress is fundamentally constrained by a dual misalignment between authentic generation trajectory and the gradient update process: (i) Process-reward misalignment. Sparse, terminal rewards are indiscriminately assigned to all intermediate steps of the generation process, failing to provide discriminative credit assignment. (ii) State-trajectory misalignment. Policy updates are often diverted toward artificial, out-of-trajectory states, squandering gradients on less informative samples. To address these limitations, we introduce Process Aligned Policy Optimization (PAPO), a novel framework that holistically aligns the RL update with the dLLM's generative trajectory via Step-Aware Process Rewards (SPR) that transform sparse terminal rewards into dense, step-wise credit, and Entropy-Guided Historical Re-enactment (EHR) that replays authentic trajectories at high-uncertainty steps. Extensive experiments on four benchmarks demonstrate that PAPO significantly outperforms baselines, achieving gains of 4.5% on GSM8K, 4.8% on MATH500, 42.2% on Countdown and 16.1% on Sudoku.
