Efficient Token-Level Policy Optimization for Masked Diffusion Models via Velocity Alignment

Chenxing Wei; Jiazhen Kang; Hong Wang; Jianqing Zhang; Hao Jiang; Xiaolong Xu; Ningyuan Sun; Ying He

cs.LGMar 2, 2026

Efficient Token-Level Policy Optimization for Masked Diffusion Models via Velocity Alignment

Authors: Chenxing Wei, Jiazhen Kang, Hong Wang, Jianqing Zhang, Hao Jiang, Xiaolong Xu, Ningyuan Sun, Ying He,

Organizations: Shenzhen University, China · Guangming Laboratory, China · ByteDance, China · Southeast University, China ·

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) successfully aligns autoregressive models in rigorous domains like mathematical reasoning and code generation. However, applying RLVR to Diffusion Large Language Models (dLLMs) is hindered by intractable sequence-level likelihoods, forcing existing methods to rely on unstable, high-variance approximations. We propose Sequence Likelihood-Free Policy Optimization (LFPO), an efficient framework that circumvents this bottleneck by formulating policy alignment as discrete velocity alignment. Inspired by continuous flow matching, LFPO eschews trajectory-level estimation and directly optimizes denoising logits via contrastive token-level updates, providing a stable, low-variance surrogate. Furthermore, predicting terminal solutions from intermediate states implicitly straightens the generation trajectory, enabling high-quality sampling with significantly fewer denoising iterations. Extensive experiments on reasoning and coding benchmarks demonstrate that LFPO matches or exceeds the accuracy of likelihood-based baselines while substantially accelerating inference.

Figures & tables

Figure 1: Overview of the LFPO framework. Step 1 Generate & Estimate : $\pi_{\text{old}}$ generates trajectories, with timesteps selected via Stratified Trajectory Sampling to minimize variance. Step 2 Block-Wise Optimization : Data is partitioned into memory-efficient blocks for parallel logit computation. Step 3 Policy Update : $\pi_{\theta}$ minimizes deviation from reward-induced implicit policies ( $\pi^{+},\pi^{-}$ ) via discrete vector field rectification. Step 4 Reference Update : Stable EMA update of $\pi_{\text{old}}$ . See Appendix C for method and Appendix E for implementation details.

Figure 2: Geometric Perspective of Discrete Lifting. Visualization on a 2-simplex (

\Delta^{2}

) where vertices represent tokens. Using a shared origin

\bm{o}_{t}

, the Target Velocity

u_{t}^{(o)}

(black) points toward the ground truth

\bm{x}_{1}

, while the Model Velocity

v_{\theta}^{(o)}

(green) points toward the model prediction

p_{\theta}

. The dashed line indicates the origin-invariant residual direction,

p_{\theta}-\bm{x}_{1}

, which corresponds exactly to the CE logit gradient

\nabla\mathcal{L}_{CE}

Model	HumanEval	HumanEval+	MBPP	MBPP+	EvalPlus	BigCodeBench	BigCodeBench	Avg.
						Full	Hard
DiffuCoder	72.0	65.2	75.1	61.9	63.6	35.7	12.2	55.1
+ diffu-GRPO	73.1 (+1.1)	66.1 (+0.9)	77.3 (+2.2)	67.3 (+5.4)	67.9 (+4.3)	39.2 (+3.5)	12.5 (+0.3)	57.6 (+2.5)
+ UniGRPO	74.5 (+2.5)	67.2 (+2.0)	78.9 (+3.8)	68.8 (+6.9)	70.3 (+6.7)	42.6 (+6.9)	12.9 (+0.7)	59.3 (+4.2)
+ SPG	74.6 (+2.6)	69.1 (+3.9)	80.2 (+5.1)	69.2 (+7.3)	68.9 (+5.3)	42.1 (+6.4)	12.9 (+0.7)	59.6 (+4.5)
+ coupled-GRPO	73.2 (+1.2)	68.3 (+3.1)	78.6 (+3.5)	67.5 (+5.6)	67.9 (+4.3)	40.4 (+4.7)	10.8 (-1.4)	58.1 (+3.0)

