Source-linked AI summary
DiRL: An Efficient Post-Training Framework for Diffusion Language Models
Ying Zhu, Jiaxin Wan, Xiaoran Liu, Siyang He, Qiqi Wang, Xu Guo, Tianyi Liang, Zengfeng Huang, Ziwei He, Xipeng Qiu
TL;DR
Post-training for diffusion language models remains underdeveloped, with computational inefficiency and training–inference mismatch among the challenges. DiRL combines efficient blockwise training and inference with online updates, while DiPO provides an unbiased GRPO implementation. DiRL-8B-Instruct is reported as the strongest dLLM on math tasks and surpasses the Qwen2.5 series.
Problem
dLLM post-training, especially reinforcement learning, remains underdeveloped, with computational inefficiency and training–inference objective mismatch posing challenges.
Source: §1 · p. 2Method
DiRL combines FlexAttention-accelerated blockwise training with LMDeploy inference and online updates, while DiPO implements unbiased GRPO for dLLMs.
Source: §1 · p. 2Results
DiRL-8B-Instruct is reported as the strongest dLLM on math tasks and as outperforming the Qwen2.5 series.
Source: §6 · p. 8Takeaways & Limitations
The authors aim to provide the community with a solid baseline for dLLM post-training.
Source: §6 · p. 8Evaluation currently focuses on math; agentic and code tasks remain for future evaluation.
Source: §5.2 · p. 8
Abstract
from arXiv · showhide
Diffusion Language Models (dLLMs) have emerged as promising alternatives to Auto-Regressive (AR) models. While recent efforts have validated their pre-training potential and accelerated inference speeds, the post-training landscape for dLLMs remains underdeveloped. Existing methods suffer from computational inefficiency and objective mismatches between training and inference, severely limiting performance on complex reasoning tasks such as mathematics. To address this, we introduce DiRL, an efficient post-training framework that tightly integrates FlexAttention-accelerated blockwise training with LMDeploy-optimized inference. This architecture enables a streamlined online model update loop, facilitating efficient two-stage post-training (Supervised Fine-Tuning followed by Reinforcement Learning). Building on this framework, we propose DiPO, the first unbiased Group Relative Policy Optimization (GRPO) implementation tailored for dLLMs. We validate our approach by training DiRL-8B-Instruct on high-quality math data. Our model achieves state-of-the-art math performance among dLLMs and surpasses comparable models in the Qwen2.5 series on several benchmarks.
1 Introduction
Post-training, especially reinforcement learning, remains underdeveloped for dLLMs, while consistent training and inference and scalable RL remain open challenges. The paper introduces DiRL and DiPO and reports strong math benchmark performance for DiRL-8B-Instruct.
1 Introduction
- Although dLLM pre-training is feasible, post-training—especially RL—remains underdeveloped, and consistent training and inference for scalable RL are underexplored.Existing methods face biased logits, training–inference mismatch, and RL computational overhead.Source: §1 · p. 2
- DiRL-8B-Instruct achieves the best math-task performance among dLLMs and surpasses the Qwen2.5 Series on AIME24, AIME25, and OlympiadBench.The model is trained from SDAR-8B-Chat on high-quality math data.Source: §1 · p. 2
- DiRL supports two-stage SFT-RL post-training with training–inference consistency, FlexAttention acceleration, fast rollouts, and online model updates.Its inference-server-based design replaces offline model loading with online policy updates.Source: §1 · p. 2
- DiPO is presented as the first unbiased GRPO implementation for dLLMs, enabled by efficient, unbiased logit computation for blockwise dLLMs.Source: §1 · p. 2
2 Preliminary
Blockwise dLLMs generate sequences through sequential blocks with parallel masked diffusion inside each block. The section also outlines GRPO’s use in reinforcement learning and the computational and objective-alignment challenges of applying it to dLLMs.
2.1 Block Diffusion Language Model
- Blockwise dLLMs combine global sequential dependence with local parallel generation, conditioning each block’s masked-diffusion distribution on preceding blocks.A sequence is partitioned into non-overlapping blocks; diffusion operates within the current block while conditioning on historical context.Source: §2.1 · p. 3
- As time increases, noising progressively masks tokens in the current block, and denoising predicts that block’s original state in parallel.The noising process acts only on the current block; the reverse process conditions on its noisy state and preceding blocks.Source: §2.1 · p. 3
- Training minimizes the Conditional Negative Evidence Lower Bound over all blocks, using a time-dependent weight and cross-entropy loss.Source: §2.1 · p. 3
- Combining intra-block autoregression with inner-block diffusion allows KV caching, improving computational efficiency and simplifying likelihood computation.Source: §2.1 · p. 3
2.2 RL and its application in dLLM
- GRPO replaces PPO’s critic with diversified sampling and within-group normalized advantages, but dLLMs lack direct token probabilities and require costly, high-variance multi-step likelihood estimates.The random masking step map used in pre-training also mismatches realistic inference decoding.Source: §2.2 · pp. 3–4
- Wang et al. (2025b)’s TraceRL reduces log-probability overhead, but lacks attention acceleration, tight training–inference integration, and GRPO; DiRL and DiPO address these gaps.Source: §2.2 · p. 4
3 DiRL Post-Training
DiRL uses SFT followed by DiPO-based RL and reports state-of-the-art dLLM performance for DiRL-8B-Instruct. The section describes the math data and training setups, then introduces the RL objective and online updates.
3 DiRL Post-Training
- DiRL-8B-Instruct achieves state-of-the-art dLLM performance after two-stage post-training with SFT followed by DiPO-based RL.The section identifies this as its reported outcome.Source: §3 · p. 4
3.1 SFT stage
- SFT uses OpenR1-Math data distilled from GLM-4.6, with reasoning capped at 8k tokens because long-reasoning evaluation for diffusion models remains scarce.The paper describes 8k tokens as the longest inference length then reported for dLLMs.Source: §3.1 · p. 4
- SFT fine-tunes SDAR-8B-Chat for 100 steps on 8×H200 GPUs, using an 8k-token maximum length and global batch size 512.The setup uses LLaMA-Factory, DeepSpeed ZeRO1, a maximum learning rate of 1e-5, zero weight decay, and cosine annealing.Source: §3.1 · p. 4
3.2 RL stage
- DiPO defines the RL objective; online updates use a stop-gradient current-policy output as the behavior policy and incorporate DAPO’s token-level policy gradient.The objective assumes trajectories are decoded over steps, with τ_i(1:t−1) denoting the prefix through step t−1.Source: §3.2 · pp. 4–5
- Figure 4 compares DiRL and TraceRL attention masks for block size 2, prompt length 2, and output length 6, with loss calculated on the repeated output.Source: Figure 4 · p. 5
- RL uses the Big-Math dataset, described as large-scale, high-quality, math-verified, and designed for RL; training runs for 40 steps on 128×H200 GPUs.The setup uses 8k-token maximum length, global batch size 128, 32 trajectories per problem, and learning rate 1e-6.Source: §3.2 · p. 5
4 Engineering Optimization
DiRL combines blockwise dLLMs and FlexAttention with an inference-server-based online update loop to streamline post-training. Its reported evaluation shows improved training efficiency and strong math benchmark performance.
4 Engineering Optimization
- DiRL integrates FlexAttention-accelerated blockwise computation with LMDeploy and distributed training to support efficient online model updates.The framework is described as combining efficiency designs at both model and system levels.Source: §4 · p. 5
4.1 Blockwise dLLM with FlexAttention
- Blockwise denoising supports parallel training, while DiRL reshapes the SFT mask for FlexAttention’s fine-grained attention masks.Unlike TraceRL’s repeated output and altered attention mask, DiRL repeats both prompt and output blockwise.Source: §4.1 · p. 5
4.2 Training and Inference Integration
- DiRL keeps the LMDeploy API server alive, loads the model once, and pushes each training-step checkpoint through an in-place update.This replaces repeated filesystem loading and saving in the existing RL loop.Source: §4.2 · p. 6, Figure 5 · p. 6
4.3 Efficiency Validation
- FlexAttention cuts training latency by nearly 6×, while in-place updates contribute to a 2.5× overall throughput improvement.The reported comparison uses 8×H200 GPUs, batch size 4, input length 1024, and output length 8192; the training gain persists across model sizes, with DiRL 8B latency below TraceRL 1.7B latency.Source: §4.3 · p. 6, Figure 6 · p. 6, Figure 7 · p. 6
4.4 Main Results
- DiRL-8B-Instruct achieves the best average score and wins on most tasks, outperforming same-size dLLMs and AR models and exceeding Qwen2.5-32B-Instruct.The comparison covers five math benchmarks and reports accuracy alongside decoding-token and output-length averages.Source: Table 1 · p. 7, §4.4 · p. 7
5 Discussion
The ablation reports that DiRL-8B-Instruct remains ahead of baselines across dynamic sampling thresholds, with τ = 0.9 optimal for average performance. Future work targets larger models, longer reasoning, faster training, and broader task evaluation.
5.1 Ablation Study
- Across varying τ values, DiRL-8B-Instruct consistently and significantly outperforms baselines on every individual task and average performance; its performance rises steadily with τ, with τ = 0.9 optimal for average performance across models.SDAR-8B-Chat plateaus across most tasks, while TraDo-8B-Instruct has significant fluctuations on AIME2024.Source: §5.1 · p. 7
5.2 Future work
- Future work will scale beyond 8B, explore test-time scaling and long CoT, use dynamic packing, and extend evaluation beyond math to agentic and code tasks.The authors note that the current 8k inference length remains shorter than that of AR models.Source: §5.2 · pp. 7–8
6 Conclusion
The authors present DiRL for efficient SFT and RL in dLLMs, combining FlexAttention, refined attention masks, LMDeploy, and online parameter updates. They propose DiPO and report that DiRL-8B-Instruct is the strongest dLLM on math tasks, outperforming the Qwen2.5 series.
6 Conclusion
- DiRL-8B-Instruct is reported as the strongest dLLM on math tasks, outperforming the Qwen2.5 series.The authors plan to validate the approach on larger models, longer reasoning chains, and broader downstream tasks.Source: §6 · p. 8
- DiRL combines FlexAttention and refined attention masks to reduce blockwise dLLM training latency, with LMDeploy and online parameter updates supporting efficient two-stage post-training.The framework targets SFT and RL and is described as achieving unbiased logits computation.Source: §6 · p. 8
- DiPO is proposed as the first unbiased GRPO implementation for dLLMs.The authors aim to refine efficiency and provide an open-source post-training toolkit for dLLMs.Source: §6 · p. 8
A More Related Work of RL in dLLMs
Related work adapts policy optimization to dLLMs through approximate likelihood estimates, variance reduction, trajectory-level updates, and methods for sparse-reward gradients. These approaches trade off approximation bias, variance, memory use, or alignment with reasoning trajectories.
A More Related Work of RL in dLLMs
- diffu-GRPO uses one-step decoding to approximate sequence log-probability at low cost but introduces substantial approximation bias; Coupled-GRPO uses complementary masks to reduce estimate variance and stabilize updates.TraceRL instead segments generation into step/block trajectories and propagates learning signals step by step.Source: A More Related Work of RL in dLLMs · p. 12
- VRPO reduces preference-score variance through more ELBO samples and antithetic sampling, while BGPO enables per-sample backpropagation and gradient accumulation, reducing memory use from O(nt) to O(1).VRPO can cause GPU memory blow-ups because it retains all Monte Carlo computation graphs at once; GDPO uses Gauss–quadrature noise levels for low-variance sequence-level importance ratios.Source: A More Related Work of RL in dLLMs · p. 12
- TraceRL replaces random masking and one-step denoising with step/block trajectories, weighted token-level ratios, stepwise learning signals, and a diffusion-based value model for step-level GAE advantages.It is presented as addressing mismatch between random sampling or masking and the actual reasoning trajectory.Source: A More Related Work of RL in dLLMs · p. 12
- IGPO uses ground-truth fragments to inpaint fully incorrect samples into learnable correct trajectories and applies entropy filtering to stabilize their influence under sparse rewards.Source: A More Related Work of RL in dLLMs · p. 12