摘要：我们提出TraceRL——一种面向扩散语言模型（DLMs）的轨迹感知强化学习框架，该框架将优选推理轨迹纳入后训练阶段，且适用于不同架构。该框架配备了基于扩散机制的价值模型，可提升训练稳定性，我们在复杂数学和编程任务上验证了其推理性能的提升。此外，它还可用于将特定模块模型适配至更大模块，从而提高采样灵活性。通过应用TraceRL，我们开发出一系列最先进的扩散语言模型，即TraDo系列。尽管模型规模小于70亿参数的自回归（AR）模型，TraDo-4B-Instruct在复杂数学推理任务中仍持续超越这些模型。在数学推理基准测试中，TraDo-8B-Instruct相较于Qwen2.5-7B-Instruct实现了6.1%的相对准确率提升，相较于Llama3.1-8B-Instruct则实现了51.3%的相对准确率提升。通过课程学习，我们还开发出首个长思维链（long-CoT）扩散语言模型，其在MATH500数据集上的表现优于Qwen2.5-7B-Instruct，相对准确率提升了18.1%。为促进可复现的研究和实际应用，我们发布了一个全面的开源框架，支持在不同架构上构建、训练和部署扩散型大语言模型（LLMs）。该框架集成了加速的键值缓存（KV-cache）技术和适用于推理与强化学习的推理引擎，并包含了针对数学、编程及通用任务的各种监督微调和强化学习方法的实现。代码与模型：https://github.com/Gen-Verse/dLLM-RL。Huggingface链接：Paper page，论文链接：2509.06949

研究背景和目的

研究背景：
随着自然语言处理技术的快速发展，大型语言模型（LLMs）在各种任务中展现出强大的能力，尤其是在需要复杂推理的任务中，如数学解题和代码生成。然而，现有的后训练框架主要针对完全注意力机制的扩散语言模型（DLMs），通过随机掩码来估计整个序列的分数，这种方法与最优推理过程存在不匹配，因为语言的顺序性和逻辑性并非完全随机。此外，尽管块注意力扩散模型（Block Diffusion）在复杂推理任务中表现出色，但其在强化学习（RL）方面的应用尚未得到充分探索。同时，现有模型在处理复杂任务时，往往缺乏对推理轨迹的精细控制，导致推理效率低下。

研究目的：
本研究旨在提出一种针对扩散语言模型的轨迹感知强化学习框架（TraceRL），通过将推理轨迹与训练目标对齐，提高模型在复杂推理任务中的表现。具体目标包括：1）设计一种能够捕捉模型推理轨迹的强化学习方法，适用于不同架构的扩散语言模型；2）通过TraceRL框架训练出具有更强推理能力的扩散语言模型，特别是在数学解题和代码生成任务中超越现有自回归模型；3）探索TraceRL在块扩散模型中的应用，提升模型处理更大块大小的能力，从而提高采样灵活性；4）发布一个完整的开源框架，支持不同架构的扩散语言模型的构建、训练和部署。

研究方法

1. 轨迹感知强化学习框架（TraceRL）：

• 轨迹捕捉：TraceRL通过记录模型在推理过程中的每一步，捕捉完整的推理轨迹，而不仅仅是最终输出。这种方法使得强化学习能够基于模型的中间推理步骤进行优化，而不仅仅是最终结果。
• 奖励机制：引入基于步骤的奖励机制，对模型在推理轨迹中的每一步进行奖励或惩罚，从而引导模型学习更优的推理策略。同时，采用扩散模型作为值模型，减少训练过程中的方差，提高稳定性。
• 轨迹缩减：为了处理长序列推理中的计算复杂性问题，引入轨迹缩减参数，将连续的推理步骤聚合，减少需要处理的轨迹长度，提高训练效率。

2. 块扩散模型的应用：

• 块注意力机制：利用块扩散模型的块注意力机制，结合半自回归微调方法，提高模型在处理长序列时的效率和准确性。
• 块大小扩展：通过TraceRL框架，探索块扩散模型在处理更大块大小时的表现，提升模型的采样灵活性。具体方法包括先在小块大小下进行训练，然后逐步增加块大小，使模型适应更大的块。

3. 开源框架实现：

