RL Post-Training Diffusion Language Models Along Their Own Generation Paths

Vaibhav Singh; Eugene Belilovsky

cs.AIOct 12, 2026

RL Post-Training Diffusion Language Models Along Their Own Generation Paths

Authors: Vaibhav Singh, Eugene Belilovsky

Organizations: Mila, Concordia University

Abstract

Diffusion language models are a promising and more efficient alternative to autoregressive models. However, most RL post-training techniques are not easily adapted to these models. One issue is that diffusion language models generate answers through many intermediate states, deciding at each step which masked positions to reveal and which tokens to place in them. Many existing RL methods used in this setting approximate intractable sequence likelihoods or weight every denoising step by the same advantage, without explicitly distinguishing how these intermediate states of the denoising trajectory affect the expected reward. To address this issue of trajectory blindness, we introduce PACT (Potential-Anchored Consistency of Transitions). Motivated by a reward-tilted reference distribution, PACT jointly learns the policy and a scalar potential that scores each intermediate state, anchored by terminal reward. For each sampled denoising transition, a consistency loss matches the policy-to-reference log-probability ratio of revealed tokens to the change in this potential. This formulation requires neither sequence-likelihood estimation nor importance-ratio clipping. On four tasks, PACT improves LLaDA-8B-Instruct by +9.9 points on GSM8K, +12.0 on MATH500, +55.7 on Countdown, and +72.6 on Sudoku, and reaches the best average accuracy across the four tasks reported for this model, +3.7 points over the next best method. Further compared to AGRPO, it reaches that baseline's best accuracy in about half the training compute on GSM8K.

Figures & tables

Figure 1: PACT has the best average accuracy across the four benchmarks. Accuracy on LLaDA-8B-Instruct at generation length 256, for three likelihood-based and three likelihood-free methods against PACT. PACT holds the highest GSM8K, MATH500 and Sudoku accuracy, on MATH500 and Sudoku by about one standard error; on Countdown, DTM stays ahead. The last panel is the average over the four benchmarks. Full details are mentioned in Table 1 .

Figure 2: Overview of PACT training. (a) PACT scores a few transitions of each sampled path; every state has a potential

f=\log h

, equal to

r/\beta

at the answer and

0

after a failed one. (b) A small potential head predicts

f

from the DLLM’s hidden state, pulled up by correct rollouts and down by failed ones. (c) On each scored step, PACT minimizes

\delta_{t}^{2}=(u_{t}-g_{t})^{2}

, where

u_{t}

is the revealed tokens’ log-ratio to the frozen reference and

g_{t}

the gain in log potential ( Equation 5 ). Numbers are illustrative.

Figure 3: PACT’s rollout reward leads AGRPO’s on all four tasks throughout training. Mean rollout reward against training step, for all four tasks against AGRPO trained on the same rollout budget, rescaled to

[0,1]

. The faint trace is the per-step value as logged, the solid line a 25-step moving average, the band

\pm 1

rolling standard deviation over the same window. On every task PACT’s moving average lies above AGRPO’s for the whole run; the gap is widest on Sudoku, where PACT reaches

0.90

by step 291 while AGRPO stands at

0.39

at that step and reaches

0.84

only at step 900.

Figure 4: PACT reaches AGRPO’s best GSM8K accuracy in about half the compute. Held-out accuracy at generation length 256 against training GPU-hours (16 H100s), at every evaluated checkpoint up to a shared budget of

150

GPU-hours. The stars mark AGRPO’s best checkpoint,

85.1

after

53.4

GPU-hours, and the point where PACT first reaches it, after

28.6

: about

2\times

saving. At matched compute, PACT leads by

2.6

4.0

and

5.1

points at

50

100

and

150

GPU-hours.

Likelihood-based methods
	GSM8K		MATH500		Countdown		Sudoku		Average
	256	512	256	512	256	512	256	512	256	512
Base	79.2	79.8	32.4	34.6	24.2	30.9	26.9	18.5	40.7	41.0
d1 ( Zhao et al., 2025 )	81.1	82.1	38.6	40.2	32.0	42.2	32.5	29.3	46.1	48.5
wd1 † ( Tang et al., 2025 )	84.4	85.4	38.6	41.4	60.9	58.6	32.1	22.5	54.0	52.0
SPG ( Wang et al., 2025a )	86.1	84.5	40.0	41.8	70.7	70.3	94.0	93.1	72.7	72.4

Table 1: PACT has the best average accuracy at both lengths. Accuracy (%) on LLaDA-8B-Instruct at lengths 256 and 512; bold marks the best value in each column , underline the second best, and the coloured offset is PACT’s margin over the best other method (blue ahead, red behind). PACT holds the largest value in every column but Countdown at 256, where DTM leads, and ties DTM on Sudoku at 512; the MATH500 and Sudoku margins are within one standard error. † Our reproduction with the authors’ code, 900 rollout rounds, decoded as in Section 4.3 . Sudoku is 3-shot on SPG’s split in every row but RSPO and DAM (zero-shot).

