TMPO：用于多样且高效扩散模型对齐的轨迹匹配策略优化

# TMPO：用于多样且高效的扩散对齐的轨迹匹配策略优化

来源: https://arxiv.org/html/2605.10983

作者: Jiaming Li\(^{1,2,*}\), Chenyu Zhu\(^{1,*}\), Nanxi Yi\(^1\), Youjun Bao\(^2\), Li Sun\(^2\), Quanying Lv\(^2\), Xiang Fang\(^3\), Daizong Liu\(^4\), Jianjun Li\(^1\), Kun He\(^1\), Bowen Zhou\(^5\), Zhiyuan Ma\(^{1,\dagger}\)

单位: \(^1\)华中科技大学, \(^2\)快手科技, \(^3\)南洋理工大学, \(^4\)武汉大学, \(^5\)清华大学

\(*\)同等贡献. \(\dagger\)通讯作者.

###### 摘要

强化学习（RL）在将扩散模型对齐到下游任务方面展现出了巨大的潜力，但大多数方法仍遭受严重的**奖励黑客（reward hacking）**问题，这通过诱导视觉模式崩溃和放大不可靠的奖励，降低了生成的多样性和质量。我们将根本原因归结为这些方法的**模式寻求（mode-seeking）**性质，即它们在不有效约束可接受轨迹上的概率分布的情况下最大化预期奖励，导致集中在少数高奖励路径上。相比之下，我们提出了**轨迹匹配策略优化（Trajectory Matching Policy Optimization, TMPO）**，用轨迹级别的奖励分布匹配替代标量奖励最大化。具体而言，TMPO 引入了一种**Softmax 轨迹平衡（Softmax-TB）**目标，将 \(K\) 条轨迹的策略概率匹配为由奖励诱导的 Boltzmann 分布。我们证明了该目标继承了正向 KL 散度的**模式覆盖（mode-covering）**特性，在优化奖励的同时保留对所有可接受轨迹的覆盖。为了进一步减少大规模流匹配模型上多轨迹训练的时间，TMPO 结合了**动态随机树采样（Dynamic Stochastic Tree Sampling）**，其中轨迹共享去噪前缀并在动态调度的步骤处分枝，从而在提高训练效果的同时减少冗余计算。在人类偏好、组合生成和文本渲染等多样化的对齐任务上的广泛实验结果表明，TMPO 比现有最先进方法的生成多样性提高了 9.1%，并在所有下游指标和效率指标上取得了具有竞争力的性能，实现了奖励与多样性之间的最佳权衡。

![Figure 1: Generative diversity comparison between TMPO (Ours) and Flow-GRPO](https://arxiv.org/html/2605.10983#S0.F1)

## 1 引言

![Figure 2: A toy experiment on a three-layer MLP diffusion model...](https://arxiv.org/html/2605.10983#S1.F2)

最近，RL 在大型语言模型（LLM）、生成扩散模型和流匹配模型的对齐方面取得了显著成功。然而，现有方法仍然遭受严重的**奖励黑客**问题，这意味着模型可能会忽略分布的多样性，而是通过过度强化一小部分被识别为高奖励的模式来走捷径。与 LLM 类似，这种困境可以称为**视觉模式崩溃**。虽然在 LLM 领域缓解奖励黑客问题已经得到了广泛研究\[49, 50\]，但在扩散或流模型中这一问题仍 largely 未被探索，这引发了两个关键问题：

*   **1) 多样性退化：** 对于语言模型，一个问题可能只有一个确定的答案，而对于生成模型（*即*，文生图模型），相同的提示可能对应许多合理的输出，这些输出在**组合**、**空间布局**和**场景语义**上有所不同。然而，奖励黑客驱动模型仅生成具有少量组合和风格的图像，从而降低生成多样性。（见图 1）
*   **2) 不可靠奖励的放大：** 对于文生图模型，基于模型的奖励（*例如*，ImageReward\[10\], HPS\[11\], PickScore\[31\]）和基于规则的奖励（*例如*，OCR 准确率和 GenEval\[44\]）本质上是人类偏好的不完美代理\[37, 12, 13, 14, 15\]。对这些代理奖励的奖励黑客放大了虚假奖励的图像属性，从而恶化生成质量。

