为扩散语言模型设计无偏强化学习方法,提升代码与逻辑推理能力。
Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
- 将去噪过程建模为马尔可夫决策过程,推导出无偏策略梯度。
- 在编码和逻辑推理任务上达到当前最优性能,数学推理表现超越现有方法。
- 通过熵引导步骤选择和单步奖励估计,降低计算成本且保持精度。
强化学习在自回归语言模型的后训练中表现良好,但扩展到扩散语言模型(DLMs)面临序列级似然不可计算的挑战。现有方法依赖代理似然或启发式近似,可能引入偏差并掩盖去噪的序列结构。本文将基于扩散的序列生成建模为有限时域马尔可夫决策过程,推导出一个在去噪步骤上可分解的精确、无偏策略梯度,其表达形式基于中间优势,无需显式评估序列似然。为获得高效实用的估计器,本文(i)通过熵引导的近似界选择用于策略更新的去噪步骤;(ii)利用扩散模型天然提供的单步去噪奖励来估计中间优势,避免代价高昂的多步回溯。在代码生成和逻辑推理基准上的实验表明,该方法达到当前最优结果,在数学推理任务上显著优于现有DLM强化学习方法。代码已开源:https://github.com/vishnutez/egspo-dllm-rl。
原文摘要 · Abstract (English)
Reinforcement learning (RL) has been effective for post-training autoregressive (AR) language models, but extending these methods to diffusion language models (DLMs) is challenging due to intractable sequence-level likelihoods. Existing approaches therefore rely on surrogate likelihoods or heuristic approximations, which can introduce bias and obscure the sequential structure of denoising. We formulate diffusion-based sequence generation as a finite-horizon Markov decision process over the denoising trajectory and derive an exact, unbiased policy gradient that decomposes over denoising steps and is expressed in terms of intermediate advantages, without requiring explicit evaluation of the sequence likelihood. To obtain a practical and compute-efficient estimator, we (i) select denoising steps for policy updates via an entropy-guided approximation bound, and (ii) estimate intermediate advantages using a one-step denoising reward naturally provided by the diffusion model, avoiding costly multi-step rollouts. Experiments on coding and logical reasoning benchmarks demonstrate state-of-the-art results, with strong competitive performance on mathematical reasoning, outperforming existing RL post-training approaches for DLMs. Code is available at https://github.com/vishnutez/egspo-dllm-rl.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。