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

• URL: http://arxiv.org/abs/2606.08501v1
• Date: Sun, 07 Jun 2026 07:59:55 GMT
• Title: Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models
• Abstract summary: Process Aligned Policy Optimization (PAPO) is a novel framework that holistically aligns the RL update with the dLLM's generative trajectory.<n>PAPO transforms sparse terminal rewards into dense, step-wise credit, and Entropy-Guided Historical Re-enactment (EHR) that replays authentic trajectories at high-uncertainty steps.
• Score: 90.09182925511317
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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 up to 4.5% on GSM8K, 4.8% on MATH500, 42.2% on Countdown and 16.1% on Sudoku.

Related papers

• REVES: REvision and VErification--Augmented Training for Test-Time Scaling [53.197756110943395]
  We propose a two-stage iterative framework that alternates between online data/prompt augmentation and policy optimization.<n>We observe gains of +6.5 points over the RL baseline and +4.0 points over standard multi-turn training.
  arXiv  Detail & Related papers  (2026-06-17T10:37:23Z)
• Process Reward Agents for Steering Knowledge-Intensive Reasoning [18.184546182516225]
  We introduce Process Reward Agents (PRA), a test-time method for providing domain-grounded, online, step-wise rewards to a frozen policy.<n>PRA consistently outperforms strong baselines, achieving 80.8% accuracy on MedQA with Qwen3-4B.<n> PRA generalizes to unseen frozen policy models ranging from 0.5B to 8B parameters, improving their accuracy by up to 25.7% without any policy model updates.
  arXiv  Detail & Related papers  (2026-04-10T16:45:44Z)
• Let's Reward Step-by-Step: Step-Aware Contrastive Alignment for Vision-Language Navigation in Continuous Environments [38.97818584066075]
  Vision-Language Navigation in Continuous Environments (VLN-CE) requires agents to learn complex reasoning from long-horizon human interactions.<n>Current training paradigms struggle to balance generalization capability, error recovery and training stability.<n>We introduce Step-Aware Contrastive Alignment (SACA), a framework designed to extract dense supervision from imperfect trajectories.
  arXiv  Detail & Related papers  (2026-03-10T14:45:50Z)
• ProRAG: Process-Supervised Reinforcement Learning for Retrieval-Augmented Generation [54.071574153853994]
  ProRAG is a process-supervised reinforcement learning framework designed to integrate learned step-level supervision into the online optimization loop.<n>Our framework consists of four stages: (1) Supervised Policy Warmup to initialize the model with a structured reasoning format; (2) construction of an MCTS-based Process Reward Model (PRM) to quantify intermediate reasoning quality; (3) PRM-Guided Reasoning Refinement to align the policy with fine-grained process preferences; and (4) Process-Supervised Reinforcement Learning with a dual-granularity advantage mechanism.
  arXiv  Detail & Related papers  (2026-01-29T16:04:59Z)
• Save the Good Prefix: Precise Error Penalization via Process-Supervised RL to Enhance LLM Reasoning [59.76691952347156]
  Reinforcement learning (RL) has emerged as a powerful framework for improving the reasoning capabilities of large language models (LLMs)<n>Most existing RL approaches rely on sparse outcome rewards, which fail to credit correct intermediate steps in partially successful solutions.<n>We propose Verifiable Prefix Policy Optimization (VPPO), which uses PRMs only to localize the first error during RL.
  arXiv  Detail & Related papers  (2026-01-26T21:38:20Z)
• Agentic Reinforcement Learning with Implicit Step Rewards [92.26560379363492]
  Large language models (LLMs) are increasingly developed as autonomous agents using reinforcement learning (agentic RL)<n>We introduce implicit step rewards for agentic RL (iStar), a general credit-assignment strategy that integrates seamlessly with standard RL algorithms.<n>We evaluate our method on three challenging agent benchmarks, including WebShop and VisualSokoban, as well as open-ended social interactions with unverifiable rewards in SOTOPIA.
  arXiv  Detail & Related papers  (2025-09-23T16:15:42Z)
• Succeed or Learn Slowly: Sample Efficient Off-Policy Reinforcement Learning for Mobile App Control [50.316067647636196]
  This paper introduces Succeed or Learn Slowly (SoLS), a novel off-policyReinforcement learning algorithm evaluated on mobile app control tasks.<n>SoLS improves sample efficiency when fine-tuning foundation models for user interface navigation via a modified off-policy actor-critic approach.<n>We augment SoLS with Successful Transition Replay (STR), which prioritises learning from successful interactions.
  arXiv  Detail & Related papers  (2025-09-01T18:55:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.