这提出了生成扩散对齐中的一个核心问题：**如何在保持模型从预训练数据中学到的输出各向异性分布的同时，鼓励模型满足具有特定偏好的奖励？**

扩散模型对齐中的奖励黑客行为可归因于当前基于扩散的 RL 方法的训练目标。现有方法使用奖励最大化作为优化目标\[4, 5, 6, 7, 26, 30\]，我们证明这在本质上是**模式寻求**的。特别是，模式寻求模型倾向于探索高奖励轨迹，而不是覆盖所有合理的对应物，如图 2 中的 GRPO 示例所示。

理想的优化目标应该是**模式覆盖**的，它鼓励模型尽可能多地覆盖合理解决方案。这一分析进一步表明，缓解奖励黑客需要优化目标发生根本性转变。现有方法主要通过两条线路试图解决奖励黑客问题：

*   **1) 约束过度优化：** KL 正则化\[38, 39, 5, 6, 15\]用于约束策略与参考模型的偏差，但它并未改变奖励最大化目标，导致生成多样性的改善有限。
*   **2) 改进训练目标：** 基于生成流网络（GFlowNet）的方法\[23, 29, 24\]试图生成概率与其奖励成正比的输出，但现有方法通常在单个去噪步骤上建立局部约束，这需要估计一个难以处理的态流函数来进行策略更新，从而引入额外的训练误差。

总体而言，现有的缓解方法仍未能提供一个直接的、针对完整去噪轨迹的模式覆盖目标。为了弥补这一差距，我们提出了**轨迹匹配策略优化（TMPO）**。受 GFlowNets\[21\]的奖励分布匹配特性启发，TMPO 从根本上将训练目标从奖励最大化转变为轨迹级别的奖励分布匹配。我们证明 TMPO 是**模式覆盖**的，可以自然地鼓励覆盖合理的解决方案，从而从根源上缓解奖励黑客问题，如图 2 所示。

具体而言，我们引入了一种**Softmax 轨迹平衡（Softmax-TB）**目标，它在完整的去噪轨迹上执行奖励分布匹配。对于从相同起点生成的 \(K\) 条轨迹，Softmax-TB 将这些轨迹的生成概率匹配到由奖励计算的目标 Boltzmann 分布（\(P_\theta(\tau) \propto \exp(\beta R(\tau))\)）。同时，通过对每组内 \(K\) 条轨迹中的轨迹概率 \(P_\theta(\tau)\) 和指数奖励 \(\exp(\beta R(\tau))\) 进行归一化，我们消除了难以处理的配分函数，得到了一个更简单且可直接优化的目标。

为了减少在 Softmax-TB\[29\]下完全采样多条轨迹的训练时间，TMPO 进一步引入了**动态随机树采样\[26, 27, 30\]**，它允许轨迹通过树结构共享去噪步骤，减少冗余计算并降低训练时间。通过动态调整分支点，树采样器还促进了不同去噪阶段的有效探索，并提高了训练性能。

我们在 FLUX.1-dev 上评估了 TMPO 的三个对齐任务：组合图像生成、视觉文本渲染和人类偏好对齐。结果表明，TMPO 在所有任务中获得了最高的多样性指标，平均生成多样性比现有最先进方法提高了 9.1%，同时实现了最佳的 GenEval 准确率和 PickScore。此外，与最先进方法相比，TMPO 将训练时间最多减少了 27%，并在奖励-多样性和奖励-效率比较中实现了最有利的权衡，如图 6 所示。

我们的贡献如下：
1.  我们确定奖励最大化是扩散基于 RL 中奖励黑客的根本原因，并将优化目标重构为奖励分布匹配。
2.  我们引入了 Softmax 轨迹平衡，这是一种无配分函数的轨迹级别分布匹配目标，并证明其模式覆盖特性自然地保持了生成多样性。
3.  我们引入了动态随机树采样，这大幅减少了我们轨迹平衡目标的训练时间，并实现了在大规模流匹配模型上的可扩展训练。