Table 1: Main Results on Code Generation Benchmarks. We compare LFPO against the base model DiffuCoder and various reinforcement learning baselines. Values in red parentheses denote absolute improvements over the base model. Standard deviations (

\pm

) across 5 random seeds are reported for the strongest baseline and our method to demonstrate statistical stability. LFPO (Pos. Only) and LFPO (Neg. Only) denote ablation variants optimized solely with the positive or negative loss, respectively, while LFPO (All Loss) utilizes the full contrastive velocity rectification objective. Best results are highlighted in bold , and second-best are underlined .

Model	GSM8K	MATH	GPQA	Hellaswag	WinoGrande	PIQA	Avg.
LLaDA 8B	69.7	30.6	25.3	70.3	75.8	71.6	57.2
+ diffu-GRPO	74.3 (+4.6)	33.7 (+3.1)	26.1 (+0.8)	80.1 (+9.8)	83.2 (+7.4)	79.3 (+7.7)	62.8 (+5.6)
+ UniGRPO	77.2 (+7.5)	34.5 (+3.9)	26.2 (+0.9)	82.3 (+12.0)	84.7 (+8.9)	81.2 (+9.6)	64.4 (+7.2)
+ SPG	76.2 (+6.5)	34.1 (+3.5)	25.9 (+0.6)	83.1 (+12.8)	84.3 (+8.5)	80.9 (+9.3)	64.1 (+6.9)
+ coupled-GRPO	75.1 (+5.4)	33.9 (+3.3)	25.7 (+0.4)	82.5 (+12.2)	85.6 (+9.8)	83.2 (+11.6)	64.3 (+7.1)
+ AGRPO	78.3 $\pm$ 1.3 (+8.6)	36.1 $\pm$ 0.9 (+5.5)	27.3 $\pm$ 1.1 (+2.0)	85.1 $\pm$ 1.0 (+14.8)	87.3 $\pm$ 1.2 (+11.5)	85.6 $\pm$ 0.8 (+14.0)	66.6 $\pm$ 1.0 (+9.4)

Table 2: Main Results on Reasoning Benchmarks. We compare LFPO against the base model LLaDA 8B and various reinforcement learning baselines. Formatting conventions, standard deviation reporting across 5 random seeds, and the nomenclature for ablation variants strictly follow those detailed in Table 1 .

Model	Code Generation				Reasoning
Model	HumanEval	MBPP	BCB-Full	Avg.	GSM8K	MATH	Hellaswag	Avg.
Base Model †	253.7	317.2	426.3	332.4	1034.2	1536.3	1036.7	1202.4
+ diffu-GRPO	254.3 (+0.6)	305.4 (-11.8)	430.7 (+4.4)	330.1 (-2.3)	1048.4 (+14.2)	1547.4 (+11.1)	1027.8 (-8.9)	1207.9 (+5.5)
+ UniGRPO	248.0 (-5.7)	319.3 (+2.1)	443.3 (+17.0)	336.9 (+4.5)	1020.3 (-13.9)	1539.3 (+3.0)	985.8 (-50.9)	1181.8 (-20.6)
+ SPG	243.7 (-10.0)	305.8 (-11.4)	416.7 (-9.6)	322.1 (-10.3)	1071.0 (+36.8)	1576.7 (+40.4)	1009.1 (-27.6)	1218.9 (+16.5)
+ coupled-GRPO	262.2 (+8.5)	326.5 (+9.3)	431.2 (+4.9)	340.0 (+7.6)	1015.9 (-18.3)	1601.6 (+65.3)	1025.8 (-10.9)	1214.4 (+12.0)

Table 3: Unified Efficiency Analysis. We report the average inference steps per problem across code generation and reasoning benchmarks. Avg. represents the mean steps within each task category. The values in red parentheses denote the reduction in inference steps (efficiency improvement) compared to the Base Model. Base Model refers to DiffuCoder for code tasks and LLaDA 8B for reasoning tasks. Lower values (

\downarrow

) indicate better efficiency.

