d-TreeRPO: Towards More Reliable Policy Optimization for Diffusion Language Models
作者: Leyi Pan, Shuchang Tao, Yunpeng Zhai, Zheyu Fu, Liancheng Fang, Minghua He, Lingzhe Zhang, Zhaoyang Liu, Bolin Ding, Aiwei Liu, Lijie Wen
分类: cs.CL
发布日期: 2025-12-10 (更新: 2026-01-06)
备注: 20 pages, 19 figures, 4 tables
💡 一句话要点
d-TreeRPO:面向扩散语言模型,提升策略优化的可靠性
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 强化学习 策略优化 奖励稀疏性 概率估计 树搜索 自蒸馏
📋 核心要点
- 现有dLLM策略优化方法面临奖励稀疏性和概率估计偏差两大挑战,影响推理能力提升。
- d-TreeRPO利用树状rollout和自底向上优势计算,提供细粒度可验证的奖励信号,提升策略优化可靠性。
- 实验表明,d-TreeRPO在Sudoku、Countdown、GSM8K和Math500等推理任务上显著优于现有基线。
📝 摘要(中文)
强化学习(RL)对于提升扩散大语言模型(dLLMs)的推理能力至关重要。然而,现有的dLLM策略优化方法存在两个关键的可靠性瓶颈:(1)奖励稀疏性,源于粗糙或无法验证的信号,阻碍了准确的优势计算;(2)它们的概率估计没有考虑到所有解码顺序上的无偏期望与单步前向传递估计之间的差距,而计算所有解码顺序是难以处理的。为了缓解这些问题,我们提出了d-TreeRPO,这是一个可靠的dLLM强化学习框架,它利用树状结构的rollout和基于可验证结果奖励的自底向上优势计算,以提供细粒度和可验证的逐步奖励信号。此外,我们提供了一个理论证明,表明增加预测置信度可以有效地最小化无偏期望预测概率与其单步前向传递估计之间的差距。在该分析的指导下,我们在训练期间引入了一个时间调度的自蒸馏损失,以增强后期训练阶段的预测置信度,从而实现更准确的概率估计和更好的性能。实验表明,d-TreeRPO优于现有的基线,并在多个推理基准测试中取得了显著的改进。具体而言,与基础模型相比,在Sudoku上实现了+86.2%,在Countdown上实现了+51.6%,在GSM8K上实现了+4.5%,在Math500上实现了+5.3%的提升。
🔬 方法详解
问题定义:论文旨在解决扩散语言模型(dLLMs)在强化学习策略优化中存在的可靠性问题。现有方法面临奖励稀疏,难以准确计算优势函数;以及概率估计存在偏差,无法有效利用所有可能的解码顺序信息。这些问题导致模型推理能力提升受限。
核心思路:论文的核心思路是通过引入树状结构的rollout和自底向上的优势计算,来提供更细粒度、可验证的奖励信号,从而缓解奖励稀疏性问题。同时,通过理论分析和自蒸馏损失,提高模型预测的置信度,减少概率估计的偏差。
技术框架:d-TreeRPO框架主要包含以下几个部分:1) 树状结构的Rollout:从当前状态出发,进行多步采样,形成树状结构。2) 自底向上优势计算:基于树状结构,从叶节点开始,逐层向上计算优势函数,利用可验证的结果奖励作为信号。3) 时间调度的自蒸馏损失:在训练后期,引入自蒸馏损失,鼓励模型产生更置信的预测。
关键创新:论文的关键创新在于:1) 提出了基于树状rollout和自底向上优势计算的奖励机制,有效缓解了奖励稀疏性问题。2) 通过理论证明和自蒸馏损失,解决了概率估计偏差问题,提高了策略优化的可靠性。3) 将时间调度机制引入自蒸馏损失,使得模型在训练后期更加关注预测置信度。
关键设计:在树状rollout中,需要确定树的深度和宽度,这会影响计算复杂度和奖励信号的丰富程度。自底向上优势计算需要设计合适的奖励函数,以反映每一步的贡献。时间调度的自蒸馏损失需要选择合适的温度系数和调度策略,以平衡探索和利用。
🖼️ 关键图片
📊 实验亮点
d-TreeRPO在多个推理基准测试中取得了显著的改进。在Sudoku上实现了+86.2%的提升,在Countdown上实现了+51.6%的提升,在GSM8K上实现了+4.5%的提升,在Math500上实现了+5.3%的提升。这些结果表明,d-TreeRPO能够有效提升扩散语言模型的推理能力。
🎯 应用场景
d-TreeRPO具有广泛的应用前景,可用于提升各种扩散语言模型在复杂推理任务中的性能,例如数学问题求解、代码生成、逻辑推理等。该方法能够提高模型在实际应用中的可靠性和准确性,具有重要的实际价值,并有望推动人工智能在推理能力方面的进一步发展。
📄 摘要(原文)
Reinforcement learning (RL) is pivotal for enhancing the reasoning capabilities of diffusion large language models (dLLMs). However, existing dLLM policy optimization methods suffer from two critical reliability bottlenecks: (1) reward sparsity, arising from coarse or unverifiable signals that impede accurate advantage calculation; and (2) their probability estimates do not account for the gap to the unbiased expectation over all decoding orders, which are intractable to compute. To mitigate these issues, we propose d-TreeRPO, a reliable RL framework for dLLMs that leverages tree-structured rollouts and bottom-up advantage computation based on verifiable outcome rewards to provide fine-grained and verifiable step-wise reward signals. Furthermore, we provide a theoretical proof demonstrating that increasing prediction confidence effectively minimizes the gap between unbiased expected prediction probabilities and its single-step forward pass estimate. Guided by this analysis, we introduce a time-scheduled self-distillation loss during training that enhances prediction confidence in later training stages, thereby enabling more accurate probability estimation and better performance. Experiments demonstrate that d-TreeRPO outperforms existing baselines and achieves significant improvements across multiple reasoning benchmarks. Specifically, it achieves +86.2% on Sudoku, +51.6% on Countdown, +4.5% on GSM8K, and +5.3% on Math500 compared to the base model.