提出新算法提升扩散语言模型的推理能力
Simple Policy Gradients for Reasoning with Diffusion Language Models
- 基于扩散模型的马尔可夫特性，优化每一步去噪过程
- 在计数和数独任务上准确率提升超59%和69%
- 无需序列级近似，适合需要精准推理的应用
扩散大语言模型（dLLMs）是自回归模型的有前景替代方案；然而，缺乏有效的后训练技术，尤其是强化学习（RL），仍是dLLMs在下游应用中的主要挑战。现有方法通常依赖序列级视角，需使用有偏的似然近似。本文提出广义组相对策略优化（AGRPO），一种利用dLLMs马尔可夫特性的策略梯度算法，通过优化单个去噪步骤而非完整序列实现改进。该方法增强了训练策略与推理过程的一致性，并通过新颖的时间步估计方案实现了高效、无偏的梯度更新。我们在多个数学与推理任务上验证了AGRPO的有效性，在Countdown和Sudoku任务上相对于基础的LLaDA模型分别取得+59.4%和+69.7%的绝对准确率提升，优于diffu-GRPO等可比方法。代码已开源：https://github.com/probablyabot/agrpo。
原文摘要 · Abstract (English)
Diffusion large language models (dLLMs) represent a promising alternative to autoregressive LLMs; however, the lack of effective post-training techniques, including reinforcement learning (RL), remains a key challenge for dLLMs, especially for downstream applications. Existing approaches often rely on a sequence-level view that requires biased likelihood approximations. In this work, we propose Amortized Group Relative Policy Optimization (AGRPO), a policy gradient algorithm that leverages the Markovian nature of dLLMs, optimizing individual denoising steps rather than full sequences. Our approach improves alignment between the trained policy and the inference process and also admits efficient, unbiased gradient updates via a novel timestep estimation scheme. We demonstrate AGRPO's effectiveness on different math and reasoning tasks, achieving absolute accuracy gains of +59.4\% and +69.7\% on Countdown and Sudoku over the base LLaDA model, exceeding comparable methods such as diffu-GRPO. Code is available at https://github.com/probablyabot/agrpo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。