拡散型大規模言語モデルのための強化学習フレームワークの革新

要旨

我々はTraceRLを提案する。これは、拡散言語モデル（DLM）向けの軌跡認識型強化学習フレームワークであり、推奨される推論軌跡をポストトレーニングに組み込み、様々なアーキテクチャに適用可能である。拡散ベースの価値モデルを備えることでトレーニングの安定性を向上させ、複雑な数学やコーディングタスクにおいて改善された推論性能を実証した。さらに、ブロック固有のモデルをより大きなブロックに適応させることも可能で、サンプリングの柔軟性が向上する。TraceRLを活用し、一連の最先端拡散言語モデルTraDoを導出した。7BスケールのARモデルよりも小規模であるにもかかわらず、TraDo-4B-Instructは複雑な数学的推論タスクにおいて一貫して優れた性能を示す。TraDo-8B-Instructは、数学的推論ベンチマークにおいてQwen2.5-7B-Instructに対して6.1%、Llama3.1-8B-Instructに対して51.3%の相対的な精度向上を達成した。カリキュラム学習を通じて、初の長文CoT DLMも導出し、MATH500においてQwen2.5-7B-Instructを18.1%の相対精度向上で上回った。再現可能な研究と実用的な応用を促進するため、様々なアーキテクチャにわたる拡散LLMの構築、トレーニング、デプロイのための包括的なオープンソースフレームワークを公開した。このフレームワークは、推論と強化学習の両方に対応した高速化KVキャッシュ技術と推論エンジンを統合し、数学、コーディング、一般タスクのための様々な教師あり微調整とRL手法の実装を含んでいる。コードとモデル: https://github.com/Gen-Verse/dLLM-RL

English

We propose TraceRL, a trajectory-aware reinforcement learning framework for diffusion language models (DLMs) that incorporates preferred inference trajectory into post-training, and is applicable across different architectures. Equipped with a diffusion-based value model that enhances training stability, we demonstrate improved reasoning performance on complex math and coding tasks. Besides, it can also be applied to adapt block-specific models to larger blocks, which improves sampling flexibility. Employing TraceRL, we derive a series of state-of-the-art diffusion language models, namely TraDo. Although smaller than 7B-scale AR models, TraDo-4B-Instruct still consistently outperforms them across complex math reasoning tasks. TraDo-8B-Instruct achieves relative accuracy improvements of 6.1% over Qwen2.5-7B-Instruct and 51.3% over Llama3.1-8B-Instruct on mathematical reasoning benchmarks. Through curriculum learning, we also derive the first long-CoT DLM, outperforming Qwen2.5-7B-Instruct on MATH500 with an 18.1% relative accuracy gain. To facilitate reproducible research and practical applications, we release a comprehensive open-source framework for building, training, and deploying diffusion LLMs across diverse architectures. The framework integrates accelerated KV-cache techniques and inference engines for both inference and reinforcement learning, and includes implementations of various supervised fine-tuning and RL methods for mathematics, coding, and general tasks. Code and Models: https://github.com/Gen-Verse/dLLM-RL

拡散型大規模言語モデルのための強化学習フレームワークの革新

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

要旨

Support