🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

本文介绍了 $f$ -轨迹平衡，这是一类损失函数族，它将均方误差方法推广至所有 $f$ -散度，从而能够利用离线策略数据训练生成模型和大语言模型，同时保留其对应的在线策略 $f$ -散度梯度所特有的优化性质（如模式覆盖）。

原作者： Jake Fawkes, Jason Hartford

发布于 2026-05-18

📖 1 分钟阅读☕ 轻松阅读

原作者： Jake Fawkes, Jason Hartford

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

想象一下，你正在教一个机器人探索一个广阔而黑暗的洞穴系统。这个洞穴有许多隐藏的洞穴室（称为“模式”），其中一些藏有黄金（高奖励），而另一些则只是空洞或危险之地。你的目标是教会机器人尽可能多地发现不同的黄金洞穴室，而不是只找到一个大金矿就忽略其他一切。

本文介绍了一套新的“教学工具”（损失函数），用于训练这些机器人，无论它们是在为新药生成分子、创作艺术，还是编写代码。

以下是用简单类比对本文思想的分解：

1. 问题所在：“淘金热”与“探险家”

过去，训练这些机器人的标准方式就像一场淘金热。机器人发现了一处有点黄金的地方，就会兴奋起来，然后把所有时间都花在那里挖掘。它会忽略洞穴中所有其他的金矿。

技术术语： 这被称为“模式寻求”。模型会坍缩到少数几个高概率的答案上，而忽略了其余的多样性。
本文的目标： 我们想要“模式覆盖”。我们希望机器人分散开来，探索所有的金矿，即使是那些较小的，从而获得洞穴的完整地图。

2. 旧工具：“平方 KL"损失

最近，研究人员发现了一个巧妙的技巧，可以使用“平方误差”方法（测量机器人认为的内容与它应该认为的内容之间的距离）来训练机器人。

类比： 想象一位老师给学生的作业打分。如果学生答错了，老师不只是说“错了”，而是计算错误的平方。
优势： 这种方法非常稳定。即使学生是从旧笔记（离线策略数据）中学习，而不是从实时课程中学习，它也能起作用。
缺陷： 尽管它很稳定，但这种特定的“平方”方法仍然像一场淘金热。它自然地推动机器人只关注少数几个答案，从而错过了洞穴的多样性。

3. 新解决方案："f-轨迹平衡”家族

作者们意识到，“平方”方法只是更大教学工具家族中的一种特定风味。他们将这个家族称为f-轨迹平衡。

把这个家族想象成收音机上的一个旋钮或音量控制。

将旋钮向一边转动（低数值）： 你会得到像超级探险家一样的工具。这些工具迫使机器人分散开来，找到每一个可能的金矿，即使是那些难以到达的。这对于药物发现非常有益，因为你想要找到任何可能有效的分子，而不仅仅是最明显的那个。
将旋钮向另一边转动（高数值）： 你会得到像淘金者一样的工具。这些工具告诉机器人忽略小矿，集中火力关注最大、最明显的金堆。如果你只想要单个最佳答案，这就很有用。
中间地带： 你可以将旋钮设置在两者之间的任何位置，以获得探索与专注的混合效果。

4. 为什么这很重要

本文证明了两个主要观点：

“魔法开关”： 他们表明，你可以用这些新的“旋钮”工具中的任何一个来替换标准的“平方”工具，数学原理依然完美适用。机器人学习得同样好，但具有不同的性格（更具探索性或更专注）。
稳定性： 就像旧工具一样，这些新工具即使机器人是从旧数据（离线策略）中学习也能起作用。这对于大型语言模型（LLM）等现实世界应用至关重要，因为在这些应用中，训练通常是异步进行的（机器人从一段时间前生成的数据中学习）。

5. 现实世界测试（洞穴地图）

作者在三个不同的“洞穴”中测试了这些工具：

