鉴于自回归大语言模型（LLM）在计算效率、吞吐量上的局限性，扩散语言模型（DLM）正受到越来越多的关注。
然而，业内目前缺少一个可适用于各种 DLM 架构（如全注意力 DLM 和块注意力 DLM）的统一且有效的强化学习（RL）框架。同时，现有研究也忽视了对齐推理轨迹与训练目标的重要性。
日前，普林斯顿大学王梦迪教授团队在一项最新研究中提出了“轨迹感知 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 作为测试基准。
表｜不同数学与编程任务的主要基准测试结果。“静态”指静态采样，“动态”指动态采样。此处采用阈值 0.9 的动态采样方法对长 CoT 模型 TraDo-8B-Instruct 进行评估
在模型训练方面，包括全注意力模型和块注意力模型，并在评估过程中同时报告静态采样和动态采样的结果。过程如下：
第一步：由 TraceRL 单独训练模型；
第二步：联合训练长 CoT 模型；
第三步：TraceRL 与其他 RL 方法进行对比实验；
第四步：TraceRL 在全注意力模型与编码任务中的验证；
第五步：进行区块大小扩展实验。
基于实验结果，他们证明了 TraceRL 的有效性与强性能，完整结果如下：
首先，他们基于 SDAR 基础模型，利用 TraceRL 开发了 TraDo-4B-Instruct 和 TraDo-8B-Instruct 两种模型。在数学、编程和 5 个推理数据集的评估中，这些模型不仅与强扩散语言模型、自回归模型形成有效对比，其生成能力也展现出显著优势。
图｜TraceRL 在数学任务中对 4B 和 8B 模型的训练曲线。红色曲线表示动态采样精度，其采样速度更快；蓝色曲线表示静态采样精度，可获得更高的精度。4B 模型使用价值模型进行训练，而 8B 模型直接使用 Jpolicy 进行训练。
TraDo-4B-Instruct 在推理任务上展现出 SOTA 水平，证明了 TraceRL 的有效性。无论是动态采样（更快）还是静态采样（更准确），模型的性能都得到了显著提升。值得注意的是，在所有数学任务中，TraDo-4B-Instruct 的表现甚至超越了 Qwen2.5-7B Instruct 等强自回归基线模型。
尽管他们在 RL 训练过程中采用了动态采样策略，但动态与静态准确率均呈现稳步提升趋势，且有趋势表明模型仍有扩展潜力。这种 RL 训练显著提升了模型的数学推理能力：
在 MATH500 测试中，TraDo-4B Instruct 模型静态准确率提升 5.4%，动态准确率提升 4.2%，优化后超越 Qwen2.5-7B-Instruct；而 TraDo-8B-Instruct 模型则分别实现静态准确率 4.2% 和动态准确率 4.8% 的提升。
图｜基于块扩散模型和数学 RL 任务的 RL 方法消融实验。红色与黄色曲线分别对应启用和禁用价值模型的 TraceRL 训练结果，蓝色曲线采用类似半自回归训练法的随机掩码目标进行块内训练，绿色曲线则通过在块内添加互补性掩码实现额外训练效果。
他们进一步将 TraceRL 与现有的 RL 方法进行了对比研究，重点关注块扩散模型。尽管当前 RL 方法主要针对全注意力模型开发，但他们将其直接适配到块结构中。对于随机掩码方法，他们限制每个块内的采样操作，使其类似于半自回归方法。对于耦合强化学习，他们在每个训练模块中引入了互补目标函数，从而获得更稳定且高效的训练效果。基于数学任务的实验结果显示，无论采用何种优化策略，TraceRL 均展现出最优性能。
图｜全注意力模型 Dream-7B-Coder-Instruct 的 RL 训练消融实验，重点聚焦编码任务；使用与不使用价值模型的对比表明，引入价值模型能有效减少训练过程中的波动。所有实验均在 4B 模型的数学任务上进行。
此外，为验证 TraceRL 的广泛适用性，他们还在编码 RL 任务中对全注意力模型进行了实验。基于 Dream-7B-Coder-Instruct，在冷启动阶段使用蒸馏数据进行微调后开展 RL 训练。为加速训练进程，将收缩参数设置为 s=8。实验表明，TraceRL 收敛速度更快且性能更优。
未来可期
基于以上实验结果，研究团队验证了 TraceRL 在不同 RL 任务中的有效性。同时，他们还展示了 TraceRL 在加速推理和扩展块大小方面的优势，这为未来的研究提供了有前景的方向。
特别地，将扩散模型的加速推理能力与其潜在的强推理能力相结合，代表了一个令人期待的研究方向。尽管当前的长 CoT LLM 在复杂任务上表现出色，但其推理时间太长。此类整合有望在大规模环境中高效地执行复杂推理任务，从而开辟新的应用可能。
研究团队还表示，他们提出的扩散价值模型能够整合过程奖励，提供比单一可验证奖励更强的监督信号。未来，他们将进一步探索基于过程奖励的 TraceRL 优化。
本文来自微信公众号 “学术头条”（ID：SciTouTiao），整理：小瑜，36氪经授权发布。
