🤖 machine learning

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

本文介绍了 TraFL，这是一种针对扩散语言模型的轨迹平衡后训练方法，它通过将策略与奖励倾斜的目标分布对齐，克服了奖励最大化方法中“轨迹锁定”的失效模式，从而在数学推理和代码生成基准测试中实现了持续的性能提升。

原作者： Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

发布于 2026-05-15

📖 1 分钟阅读☕ 轻松阅读

原作者： Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用通俗易懂的语言和富有创意的类比，对论文《超越模式寻求型强化学习：扩散语言模型的轨迹平衡后训练》的解释。

宏观图景：一种训练 AI 的新方法

想象你有一位才华横溢的艺术家（AI 模型），他可以通过从一块充满噪点和杂讯的画布开始，慢慢将其清理，直到一幅清晰的图像显现出来。这就是扩散语言模型的工作原理：它们不像旧式 AI 模型那样像打字机一样逐字逐句地写作；相反，它们从一堆杂乱的字母开始，逐渐将这些字母“去噪”，形成连贯的句子。

这篇论文指出，目前我们训练这些艺术家去解决难题（如数学或编程）的方法存在缺陷。作者提出了一种名为TraFL（轨迹流平衡，Trajectory Flow baLancing）的新方法，旨在修复一种被称为**“轨迹锁定”**（Trajectory Locking）的具体缺陷。

问题所在：“单一路径”陷阱（轨迹锁定）

要理解这个问题，请想象一个迷宫。

目标：迷宫有许多不同的正确出口（对应数学问题的不同有效解）。
旧方法（奖励最大化强化学习）：想象你在训练一只狗寻找出口。每次狗找到一个出口时，你就给它一个奖励。
- 缺陷：狗并不关心它走了哪条路才得到奖励，它只在乎是否得到了奖励。
- 结果：如果狗碰巧在早期发现了一条通往奖励的特定路径，它就会得到奖励。它记住了那条路。下次，它会尝试那条路。它再次得到奖励。它对那一条路变得更加自信。
- 轨迹锁定：最终，狗完全停止探索迷宫。它只沿着那条狭窄的路径奔跑，即使还有 50 个其他有效的出口它本可以找到。它已经“锁定”在单一解决方案上，忘记了如何找到其他路径。

在论文中，作者将这种现象称为轨迹锁定。因为 AI 仅因最终答案而获得“奖励”（分数），它忽略了通往那里有许多不同方式（路径）的事实。它将所有的创造力坍缩到一条狭窄的路径上，导致当你要求它再次尝试时，它难以找到替代性的正确答案。

解决方案：“指南书”方法（TraFL）

作者提出了TraFL，它改变了训练规则。不再仅仅因为找到任何出口就给予奖励，而是给 AI 一本指南书（一个冻结的参考模型）和一张地图。

指南书（参考模型）：这是 AI 的一个版本，尚未针对特定奖励进行训练。它代表了一种“健康”且多样化的思维方式。它知道解决问题有多种方法。
地图（奖励倾斜目标）：我们告诉 AI：“你想要找到能获得奖励的出口（正确答案），但是你必须保持你的移动模式与指南书相似。”

类比：
想象你在教一名学生解决数学问题。

旧方法：你说：“用任何你能想到的方法得到正确答案！”学生发现了一个奏效的窍门，将其死记硬背，并拒绝学习任何其他方法。
TraFL 方法：你说：“找到正确答案，但保持你的思维过程多样且灵活，就像一位知道多种不同策略的优等生那样。”

TraFL 迫使 AI 平衡两件事：

获得奖励：找到正确答案。
保持多样性：不坍缩成单一、重复的路径。它将“概率质量”（选择某条路径的可能性）分散在多个不同的有效解上，并以指南书为锚点。

他们是如何实现的

论文指出，扩散模型很棘手，因为它们不像旧模型那样逐步展示其“思维过程”。为了解决这个问题，作者发明了两个工具：

代理分数：由于他们无法看到每一步的确切数学计算，他们创建了一个“代理”分数来估算完整答案的好坏，从而允许他们在不需要完美观察每一个微小步骤的情况下训练模型。
学习到的归一化器：他们教会 AI 一个特殊的“计算器”，根据具体问题（提示词）调整任务的难度，确保训练保持公平和平衡。

结果：它真的有效吗？

作者在数学（解决文字题）和代码（编写计算机程序）上测试了这种新方法。

“单一路径”陷阱被打破：与其他有时在找到一个解决方案时变好、却在找到不同解决方案时变差的方法不同，TraFL 在每一项测试中都有提升。
更多样本 = 更好结果：当他们要求 AI 生成 16 个不同答案而不是仅仅 1 个时，TraFL 的表现显著提高。这证明它实际上找到了更多不同的正确答案，而不仅仅是一个幸运的猜测。
它适用于新事物：AI 不仅仅是死记硬背训练问题。当在从未见过的全新数学问题（Minerva Math）和新的编程挑战（LiveCodeBench）上进行测试时，TraFL 表现最强。
多样性检查：他们使用了一个“裁判”（另一个 AI）来审视答案。裁判确认，TraFL 不仅仅是用不同的措辞给出相同的答案；它实际上使用了不同的推理策略和不同的算法来解决相同的问题。

总结

这篇论文指出了一个缺陷，即 AI 模型会“卡”在解决某种问题的单一方式上，而忽略了其他有效的解决方案。他们通过TraFL修复了这一问题，该方法教导 AI 在寻求正确答案的同时，保持对不同路径的多样化“探索”，并由参考模型引导。其结果是，AI 不仅更聪明，而且在被要求生成多个解决方案时，也更具创造力和可靠性。

技术摘要：扩散语言模型的轨迹平衡后训练

