Source-linked AI summary
RL Post-Training Diffusion Language Models Along Their Own Generation Paths
Vaibhav Singh, Eugene Belilovsky
TL;DR
DLLM RL adaptation is challenging because exact sequence likelihoods are intractable and many methods do not distinguish transition-specific reward changes. PACT trains a reward-anchored potential alongside the policy to constrain sampled transitions, and achieves the best average accuracy across four benchmarks.
Problem
DLLMs require marginalizing over many unmasking paths for exact sequence likelihoods, and completion-level advantages do not explicitly estimate transition-specific changes in expected reward.
Source: §1 · p. 1Method
PACT jointly learns a policy and terminal-reward-anchored scalar potential, matching each sampled transition’s policy-to-reference log-probability ratio to the change in log potential.
Source: §1 · p. 2Results
PACT achieves the best average accuracy among DLLM post-training methods across four benchmarks, +3.7 points over the next best method at length 256.
Source: §1 · p. 2Takeaways & Limitations
With the ideal potential and all unmasking decisions scored, PACT’s transition condition reproduces the reward-tilted target distribution.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
DLLM RL adaptation is challenging because exact response likelihoods require marginalizing over unmasking paths, while common approaches do not track transition-specific reward changes. PACT links sampled transitions to learned, terminal-reward-anchored potentials and reports the best average accuracy among DLLM post-training methods across four benchmarks.
1 INTRODUCTION
- Exact sequence likelihoods are intractable because DLLMs require marginalizing over many unmasking paths, and completion-level advantages weight denoising steps alike.Those methods do not explicitly estimate how each transition changes expected reward.Source: §1 · p. 1
- Distribution-matching methods fit reward-tilted targets on completed responses or constructed intermediate states, but leave open whether unfinished-answer value can guide completion.The supplied text frames the missing link as connecting each sampler-path transition to changes in intermediate-state value.Source: §1 · pp. 1–2
- +3.7 points over the next best method at length 256, PACT has the best average accuracy across four benchmarks.PACT also improves over the base model by +9.9 to +72.6 points and matches AGRPO’s best GSM8K accuracy in about half the training compute.Source: Figure 1 · p. 2, §1 · p. 2
- PACT learns a scalar potential for each partial response, anchors it with terminal rewards, and matches each sampled transition’s policy-to-reference log-probability ratio to the change in log potential.This provides local training targets without sequence-likelihood estimation or importance-ratio clipping.Source: §1 · p. 2
2 RELATED WORK
Prior DLLM post-training methods use likelihood surrogates, transition-based objectives, or reward-guided distribution matching. PACT’s distinction is an explicit reference-relative constraint connecting each sampled transition to learned potential changes along the model’s own denoising trajectory.
Likelihood-based reward maximization
- Trajectory-based methods such as AGRPO, LLaDOU, and MAPG already optimize intermediate decisions; PACT adds a constraint linking each sampled transition to learned potential change.Sequence-likelihood surrogates include d1, wd1, SPG, and RSPO; wd1++ scores intermediate completions with group advantage but does not learn potential changes between successive states.Source: §2 · pp. 2–3
Likelihood-free distribution matching
- DMPO and DTM incorporate reward into local predictions, while PACT learns a scalar potential and constrains its change across recorded denoising transitions.This directly relates the policy’s preference for a sampled transition to predicted improvement in the partial response.Source: §2 · p. 3
Flow balance, relative to a reference
- PACT’s reference-relative detailed-balance condition equates each sampled transition’s policy-to-reference probability ratio with its successor-to-current potential ratio.At the residual’s fixed point, the fully masked state’s potential equals log Zq, so no separate normalizer network is fitted.Source: §2 · p. 3
3 METHOD: POTENTIAL-ANCHORED CONSISTENCY
PACT trains a policy and a reward-anchored potential over partial responses, matching each sampled transition’s revealed-token log-ratio to the change in potential. The method derives this local target from reward-tilted reference transitions and uses sampled denoising paths without estimating sequence likelihoods.
Denoising trajectories
- A trajectory runs from a masked answer window to a completed response, revealing selected masked slots and committing tokens that are conditionally independent given slot selection.Slot selection uses the policy’s confidence scores; the policy learns the order only implicitly, and the loss does not train it.Source: §3 · p. 3
Reward-tilted target
- The ideal potential measures expected exponential reward under reference-model completion, with h(x0) = e^r(y)/β at the answer and h(xT) = Zq at the masked start.Reward tilting reweights each reference transition according to the successor-to-current potential ratio.Source: §3 · pp. 3–4
PACT objective
- PACT matches each sampled transition’s revealed-token log-ratio to its change in log potential, with reward anchors tying intermediate scores to terminal rewards.The anchor uses scaled reward r(y)/β as a practical regression target rather than an unbiased estimate of the ideal log potential.Source: §3 · p. 4, Figure 2 · p. 4
Implementation
- A two-layer potential head reads the mean-pooled answer-window hidden state, while K transitions stratified across each rollout estimate the transition expectation.The head shares the policy’s forward pass, gradients flow through it into the policy adapters, and decoding uses the policy alone.Source: §3 · p. 5
Relation to trajectory balance
- Summing full-path residuals cancels intermediate potentials, so a small trajectory-level error can conceal large local inconsistencies that PACT penalizes separately.The residual also gives a consistency gradient on prompts whose completions all receive the same reward, where group-relative advantage is zero.Source: §3 · p. 5
4 EXPERIMENTS
PACT is evaluated across four tasks against matched-budget and published baselines; it has the best average accuracy at both tested generation lengths, while its rollout reward stays above AGRPO’s throughout training.
4.1 TRAINING DETAILS
- PACT is trained on LLaDA-8B-Instruct without supervised fine-tuning, using LoRA with r = 128 and 900 rollout batches of 128 completions.Training uses 128 denoising steps, generation length 256, and semi-autoregressive blocks of 32.Source: §4.1 · pp. 5–6
- Across all four tasks, PACT’s 25-step moving-average rollout reward stays above AGRPO’s throughout matched-budget training, with the widest gap on Sudoku.At step 291 on Sudoku, PACT reaches 0.90 while AGRPO is at 0.39; AGRPO reaches 0.84 only at step 900.Source: Figure 3 · p. 6
4.2 DATASETS
- The four tasks span grade-school and competition math, arithmetic-expression construction, and 4×4 Sudoku, with task-specific rewards rescaled to [0, 1].GSM8K and MATH500 use format plus exact-match rewards; Countdown checks arithmetic and number use, while Sudoku scores the fraction of empty cells filled correctly.Source: §4.2 · p. 6
4.3 EVALUATION
- Evaluation uses greedy Pass@1 at lengths 256 and 512, reports each method’s best checkpoint, and retrains AGRPO, wd1, and DMPO for matched 900-round comparisons.Other baseline rows use each paper’s self-reported results.Source: §4.3 · p. 6
- PACT has the best average accuracy on LLaDA-8B-Instruct at lengths 256 and 512, leading every column except Countdown at 256 and tying DTM on Sudoku at 512.The MATH500 and Sudoku margins are within one standard error.Source: Table 1 · p. 7
5 RESULTS
PACT improves the base model on all four benchmarks and has the best reported average accuracy across published methods. Its compute-matched GSM8K results favor PACT over AGRPO, while ablations highlight the roles of reward tilt, order gradients, and transition-level constraints.
5 RESULTS
- +9.9 on GSM8K, +12.0 on MATH500, +55.7 on Countdown, and +72.6 on Sudoku versus the base model under the same rollout budget.The largest gains are on Countdown and Sudoku, where base-model accuracy is 24.2 and 26.9, respectively.Source: §5 · p. 7
- 78.2 versus DTM’s 74.5 is PACT’s best average accuracy across published methods at length 256; DTM leads PACT on Countdown at that length, 81.3 to 79.9.PACT scores 80.7 on Countdown at length 512, the highest reported there.Source: §5 · p. 7
- 2.6, 4.0, and 5.1 points are PACT’s GSM8K leads over AGRPO at matched compute of 50, 100, and 150 GPU-hours.PACT reaches AGRPO’s best accuracy after 28.6 GPU-hours, compared with AGRPO’s 53.4 GPU-hours; the crossing depends on checkpoint noise.Source: Figure 4 · p. 7, §5 · p. 7
5.1 ABLATIONS
- PACT outperforms TraFL under the shared rollout configuration: all three TraFL settings peak within 300 steps and decline, while PACT is still improving at step 775.TraFL uses one residual per completion, leaving errors across its 128 transitions able to cancel; its comparison uses the authors’ setup with greedy Pass@1, not their Pass@5.Source: Figure 6 · p. 8, §5.1 · p. 9
5 RESULTS
- PACT evaluates the policy-to-reference ratio on the sampled transition, avoiding a marginal over unmasking orders, and assigns credit according to each step’s potential change.Unlike a policy gradient that weights every token by one scalar advantage, the residual raises steps that increase potential and keeps steps with no gain close to the reference.Source: §5 · p. 8
5.1 ABLATIONS
- Explicitly training the unmasking order with γord=1 lowers accuracy in both tested implementations, so PACT keeps the order stochastic but ungraded at τo=0.2.Sweeping τo from near-greedy to near-random leaves accuracy flat; the near-random arm rolls out worst.Source: §5.1 · pp. 8–9, Figure 5 · p. 8
- β has the largest effect on held-out accuracy among the tested knobs: tilts of 0.2 and 0.5 score 87.3 and 86.0, while β=0.1 is retained as the default.Because β both sets the target through r/β and scales the gradient, a fixed learning rate cannot separate those effects; the authors do not anneal β.Source: §5.1 · p. 9
6 CONCLUSION
PACT trains diffusion policies with transition-level consistency between revealed-token policy/reference ratios and changes in a reward-anchored potential. It improves accuracy across four reasoning tasks and reaches AGRPO’s best GSM8K accuracy in about half the training compute.
6 CONCLUSION
- PACT improves on all four reasoning tasks and achieves the best average accuracy across them.It reaches AGRPO’s best GSM8K accuracy in about half the training compute.Source: §6 · p. 9
- PACT matches each sampled transition’s policy-to-reference probability ratio to the potential gained at that step.A learned potential head carries reward from the completed answer to earlier states; the ideal condition is the Doob h-transform of the reference sampler.Source: §6 · p. 9
- The fully masked state’s potential acts as a normalizer, so PACT needs neither sequence-likelihood estimates nor clipped importance ratios.Source: §6 · p. 9
- PACT’s training signal can remain nonzero when all prompt rollouts receive the same reward, but then it is a consistency correction, not reward information.Source: §6 · p. 9
A TRAINING CONFIGURATIONS
The experiments use a shared PACT recipe across four tasks, with task-specific learning rates, and evaluate checkpoints under a common greedy Pass@1 protocol. AGRPO is run with its released configuration and a matched training budget.
A TRAINING CONFIGURATIONS
- PACT uses one recipe across four tasks except for learning rate, with β=0.1, τo=0.2, K=16 transitions, and a token-only objective.Its potential head is a zero-initialized two-layer MLP of width 256; the reference is the same network with adapters disabled.Source: A TRAINING CONFIGURATIONS · p. 13, Table 3 · p. 13
- AGRPO uses its released configuration, while the rerun uses the paper’s current code and the upper end of its 400–900 training-step range.The paper reports an unresolved discrepancy between its prior Countdown result and the rerun, and reports the rerun as the only row trained on the matched budget and decoded under its protocol.Source: A TRAINING CONFIGURATIONS · p. 13
- Evaluation uses greedy decoding at lengths 256 and 512, while PACT’s training rollouts use 256 tokens on every task.Evaluation uses low-confidence order, block length 32, and denoising steps equal to half the generation length.Source: A TRAINING CONFIGURATIONS · p. 13
B A MATCHED COMPARISON WITH TRAJECTORY BALANCE
In the matched trajectory-balance comparison, PACT leads at shared checkpoints, while completion-level residuals constrain the trajectory only through one scalar. The comparison is bounded by the reproduced setup and a sparse sweep of β values.
B A MATCHED COMPARISON WITH TRAJECTORY BALANCE
- TraFL uses one residual per completion, based on a masked-reconstruction likelihood surrogate, leaving all 128 denoising steps constrained through one scalar.Its objective also uses a learned per-prompt normalizer.Source: B A MATCHED COMPARISON WITH TRAJECTORY… · p. 13
- 12.8 points at β=1 and 8.6 points at β=0.05: PACT leads TraFL past step 600 and is ahead at every shared checkpoint.Source: B A MATCHED COMPARISON WITH TRAJECTORY… · p. 14, Table 4 · p. 14
- Trajectory balance’s completion-level quantity moves only when all transitions move together, and its usable β range is narrow; PACT’s per-transition residual constrains each step.Under the comparison configuration, every β in Table 2 lands between 86.0 and 89.1; trajectory balance peaks within 300 steps and declines, while PACT is still improving at step 775.Source: B A MATCHED COMPARISON WITH TRAJECTORY… · p. 14
- The comparison uses the authors’ implementation under their sampler, reward, and metric, tests only three widely spaced β values, and does not reproduce Pass@5.Source: B A MATCHED COMPARISON WITH TRAJECTORY… · pp. 13–14