Révolutionner le cadre d'apprentissage par renforcement pour les modèles de langage massifs à diffusion

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

September 8, 2025

papers.authors: Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang

cs.AI

papers.abstract

Nous proposons TraceRL, un cadre d'apprentissage par renforcement conscient des trajectoires pour les modèles de langage à diffusion (DLM) qui intègre une trajectoire d'inférence privilégiée dans le post-entraînement, et qui est applicable à différentes architectures. Doté d'un modèle de valeur basé sur la diffusion qui améliore la stabilité de l'entraînement, nous démontrons une amélioration des performances de raisonnement sur des tâches complexes de mathématiques et de codage. De plus, il peut également être appliqué pour adapter des modèles spécifiques à des blocs à des blocs plus grands, ce qui améliore la flexibilité de l'échantillonnage. En utilisant TraceRL, nous dérivons une série de modèles de langage à diffusion de pointe, nommés TraDo. Bien que plus petits que les modèles AR de 7 milliards de paramètres, TraDo-4B-Instruct surpasse systématiquement ces derniers sur des tâches complexes de raisonnement mathématique. TraDo-8B-Instruct obtient des améliorations relatives de précision de 6,1 % par rapport à Qwen2.5-7B-Instruct et de 51,3 % par rapport à Llama3.1-8B-Instruct sur des benchmarks de raisonnement mathématique. Grâce à l'apprentissage curriculaire, nous dérivons également le premier DLM à long-CoT, surpassant Qwen2.5-7B-Instruct sur MATH500 avec un gain de précision relative de 18,1 %. Pour faciliter la recherche reproductible et les applications pratiques, nous publions un cadre open-source complet pour la construction, l'entraînement et le déploiement de LLM à diffusion sur diverses architectures. Le cadre intègre des techniques accélérées de cache KV et des moteurs d'inférence pour l'inférence et l'apprentissage par renforcement, et inclut des implémentations de diverses méthodes de fine-tuning supervisé et de RL pour les tâches de mathématiques, de codage et générales. Code et Modèles : 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