超越Mode-Seeking RL：扩散语言模型的轨迹平衡后训练

# 超越追求模式强化学习：面向扩散语言模型的轨迹平衡后训练技术 来源：https://arxiv.org/html/2605.13935 Saba Ahmadi Prasanna Parthasarathi¹ Yufei Cui Noah’s Ark Lab ###### 摘要 扩散语言模型是自回归模型的一种有前景的替代方案，然而针对它们的后训练方法主要采用奖励最大化目标。我们在此场景中发现了一个核心失效模式，称为 **轨迹锁定**：基于采样的奖励驱动更新会将概率质量过度集中于一个狭窄的去噪路径集合上，从而在重复采样时减少对替代正确解的覆盖。为了解决这个问题，我们提出了 TraFL（TrajectoryFlow baLancing，轨迹流平衡），一种轨迹平衡目标，它将策略训练向一个以冻结参考模型为锚定的奖励倾斜目标分布。我们通过两个要素使其对扩散语言模型实用：一个与扩散兼容的序列级替代方法，以及一个学习的、依赖于提示的归一化项。在数学推理和代码生成基准测试中，TraFL 是唯一一个在所有基准-长度设置下均优于基础模型的经过评估的后训练方法，并且随着采样预算的增加，其增益持续存在。这些改进迁移到了留出评估中：TraFL 在 Minerva Math (Lewkowycz 等人, 2022) 上保持高于基础模型，并且在每个 LiveCodeBench (Jain 等人, 2025) 难度划分上都是最强方法。

## 1 引言

近期的扩散语言模型（dLLMs）已成为自回归模型的一个引人注目的替代方案，在推理和代码生成方面显示出早期潜力 (Zhu 等人, 2025；Ye 等人, 2025)。进一步推进这些模型的一个关键开放问题是如何有效地对它们进行后训练。从能力角度来看，当前的后训练方法提高了单样本准确率，但在重复采样下难以产生 **多样化** 的正确解集合——这在存在多个有效答案时是一个关键要求。设计这样的目标颇具挑战性，因为扩散语言模型不像自回归模型那样暴露词元级别的条件对数概率。因此，现有方法通过替代似然、自回归简化或逐步近似，将奖励最大化强化学习适应到扩散生成 (Zhao 等人, 2025；Tang 等人, 2025；Wang 等人, 2025a；Kunde 等人, 2026；Ni 等人, 2026)。在这些方法中，我们观察到一个反复出现的失效模式，我们称之为 **轨迹锁定**：由于奖励仅取决于最终完成，并且对产生它的具体去噪路径不敏感，基于采样的策略梯度更新强化了已经被偏好的路径，逐渐将概率质量坍缩到一个狭窄的轨迹子集上——随之坍缩的还有对替代正确解的覆盖范围。

旨在向奖励倾斜目标分布的分布匹配方法 (Zhu 等人, 2026a, b) 提供了一种原理性的替代方案，然而，正如我们所展示的，归一化项的处理方式可能会重新引入相同的集中失效模式。这引发了一个基本问题：*对于扩散语言模型来说，什么是既能避免轨迹锁定又保持实用的正确后训练目标？*

为此，我们提出了 TraFL（TrajectoryFlow baLancing，轨迹流平衡），一种基于生成流网络 (GFlowNets) 轨迹平衡原理 (Bengio 等人, 2023) 的后训练目标。TraFL 不是将奖励用作采样轨迹的无约束放大器，而是将策略训练向一个以冻结参考模型为锚定的奖励倾斜目标分布。在这个目标中，奖励增加了分配给成功完成的概率质量，而参考模型则规范了如何在合理生成上分配概率质量。我们通过两个要素使其对扩散语言模型实用：(i) 一个与扩散兼容的序列级替代方法，用于比较当前模型和参考模型下完全去噪的完成结果；(ii) 一个学习的、依赖于提示的归一化项，与策略一起联合训练，并且在每个训练步骤都能接收梯度信号，无论采样到哪些完成结果。

在数学推理 (GSM8K (Cobbe 等人, 2021)、MATH-500 (Lightman 等人, 2023)) 和代码生成 (HumanEval (Chen 等人, 2021)、MBPP (Austin 等人, 2021)) 基准测试中，TraFL 是唯一一个在所有基准-长度设置下均优于基础模型的直接评估的后训练方法。这些改进转移到了留出的评估中：在 Minerva Math (Lewkowycz 等人, 2022) 和 LiveCodeBench (Jain 等人, 2025) 上，TraFL 在每个 LiveCodeBench 难度划分上都是最强方法。

我们的贡献如下：

