单人一作 ICML 论文预警！是否曾想过用传统的策略梯度（Policy Gradients）对扩散语言模型（dLLM）进行后训练，而无需处理 ELBO 或代理目标？在《Simple Policy Gradients for Reasoning with Diffusion Language Models》中，我们展示了一种直接且有效的方法来实现这一点。我们的框架 Amortized GRPO (AGRPO) 允许模型通过时间步估计从无偏策略梯度更新中学习，既自然契合 dLLM 的推理过程，又保持了高效与可扩展性。Paper: https://arxiv.org/abs/2510.04019 Code: https://github.com/probablyabot/agrpo… 1/n —

Simple Policy Gradients for Reasoning with Diffusion Language Models

Source: https://arxiv.org/html/2510.04019

Abstract

扩散大语言模型（dLLM）作为传统自回归 LLM 的一个极具前景的替代方案，近期在预训练阶段已展现出强劲的性能。然而，由于其缺乏可处理的序列级似然函数，它们尚未能从强化学习（RL）等现代 LLM 后训练技术中受益，限制了其在实际应用中的潜力。现有的 dLLM 后训练尝试多依赖于启发式近似或真实似然的下界。本文提出 Amortized Group Relative Policy Optimization（AGRPO），一种利用 dLLM 生成分步马尔可夫特性的策略梯度算法。该算法优化单个去噪步骤而非整个序列。我们在不同的数学与推理任务上验证了 AGRPO 的有效性：在基础 LLaDA 模型上，GSM8K 绝对提升 +9.9%，MATH-500 提升 +4.6%，Countdown 提升 +59.4%，Sudoku 提升 +69.7%，性能优于 diffu-GRPO 等同类 dLLM RL 方法。此外，我们分析了后训练收益在不同推理配置下的保持情况，结果表明，使用 AGRPO 训练的模型在采样速度提升 4 倍的情况下，性能损失微乎其微。 Diffusion language model, dLLM, Reinforcement learning, Post-training, Reasoning, Policy gradient

1 Introduction

近期 LLM 研究的许多工作都集中在强化学习上，特别是在可验证奖励（RLVR）设定下。在典型设置中，基础模型会在数学或编程问题上进行训练，并被激励逐步推理出解决方案，若最终答案正确则获得奖励。RLVR 的主要目标是激发数学思维/推理能力，使模型能够解决复杂的现实世界任务。这股对 RL 和推理的兴趣浪潮，最初由 OpenAI 的 o1 和 DeepSeek 的 R1 等模型引发，现已催生了大量专为基于 Transformer 的自回归（AR）LLM 设计的后训练算法。随着这些算法的成功，尤其是 Group Relative Policy Optimization（GRPO）的脱颖而出，AR LLM 在解题基准测试上变得极其强大，部分闭源模型甚至在 IMO 和 IOI 等竞赛中夺得金牌，取得了令人瞩目的成就。

在另一条并行研究线上，扩散语言模型近期作为传统自回归范式的替代方案崭露头角。连续扩散模型长期以来一直是图像和视频生成的主导框架，依赖于去噪/分数匹配目标。D3PM 和 SEDD 等研究成功将这一扩散框架迁移到离散设置（包括语言）中。随后的 MDLM 和 RADD 等工作进一步简化了理论框架，最新研究多采用“吸收”或“掩码”扩散框架。下文我们将此类掩码扩散模型统称为 dLLM。当前最先进的 dLLM（如 LLaDA 和 MMaDA）在常见 NLP 基准测试上的表现已接近或媲美 LLaMA3-8B 和 Qwen2.5-7B 等开源 AR 模型。训练完成后，这些模型可以从部分掩码序列开始，超越标准的从左到右生成方式，并且可以通过减少采样步数（即并行采样更多 token）来权衡计算量与生成质量。然而，在需要长篇思考与推理的下游任务中，这些模型仍难以匹敌 AR 模型。这种后训练上的差距源于为 dLLM 设计训练目标时的根本性挑战：AR 模型可通过自回归分解轻松获取序列级似然，而扩散模型则不得不依赖近似或类似 ELBO 的似然下界。解锁真正的推理能力将是 dLLM 向前迈出的一大步，使其真正具备与 AR LLM 抗衡的实力。我们的工作通过提出一种专为 dLLM 设计的原理性策略梯度算法——Amortized Group Relative Policy Optimization（AGRPO）——来帮助 dLLM 弥合这一差距。与传统的一步法不同，我们首先将后训练问题建立为多步 MDP 公式，并将其与 dLLM 使用的迭代解掩码过程联系起来。然后，通过对策略梯度目标进行简单修改——将所有 token 的内部求和视为对时间步的期望——我们展示了如何使长篇推理任务的训练变得可处理。我们的主要贡献如下：