• 模型支持：框架支持多种架构的扩散语言模型，包括预训练的全注意力模型、适应后的全注意力模型和块扩散模型。
• 加速推理技术：集成多种加速推理技术，如KV-cache技术，提高模型在推理过程中的效率。
• 训练方法实现：提供多种后训练方法的实现，包括监督微调（SFT）、随机掩码强化学习（RL）和TraceRL，支持多节点训练。

研究结果

1. 模型性能提升：

• 在数学推理任务中，TraDo-4B-Instruct和TraDo-8B-Instruct模型在MATH500和AIME2024等基准测试上取得了显著提升，超越了多个强大的自回归模型，如Qwen2.5-7B-Instruct和Llama3.1-8B-Instruct。特别是在MATH500上，TraDo-8B-Instruct的准确率达到了87.4%，相对准确率提升了13.1%。
• 在代码生成任务中，TraDo模型在LiveCodeBench-V2等基准测试上也展现了优越性能，特别是在动态采样设置下，取得了新的最高水平。

2. 长链推理能力：

• 通过结合TraceRL和长链推理监督微调（long-CoT SFT），训练出了第一个长链推理扩散语言模型TraDo-8B-Thinking。该模型在MATH500和LiveCodeBench-V2等基准测试上展现了强大的长链推理能力，相对准确率分别提升了18.1%和28.6%。

3. 块大小扩展效果：

• 通过TraceRL框架，成功地将块扩散模型的块大小从4扩展到8，提高了模型的采样灵活性。实验结果表明，扩展后的模型在保持准确性的同时，能够更高效地处理长序列推理任务。

4. 开源框架贡献：

• 发布的开源框架为研究人员和开发者提供了一个完整的工具链，支持不同架构的扩散语言模型的构建、训练和部署。框架的模块化设计和高效实现，促进了扩散语言模型在复杂推理任务中的广泛应用。

研究局限

1. 数据依赖性问题：

• 尽管TraceRL框架在多个基准测试上取得了显著提升，但其性能仍然高度依赖于训练数据的质量和多样性。如果训练数据存在偏差或不足，可能会影响模型的泛化能力和鲁棒性。

2. 模型规模限制：

• 当前研究主要关注8B参数规模的模型，虽然这些模型在复杂推理任务中展现了优越性能，的模型可能会进一步提升性能。然而，更大规模模型的训练和推理成本也会显著增加，限制了其在实际应用中的广泛部署。

3. 评估方法局限性：

• 当前评估主要依赖于有限的基准测试，这些测试可能无法全面反映模型在所有实际场景中的性能。此外，评估过程中可能存在主观性和偏见，影响评估结果的准确性。

4. 训练效率问题：

• 尽管TraceRL框架通过轨迹缩减等技术提高了训练效率，但在处理极长序列推理任务时，训练过程仍然可能非常耗时。如何进一步提高训练效率，是未来研究需要解决的问题之一。

未来研究方向

1. 数据增强与多样性提升：

• 探索更有效的数据增强方法，如结合人类反馈的强化学习（RLHF），以生成更多样化、更具挑战性的训练数据。同时，扩展数据集的覆盖范围，涵盖更多领域和任务类型，提高模型的泛化能力。

2. 更大规模模型研究：

• 研究更大规模（如32B、100B参数）的扩散语言模型在TraceRL框架下的性能表现。探索参数规模与模型能力之间的关系，为实际应用提供更强大的模型支持。

3. 评估方法创新：

• 开发更客观、全面的评估体系，结合人类评估和自动化评估方法，确保评估结果的准确性和可靠性。同时，探索新的基准测试，覆盖更多实际场景中的复杂任务，以更全面地评估模型的性能。

4. 实际应用探索：

• 将TraceRL框架和训练出的高性能扩散语言模型应用于实际场景中，如智能助手、自动驾驶、医疗诊断等领域。探索模型在实际应用中的表现和挑战，为模型的进一步优化和改进提供反馈。

5. 跨模态学习与融合：

• 探索扩散语言模型与其他模态（如图像、音频）的融合方法，开发跨模态的扩散模型。通过结合多模态信息，进一步提升模型在复杂任务中的推理能力和泛化能力。