$f$-轨迹平衡：一种用于离策略和在线策略数据调优GFlowNet、生成模型和LLM的损失函数族

# 用于调优 GFlowNet、生成模型和 LLM 的损失函数家族：支持离策略和在策略数据

来源：https://arxiv.org/html/2605.15417

###### 摘要

在 GFlowNet 和变分推断中，已证明目标对数概率与模型对数概率之间的均方误差是一种有效、低方差的代理损失函数，可用于训练生成模型。该损失函数的一个特性是：当在**在策略**数据上评估时，其梯度对应于 KL 散度的梯度；而**离策略**时，它仍然是一个有效的损失函数，且具有相同的全局最小值。本文中，我们证明这一构造可推广至整个 **f**-散度家族，从而导出一族损失函数，其在策略梯度对应相应 **f**-散度的梯度，而离策略时仍保持相同的全局最小值。具体来说，我们证明在策略梯度在目标与模型对数概率上的平移不变损失函数与 **f**-散度之间存在一一对应关系。这一等价性使我们能够为调优广泛的生成模型设计新的代理损失函数，这些函数继承了对应 **f**-散度的性质（例如更具模式覆盖性），同时适用于离策略数据。我们将所提出的损失函数应用于一系列任务，包括经典合成示例、用于分子发现的 SynFlowNet，以及异步大型语言模型 (LLM) 调优，结果表明我们的模型在在策略和离策略情况下均能在广泛的生成模型类别中保持其预测的性质。

机器学习，ICML

## 1 引言

参照图注

图 1：**f**-散度包含 α\\-散度族，其中正向 KL (α=0) 和反向 KL / 标准轨迹平衡 (α=1) 是特例。较低的 α 产生更具模式覆盖性的损失，较高的 α 则产生更具模式寻找性的损失。

生成模型（无论是生成文本的大型语言模型 (LLM) 还是生成分子图的生成流网络 (GFlowNet)）的对齐与微调，都依赖于最小化模型分布之间的差异。在强化学习 (RL) 微调中，这一目标通常被表述为在 KL 散度惩罚下最大化期望奖励，或等价地，最小化策略与目标分布之间的散度。传统上，优化这些目标需要高方差的梯度估计器（如 REINFORCE (Williams, 1992)）或复杂的策略梯度方法（例如 Schulman 等人, 2017）。然而，近期 GFlowNet 文献和 Kimi (Team 等人, 2025a) 中的大规模 LLM 调优采用了一种更简单的方法：最小化模型对数概率与目标对数概率之间的均方误差 (MSE)。这与 Schulman (2020) 提出的 KL 散度 K2 估计器类似，我们称之为 KL_sq。Tang 和 Munos (2025) 指出，这种“平方 KL”损失具有一个令人惊讶的性质：虽然它是 KL 散度的有偏估计器，但其在在策略数据上的梯度却是真实 KL 梯度的无偏估计。此外，由于它基于 MSE，当应用于离策略数据时，它仍然是具有相同全局最小值的有效损失函数 (Bartoldson 等人, 2025; Tang 等人, 2025)。虽然 KL_sq 估计器提供了稳定性和离策略兼容性，但根据其构造，它继承了反向 KL 散度的性质。特别是，在此目标下训练的模型通常会表现出“模式寻找”行为，即模型坍缩到少数高概率模式上，而不是覆盖目标分布的完整多样性。在许多生成任务中，例如药物发现或探索性智能体，通常更偏好“模式覆盖”行为，因为生成模型用于探索高奖励候选空间（例如所有可能的药物靶点），而不是寻找一个（或少数几个）特别好的候选。本文中，我们证明 KL_sq 损失的有效性并非 KL 散度独有的现象，我们为整个 **f**-散度家族推导出类似的损失函数。我们建立了一个理论等价性：对数概率上的**任何**平移不变损失函数都对应一个特定的 **f**-散度。这使我们能够推导出一族新的代理损失函数，这些函数继承了其母散度的特定性质（例如正向 KL 或 Hellinger 距离的模式覆盖行为），同时保留了 KL_sq 估计器的优化优势：在策略上低方差梯度，离策略上仍有效。我们在各种任务上展示了这族损失函数的行为和效用。在超网格任务 (Bengio 等人, 2021)（一个专门设计用于测试模式覆盖的合成任务）上，我们展示了 Hellinger 和正向 KL 损失发现了比轨迹平衡损失 (Malkin 等人, 2022b) 所隐含的反向 KL 损失更多的模式。然后我们将这些损失应用于分子生成（通过修改 SynFlowNet (Cretu 等人, 2025) 中的损失）、扩散模型中的类条件采样以及语言模型上的异步强化学习。在这两种情况下，我们发现能够利用我们代理损失函数家族中更具模式覆盖性的实例来学习更高熵的策略。总之，我们的贡献如下：

- • 我们推导出一族通用的平移不变代理损失函数 L_f，其期望的自动微分梯度在在策略数据上与对应 **f**-散度的梯度相匹配。
- • 我们证明逆关系：对数概率上的任何凸的、平移不变的损失函数都对应于最小化一个特定的 **f**-散度。
- • 我们将“Vargrad”(Richter 等人, 2020) 和批量归一化技术推广到任意 **f**-散度，解决了 RL 微调中难处理的配分函数问题。
- • 我们将此框架应用于 GFlowNet，表明我们的损失函数推广了标准的轨迹平衡目标 (Malkin 等人, 2022a)，并在合成网格和 SynFlowNet 分子发现任务上通过实验证明，仅通过改变代理损失函数，我们就可以控制探索-利用权衡（模式寻找 vs. 模式覆盖）。

