Revolucionando el Marco de Aprendizaje por Refuerzo para Modelos de Lenguaje de Difusión a Gran Escala

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

September 8, 2025

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

cs.AI

Resumen

Proponemos TraceRL, un marco de aprendizaje por refuerzo consciente de trayectorias para modelos de lenguaje basados en difusión (DLMs) que incorpora trayectorias de inferencia preferidas en el post-entrenamiento, siendo aplicable a diferentes arquitecturas. Equipado con un modelo de valor basado en difusión que mejora la estabilidad del entrenamiento, demostramos un mejor rendimiento en tareas complejas de razonamiento matemático y codificación. Además, también puede aplicarse para adaptar modelos específicos de bloques a bloques más grandes, lo que mejora la flexibilidad en el muestreo. Utilizando TraceRL, derivamos una serie de modelos de lenguaje basados en difusión de vanguardia, denominados TraDo. Aunque más pequeños que los modelos AR de escala 7B, TraDo-4B-Instruct supera consistentemente a estos en tareas complejas de razonamiento matemático. TraDo-8B-Instruct logra mejoras relativas en precisión del 6.1% sobre Qwen2.5-7B-Instruct y del 51.3% sobre Llama3.1-8B-Instruct en benchmarks de razonamiento matemático. Mediante aprendizaje curricular, también derivamos el primer DLM de cadena de pensamiento larga (long-CoT), superando a Qwen2.5-7B-Instruct en MATH500 con una ganancia relativa en precisión del 18.1%. Para facilitar la investigación reproducible y aplicaciones prácticas, publicamos un marco de código abierto integral para construir, entrenar y desplegar modelos de lenguaje grandes (LLMs) basados en difusión en diversas arquitecturas. El marco integra técnicas aceleradas de caché KV y motores de inferencia tanto para inferencia como para aprendizaje por refuerzo, e incluye implementaciones de varios métodos de ajuste fino supervisado y RL para tareas de matemáticas, codificación y generales. Código y Modelos: 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