Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
作者: Vishnu Teja Kunde, Fatemeh Doudi, Mahdi Farahbakhsh, Dileep Kalathil, Krishna Narayanan, Jean-Francois Chamberland
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-03-13
🔗 代码/项目: GITHUB
💡 一句话要点
提出基于熵引导步选择和逐步优势的强化学习方法，用于扩散语言模型后训练。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 强化学习 扩散语言模型 策略梯度 序列生成 后训练
📋 核心要点
- 扩散语言模型的强化学习后训练面临序列似然难以处理的问题，现有方法依赖近似，引入偏差。
- 论文提出一种精确、无偏的策略梯度方法，将扩散过程建模为马尔可夫决策过程，避免直接评估序列似然。
- 实验表明，该方法在编码、逻辑推理和数学推理任务上表现出色，优于现有扩散语言模型的强化学习方法。
📝 摘要（中文）
强化学习在自回归语言模型后训练中表现出色，但将其扩展到扩散语言模型（DLMs）面临挑战，因为序列级似然难以处理。现有方法依赖于替代似然或启发式近似，这会引入偏差并掩盖去噪的序列结构。本文将基于扩散的序列生成建模为去噪轨迹上的有限horizon马尔可夫决策过程，并推导出精确、无偏的策略梯度，该梯度分解为去噪步骤，并以中间优势表示，无需显式评估序列似然。为了获得实用且计算高效的估计器，我们（i）通过熵引导的近似界选择用于策略更新的去噪步骤，以及（ii）使用扩散模型自然提供的一步去噪奖励来估计中间优势，避免了代价高昂的多步rollout。在编码和逻辑推理基准上的实验表明，该方法取得了最先进的结果，并在数学推理方面表现出强大的竞争力，优于现有的DLM强化学习后训练方法。
🔬 方法详解
问题定义：现有扩散语言模型的强化学习方法，由于难以处理序列级别的似然函数，通常采用替代似然或启发式近似，这会导致偏差，并且无法充分利用去噪过程的序列结构信息。因此，如何设计一种无偏且高效的强化学习方法，充分利用扩散模型的特性，成为一个关键问题。
核心思路：论文的核心思路是将扩散模型的序列生成过程视为一个有限horizon的马尔可夫决策过程（MDP），并在此基础上推导出精确的策略梯度。通过将策略梯度分解为每个去噪步骤的中间优势函数，避免了直接计算序列似然，从而消除了偏差。同时，利用扩散模型本身提供的去噪奖励来估计中间优势，避免了多步rollout，提高了计算效率。
技术框架：该方法的技术框架主要包含以下几个步骤：1. 将扩散模型的去噪过程建模为MDP，其中状态是去噪过程中的中间状态，动作是去噪步骤的选择，奖励是扩散模型提供的去噪奖励。2. 推导出精确的策略梯度，该梯度分解为每个去噪步骤的中间优势函数。3. 使用熵引导的近似界来选择用于策略更新的去噪步骤，以提高采样效率。4. 使用一步去噪奖励来估计中间优势，避免多步rollout。5. 使用策略梯度算法更新策略，优化扩散模型的生成能力。
关键创新：该方法最重要的技术创新点在于提出了一个精确且无偏的策略梯度估计器，该估计器避免了直接计算序列似然，并且能够充分利用扩散模型的特性。此外，熵引导的步选择和一步优势估计也提高了算法的效率。与现有方法的本质区别在于，该方法避免了使用替代似然或启发式近似，从而消除了偏差。
关键设计：在关键设计方面，论文采用了熵引导的步选择策略，通过最大化选择步骤的熵，保证了探索的多样性。同时，使用一步去噪奖励作为优势函数的估计，避免了多步rollout的计算开销。具体的损失函数是基于策略梯度的标准形式，通过优化策略参数来最大化期望奖励。网络结构方面，可以使用现有的扩散模型结构，并在此基础上添加策略网络，用于选择去噪步骤。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在编码和逻辑推理任务上取得了最先进的结果，并在数学推理任务上表现出强大的竞争力，优于现有的扩散语言模型强化学习后训练方法。具体而言，在某些任务上，该方法能够显著提高生成结果的准确率和流畅度，证明了其有效性和优越性。
🎯 应用场景
该研究成果可广泛应用于各种序列生成任务，例如文本生成、代码生成、图像生成等。通过强化学习的后训练，可以显著提升扩散语言模型在特定任务上的性能，例如提高代码生成的正确率、提升数学推理的准确性。该方法具有重要的实际价值，可以促进人工智能在各个领域的应用。
📄 摘要（原文）
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.