Figure 3: Convergence Analysis on Code and Reasoning Tasks. We plot accuracy progression against training time under the matched hardware setting. The red and blue curves denote our proposed LFPO and the AGRPO baseline, respectively. The horizontal dashed line marks the baseline’s final converged accuracy. Notably, LFPO requires less time to match or surpass this peak performance, underscoring LFPO sample efficiency and rapid convergence.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Component	DiffusionNFT / Flow-GRPO	AGRPO	LFPO
Positive–negative or reward-induced target direction	Inspired by this line of diffusion RL/FM work Zheng et al. (2026) ; Liu et al. (2025)	Not the central design	Adopted as a detached contrastive target in discrete logit space
Discrete masked-dLLM token-level formulation	Continuous/image or flow-matching setting	Likelihood-based dLLM policy-gradient estimator	Token-level CE/logit surrogate for masked dLLMs
Sequence/trajectory likelihood handling	Not designed for masked dLLM sequence likelihood	Estimates or approximates policy-gradient quantities	Avoids sequence/trajectory likelihood estimation; retains token-level CE/logit supervision
Perturbation space	Model/flow-space target construction	Policy-gradient style update	Logit-space perturbation to keep simplex constraints smooth
Temporal credit assignment	Method-specific continuous-time sampling	Trajectory-level PG sampling	Stratified timestep sampling over denoising trajectories
Scalable implementation	Not focused on dLLM token-logit memory	Monte Carlo estimator for dLLMs	Block-wise logit parallelization for $B\times N\times K$ training samples

Appendix

Table 4: Positioning of LFPO relative to closely related RL/FM methods. The positive–negative implicit target idea is inspired by prior diffusion RL/FM work. The main contribution of LFPO is the discrete dLLM instantiation: token-level CE/logit supervision that avoids sequence-likelihood estimation, together with scalable sampling/optimization and standardized dLLM RLVR evaluation.

Aspect	Naive reward-weighted CE	LFPO
Target construction	Pull toward sampled high-reward outputs	Detached positive/negative implicit policies $(\pi^{+},\pi^{-})$
Geometric view	Supervised attraction heuristic	Shared-origin residual rectification
Update effect	Attraction-only signal	Pull high-reward directions and push away low-reward directions
Empirical check	Pos. Only variant	All Loss variant in Tables 1 – 2

Appendix

Table 5: Why the FM view changes the implementation. The algebraic CE gradient gives a tractable token-level update, while the shared-origin velocity view motivates the contrastive target construction used by LFPO.

Property	Probability Space	Log-Prob. Space	Logit Space (Ours)
Domain	$p_{i}\in[0,1]$	$\log p_{i}\leq 0$	$z_{i}\in\mathbb{R}$
Linear perturbation	Requires clipping	Shift-normalized	Unconstrained before Softmax
Simplex constraint	Broken before truncation	Recovered by normalization	Recovered by Softmax
Gradient continuity	Fragmented by clipping	Mostly smooth	Smooth relative update
HumanEval check	73.9 with clipping	–	75.6 with full LFPO

Appendix

Table 6: Properties of different spaces for policy perturbation. Probability-space perturbation uses clipping/truncation after a linear update, while LFPO constructs targets in logit space and maps them back to the simplex.

Hyperparameter	Value	Hyperparameter	Value
Contrastive Coefficient ( $\beta$ )	1.0	Block-wise Partition Size	8
Learning Rate (Policy)	$1\times 10^{-6}$	Trajectories per Prompt ( $N$ )	10
EMA Decay Rate ( $\alpha$ )	$\min(10^{-3}\times\text{step},0.5)$	Max Generation Length	512
Stratified Timesteps ( $K$ )	8	Max Diffusion Steps	2048
Optimizer	AdamW	Warm-up Ratio	0.1
Weight Decay	0.01	Gradient Clipping	1.0

