Революционизация фреймворка обучения с подкреплением для диффузионных крупномасштабных языковых моделей

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

September 8, 2025

Авторы: Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang

cs.AI

Аннотация

Мы представляем TraceRL — фреймворк для обучения с подкреплением, учитывающий траектории, предназначенный для диффузионных языковых моделей (DLMs), который интегрирует предпочтительные траектории вывода в процесс пост-обучения и применим к различным архитектурам. Оснащённый диффузионной моделью ценности, которая повышает стабильность обучения, мы демонстрируем улучшенную производительность в сложных задачах математики и программирования. Кроме того, он также может быть использован для адаптации моделей, специфичных для блоков, к более крупным блокам, что повышает гибкость сэмплирования. Используя TraceRL, мы создаём серию передовых диффузионных языковых моделей, названных TraDo. Хотя TraDo-4B-Instruct меньше, чем AR-модели масштаба 7B, она стабильно превосходит их в сложных задачах математического рассуждения. TraDo-8B-Instruct достигает относительного улучшения точности на 6,1% по сравнению с Qwen2.5-7B-Instruct и на 51,3% по сравнению с Llama3.1-8B-Instruct в тестах на математическое рассуждение. С помощью обучения по учебному плану мы также создаём первую DLM с длинной цепочкой рассуждений (long-CoT), превосходящую Qwen2.5-7B-Instruct на MATH500 с относительным улучшением точности на 18,1%. Для обеспечения воспроизводимости исследований и практического применения мы выпускаем комплексный открытый фреймворк для создания, обучения и развёртывания диффузионных языковых моделей на различных архитектурах. Фреймворк интегрирует ускоренные техники KV-кэша и механизмы вывода как для инференса, так и для обучения с подкреплением, а также включает реализации различных методов тонкой настройки с учителем и RL для задач математики, программирования и общего назначения. Код и модели: 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