• •严谨性：我们从 dLLM 生成过程的多步视角推导出了无偏的策略梯度目标，阐明了我们的方法如何在保持理论严谨性的同时，避开对启发式似然近似或 ELBO 类下界的依赖。
• •高效性：利用统计技术，我们展示了如何以稳定且节省内存的方式实现所提出的算法，并讨论了各种实际权衡。
• •有效性：我们在四个推理任务（GSM8K、MATH、Countdown 和 Sudoku）上训练模型，结果表明 AGRPO 超越了所有先前基于近似的方法。此外，我们还证明，使用 AGRPO 后训练的模型在大幅减少采样步数进行评估时仍能保持高精度，这是预训练基础 dLLM 所不具备的能力。

2 Preliminaries

2.1 dLLM Pretraining

语言领域最常见的离散扩散形式是掩码（或“吸收”）方法，即模型通过随机掩码破坏数据，并学习将其恢复。具体而言，给定离散 token 序列的分布 (p) 上的序列 (x=(x_1,\dots,x_n))，模型被训练以最大化以下对数似然的证据下界（ELBO）： (\mathcal{L}(\theta)=\mathbb{E}{\begin{subarray}{c}t\sim U[0,1]\ x,x^{t}\sim p^{t}\end{subarray}}\left[\frac{1}{t}\sum{x^{t}{i}=\blacksquare}\log p{\theta}(x_{i}\mid x^{t})\right]) (1) 其中 (x,x^{t}\sim p^{t}) 表示 (x) 从 (p) 中采样，而 (x^{t}) 是通过以概率 (t) 独立地将每个 token (x_i) 设置为掩码 token ■ 从 (x) 中获得的。与 BERT 类似，目标是让模型学习在给定上下文条件下掩码 token 的边缘分布。关键点在于，与经典的自回归目标相比，这种掩码 token 预测目标更难且更通用，因为解掩码顺序可以是任意的，且模型必须预测多个掩码 token。相比之下，自回归模型将其自身限制为按从左到右的顺序对下一个token进行建模。施加此约束的好处是双重的：它使自回归模型能够通过链式法则最大化精确的序列似然，并且当与仅解码器 Transformer 结合使用时，还能通过因果自注意力机制实现训练并行化。

2.2 dLLM Inference

为了生成文本，dLLM 从一个全掩码或部分掩码序列开始，获取每个掩码 token 的边缘分布，然后通过从其边缘分布中采样来解掩码其中的一部分。解掩码的位置可以选择随机选择（遵循理论上的“逆向过程”），或者保留概率最高的 token（如 Nie et al. (2025) 所提出）。（我们分别称之为“随机”和“基于置信度”的解掩码）。其余 token 保持不变——掩码 token 保持掩码状态，未掩码 token 保持未掩码状态——并将这个新序列反馈回模型。重复此过程直到所有 token 都被解掩码。在本文中，我们用 (m) 表示采样步数，用 (n) 表示序列长度。扩散模型的一个显著优势是能够动态调整每步解掩码的 token 数量 (n/m)。通常，(n/m) 的比例选择得相对较小（(\leq 8)）——在每一步解掩码更多 token 会严重降低以困惑度或准确率衡量的质量。我们在后续章节将展示，对于特定任务，与预训练模型不同，后训练实际上允许使用高得多的 (n/m) 值而不会出现明显的性能下降。有关 dLLM 推理的更多细节，请参阅附录 F。

2.3 Reinforcement Learning and MDPs

