Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization

TL;DR

This paper tackles the challenge of reinforcement learning fine-tuning for diffusion language models by grounding training in sequence-level likelihood via the ELBO. It analyzes the variance sources in ELBO estimation for DLMs and shows that random time contributes most to variance, motivating a Semi-deterministic Monte Carlo (SDMC) approach that applies deterministic time integration with Gaussian quadrature. Building on this, the authors introduce Group Diffusion Policy Optimization (GDPO), a policy-gradient method that uses sequence-level ELBO with SDMC to achieve provably lower-variance estimates and practical efficiency. Empirically, GDPO yields consistent gains over pretrained checkpoints and outperforms diffu-GRPO on mathematics, reasoning, and coding benchmarks, with good performance under limited compute budgets, demonstrating the method’s practicality and effectiveness for aligning DLMs.

Abstract

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.

Figures (5)

• Figure 1: Across reasoning, planning, and coding tasks, our GDPO algorithm for diffusion language models — using the best of 128/256/512 generations — significantly outperforms the LLaDA baseline and prior RL methods such as diffu-GRPO.
• Figure 2: We plot the mean and variance of the loss functions as a function of the noise level $t$. (a) We observe that most of the variance comes from picking the random time (b) The loss function follows a simple, predictable shape across many prompts. (c) The loss variance varies highly at the end but stabilizes for most times.
• Figure 3: Estimation error and variance for Double Monte Carlo vs our Semi-deterministic Monte Carlo method. SD-MC achieves lower bias and variance, with most benefits obtained using only $2$–$3$ points.
• Figure 4: Test accuracy with different training iterations and ELBO estimators on the Countdown dataset.
• Figure 5: Reward curves during RL training for the models reported in Table \ref{['tab:math-planning']}.

Theorems & Definitions (14)

• Lemma B.1
• proof
• Remark B.1
• Proposition B.1
• Remark B.2
• Proposition B.2
• proof
• Lemma B.2
• proof
• Remark B.3