TraceRL框架：解锁扩散语言模型的高效优化

鉴于自回归大语言模型（LLM）在计算效率和吞吐量上的局限性，扩散语言模型（DLM）逐渐受到广泛关注。

然而，当前业界缺乏一个统一且有效的强化学习（RL）框架，适用于各种DLM架构，如全注意力DLM和块注意力DLM。同时，现有研究也未能充分重视推理轨迹与训练目标的一致性。

近日，普林斯顿大学王梦迪教授团队在最新研究中提出了“轨迹感知RL”框架——TraceRL，该框架可应用于全注意力和块注意力模型，实现快速优化。

论文链接：https://arxiv.org/abs/2509.06949

值得一提的是，使用TraceRL训练的4B DLM模型，在多项复杂数学推理任务上的表现超越了7B自回归模型。

通过课程学习，他们还推出了首个长思维链扩散语言模型（long-CoT DLM），在MATH500上的相对准确率比Qwen2.5-7B-Instruct提高了18.1%。

图｜左：不同方法的RL训练动态，其中TraceRL达到最佳优化。右：基于KV缓存评估的复杂数学推理任务与LiveCodeBench-V2基准测试结果对比。

同时，他们提出了一个基于扩散的价值模型，以降低方差并提高训练稳定性。此外，他们还探索了TraceRL的更多潜在应用，如增加模型块大小和加速推理。

此外，他们还开源了一个完整集成框架——dLLM-RL，用于跨不同架构构建、训练和部署DLM。该框架包含各种后训练方法和加速KV缓存技术的实现，支持可复现的研究和实际应用。

地址：https://github.com/Gen-Verse/dLLM-RL

亟需解决DLM“不匹配”难题

研究团队强调，DLM在后训练阶段所采用的目标与实际推理时遵循的轨迹之间存在显著不匹配。标准训练方法如完全随机掩码虽能实现并行解码，但忽略了语言本身的有序逻辑。这种训练与推理的脱节导致模型优化效率低下。

为了说明这一差异，他们通过实验证明，采用半自回归式微调方法（即根据先前上下文生成后续内容）的优化性能显著优于全随机掩码方法，即使计算负载相同也是如此。这表明，将训练目标与推理模式对齐至关重要。

为了进一步验证对齐的重要性，他们收集了模型自身的“偏好推理轨迹”，即模型生成内容时实际遵循的步骤顺序。实验结果显示，使用这些真实的推理轨迹进行微调所取得的性能超过了其他基线方法，且计算成本更低或相当。

最后，尽管使用偏好轨迹进行微调效果很好，但收集这些轨迹需要大量额外工作。相比之下，RL在其“rollouts”（即模型生成样本）过程中会自然地产生这些推理轨迹。因此，RL是一种更实用、更有效的训练后策略，能够天然地利用这些轨迹来优化模型。

TraceRL：“小扩散语言”>“大自回归”

在这项工作中，TraceRL关注由DLM生成的中间轨迹且可跨架构应用。

图｜TraceRL概述。本示例采用s=2、L=6和B=3的参数配置。通过聚合每s个相邻步骤实现轨迹感知RL。方框内数字对应策略推理过程的执行顺序。

在数据方面，研究团队采用了不同的数据源：

（1）选用Math训练数据集生成8000个硬核任务；

（2）选择GSM8K、MATH500、AIME2024作为测试基准以评估数学与编程领域的推理任务；

（3）在编码强化学习场景中采用PrimeIntellect平台提供的6000道已验证题目进行验证；

（4）针对编程测试在评估时选用LiveCodeBench-V2、LiveBench作为测试基准。

表｜不同数学与编程任务的主要基准测试结果。“静态”指静态采样，“动态”指动态采样。此处采用阈值