马尔可夫决策过程（MDP）是对序列决策问题的形式化，由状态空间 (\mathcal{S})、动作空间 (\mathcal{A})、转移核 (p(\cdot \mid s,a)) 和奖励函数 (r(s,a)) 组成。广义而言，强化学习（RL）的目标是学习一个策略，即最大化期望奖励总和的分布 (\pi(a \mid s))： (\mathbb{E}{\tau\sim\pi}\left[\sum{t=0}^{T}r(s_{t},a_{t})\right]) (2) 其中 (\tau) 表示一条轨迹（或“采样路径”），即状态和动作的序列 ((s_{0},a_{0},s_{1},a_{2},\dots,s_{T},a_{T}))，其中 (a_{i}\sim\pi(\cdot \mid s_{i})) 且 (s_{i+1}\sim p(\cdot \mid s_{i},a_{i}))。

3 Policy Gradients for Diffusion Models

Refer to caption Figure 1: Existing RL post-training algorithms focus on sequence-level likelihoods and require either autoregressive factorization or ELBO-like bounds, which result in biased policy gradients. Our proposed algorithm instead focuses on individual unmasking steps, aligning more naturally with the dLLM generation process.

在本节中，我们将 dLLM 后训练框定为一个多步 RL 问题，这与 LLM 后训练的标准序列级框定方式不同。为了说明这一动机，我们首先概述策略梯度。

3.1 Policy Gradient Methods

策略梯度（PG）是一类广受欢迎的算法，用于训练神经网络参数化的策略 (\pi_\theta) 以最大化期望奖励。最简单的策略梯度形式 REINFORCE 涉及以下梯度更新： (\nabla_\theta \mathcal{J}{PG}(\theta)=\mathbb{E}{\tau\sim\pi}\left[\left(\sum_{t=0}^T\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right)\left(\sum_{t=0}^Tr(s_t,a_t)\right)\right].) (3) 更复杂的算法（如近端策略优化 PPO）通过减去可学习的奖励基线（即价值函数）并允许通过重要性采样进行 off-policy 更新来改进这一公式。然而，所有 PG 方法的底层结构保持不变：计算轨迹上所有动作的(精确)似然，并以奖励作为权重。基于这一共同结构，我们展示了如何为 dLLM 开发一种原理性的 PG 方法。

3.2 LLM Post-Training with RL

在 LLM 后训练的语境下，特别是使用可验证奖励的 RL（RLVR），状态 (s) 对应上下文（即提示词），动作 (a) 对应模型输出，奖励 (r) 通过真实答案提供。（由于转移是确定性的，我们可以放心忽略转移核 (p)。）得益于自回归分解，人们可以将 AR LLM 视为在整个序列上生成分布的策略，即一个单步MDP。在此视角下，公式 (2) 中的轨迹仅由初始提示 (s_0) 和模型输出 (a_0) 组成。这使得 PG 目标 (3) 非常便利，并且对许多后训练任务十分有效；我们将在第 6 节讨论大量致力于此设定的研究工作。

3.3 Text Generation as a Multi-Step MDP

扩散模型不具备清晰的从左到右分解特性，这使得序列似然难以处理。特别是对于 LLaDA 等 dLLM，计算给定长度 (n) 序列的精确似然需要对 (O(n!)) 种解掩码顺序进行边缘化，这对于较大的 (n) 来说根本不可行。在预训练中最大化似然下界已经足够，但这与需要精确序列似然的后训练单步 MDP 视图不兼容。虽然现有的 dLLM 后训练方法简单地使用 ELBO 替代真实似然，但我们希望在不作此类妥协的情况下开发一种更具原理性的方法。这促使我们考虑生成过程的多步MDP 公式，其中：

• •状态 (s) 是当前部分掩码的序列，
• •动作 (a) 是待解掩码的 (n/m) 个 token¹，
• •奖励 (r) 在序列完全解掩码后提供，即最后一步。 ¹ The positions to be unmasked are determined by the unmasking strategy (random or confidence-based) and are not explicitly optimized, following theoretical assumptions of diffusion models.

换言之，此 MDP 中的单步被视为一个单独的解掩码步骤，而非整个生成过程。该公式与扩散模型的参数化方式更自然地契合，已被 DDPO 和 Flow-GRPO 等连续扩散 RL 的先前工作采用。关键在于，将去噪视为多步 MDP 的这一视角允许在 PG 目标中使用精确的动作似然，从而避开了对下 b