网格游戏： 一个拥有四个满是黄金角落的简单计算机迷宫。标准工具只找到了一个角落。而新的“探险家”工具则迅速找到了所有四个角落。
分子发现（SynFlowNet）： 他们尝试生成新的化学分子。通过将旋钮调至“探险家”模式，他们生成了种类更广泛、独特的潜在药物分子，而不仅仅是同几种化学物质的变体。
语言模型（LLMs）： 他们调整了 AI 模型以解决数学问题。新工具使 AI 能够探索解决问题的不同方法，而不会陷入重复同一答案的循环，即使训练数据是延迟的（异步的）。

总结

本文并没有发明一种新型机器人。相反，它发明了一套新的指令，用于如何训练它们。

旧方法： “找到最大的金矿并在那里挖掘。”（稳定，但乏味）。
新方法： “这里有一个旋钮。你可以选择成为淘金者、超级探险家，或者介于两者之间的任何角色。无论你选择哪一个，训练过程都将保持稳定，并能利用旧数据。”

这为工程师提供了一个简单的旋钮，用于控制其 AI 模型应该具有多大的创造力或专注度，而不会破坏训练过程。

技术摘要：f-轨迹平衡

问题陈述

生成模型（包括大型语言模型（LLM）和生成流网络（GFlowNets））的对齐与微调通常依赖于最小化模型分布与目标分布之间的差异。在强化学习（RL）微调中，这通常被表述为在 Kullback-Leibler（KL）散度惩罚的约束下最大化预期奖励。

该领域的一个重大挑战是寻求模式（mode-seeking）与覆盖模式（mode-covering）行为之间的权衡。标准目标（如反向 KL 散度，通常通过轨迹平衡或平方 KL 损失实现）往往倾向于寻求模式，导致模型坍缩到少数几个高概率模式上。这对于需要探索的任务（如药物发现或多样化候选生成）而言是次优的，在这些任务中，“覆盖模式”（即捕捉目标分布的完整多样性）更为可取。

此外，现有方法往往难以处理离线策略有效性（off-policy validity）。虽然 REINFORCE 或 PPO 等估计器很常见，但它们存在高方差问题。最近使用平方 KL 损失（ $KL_{sq}$ ）的方法提供了在线策略下的低方差梯度和离线策略下的有效性，但它们本质上与反向 KL 散度的属性绑定，限制了控制探索 - 利用权衡的能力。

方法论

作者提出了f-轨迹平衡（f-Trajectory Balance），这是一族代理损失函数，将平方 KL 损失推广到整个f-散度族。

理论基础

核心理论贡献在于建立了对数概率上的平移不变损失函数与f-散度之间的一一对应关系。

从散度到损失：对于定义 f-散度 $D_f(p_\theta \| p^*)$ 的任意凸函数 $f$ ，作者推导出了定义在对数概率差 $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ 上的损失函数 $L_f$ ：
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
该损失具有两个关键属性：
- 在线策略梯度等价性：当样本从模型策略（ $\mu = p_\theta$ ）中抽取时， $L_f$ 的期望自动微分梯度与相应 f-散度 $D_f$ 的梯度相匹配。
- 离线策略有效性：当在离线策略数据（从行为分布 $\mu$ 采样）上评估时，只要目标分布的支撑集包含在行为策略的支撑集中，该损失仍然是一个有效的目标，且具有相同的全局最小值（ $p_\theta = p^*$ ）。
从损失到散度：反之，对数概率上的任何严格凸且平移不变的损失 $\ell$ 都对应于最小化特定的 f-散度。论文证明了这些映射互为逆映射。

实际实现

为了解决训练中的实际挑战，作者引入了几种扩展：

