Yinjie Wang Ling Yang Bowen Li Ye Tian Ke Shen Mengdi Wang

摘要

我们提出了一种名为TraceRL的轨迹感知强化学习框架，用于扩散语言模型（Diffusion Language Models, DLMs），该框架在后训练阶段引入了偏好推理轨迹，适用于多种模型架构。通过配备基于扩散的值函数模型以增强训练稳定性，我们在复杂数学与编程任务上显著提升了模型的推理性能。此外，该方法还可用于将针对特定块（block）设计的模型适配至更大块结构，从而提升采样灵活性。基于TraceRL，我们构建了一系列当前领先的扩散语言模型，统称为TraDo。尽管TraDo-4B-Instruct的规模小于7B量级的自回归（AR）模型，但在复杂数学推理任务上仍持续优于后者。TraDo-8B-Instruct在数学推理基准测试中，相较Qwen2.5-7B-Instruct实现了6.1%的相对准确率提升，相较Llama3.1-8B-Instruct更是提升了51.3%。通过课程学习（curriculum learning）策略，我们进一步构建出首个长链思维（long-CoT）扩散语言模型，在MATH500基准上相较Qwen2.5-7B-Instruct实现了18.1%的相对准确率提升。为促进可复现研究与实际应用，我们开源了一个全面的框架，支持在多种架构下构建、训练与部署扩散大语言模型。该框架集成了加速的KV缓存技术与推理引擎，适用于推理与强化学习双重场景，并包含针对数学、编程及通用任务的多种监督微调与强化学习方法的实现。代码与模型地址：https://github.com/Gen-Verse/dLLM-RL