Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

我们提出了TraceRL，这是一种轨迹感知的强化学习框架，适用于扩散语言模型（DLMs），它将优选的推理路径纳入后训练过程，并可应用于不同架构。通过配备基于扩散的价值模型来提升训练稳定性，我们在复杂的数学和编程任务中展示了更优的推理性能。此外，该方法还可用于将块特定模型适配到更大的块上，从而提高采样灵活性。借助TraceRL，我们开发了一系列先进的扩散语言模型，统称为TraDo。尽管参数规模小于70亿的自回归模型，TraDo-4B-Instruct在多项复杂数学推理任务中仍持续优于这些模型。TraDo-8B-Instruct在数学推理基准测试中相较Qwen2.5-7B-Instruct提升了6.1%，相较Llama3.1-8B-Instruct则提升了51.3%。通过课程学习方法，我们还开发出了首个支持长思维链（long-CoT）的DLM，在MATH500数据集上相较Qwen2.5-7B-Instruct实现了18.1%的相对准确率提升。为了促进可重复研究和实际应用，我们发布了一个全面的开源框架，支持在多种架构上构建、训练和部署扩散大语言模型。该框架整合了加速KV缓存技术，以及适用于推理和强化学习的推理引擎，并包含了多种针对数学、编程和通用任务的监督微调与强化学习方法的实现。代码与模型请访问：https://github.com/Gen-Verse/dLLM-RL

论文旨在解决扩散语言模型（DLMs）在复杂推理任务（如数学和编程）中的性能不足问题，并探索如何提高其训练稳定性与推理灵活性。这属于当前生成模型研究的热点问题，但针对DLMs结合强化学习进行优化仍是一个较新的方向。

提出了一种轨迹感知的强化学习框架TraceRL，通过在后训练过程中引入偏好的推理轨迹，提升扩散语言模型在复杂任务上的表现。其核心创新在于使用基于扩散的价值模型来增强训练稳定性，并支持模型适应不同块结构，提升采样灵活性。

1. 通过TraceRL训练的TraDo模型在数学推理任务上显著优于同规模的自回归模型（如Qwen2.5-7B-Instruct和Llama3.1-8B-Instruct） 2. 首次实现了支持长链思维（long-CoT）的扩散语言模型，并在MATH500数据集上取得了显著的准确率提升 3. 提出的方法适用于不同架构的DLMs，并通过课程学习进一步优化模型推理能力 4. 论文作者开源了完整的训练、推理与强化学习框架，包含加速KV-cache技术与多种优化方法 5. 实验设计覆盖了数学、编码和通用任务，使用了MATH500等权威基准数据集

1. Diffusion Models for Language Generation (2022) 2. Training-Free Diffusion Model Inversion for Text Editing (2023) 3. Flow Matching for Generative Modeling (2023) 4. AR-LLMs: Advancements in Autoregressive Language Models (2023) 5. Curriculum Learning in Large Language Models: A Survey (2024)