wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models

📄 arXiv: 2507.08838v1 📥 PDF

作者: Xiaohang Tang, Rares Dolga, Sangwoong Yoon, Ilija Bogunovic

分类: cs.LG, cs.AI, stat.ML

发布日期: 2025-07-07

备注: Preprint

💡 一句话要点

提出wd1以提升扩散语言模型的推理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散语言模型 强化学习 推理能力 加权似然 策略优化 计算效率 自然语言处理

📋 核心要点

  1. 现有方法在扩散语言模型的推理能力提升上面临计算开销大和偏差大的挑战。
  2. 本文提出的wd1方法通过将优化目标重构为加权似然，简化了策略优化过程。
  3. 实验结果显示wd1在推理基准上准确率提升高达16%，且训练时间和函数评估次数显著减少。

📝 摘要（中文）

通过强化学习（RL）提升扩散基础的大型语言模型（dLLMs）的推理能力仍然是一个未解决的问题。dLLMs的似然函数的复杂性要求在每个策略优化步骤中近似当前、旧的和参考策略的似然。这种依赖引入了额外的计算开销，并可能导致较大的偏差，尤其是在重要性采样中策略比率的分母出现近似误差时。为了解决这些问题，本文提出了wd1，一种新的策略优化方法，将目标重新表述为加权似然，仅需对当前参数化策略的似然进行一次近似。实验表明，wd1在广泛使用的推理基准上，无需监督微调（SFT）或任何监督数据，超越了现有的RL方法，准确率提高了16%。wd1还带来了额外的计算收益，包括减少训练时间和每个梯度步骤的函数评估次数（NFE）。这些发现结合方法实现的简便性和类似R1-Zero的训练（无SFT），使wd1成为应用RL于dLLMs推理的更有效和高效的方法。

🔬 方法详解

问题定义：本文旨在解决扩散语言模型在推理能力提升中的计算复杂性和偏差问题。现有方法在策略优化中需要多次近似，导致计算开销大且可能引入偏差。

核心思路：wd1方法的核心在于将策略优化目标重新表述为加权似然，这样只需对当前策略的似然进行一次近似，从而降低了计算复杂性和偏差风险。

技术框架：wd1的整体架构包括策略优化的重构、加权似然的计算以及单次近似的实现。主要模块包括策略评估、加权计算和优化步骤。

关键创新：wd1的主要创新在于其将策略优化目标转化为加权似然的方式，这与传统方法需要多次近似的做法本质上不同，显著提高了效率和准确性。

关键设计：wd1在设计上采用了简化的损失函数，减少了对历史策略的依赖，并通过优化算法的选择来降低每个梯度步骤的函数评估次数（NFE）。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，wd1在多个推理基准上表现优异，准确率比现有的强化学习方法提高了高达16%。此外，wd1在训练时间和每个梯度步骤的函数评估次数上也显著减少，展示了其高效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、智能对话系统和自动推理等。wd1方法的高效性和准确性使其在需要快速推理和决策的场景中具有实际价值，未来可能推动更多基于扩散语言模型的智能应用的发展。

📄 摘要（原文）

Improving the reasoning capabilities of diffusion-based large language models (dLLMs) through reinforcement learning (RL) remains an open problem. The intractability of dLLMs likelihood function necessitates approximating the current, old, and reference policy likelihoods at each policy optimization step. This reliance introduces additional computational overhead and lead to potentially large bias -- particularly when approximation errors occur in the denominator of policy ratios used for importance sampling. To mitigate these issues, we introduce $\mathtt{wd1}$, a novel policy optimization approach that reformulates the objective as a weighted likelihood, requiring only a single approximation for the current parametrized policy likelihood. Experiments on widely used reasoning benchmarks demonstrate that $\mathtt{wd1}$, without supervised fine-tuning (SFT) or any supervised data, outperforms existing RL methods for dLLMs, achieving up to 16% higher accuracy. $\mathtt{wd1}$ delivers additional computational gains, including reduced training time and fewer function evaluations (NFEs) per gradient step. These findings, combined with the simplicity of method's implementation and R1-Zero-like training (no SFT), position $\mathtt{wd1}$ as a more effective and efficient method for applying RL to dLLMs reasoning.