Revolutionierung des Reinforcement-Learning-Frameworks für Diffusion Large Language Models

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

Wir stellen TraceRL vor, ein trajektorienbewusstes Reinforcement-Learning-Framework für Diffusionssprachmodelle (DLMs), das bevorzugte Inferenztrajektorien in das Post-Training integriert und auf verschiedene Architekturen anwendbar ist. Ausgestattet mit einem diffusionsbasierten Wertmodell, das die Trainingsstabilität verbessert, demonstrieren wir verbesserte Leistungen bei komplexen mathematischen und Programmieraufgaben. Darüber hinaus kann es auch angewendet werden, um block-spezifische Modelle an größere Blöcke anzupassen, was die Sampling-Flexibilität erhöht. Mit TraceRL leiten wir eine Reihe von state-of-the-art Diffusionssprachmodellen ab, nämlich TraDo. Obwohl kleiner als 7B-skalierte AR-Modelle, übertrifft TraDo-4B-Instruct diese dennoch konsistent bei komplexen mathematischen Denkaufgaben. TraDo-8B-Instruct erzielt relative Genauigkeitsverbesserungen von 6,1 % gegenüber Qwen2.5-7B-Instruct und 51,3 % gegenüber Llama3.1-8B-Instruct auf mathematischen Denkbenchmarks. Durch Curriculum Learning leiten wir auch das erste lang-CoT DLM ab, das Qwen2.5-7B-Instruct auf MATH500 mit einer relativen Genauigkeitssteigerung von 18,1 % übertrifft. Um reproduzierbare Forschung und praktische Anwendungen zu erleichtern, veröffentlichen wir ein umfassendes Open-Source-Framework für den Aufbau, das Training und die Bereitstellung von Diffusions-LLMs über verschiedene Architekturen hinweg. Das Framework integriert beschleunigte KV-Cache-Techniken und Inferenz-Engines sowohl für die Inferenz als auch für das Reinforcement Learning und umfasst Implementierungen verschiedener überwachter Feinabstimmungs- und RL-Methoden für Mathematik, Programmierung und allgemeine Aufgaben. Code und Modelle: 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