Rivoluzionare il Framework di Apprendimento per Rinforzo per Modelli Linguistici di Diffusione su Larga Scala

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

September 8, 2025

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

cs.AI

Abstract

Proponiamo TraceRL, un framework di apprendimento per rinforzo basato su traiettorie per modelli linguistici di diffusione (DLM) che incorpora traiettorie di inferenza preferite nel post-addestramento, applicabile a diverse architetture. Dotato di un modello di valore basato su diffusione che migliora la stabilità dell'addestramento, dimostriamo prestazioni migliorate in compiti complessi di matematica e programmazione. Inoltre, può essere applicato per adattare modelli specifici per blocchi a blocchi più grandi, migliorando la flessibilità del campionamento. Utilizzando TraceRL, deriviamo una serie di modelli linguistici di diffusione all'avanguardia, denominati TraDo. Sebbene più piccoli dei modelli AR su scala 7B, TraDo-4B-Instruct supera costantemente questi ultimi in compiti complessi di ragionamento matematico. TraDo-8B-Instruct ottiene miglioramenti relativi di accuratezza del 6,1% rispetto a Qwen2.5-7B-Instruct e del 51,3% rispetto a Llama3.1-8B-Instruct su benchmark di ragionamento matematico. Attraverso l'apprendimento curriculare, deriviamo anche il primo DLM a lunga catena di pensiero (long-CoT), superando Qwen2.5-7B-Instruct su MATH500 con un guadagno relativo di accuratezza del 18,1%. Per favorire la ricerca riproducibile e le applicazioni pratiche, rilasciamo un framework open-source completo per la costruzione, l'addestramento e il deployment di LLM di diffusione su diverse architetture. Il framework integra tecniche accelerate di KV-cache e motori di inferenza sia per l'inferenza che per l'apprendimento per rinforzo, e include implementazioni di vari metodi di fine-tuning supervisionato e RL per compiti di matematica, programmazione e generali. Codice e Modelli: 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