DevGrad 损失（批次归一化）：为了处理不可处理的配分函数（在 GFlowNets 和非归一化目标中很常见），作者推广了VarGrad估计器。他们不使用批次方差，而是使用对数概率差的批次广义偏差。这涉及通过在批次上最小化损失来估计归一化常数 $\log Z$ ，有效地对得分函数系数进行中心化并降低方差。
平滑损失（Tempered Loss）：为了防止在处理小温度参数 $\beta$ （此时目标分布为 $p^* \propto \pi_{ref} \exp(r/\beta)$ ）时出现数值溢出，作者提出了一种平滑损失。该损失将损失缩放 $1/\beta$ 并应用于平滑分布，确保梯度保持有界且稳定。
GFlowNets 的 f-轨迹平衡：该框架通过利用 $L_f$ 损失强制执行轨迹平衡约束，应用于 GFlowNets。这将标准的轨迹平衡（对应于反向 KL）推广到了任意 f-散度。

主要贡献

广义损失族：推导出一族通用的平移不变代理损失（ $L_f$ ），其在线策略梯度与相应 f-散度的梯度相匹配，同时保留离线策略有效性。
理论等价性：证明了平移不变的对数概率损失与 f-散度之间的逆关系，使得能够系统地设计具有所需属性（例如覆盖模式）的新损失函数。
扩展到非归一化目标：通过DevGrad损失将 VarGrad 技术推广到任意 f-散度，解决了 RL 微调中配分函数不可处理的问题。
探索控制：证明了 f-散度的选择（在 $\alpha$ -散度中由 $\alpha$ 参数化）直接控制寻求模式与覆盖模式的行为。较低的 $\alpha$ 值（如前向 KL、Hellinger）促进覆盖模式，而较高的值（如 Pearson $\chi^2$ ）促进寻求模式。

实验结果

作者在四个不同的任务中评估了该损失族：

合成网格（GFlowNets）：在一个旨在测试模式覆盖的 2D 网格任务中，具有较低 $\alpha$ 值（前向 KL、Hellinger）的损失比标准轨迹平衡（反向 KL）更快、更可靠地发现了目标分布的所有四个模式。相反，Pearson $\chi^2$ （ $\alpha=2$ ）表现出强烈的寻求模式行为，被困在第一个发现的模式中。
分子发现（SynFlowNet）：应用于 SynFlowNet 框架以生成合成可及的分子。作者发现，调节 $\alpha$ 可以控制生成分子的多样性。退火调度（从较低的 $\alpha$ 开始用于探索，然后增加到 $\alpha=1$ 用于利用）产生了最佳权衡，与标准轨迹平衡相比，生成了具有更高奖励和更大多样性的分子。
扩散模型：在 MNIST 数字的条件采样（奇数与偶数）中，标准轨迹平衡未能均匀地在目标模式之间采样（过度采样特定数字如 0 和 6）。f-轨迹平衡损失，特别是那些具有较低 $\alpha$ 的损失，实现了在目标模式之间更均匀的采样。
异步 LLM 微调：这些损失被应用于微调 LLM（Qwen 和 OLMo 系列）以解决数学问题（GSM8k、MATH），并带有 50 步的异步延迟。f-散度损失表现出稳定的训练和一致的熵 - 奖励权衡，而标准的 PPO 实现在相同的异步条件下表现出不稳定性。

意义与主张

论文声称，平方 KL 损失的有效性并非孤立现象，而是更广泛损失族的一个具体实例。通过将“在线策略梯度匹配”和“离线策略有效性”的属性扩展到整个 f-散度族，作者提供了一个统一的框架来微调生成模型。

其主要意义在于能够显式控制探索 - 利用权衡，只需选择不同的代理损失函数（或调节 $\alpha$ 参数），而无需更改底层训练算法。这使得从业者可以在寻求模式（用于奖励最大化）和覆盖模式（用于多样性和探索）行为之间进行选择，同时保持现代大规模异步训练 LLM 和 GFlowNets 所需的稳定性和离线策略兼容性。这项工作表明，该框架对于药物发现和探索性智能体等应用特别有价值，在这些应用中，覆盖高奖励候选者的空间至关重要。