Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models
作者: Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang
分类: cs.CL
发布日期: 2025-09-08
备注: Code and Models: https://github.com/Gen-Verse/dLLM-RL
🔗 代码/项目: GITHUB
💡 一句话要点
提出TraceRL，一种轨迹感知的扩散语言模型强化学习框架，提升推理性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 强化学习 轨迹感知 推理性能 数学推理 课程学习 思维链 值模型
📋 核心要点
- 现有扩散语言模型在复杂推理任务中面临挑战，需要更有效的训练方法来提升性能。
- TraceRL通过将偏好推理轨迹融入后训练，并结合扩散值模型，提升训练稳定性和推理能力。
- TraDo系列模型在数学推理任务上超越了同等规模的自回归模型，并在长程CoT任务上取得显著提升。
📝 摘要（中文）
本文提出TraceRL，一种轨迹感知的扩散语言模型（DLM）强化学习框架，该框架将偏好的推理轨迹融入到后训练中，并且适用于不同的模型架构。通过使用基于扩散的值模型来增强训练稳定性，该框架在复杂的数学和编程任务上表现出改进的推理性能。此外，TraceRL还可以应用于将特定块的模型适配到更大的块，从而提高采样灵活性。利用TraceRL，作者们得到了一系列最先进的扩散语言模型，即TraDo。尽管TraDo-4B-Instruct小于7B规模的自回归模型，但它在复杂的数学推理任务中仍然始终优于它们。TraDo-8B-Instruct在数学推理基准测试中，相对于Qwen2.5-7B-Instruct和Llama3.1-8B-Instruct分别实现了6.1%和51.3%的相对准确率提升。通过课程学习，作者们还得到了第一个长程思维链（CoT）DLM，在MATH500上优于Qwen2.5-7B-Instruct，相对准确率提升了18.1%。为了方便可复现的研究和实际应用，作者们发布了一个全面的开源框架，用于构建、训练和部署各种架构的扩散LLM。该框架集成了加速KV-cache技术和推理引擎，用于推理和强化学习，并包括各种监督微调和RL方法的实现，适用于数学、编码和通用任务。
🔬 方法详解
问题定义：论文旨在解决扩散语言模型（DLMs）在复杂推理任务（如数学和编程）中的性能瓶颈。现有的DLM训练方法可能无法充分利用有效的推理轨迹，导致性能受限，尤其是在需要长程推理的任务中。此外，训练过程的稳定性也是一个挑战。
核心思路：TraceRL的核心思路是将强化学习与扩散模型相结合，通过奖励模型引导DLM学习更优的推理轨迹。具体来说，它利用一个扩散值模型来估计轨迹的价值，并使用强化学习算法来优化策略，从而鼓励模型生成更符合人类偏好的推理过程。这种方法能够有效地利用数据中的信息，并提高模型的推理能力。
技术框架：TraceRL框架主要包含以下几个模块：1) 扩散语言模型（DLM）：作为生成模型，负责生成推理轨迹。2) 奖励模型：评估生成轨迹的质量，提供强化学习的奖励信号。3) 扩散值模型：估计轨迹的价值，用于指导策略优化，增强训练稳定性。4) 强化学习算法：根据奖励信号和价值估计，优化DLM的策略，使其生成更优的推理轨迹。整个流程通过迭代训练，不断提升DLM的推理性能。
关键创新：TraceRL的关键创新在于将轨迹信息融入到扩散语言模型的强化学习训练中。与传统的强化学习方法不同，TraceRL不仅关注最终结果，还关注整个推理过程，从而能够更有效地学习到有效的推理策略。此外，扩散值模型的使用增强了训练的稳定性，使得模型能够更好地收敛。
关键设计：TraceRL的关键设计包括：1) 奖励函数的设计：奖励函数需要能够准确地反映推理轨迹的质量，例如，对于数学问题，可以使用答案的正确性作为奖励信号。2) 扩散值模型的选择：扩散值模型需要能够准确地估计轨迹的价值，可以使用基于Transformer的模型。3) 强化学习算法的选择：可以使用PPO等常用的强化学习算法。4) 课程学习策略：为了提高训练效率，可以使用课程学习策略，先训练简单的任务，再逐步增加任务的难度。
🖼️ 关键图片
📊 实验亮点
TraDo-4B-Instruct在数学推理任务中优于7B规模的自回归模型。TraDo-8B-Instruct在数学推理基准测试中，相对于Qwen2.5-7B-Instruct和Llama3.1-8B-Instruct分别实现了6.1%和51.3%的相对准确率提升。在MATH500上，长程CoT DLM优于Qwen2.5-7B-Instruct，相对准确率提升了18.1%。
🎯 应用场景
TraceRL具有广泛的应用前景，可用于提升各种扩散语言模型在复杂推理任务中的性能，例如数学问题求解、代码生成、逻辑推理等。该框架还可以应用于教育、金融、医疗等领域，帮助人们更好地解决实际问题。此外，TraceRL的开源框架可以促进扩散语言模型的研究和应用。
📄 摘要（原文）
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