Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization

About

Diffusion language models (DLMs) enable parallel, order-agnostic generation with iterative refinement, offering a flexible alternative to autoregressive large language models (LLMs). However, adapting reinforcement learning (RL) fine-tuning to DLMs remains an open challenge because of the intractable likelihood. Pioneering work such as diffu-GRPO estimated token-level likelihoods via one-step unmasking. While computationally efficient, this approach is severely biased. A more principled foundation lies in sequence-level likelihoods, where the evidence lower bound (ELBO) serves as a surrogate. Yet, despite this clean mathematical connection, ELBO-based methods have seen limited adoption due to the prohibitive cost of likelihood evaluation. In this work, we revisit ELBO estimation and disentangle its sources of variance. This decomposition motivates reducing variance through fast, deterministic integral approximations along a few pivotal dimensions. Building on this insight, we introduce Group Diffusion Policy Optimization (GDPO), a new RL algorithm tailored for DLMs. GDPO leverages simple yet effective Semi-deterministic Monte Carlo schemes to mitigate the variance explosion of ELBO estimators under vanilla double Monte Carlo sampling, yielding a provably lower-variance estimator under tight evaluation budgets. Empirically, GDPO achieves consistent gains over pretrained checkpoints and outperforms diffu-GRPO, one of the state-of-the-art baselines, on the majority of math, reasoning, and coding benchmarks.

Kevin Rojas, Jiahe Lin, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, Molei Tao, Wei Deng• 2025

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningGSM8K (train)
Accuracy84.5
22
Mathematical ReasoningMATH-500 (train)
Accuracy41.4
21
Grade School Math Word ProblemsGSM8k 256 tokens
Pass@181.1
15
Arithmetic ReasoningCountdown 256 tokens
Pass@164.1
15
Arithmetic ReasoningCountdown 512 tokens
Pass@160.2
15
Sudoku SolvingSudoku 256 tokens
Pass@131.2
15
Sudoku SolvingSudoku 512 tokens
Pass@130.6
15
Advanced Mathematical ReasoningMath500 256 tokens
Pass@1 Accuracy38.4
15
Grade School Math Word ProblemsGSM8k 512 tokens
Pass@182
15
Advanced Mathematical ReasoningMath500 512 tokens
Pass@1 Accuracy41.2
15
Showing 10 of 12 rows

Other info

Follow for update