1. 我们识别出 **轨迹锁定** 是 dLLM 后训练的一个核心失效模式，其中基于采样的奖励驱动更新将概率质量坍缩到一个狭窄的去噪路径集合上。
2. 我们提出了 TraFL（TrajectoryFlow baLancing，轨迹流平衡），一种用于扩散语言模型后训练的轨迹平衡目标，它向以参考模型为锚定的奖励倾斜目标分布进行训练，并通过一个与扩散兼容的序列级替代方法和一个学习的、依赖于提示的归一化项使其变得实用。
3. TraFL 是唯一一个在 GSM8K、MATH-500、HumanEval 和 MBPP 上所有基准-长度设置下均优于基础模型的评估后训练方法，平均性能优于 ESPO (Ou 等人, 2026) 和 JustGRPO (Ni 等人, 2026)。
4. 这些改进转移到了留出的 Minerva Math (Lewkowycz 等人, 2022) 和 LiveCodeBench (Jain 等人, 2025) 上，其中 TraFL 在每个 LiveCodeBench 难度划分上都是最强方法，并且 LLM-as-judge 分析提供了证据表明改进与更广泛的正确解覆盖相关，而不仅仅是单一模式的更尖锐评分。

## 2 相关工作

#### 扩散语言模型的强化学习后训练。
近期的扩散语言模型，如 LLaDA (Zhu 等人, 2025) 和 Dream (Ye 等人, 2025)，已经使基于扩散的文本生成成为自回归语言模型的一个可行替代方案，但这些模型的强化学习仍然具有挑战性，因为扩散生成并不暴露自回归 LLM 中 PPO 或 GRPO 风格训练所使用的相同从左到右的词元级别条件分解。早期工作通过引入似然替代方法，将自回归强化学习目标适应到扩散模型。特别地，d1 (Zhao 等人, 2025) 提出了 diffu-GRPO，这是一种针对掩码 dLLM 的 GRPO 风格方法，建立在每步每词元对数概率估计以及序列似然的平均场近似之上，并在其完整方案中结合了前置的监督微调阶段。wd1 (Tang 等人, 2025) 去除了显式的策略比率估计，而是优化一个从逆 KL 正则化策略优化推导出的无比率加权对数似然目标。SPG (Wang 等人, 2025a) 通过最大化正优势样本的下界和最小化负优势样本的证据上界，结合用于更稳定蒙特卡洛估计的块级掩码策略，解决了一侧似然替代方法引入的偏差。

另一条互补的工作线认为，核心问题不仅在于替代方法的质量，还在于强化学习目标所使用的动作粒度。ESPO (Ou 等人, 2026) 最明确地形式化了这一观点：它将整个序列生成视为一个单一动作，并使用 ELBO 作为可处理的序列级代理，结合稳定的比率和 KL 估计量，认为词元级目标从根本上与非自回归扩散生成不匹配。TraceRL (Wang 等人, 2025b) 则采取了另一种视角，强调训练目标与模型偏好的推理轨迹之间的一致性。它通过去噪轨迹执行轨迹感知优化，并引入一个基于扩散的价值模型用于方差缩减。相比之下，JustGRPO (Ni 等人, 2026) 认为在强化学习期间保留任意顺序的生成本身可能对推理适得其反，而是将训练约束为自回归顺序，以便标准 GRPO 可以直接应用，同时仍在推理时保留 dLLM 的并行解码优势。

#### 分布匹配与奖励倾斜目标分布。
与我们工作精神最接近的是那些超越纯奖励最大化，转而优化向奖励倾斜目标分布的方法。在 dLLM 方面，DMPO (Zhu 等人, 2026b) 将后训练公式化为策略分布匹配向一个奖励倾斜目标分布，并通过重要性采样和加权去噪交叉熵（即可扩展的前向 KL 风格分布匹配近似）来实现这一点。在自回归方面，FlowRL (Zhu 等人, 2026a) 同样倡导匹配奖励倾斜分布而不是仅最大化奖励，但通过一种适用于自回归推理模型而非扩散语言模型的 GFlowNet 风格公式来实现。我们的方法与这种更广泛的分布匹配观点一致，但在针对 dLLM 的具体实现上有所不同：我们采用了一种适合扩散生成的轨迹平衡视角，将其与一个与扩散兼容的序列级替代方法相结合，并学习一个捕获奖励倾斜目标归一化的、依赖于提示的配分函数。

## 3 轨迹流平衡

我们介绍 TraFL（TrajectoryFlow baLancing，轨迹流平衡），一种用于扩散语言模型后训练的、以参考模型为锚定的轨迹平衡目标。我们首先形式化为什么终端奖励可能诱发轨迹锁定，然后定义 TraFL 所基于的奖励倾斜目标。

