提出d2框架，用轨迹似然估计提升扩散语言模型的推理能力。
d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation
- 设计新策略梯度算法，精准估计采样轨迹似然。
- d2-AnyOrder和d2-StepMerge分别实现精确与可分析近似计算。
- 在数理推理任务上超越现有强化学习基线，刷新性能纪录。
尽管扩散语言模型（DLMs）在文本生成上表现优异，但利用强化学习提升其推理能力仍是研究热点。本文提出d2，一种专为掩码型DLM设计的推理框架。核心是基于准确轨迹似然估计的新策略梯度算法。由于掩码DLM的似然计算成本高，我们针对不同模型类型设计了相应估计器：对支持任意顺序解码的DLM，提出d2-AnyOrder，仅需一次模型前向即可获得精确轨迹似然；针对标准掩码扩散模型，提出d2-StepMerge，以可分析方式在计算开销与近似精度间权衡。实验表明，d2在多个主流DLM上显著优于现有RL基线，在逻辑推理任务（Countdown、Sudoku）和数学推理基准（GSM8K、MATH500）上达到新SOTA。代码与项目博客已公开。
原文摘要 · Abstract (English)
While diffusion language models (DLMs) have achieved competitive performance in text generation, improving their reasoning ability with reinforcement learning remains an active research area. Here, we introduce d2, a reasoning framework tailored for masked DLMs. Central to our framework is a new policy gradient algorithm that relies on accurate estimates of the sampling trajectory likelihoods. Because computing these likelihoods naively is computationally expensive for masked DLMs, we develop a family of estimators tailored to distinct model classes. For DLMs that support a sampling algorithm called any-order decoding, we propose d2-AnyOrder, which achieves exact trajectory likelihood with a single model pass. Through an empirical study of widely used DLMs, we show that any-order decoding is not universally supported in practice. For standard masked diffusion models, we propose d2-StepMerge, which approximates the trajectory likelihood, trading off compute for approximation accuracy in an analytically tractable manner. Empirically, d2 significantly outperforms widely-used RL baselines when applied to popular DLMs, and sets a new state-of-the-art performance for DLMs on logical reasoning tasks (Countdown and Sudoku) and math reasoning benchmarks (GSM8K and MATH500). We provide the code along with a blog post on the project page: https://guanghanwang.com/d2
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。