Revolutionair Reinforcement Learning Framework voor Diffusion Grote Taalmodellen

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

September 8, 2025

Auteurs: Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang

cs.AI

Samenvatting

We stellen TraceRL voor, een trajectbewust reinforcement learning-framework voor diffusie-taalmodelen (DLMs) dat de voorkeursinferentietrajectorie integreert in post-training en toepasbaar is op verschillende architecturen. Uitgerust met een diffusiegebaseerd waardemodel dat de trainingsstabiliteit verbetert, tonen we verbeterde redeneerprestaties aan op complexe wiskunde- en coderingstaken. Daarnaast kan het ook worden toegepast om blokspecifieke modellen aan te passen aan grotere blokken, wat de steekproefflexibiliteit verbetert. Met behulp van TraceRL leiden we een reeks state-of-the-art diffusie-taalmodelen af, genaamd TraDo. Hoewel kleiner dan 7B-schaal AR-modellen, presteert TraDo-4B-Instruct nog steeds consistent beter op complexe wiskundige redeneertaken. TraDo-8B-Instruct behaalt relatieve nauwkeurigheidsverbeteringen van 6,1% ten opzichte van Qwen2.5-7B-Instruct en 51,3% ten opzichte van Llama3.1-8B-Instruct op wiskundige redeneerbenchmarks. Door middel van curriculum learning leiden we ook het eerste lange-CoT DLM af, dat Qwen2.5-7B-Instruct overtreft op MATH500 met een relatieve nauwkeurigheidsverbetering van 18,1%. Om reproduceerbaar onderzoek en praktische toepassingen te faciliteren, geven we een uitgebreid open-source framework vrij voor het bouwen, trainen en implementeren van diffusie-LLM's over diverse architecturen. Het framework integreert versnelde KV-cache-technieken en inferentie-engines voor zowel inferentie als reinforcement learning, en bevat implementaties van verschillende supervised fine-tuning en RL-methoden voor wiskunde, codering en algemene taken. Code en Modellen: 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