Figure 5: Ablations: training reward. Mean rollout reward on GSM8K, one panel per knob of Table 2 , as a 100-step moving average. The arms separate more here than in their held-out accuracy: the near-random orders and the weaker or stronger tilts roll out worst, yet of the four knobs

\beta

moves the final accuracy the most.

$K$	%	$\tau_{o}$	%	$\gamma_{\mathrm{ord}}$	%	$\beta$	%
$8$	$88.4$	$0.05$	$88.8$	off	89.1	$0.05$	$89.1$
$\mathbf{16}$	$\mathbf{89.1}$	$\mathbf{0.2}$	$\mathbf{89.1}$	shared	$87.7$	$\mathbf{0.1}$	$\mathbf{89.1}$
$32$	$88.6$	$1$	$88.3$	head	$88.6$	$0.2$	$87.3$
$64$	$89.4$	$5$	$89.3$			$0.5$	$86.0$

Table 2: Ablations. Best held-out GSM8K accuracy (%) at length 256 per knob value; default in bold . Seconds per step for

K{=}8/16/32/64

32.8/41.1/58.3/92.0

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
base model	LLaDA-8B-Instruct
LoRA rank / $\alpha$ / dropout	128 / 128 / 0.05
LoRA targets	all attention + MLP projections
backbone precision	bfloat16 (unquantized)
adapters trained	LoRA only, no SFT stage
optimizer	AdamW ( $\beta_{1}{=}0.9$ , $\beta_{2}{=}0.99$ )

Appendix

Table 3: PACT training configuration , shared across all four tasks except the learning rate.

Method	GSM8K (%)
TraFL, $\beta=0.05$	81.5
TraFL, $\beta=1$	83.4
TraFL, $\beta=10$	72.0
PACT (ours)	89.1

Appendix

Table 4: Trajectory balance under PACT’s configuration. Best held-out GSM8K accuracy at length 256, one decode protocol for every row. The untuned model scores

79.2

Language Model Post-Training

Language model post-training adapts pretrained models to follow instructions, reason, use tools, and meet task-specific or human-defined goals.

See all 197 papers

May 13, 2026

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

Diffusion language models are a promising alternative to autoregressive models, yet post-training methods for them largely adapt reward-maximizing objectives. We identify a central failure mode in this setting we call trajectory locking: sampled reward-driven updates over-concentrate probability mass onto a narrow set of denoising paths, reducing coverage of alternative correct solutions under repeated sampling. To address this, we propose TraFL (Trajectory Flow baLancing), a trajectory-balance objective that trains the policy toward a reward-tilted target distribution anchored to a frozen reference model. We make this practical for diffusion language models with a diffusion-compatible sequence-level surrogate and a learned prompt-dependent normalization. Across mathematical reasoning and code generation benchmarks, TraFL is the only evaluated post-training method that improves over the base model in every benchmark-length setting, with gains that persist as the sampling budget increases. The improvements transfer to held-out evaluations: TraFL stays above the base model on Minerva Math and is the strongest method on every LiveCodeBench difficulty split.

Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

Noah’s Ark Lab

RL for Language Model Reasoning RL for Diffusion Models Diffusion Language Models Language Model Post-Training Diffusion Language Model Training

Ask PDF

May 8, 2026

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

Diffusion large language models are a compelling alternative to autoregressive models, yet existing RL methods for diffusion treat all denoising steps as equally important and rely on biased, high-variance likelihood estimates. We identify two fundamental weaknesses: the absence of temporal credit assignment across the denoising trajectory, and the systematic bias of mean-field likelihood estimates used for policy optimization. To address these, we propose Denoising-Aware Credit Assignment for GRPO (DACA-GRPO), a lightweight, plug-and-play enhancement for any GRPO-style trainer. DACA-GRPO introduces two complementary mechanisms: Denoising Progress Scores, which extract per-token importance weights from intermediate predictions at no additional forward cost, and Stratified Masking Likelihood, which partitions token positions into strata so that each token is predicted with most of the sequence as context, reducing the mean-field bias. Applied on top of three GRPO base methods, DACA-GRPO achieves consistent improvements across seven benchmarks spanning mathematical reasoning, code generation, constraint satisfaction, and constrained generation, with gains of up to 5.6pp on math reasoning, 7.4pp on code generation, 36.3pp on constraint satisfaction, and 5.9pp on JSON schema adherence.

Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade +4

The Ohio State University · Apple

Group Relative Policy Optimization Credit Assignment in RL Step-Level Credit Assignment in RL Step-Level Credit Assignment RL for Diffusion Models Diffusion Language Models+1

Ask PDF

Jun 7, 2026

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

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.

Yawen Shao, Jie Xiao, Kai Zhu +6

University of Science and Technology of China · Tongyi Lab · Northeastern University

RL for Language Model Reasoning Credit Assignment in RL RL for Diffusion Models Diffusion Language Models Reward Design for RL

Ask PDF

Abstract

Language Model Post-Training

Explore similar work

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

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