Appendix

Table 7: Unified Hyperparameters for LFPO. This single configuration applies uniformly to all tasks and model architectures evaluated in our study.

Parameter	Value
Max Diffusion Steps ( $T$ )	2048
Noise Schedule	Cosine
Confidence Metric	Max softmax probability
Dynamic Parallel Update	$n_{t}=\lfloor L\gamma_{t}\rfloor$ revealed positions
Confidence Threshold ( $\tau$ )	0.99
Patience (Early Stopping)	3 identical sequence predictions

Appendix

Table 8: Inference hyperparameters and efficiency protocol. The parallel update size is dynamically determined by the cosine schedule, rather than by a separately tuned fixed top-

k

Model	Hardware / Decode Setting	Avg. Steps	Examples/s	Tokens/s	Generated Tokens/Forward
LLaDA 8B Base	8 $\times$ H20, max answer length 512, CoT prompt	1034.2	0.051	23	0.436
+ AGRPO	8 $\times$ H20, max answer length 512, CoT prompt	1004.8	0.053	24	0.451
+ LFPO	8 $\times$ H20, max answer length 512, CoT prompt	891.2	0.057	29	0.571

Appendix

Table 9: Practical decoding throughput on GSM8K. Tokens/s counts only generated answer tokens and excludes prompt tokens. Generated tokens/forward is computed as

(\text{Tokens/s}/\text{Examples/s})/\text{Avg. Steps}

Method	Checkpoint	Training Strategy	Hardware	Wall-clock Hours/Seed	GSM8K Acc. (%)
AGRPO	Converged	Diffusion PG estimator	8 $\times$ H20	152	$78.3\pm 1.3$
JustGRPO	Converged (step 253)	AR-style GRPO training	8 $\times$ H20	$\approx 153$	$79.4\pm 3.1$
LFPO	Converged	Token-level velocity alignment	8 $\times$ H20	$\approx 42$	$\mathbf{79.6\pm 0.3}$

Appendix

Table 10: Training comparison with JustGRPO on GSM8K using the LLaDA 8B base checkpoint. We include both wall-clock cost and accuracy because JustGRPO uses an AR-style GRPO training strategy, which is substantially more expensive in our reproduction. The JustGRPO wall-clock time is extrapolated from the measured step-200 runtime (121 hours).

Domain	Mean Difference	Approx. 95% CI	Performance Interpretation
Code Avg.	+0.2	$[-1.1,\,1.5]$	Comparable Mean, Superior Stability
Reasoning Avg.	+0.5	$[-0.8,\,1.8]$	Comparable Mean, Superior Stability

Appendix

Table 11: Approximate 95% confidence intervals for aggregate accuracy differences between LFPO and the strongest baseline (AGRPO). While the mean performance is strictly comparable, LFPO achieves this with significantly lower variance across seeds.

Control	Purpose	Observed Outcome
No stratified sampling	Tests whether dense temporal coverage matters	Uniform timestep sampling yields higher gradient variance and lower GSM8K accuracy (Table 15 )
No block-wise accumulation	Tests implementation feasibility of the full $B\times N\times K$ training tensor	Default configuration is not memory-feasible in our implementation; block-wise accumulation bounds peak VRAM to approximately 92.7 GB (Table 14 )
Probability space perturbation	Tests whether targets can be constructed directly in probability space	Requires clipping or truncation at simplex boundaries, which fragments gradients; we therefore enforce logit space perturbation (Table 6 )
Fixed EMA	Tests whether dynamic reference updates matter	Fixed $\alpha=0.9$ drops HumanEval accuracy to 70.8% (Table 14 )
Hard reference update	Tests whether replacing EMA with periodic old-policy replacement is stable	Hard update is substantially less stable in dLLM training and drops HumanEval accuracy to 60.1% (Table 13 )
Low $N$ , low $K$ , high $\beta$	Tests sensitivity to trajectories, timesteps, and perturbation scale	Reduces computational cost but lowers accuracy, or destabilizes training due to overly aggressive targets (Table 14 )