## 2 相关工作

##### GFlowNet 和 GFlowNet 损失

Bengio 等人 (2021) 引入 GFlowNet 并通过流匹配损失进行训练，该损失直接应用于神经网络，旨在学习给定 GFlowNet 的最优马尔可夫流，并在动作级别上应用。详细平衡目标由 Bengio 等人 (2023) 提出，同样应用于动作级别，但参数化了前向和后向策略分布。Malkin 等人 (2022a) 提出了轨迹平衡，它在轨迹级别上强制前向策略、后向策略和归一化常数之间的等式约束，并用均方误差进行惩罚。Malkin 等人 (2023) 表明，在此损失下进行在策略训练等价于一种分层变分推断 (Ranganath 等人, 2016) 形式，且使用 KL 散度。我们的结果将其扩展到所有 **f**-散度。Vargrad 目标 (Richter 等人, 2020) 常用于 GFlowNet 训练，它使用与 Malkin 等人 (2022a) 相同的平方误差形式，但以批量方式估计归一化常数。这个 Vargrad 损失来源于 Nüsken 和 Richter (2021) 引入的对数方差散度，他们还提出了一个基于方差的估计器，建立在 Pearson χ² 散度之上，与我们的工作最为相似。然而，正如我们在本文中所示，这一方法不会共享相同的在策略梯度匹配性质，因为正确扩展这一性质的方法是使用广义偏差 (Rockafellar 等人, 2006) 而非方差。最后，与我们的工作最相关的是 Silva 等人 (2024)，他们讨论了用于训练 GFlowNet 的替代散度度量，但他们的工作侧重于在策略设置，而 GFlowNet 的主要探索优势来自于离策略训练。

##### **f**-散度和生成模型

**f**-散度 (Ali 和 Silvey, 1966; Morimoto, 1963) 在概率模型调优中有着悠久的历史 (Minka 等人, 2005)。在生成模型的背景下，它们被广泛应用于变分推断 (Wang 等人, 2018) 以及 GAN (Nowozin 等人, 2016)，近期也有一些工作将 **f**-散度应用于调优扩散模型 (Tang, 2024; Novello 等人, 2025)。**f**-散度已被应用于 LLM 中，或直接作为目标 (Go 等人, 2023; Han 等人, 2024)，或作为正则化项 (Huang 等人)。然而据我们所知，这些工作均未展示离策略的有效性，而这一点很重要，因为许多 LLM 调优是异步进行的 (Intellect, 2025) 或完全离策略的。**f**-散度也被广泛应用于模仿学习 (Ke 等人, 2020; Ghasemipour 等人, 2020)，这表明我们的工作可以应用于 LLM 的异步蒸馏 (Lu 和 Lab, 2025)。

## 3 动机

##### RL 调优 LLM 和 KL 梯度

我们用 π_θ 表示 LLM，它被视为给定提示 x 下关于 token 序列 y 的策略。然后使用强化学习通过最大化以下奖励来调优该模型，其中 P(x) 是提示/状态 x 的分布：

J(θ) = E_{x∼P(x)}[E_{y∼π_θ(y|x)}[r(x,y)]] - β KL(π_θ(·|x), π_ref(·|x)).

此目标通常通过 PPO (Schulman 等人, 2017) 或 REINFORCE (Williams, 1992) 进行优化。然而，正如 Tang 和 Munos (2025) 所指出的，许多当时的开源实现无法通过 KL 散度的蒙特卡洛估计器的自动微分获得 KL 项的梯度。这是因为 KL 散度在采样和目标两方面都依赖于 θ，而自动微分只考虑了后者。正确的梯度是：

∇_θ KL = E_{π_θ(y|x)}[ log(π_θ(y|x)/π(y|x)) ∇_θ log_θ π_θ(y|x) ]，

而朴素的自动微分会给出梯度：

E_{π_θ}[ ∇_θ log(π_θ(y|x)/π(y|x)) ] = E_{π_θ}[ ∇_θ log π_θ(y|x) ]，

该表达式的期望为零。Tang 和 Munos (2025) 建议使用标准的 REINFORCE 梯度估计器，但也讨论了 Schulman (2020) 的平方 KL 估计器。其记作 KL_sq，定义为：

KL_sq(π_θ ∥ π_ref)(x) = 1/2 E_{y∼π_θ}[ (log(π_θ(y|x)/π_ref(y|x)))^2 ]，

这里我们将其写为条件 x 的函数。它本身就是 KL 散度的有偏估计器，但具有以下性质：朴素的自动微分梯度在期望上是真实 KL 的梯度。具体地，

E_{y∼π_θ}[ ∇_θ 1/2 (log(π_θ(y|x)/π_ref(y|x)))^2 ] = ∇_θ KL(π_θ ∥ π_ref)(x).

##### 从 KL 散度出发用于离策略和在策略 RL 的单一损失

我们现在描述 KL_sq 估计器如何导出一个可用于优化方程 (1) 中 J(θ) 目标的损失函数，无论使用离策略还是在策略数据，正如 Tang 等人 (2025) 和 Bartoldson 等人 (2025) 所做的那样。首先，J(θ) 可以写为：

J(θ) = β E_{x∼P(x)}[ E_{y∼π_θ(y|x)}[ log e^{β^{-1} r(x,y)} ] ]，