DiRL: 拡散言語モデルのための効率的な学習後フレームワーク

要旨

拡散言語モデル（dLLM）は、自己回帰（AR）モデルに代わる有望なアプローチとして登場した。近年の研究ではその事前学習の可能性が検証され、推論速度の高速化が進められているが、dLLMの学習後調整（post-training）に関する技術基盤は未だ未成熟である。既存手法は計算効率の低さや、学習と推論の目的関数の不一致に悩まされており、数学のような複雑な推論タスクにおける性能が大幅に制限されている。この問題を解決するため、我々は効率的な学習後調整フレームワークであるDiRLを提案する。DiRLは、FlexAttentionにより高速化されたブロック単位学習とLMDeployによって最適化された推論を緊密に統合する。このアーキテクチャは、効率的なオンラインモデル更新ループを実現し、監督ファインチューニングと強化学習からなる二段階の学習後調整を効率的に実行する。このフレームワークに基づき、我々はdLLMに特化した初の不偏なグループ相対方策最適化（GRPO）実装であるDiPOを提案する。高品質な数学データを用いてDiRL-8B-Instructを学習し、本アプローチを検証した。その結果、本モデルはdLLMの中で数学タスクにおいて最先端の性能を達成し、いくつかのベンチマークではQwen2.5シリーズの同規模モデルを上回った。

English

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: 拡散言語モデルのための効率的な学習後フレームワーク

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

要旨

Support