## 2 相关工作

##### 缓解扩散模型对齐中的奖励黑客

基于扩散的 RL 对齐在基于模型或基于规则的奖励下极易受到奖励黑客的影响\[13, 14, 16, 18\]。现有缓解措施，包括 KL 正则化\[5, 6\]、成对偏好奖励\[20\]和感知多样性的正则化\[15\]，仍然优化标量奖励，因此继承了模式崩溃的风险\[17\]。TMPO 通过将标量奖励重构为需要匹配的偏好分布来解决这一局限性。

##### 分布匹配和基于 GFlowNet 的对齐

GFlowNets 训练随机策略，使其终端分布与奖励成正比\[21\]。DGFS\[23\]、DAG\[29\]和 \(\nabla\)-GFlowNet\[24\]使用局部详细平衡目标将这一思想适应于扩散对齐，其中 \(\nabla\)-GFlowNet 进一步结合了奖励梯度。然而，它们的局部或部分轨迹约束需要难以处理的态流估计，并且可能在长去噪视界上累积误差。TMPO 则直接将完整轨迹概率匹配到由奖励诱导的 Boltzmann 分布，避免了每次转换的流估计和显式的配分计算。

##### 扩散模型对齐中的高效轨迹采样

大规模流模型的 RL 后训练受限于冗余的回滚（rollouts），而完整轨迹目标通过要求多个轨迹上的策略概率放大了这一成本\[29\]。MixGRPO 将随机优化减少到选定的 SDE 窗口\[26\]，而 Dynamic-TreeRPO 和 TreeGRPO 通过共享去噪前缀摊销早期计算\[27, 30\]。TMPO 采用前缀共享，并进一步引入动态分支计划，以在训练阶段保持有效的探索。

## 3 预备知识

### 3.1 轨迹平衡与奖励匹配

生成流网络（GFlowNets）训练随机策略以按比例采样与非负奖励成比例的终端对象\[21\]。对于以 \(x=s_n\) 结束的轨迹 \(\tau=(s_0 \to \dots \to s_n)\)，其奖励为 \(R(x)\)，轨迹平衡（TB）施加如下条件\[22\]：

$$
Z_\theta \prod_{t=1}^n P_F(s_t \mid s_{t-1}; \theta) = R(x) \prod_{t=1}^n P_B(s_{t-1} \mid s_t; \theta), \quad (1)
$$

其中 \(P_F\)、\(P_B\) 和 \(Z_\theta\) 分别是正向策略、反向策略和总流。这种奖励比例视图保留了多个高奖励模式。扩散 GFlowNet 方法和 Flow-GRPO 同样以用奖励分布匹配替代标量奖励最大化为动机\[23, 24, 25\]。

### 3.2 随机扩散策略的 ODE 到 SDE 转换

现代文生图生成器通常建立在扩散或流匹配模型之上\[32, 1, 34, 35, 9, 2, 3\]。在整流流中，干净样本 \(x_0\) 和噪声 \(x_1 \sim \mathcal{N}(0, \mathbf{I})\) 通过以下公式连接：

$$
x_t = (1-t)x_0 + tx_1, \quad dx_t = v_\theta(x_t, t, c) \, dt, \quad (2)
$$

确定性 ODE 求解器\[33, 28\]缺乏策略梯度 RL 所需的转移概率。遵循先前工作\[8, 6, 7\]，我们将概率流 ODE 转换为等价的 SDE，该 SDE 允许进行信用分配的易处理似然度，同时保留边际分布：

$$
dx_t = \left[ v_\theta(x_t, t, c) + \frac{1}{2}\sigma^2(t) \nabla_x \log p_\theta(x_t \mid c, t) \right] dt + \sigma(t) dW_t. \quad (3)
$$

这里 \(\sigma(t)\) 控制噪声尺度，当 \(\sigma(t) \equiv 0\) 时恢复确定性 ODE。

## 4 轨迹匹配策略优化

在本节中，我们介绍轨迹匹配策略优化...