dTRPO: Trajectory Reduction in Policy Optimization of Diffusion Large Language Models

📄 arXiv: 2603.18806v1 📥 PDF

作者: Wenxuan Zhang, Lemeng Wu, Changsheng Zhao, Ernie Chang, Mingchen Zhuge, Zechun Liu, Andy Su, Hanxian Huang, Jun Chen, Chong Zhou, Raghuraman Krishnamoorthi, Vikas Chandra, Mohamed Elhoseiny, Wei Wen

分类: cs.AI

发布日期: 2026-03-19

💡 一句话要点

dTRPO：通过轨迹缩减优化扩散大语言模型的策略

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散语言模型 策略优化 轨迹缩减 离线训练 人类偏好对齐

📋 核心要点

  1. 扩散语言模型对齐人类偏好面临挑战，现有方法计算轨迹概率成本高昂，限制了离线策略训练的规模。
  2. dTRPO通过轨迹缩减策略，利用参考策略正则化和重掩码技术，降低轨迹概率计算的复杂度，实现高效策略优化。
  3. 实验表明，dTRPO显著提升了dLLMs在STEM、编码和指令跟随任务上的性能，并展现出强大的训练和生成效率。

📝 摘要（中文）

扩散大语言模型（dLLMs）为语言生成引入了一种新的范式，同时也为将其与人类偏好对齐带来了新的挑战。本文旨在通过降低轨迹概率计算的成本来改进dLLMs的策略优化，从而实现大规模的离线策略训练。我们证明了：（i）在参考策略正则化下，新近未掩码的token的概率比是中间扩散状态的无偏估计；（ii）可以通过重新掩码的最终状态的单次前向传播有效地估计完整轨迹的概率。通过将这两种轨迹缩减策略集成到策略优化目标中，我们提出了轨迹缩减策略优化（dTRPO）。我们在7B dLLMs上，针对指令跟随和推理基准评估了dTRPO。结果表明，它显著提高了最先进dLLMs的核心性能，在STEM任务上实现了高达9.6%的收益，在编码任务上实现了高达4.3%的收益，在指令跟随任务上实现了高达3.0%的收益。此外，由于其离线、单次前向的特性，dTRPO表现出强大的训练效率，并通过高质量的输出实现了更高的生成效率。

🔬 方法详解

问题定义：论文旨在解决扩散大语言模型（dLLMs）与人类偏好对齐的问题，具体来说，是如何高效地进行策略优化。现有方法在计算轨迹概率时成本过高，这限制了离线策略训练的规模，使得难以充分利用大规模离线数据来提升模型性能。

核心思路：论文的核心思路是通过轨迹缩减来降低策略优化的计算成本。具体而言，论文提出了两种策略：一是利用参考策略正则化，证明新近未掩码的token的概率比是中间扩散状态的无偏估计；二是利用重掩码技术，通过单次前向传播来估计完整轨迹的概率。

技术框架：dTRPO的整体框架是将上述两种轨迹缩减策略集成到策略优化目标中。该框架主要包含以下几个阶段：1) 使用扩散模型生成文本轨迹；2) 利用参考策略正则化和重掩码技术进行轨迹缩减，降低计算复杂度；3) 基于缩减后的轨迹，优化策略模型，使其与人类偏好对齐。整个过程是离线的，只需要单次前向传播，因此具有较高的效率。

关键创新：论文的关键创新在于提出了两种有效的轨迹缩减策略，并将其成功应用于扩散大语言模型的策略优化中。与现有方法相比，dTRPO能够在保证性能的同时，显著降低计算成本，从而实现更大规模的离线策略训练。

关键设计：论文的关键设计包括：1) 参考策略正则化的具体形式，如何选择合适的参考策略；2) 重掩码技术的实现细节，如何选择合适的掩码策略；3) 策略优化目标的构建，如何平衡策略改进和策略稳定性；4) 实验中使用的具体参数设置，例如学习率、batch size等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

dTRPO在7B dLLMs上进行了广泛的实验评估，结果表明其性能显著优于现有方法。具体而言，在STEM任务上，dTRPO实现了高达9.6%的收益；在编码任务上，实现了高达4.3%的收益；在指令跟随任务上，实现了高达3.0%的收益。此外，dTRPO还展现出强大的训练效率和生成效率，证明了其在扩散大语言模型策略优化方面的有效性。

🎯 应用场景

dTRPO具有广泛的应用前景，可用于提升各种扩散大语言模型在文本生成、代码生成、问题求解等任务上的性能。通过与人类偏好对齐，可以生成更符合人类意图和价值观的文本内容。此外，dTRPO的高效性使其能够应用于资源受限的场景，例如移动设备或边缘计算平台。未来，dTRPO有望推动扩散模型在自然语言处理领域的更广泛应用。

📄 摘要（原文）

Diffusion Large Language Models (dLLMs) introduce a new paradigm for language generation, which in turn presents new challenges for aligning them with human preferences. In this work, we aim to improve the policy optimization for dLLMs by reducing the cost of the trajectory probability calculation, thereby enabling scaled-up offline policy training. We prove that: (i) under reference policy regularization, the probability ratio of the newly unmasked tokens is an unbiased estimate of that of intermediate diffusion states, and (ii) the probability of the full trajectory can be effectively estimated with a single forward pass of a re-masked final state. By integrating these two trajectory reduction strategies into a policy optimization objective, we propose Trajectory Reduction Policy Optimization (dTRPO). We evaluate dTRPO on 7B dLLMs across instruction-following and reasoning benchmarks. Results show that it substantially improves the core performance of state-of-the-art dLLMs, achieving gains of up to 9.6% on STEM tasks, up to 4.3% on coding tasks, and up to 3.0% on instruction-following tasks. Moreover, dTRPO exhibits strong training efficiency due to its offline, single-forward nature, and achieves improved generation efficiency through high-quality outputs.