Simple Policy Gradients for Reasoning with Diffusion Language Models

2026-07-25 · Source: cs.CL updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Natural Language Processing · Depth: Expert, extended

Summary

Amortized Group Relative Policy Optimization (AGRPO) is a new, principled on-policy reinforcement learning algorithm specifically designed for diffusion large language models (dLLMs). Traditional RL techniques for autoregressive LLMs are incompatible with dLLMs due to their iterative unmasking generation process and parallel token prediction. AGRPO addresses this by employing Monte Carlo sampling to compute an unbiased policy gradient estimate, making it the first tractable and theoretically sound adaptation of policy gradient methods for dLLMs. Experiments demonstrate AGRPO's effectiveness on reasoning tasks, achieving up to a +7.6% absolute gain on GSM8K and 3.8x performance on the Countdown task over the LLaDA-8B-Instruct baseline. It also outperforms comparable RL methods like diffu-GRPO by 1.3x and improves inference compute/quality tradeoffs.

Key takeaway

For AI Scientists and Machine Learning Engineers aiming to enhance diffusion LLM reasoning and optimize inference, AGRPO presents a theoretically sound and effective post-training solution. You can achieve significant performance gains on tasks like GSM8K and Countdown, while also improving the compute-quality frontier by maintaining accuracy with fewer sampling steps. Be mindful of AGRPO's compute-intensive nature and the potential for entropy collapse, which may require careful hyperparameter tuning.

Key insights

AGRPO provides a principled, tractable reinforcement learning method for diffusion LLMs using unbiased policy gradients.

Principles

• RL algorithms for autoregressive LLMs are incompatible with diffusion LLMs.
• Unbiased policy gradient estimates are critical for dLLM theoretical soundness.
• Monte Carlo sampling can make complex RL objectives tractable for dLLMs.

Method

AGRPO reinterprets the GRPO objective's inner sum as an expectation over timesteps, then estimates this via Monte Carlo sampling, computing exact token probabilities at sampled timesteps. Low-discrepancy sampling reduces variance.

In practice

• Cache unmasking order to efficiently recreate partially masked states.
• Use "float64" for Gumbel-based categorical sampling to prevent numerical instability.
• Limit max sampling timestep to the last non-EOS token for meaningful gradient updates.

Topics

• Diffusion LLMs
• Reinforcement Learning
• Policy Gradients
• Mathematical Reasoning
• Monte Carlo Sampling
• LLaDA-8B-Instruct

Best for: Research Scientist, AI Scientist, Machine Learning Engineer, NLP Engineer

Related on AIssential

• Paper: DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning — DeepSeek R1 enhances LLM reasoning via pure reinforcement learning and rule-based rewards, fostering autonomous problem-solving without supervised data.
• Roll Out and Roll Back: Diffusion LLMs are Their Own Efficiency Teachers — Revokable parallel decoding and trajectory-guided training overcome DLLM quality-speed trade-offs.
• Zone of Proximal Policy Optimization: Teacher in Prompts, Not Gradients — ZPPO enhances knowledge distillation for small student models by embedding teacher guidance directly into prompts, avoiding policy gradient issues.
• Learning Unmasking Policies for Diffusion Language Models — Reinforcement learning can optimize dLLM unmasking policies, outperforming heuristics and improving inference efficiency.
• Learning from the Self-future: On-policy Self-distillation for dLLMs — d-OPSD adapts self-distillation for dLLMs by using suffix conditioning and step-level supervision, significantly boosting sample efficiency.
• Diffusion-GR2: Diffusion Generative Reasoning Re-ranker — Diffusion-GR2 enables fast, accurate generative reasoning re-ranking by converting autoregressive models to block-diffusion via targeted fine-tuning.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.CL updates on arXiv.org.