Appendix

Table 12: Ablation and negative control coverage. Numeric ablations are detailed in subsequent tables; feasibility controls are reported as implementation outcomes rather than quantitative accuracy comparisons.

Reference Update	Update Rule	HumanEval Acc. (%)
Dynamic EMA (LFPO)	$\alpha^{*}$	75.6
Fixed EMA	$\alpha=0.9$	70.8
Hard update	Periodic replacement	60.1

Appendix

Table 13: Reference update ablation on HumanEval. Hard update denotes periodic replacement of the old/reference policy by the current policy, while dynamic EMA uses

\alpha^{*}=\min(10^{-3}\times\text{step},0.5)

Config	$\beta$	$K$	$N$	$\alpha$	Peak VRAM (GB)	GPU Hours	HumanEval Acc. (%)
Default	1.0	8	10	$\alpha^{*}$	92.7	96	75.6
Low $N$	1.0	8	5	$\alpha^{*}$	92.6	49	74.1
Low $K$	1.0	4	10	$\alpha^{*}$	92.5	52	74.2
High $\beta$	1.5	8	10	$\alpha^{*}$	92.7	96	73.9
Fixed EMA	1.0	8	10	0.9	92.7	96	70.8

Appendix

Table 14: Core Hyperparameters vs. Training Costs. Evaluated on HumanEval. The dynamic EMA rate is defined as

\alpha^{*}=\min(10^{-3}\times\text{step},0.5)

Sampling Method	Gradient Variance	Accuracy (%)
Uniform	High	$72.3\pm 5.3$
Stratified (Ours)	Low	$79.6\pm 1.2$

Appendix

Table 15: Impact of Sampling Strategy on GSM8K (reported across 5 random seeds).

Masking Schedule	Threshold ( $\tau$ )	Accuracy (%)	Avg. Steps	Step Trade-off
Linear	0.99	76.3	229.6	Suboptimal step allocation.
Cosine (Ours)	1.00 (No Early Stop)	76.9	2048 (Max)	Upper bound accuracy, max cost.
Cosine (Ours)	0.99 (Default)	76.5	218.3	Favorable accuracy and step balance.
Cosine (Ours)	0.95	74.8	197.5	Fewer steps, minor accuracy drop.
Cosine (Ours)	0.90	73.5	173.2	Fewest steps, noticeable degradation.

Appendix

Table 16: Sensitivity Analysis of Inference Decoding Parameters on HumanEval.

Method	HumanEval (%)	MATH (%)
Base Model (Dream)	79.3	32.7
+ AGRPO	80.7	34.9
+ LFPO	81.9 (+1.2)	36.1 (+1.5)

Appendix

Table 17: Generalization capabilities of LFPO applied to Dream Models.

RL for Diffusion Models

Reinforcement learning for diffusion models adapts denoising and sampling to rewards, preferences, and task outcomes.

See all 179 papers

Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models
Jun 7, 2026Yawen Shao, Jie Xiao, Kai Zhu +6RL for Language Model Reasoning Credit Assignment in RL
University of Science and Technology of China · Tongyi Lab · Northeastern University
Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
Mar 13, 2026Vishnu Teja Kunde, Fatemeh Doudi, Mahdi Farahbakhsh +3RL for Diffusion Models Diffusion Language Model Training
Department of Electrical & Computer Engineering, Texas A&M University, College Station, TX 77843
Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Jul 18, 2026Haolin Ren, Ziyang Huang, Chenhao Yuan +2Masked Diffusion Models Diffusion Model Distillation
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

Jun 7, 2026J/K move · Enter open · S save

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

Yawen Shao, Jie Xiao, Kai Zhu, Yu Liu+5

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

Open paper AskPDF arXiv

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.

Abstract

RL for Diffusion Models

Explore similar work

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

Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages

Trace-Based On-Policy Distillation for Masked Diffusion Language Models