问题：扩散强化学习中的轨迹锁定

扩散语言模型（dLLMs）为推理和代码生成提供了一种有前景的替代方案，可取代自回归模型。然而，现有的 dLLMs 后训练方法大多借鉴了专为自回归模型设计的奖励最大化强化学习（RL）目标（例如 PPO、GRPO）。作者指出了这种适配中存在的一个核心失效模式，称为轨迹锁定。

在扩散模型中，奖励仅分配给最终的去噪完成结果（ $y$ ），而模型分布则诱导在整个去噪轨迹（ $\tau$ ）上。由于终端奖励是“路径无差别”的——即任何到达相同正确 $y$ 的轨迹都会获得相同的奖励——采样的策略梯度更新倾向于强化在训练早期恰好被采样的特定轨迹。这形成了一个自我强化的反馈循环，导致概率质量坍缩到一组狭窄的去噪路径上。因此，即使模型学会了生成正确答案，它在重复采样下也会失去生成多样化有效解（模式）的能力。作者通过定理 1对此进行了形式化证明，表明终端模式覆盖率受轨迹多样性的上限约束；因此，轨迹锁定直接限制了模型的解覆盖率。

方法论：TraFL（轨迹流平衡）

为了解决轨迹锁定问题，作者提出了TraFL，这是一种基于生成流网络（GFlowNets）中的轨迹平衡原则的后训练目标。TraFL 不直接最大化奖励，而是训练策略以匹配一个锚定在冻结参考模型上的奖励倾斜目标分布。

目标函数

目标分布 $p^*(\tau | x)$ 定义为：
$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$
其中：

$p_{\text{ref}}$ 是冻结的参考模型。
$r(x, y)$ 是终端奖励。
$\beta$ 控制奖励的锐度。
$Z(x)$ 是依赖于提示词的归一化项（配分函数）。

训练目标是最小化当前模型的对数概率与目标之间的平方残差：
$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$
$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$

扩散模型的实际实现

由于 dLLMs 不暴露精确的轨迹对数概率，作者引入了两个关键组件以使 TraFL 具有实用性：

扩散兼容的序列级代理：他们基于掩码重建下界推导出了一个可处理的对数概率代理（ $\log \hat{p}_\theta(y|x)$ ）。这涉及采样一个掩码层级 $l$ ，破坏 $l$ 个令牌，并对重建令牌的 log 概率取平均值。使用 $1/l$ 因子来归一化分数，使其独立于被掩码令牌的數量。
学习的提示词依赖归一化：与通过批次上的 softmax 估计 $Z(x)$ 的方法（这可能导致集中）不同，TraFL 通过一个小 MLP 学习配分函数 $Z_\phi(x)$ 。无论采样到哪些完成结果，该术语在每个步骤都会接收梯度信号，从而避免了其他分布匹配方法中出现的“批次局部”归一化失效问题。

主要贡献

轨迹锁定的识别：该论文正式识别并证明，dLLMs 中的奖励最大化目标会导致轨迹多样性的坍缩，进而限制不同正确解模式的覆盖率。
TraFL 目标：提出了一种轨迹平衡目标，将策略锚定在参考模型上，将奖励的作用（倾斜终端完成结果）与参考模型的作用（在合理轨迹上分配概率质量）分离开来。
实用的代理与归一化：开发了一种扩散兼容的序列级代理和一种学习的归一化项，使得在没有精确轨迹概率的情况下也能进行稳定训练。
实证优越性：TraFL 是唯一在所有基准长度设置（GSM8K、MATH-500、HumanEval、MBPP）和采样预算下均一致优于基础模型的后训练方法。

结果

作者在数学推理和代码生成基准测试中评估了 TraFL，对比了基础模型（LLaDA-8B-Instruct）和强基线模型（ESPO、JUSTGRPO）。

分布内性能：TraFL 在所有采样预算（ $k \in \{1, 2, 4, 8, 16\}$ ）和温度下均提升了 Pass@k。值得注意的是，它在提升 Pass@1（单次采样准确率）的同时，随着采样预算的增加也能更好地扩展，而基线模型往往为了多采样增益而牺牲单次采样准确率。
未见数据泛化：增益转移到了分布外基准测试上。在Minerva Math上，TraFL 保持在基础模型之上。在LiveCodeBench上，TraFL 在所有难度划分（简单、中等、困难）中都是最强的方法，在中等和困难问题上显著优于基线模型。
多样性分析：使用“LLM 作为裁判”的协议，作者表明 TraFL 的增益与更广泛的独特解决方案策略（例如不同的证明方法或算法分解）的覆盖率相关，而不仅仅是锐化单一模式的分数。像 ESPO 这样的基线虽然提高了准确率，但并未显示出相对于基础模型的显著多样性增益，这与轨迹锁定假设一致。

意义与主张

该论文声称，TraFL 为扩散模型标准 RL 后训练的“模式寻求”失效提供了一种原则性解决方案。通过从奖励最大化转向轨迹平衡，该方法保留了覆盖多个有效解所需的去噪路径多样性。

作者指出，TraFL 是唯一在所测试的每个基准长度设置中均优于基础模型的方法。他们强调，这些改进对采样预算的增加具有鲁棒性，并泛化到未见数据集，这表明轨迹平衡目标有效地将奖励信号与路径分配机制解耦。该工作主张在 dLLM 后训练中采用分布匹配目标而非纯粹的奖励最大化，提供了一个实用的框架，避免了标准策略梯度方法固有的集中失效模式。

该论文在局限性方面保持适度，指出该方法依赖于近似的序列级代理而非精确的轨迹概率，并使用粗略的二元奖励（完全匹配或测试通过），这可能会在成功罕见的场景中限制学习效率。