Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DiRL: An Efficient Post-Training Framework for Diffusion Language Models

Ying Zhu, Jiaxin Wan, Xiaoran Liu, Siyanag He, Qiqi Wang, Xu Guo, Tianyi Liang, Zengfeng Huang, Ziwei He, Xipeng Qiu

TL;DR

This work targets the inefficient post-training of diffusion language models (dLLMs) by introducing DiRL, a two-stage framework (SFT followed by RL) that tightly couples training with an inference backend for online updates. It also presents DiPO, the first unbiased Group Relative Policy Optimization (GRPO) implementation for dLLMs, enabled by blockwise diffusion and FlexAttention to compute logits efficiently. Trained on OpenR1-Math, the resulting DiRL-8B-Instruct achieves state-of-the-art math performance among dLLMs and often surpasses larger autoregressive baselines like Qwen2.5-32B-Instruct on multiple benchmarks, demonstrating the practical impact of efficient post-training and unbiased RL for complex reasoning tasks. The approach combines a blockwise BDLM foundation, efficient training–inference co-design, and targeted RL data to deliver a scalable, open-source post-training toolkit for diffusion-based language models with strong math reasoning capabilities.

Abstract

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.

DiRL: An Efficient Post-Training Framework for Diffusion Language Models

TL;DR

This work targets the inefficient post-training of diffusion language models (dLLMs) by introducing DiRL, a two-stage framework (SFT followed by RL) that tightly couples training with an inference backend for online updates. It also presents DiPO, the first unbiased Group Relative Policy Optimization (GRPO) implementation for dLLMs, enabled by blockwise diffusion and FlexAttention to compute logits efficiently. Trained on OpenR1-Math, the resulting DiRL-8B-Instruct achieves state-of-the-art math performance among dLLMs and often surpasses larger autoregressive baselines like Qwen2.5-32B-Instruct on multiple benchmarks, demonstrating the practical impact of efficient post-training and unbiased RL for complex reasoning tasks. The approach combines a blockwise BDLM foundation, efficient training–inference co-design, and targeted RL data to deliver a scalable, open-source post-training toolkit for diffusion-based language models with strong math reasoning capabilities.

Abstract

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.
Paper Structure (24 sections, 8 equations, 8 figures, 1 table)

This paper contains 24 sections, 8 equations, 8 figures, 1 table.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Block Diffusion Language Model
  4. Forward Process (Noising Process)
  5. Reverse Process (Denoising Process)
  6. RL and its application in dLLM
  7. DiRL Post-Training
  8. SFT stage
  9. SFT Data
  10. SFT setup
  11. RL stage
  12. DiPO
  13. RL Data
  14. RL setup
  15. Engineering Optimization
  16. ...and 9 more sections

Figures (8)

  • Figure 1: Performance of DiRL-8B-Instruct.
  • Figure 2: Features of our DiRL framework.
  • Figure 3: Overview of the RL pipeline in our DiRL framework.
  • Figure 4: The visualized comparison of the attention mask between our DiRL framework and TraceRL, where block size is 2, the length of prompt colored in green is 2, and the length of output colored in blue is 6. The loss is calculated based on the repeated output part colored in full purple.
  • Figure 5: The training and inference integration of our DiRL compared to RL in current dLLMs.
  • ...and 3 more figures