Revolucionando o Framework de Aprendizado por Reforço para Modelos de Linguagem de Grande Escala Baseados em Difusão

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

Resumo

Propomos o TraceRL, um framework de aprendizado por reforço consciente da trajetória para modelos de linguagem de difusão (DLMs) que incorpora trajetórias de inferência preferenciais no pós-treinamento, sendo aplicável a diferentes arquiteturas. Equipado com um modelo de valor baseado em difusão que aumenta a estabilidade do treinamento, demonstramos melhorias no desempenho de raciocínio em tarefas complexas de matemática e codificação. Além disso, ele também pode ser aplicado para adaptar modelos específicos de blocos a blocos maiores, o que melhora a flexibilidade de amostragem. Utilizando o TraceRL, derivamos uma série de modelos de linguagem de difusão de última geração, denominados TraDo. Embora menores do que modelos AR de escala 7B, o TraDo-4B-Instruct ainda supera consistentemente esses modelos em tarefas complexas de raciocínio matemático. O TraDo-8B-Instruct alcança melhorias relativas de precisão de 6,1% sobre o Qwen2.5-7B-Instruct e 51,3% sobre o Llama3.1-8B-Instruct em benchmarks de raciocínio matemático. Por meio de aprendizado curricular, também derivamos o primeiro DLM de long-CoT, superando o Qwen2.5-7B-Instruct no MATH500 com um ganho relativo de precisão de 18,1%. Para facilitar pesquisas reproduzíveis e aplicações práticas, disponibilizamos um framework open-source abrangente para construção, treinamento e implantação de LLMs de difusão em diversas arquiteturas. O framework integra técnicas aceleradas de cache KV e motores de inferência tanto para inferência quanto para aprendizado por reforço, e inclui implementações de vários métodos de ajuste fino supervisionado e RL para tarefas de matemática, codificação e gerais. Código e 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