#### 设置：扩散轨迹与终端奖励
给定一个提示 \(\mathbf{x}\)，一个离散扩散语言模型定义了一个关于去噪轨迹 \(\tau^{(i)}\) 的分布，其中 \(\tau^{(i)} = (z_T^{(i)}, z_{T-1}^{(i)}, \ldots, z_0^{(i)})\)，\(\tau^{(i)} \sim p_\theta(\tau \mid \mathbf{x})\)。这里，\(z_T\) 表示完全加噪状态，\(z_0\) 表示最终完成 \(\mathbf{y}\)。奖励 \(r(\mathbf{x}, \mathbf{y})\) 仅分配到这个终端完成，而模型分布是通过整个去噪轨迹诱导的。这个区别很重要，因为多个去噪轨迹可以终止于同一个完成 \(\mathbf{y}\)，而不同的完成可能对应不同的有效解模式。因此，后训练可以同时影响概率在到达固定答案的路径上的分配以及对不同终端解模式的覆盖。

### 3.1 轨迹锁定：为什么仅有终端奖励是不够的

仅最大化终端奖励对于概率质量如何在到达同一完成的去噪路径之间分配是盲目的。该目标取决于分配给最终完成 \(\mathbf{y}\) 的总概率，但不取决于这个概率在产生 \(\mathbf{y}\) 的去噪轨迹之间如何分割。我们在下面形式化了这种路径无关性。

###### 命题 1（终端奖励的路径无关性）。
固定一个提示 \(\mathbf{x}\) 和一个最终完成 \(\mathbf{y}\)。令 \(\mathcal{T}(\mathbf{y})\) 表示终止于 \(\mathbf{y}\) 的去噪轨迹集合。任何在 \(\mathcal{T}(\mathbf{y})\) 内的轨迹之间保持 \(p_\theta(\mathbf{y} \mid \mathbf{x})\) 不变的概率质量重新分配，都不会改变终端奖励目标。

###### 证明。
终端奖励目标分解为：
\[
J(\mathbf{x}) = \sum_{\mathbf{y}} p_\theta(\mathbf{y} \mid \mathbf{x}) \, r(\mathbf{x}, \mathbf{y}), \quad p_\theta(\mathbf{y} \mid \mathbf{x}) = \sum_{\tau \in \mathcal{T}(\mathbf{y})} p_\theta(\tau \mid \mathbf{x}).
\]
\(\mathbf{y}\) 对 \(J\) 的贡献仅取决于标量 \(p_\theta(\mathbf{y} \mid \mathbf{x})\)，而不取决于这个质量如何分布在 \(\mathcal{T}(\mathbf{y})\) 中的轨迹上。因此，任何在 \(\mathcal{T}(\mathbf{y})\) 内部保持 \(p_\theta(\mathbf{y} \mid \mathbf{x})\) 不变的质量重新分配都会使 \(J(\mathbf{x})\) 保持不变。∎

命题 1 是一个目标层面的陈述：奖励最大化并不区分到达同一终端完成的不同路径。然而，在采样优化中，这种平坦性可能变得不稳定。只有采样到的轨迹会收到直接的梯度信号。因此，如果两条轨迹到达相同的奖励完成，但一条在训练早期被稍频繁地采样，它会收到更多正向更新，变得更可能被采样，从而在后续更新中被更频繁地采样。我们将这种自我强化的集中称为 **轨迹锁定**。关于这种反馈效应的正式分析见附录 C.2 和附录 C.3。

轨迹锁定之所以重要，是因为轨迹多样性是终端解覆盖的上限。我们接下来明确这一联系。

###### 定理 1（轨迹多样性是模式覆盖的必要条件）。
令 \(\mathcal{T}\) 为一个随机去噪轨迹，取值于可数集 \(\Omega_\mathcal{T}\)，并令 \(M = g(\mathcal{T})\) 为在确定性映射 \(g: \Omega_\mathcal{T} \to \Omega_M\) 下的终端解模式。则
\[
|\mathrm{supp}(M)| \leq |\mathrm{supp}(\mathcal{T})| \quad \text{且} \quad H(M) \leq H(\mathcal{T}),
\]
其中 \(\mathrm{supp}(\cdot)\) 表示分布的支持集，\(H(\cdot)\) 表示香农熵。特别地，终端模式覆盖不能超过轨迹多样性。

###### 推论 1（轨迹锁定限制了模式覆盖）。
如果轨迹锁定将轨迹分布减少到所有可能路径的子集 \(\mathcal{S}\)，那么可以覆盖的终端模式集合最多为 \(g(\mathcal{S})\)。