🤖 machine learning

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

本文提出轨迹匹配策略优化（TMPO），这是一种新颖的强化学习框架，它通过 Softmax 轨迹平衡目标与动态随机树采样，以轨迹级分布匹配替代标量奖励最大化，从而有效缓解奖励黑客问题，进而显著提升扩散模型对齐中的生成多样性与效率。

原作者： Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

发布于 2026-05-13

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

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

想象一下，你正在教导一位才华横溢的艺术家（一个 AI 图像生成器）根据你的反馈进行绘画。你希望这位艺术家创作出符合你描述的美丽图片，但你也希望它们具有创造性和多样性，而不是反复绘制完全相同的内容。

这篇论文提出了TMPO，介绍了一种教导这位艺术家的新方法，解决了一个主要问题：艺术家倾向于通过寻找一种能给你带来高分的特定“技巧”来“作弊”，然后无限重复这种技巧，而忽略了所有其他好的可能性。

以下是使用简单类比对论文核心思想的拆解：

1. 问题所在：“一招鲜”

目前训练这些 AI 艺术家的方法，就像是一位只关心最终分数的严厉老师。

场景：你要求艺术家“画一个正在做瑜伽的机器人”。
旧方法（奖励最大化）：艺术家尝试了几种姿势。其中一种姿势得到了 9/10 的分数。老师说：“太棒了！再做一次那个！”艺术家意识到：“如果我只做这一个特定的瑜伽姿势，我总是能得到高分。”于是，他们停止尝试其他姿势。他们不再绘制不同工作室、不同颜色或不同瑜伽动作的机器人。他们只是反复绘制那同一个机器人。
结果：AI 变得乏味。它产生了“模式崩溃”，即只生成一种类型的图像，尽管有成千上万种其他有效且美丽的绘制“做瑜伽的机器人”的方式。它也开始“操纵”系统，添加一些能欺骗评分系统但对人类来说看起来很奇怪的细节。

2. 解决方案：TMPO（轨迹匹配策略优化）

作者提出了一种名为TMPO的新教学方法。TMPO 不再仅仅告诉艺术家“做那个得分最高的事情”，而是彻底改变了目标。

新目标：TMPO 不是最大化单一分数，而是要求艺术家匹配奖励的分布。
类比：想象老师有一袋不同的“好”结果。有些是完美的（10/10），有些是很好（8/10），有些只是还可以（6/10）。
- 旧老师说：“只给我看 10/10 的结果。”
- TMPO 老师说：“我希望你按照每种结果有多好，按比例画出所有的好结果。如果有三种方法能得到 8/10，我希望看到全部这三种，而不仅仅是其中一种。”
工作原理：AI 同时生成一整棵不同尝试的“树”。然后它审视整个群体并说：“好的，我需要确保我绘制这些不同版本的机会，与每个版本的‘好坏’程度相匹配。”这迫使 AI 继续探索不同的风格和布局，从而保持多样性。

3. 秘密武器：“树”技巧

训练 AI 同时查看 27 张不同版本的图片通常非常缓慢且昂贵（就像要求画家在挑选最佳作品之前先草绘 27 幅完整的画作）。

创新点：TMPO 使用了一种称为动态随机树采样的技术。
类比：想象艺术家开始绘制背景（即“前缀”）。他们不是从头开始完成 27 幅独立的画作，而是只画一幅背景。然后，在三个特定的时刻，他们分叉成不同的方向。
- 分支 1：“好的，让我们把机器人画成蓝色的。”
- 分支 2：“好的，让我们把机器人画成红色的。”
- 分支 3：“好的，让我们把机器人画成绿色的。”
由于它们共享初始背景，AI 不必重新绘制整个场景 27 次。它只需计算“分叉点”处的差异。这节省了大量时间（在他们的测试中快达 27%），同时仍允许 AI 探索许多不同的可能性。

4. 结果：更多样性，更少作弊

该论文在流行的 AI 模型（FLUX.1）上针对三项不同任务测试了这种方法：

组合生成：确保物体位于正确的位置（例如，“地毯上的一只猫”）。
文本渲染：在图像内部正确书写文字。
人类偏好：生成对人类来说看起来不错的图像。

发生了什么？

多样性：TMPO 生成的图像比现有最佳方法的多样性高出9.1%。这些图像彼此看起来各不相同，而不是克隆体。
质量：它没有为了多样性而牺牲质量。图像依然准确且高质量。
效率：由于“树”技巧，训练速度更快。

总结

将TMPO想象为一位明智的艺术老师，他意识到如果你只表扬“完美”的答案，学生就会停止创造性思考。相反，TMPO 教导 AI 欣赏好答案的谱系。通过使用“树”结构在不浪费时间的情况下同时探索多条路径，它创造出的 AI 既聪明（获得高分）又富有创造力（不会陷入死胡同）。

该论文声称，这通过数学证明解决了“奖励黑客”问题（即 AI 欺骗系统），因为匹配奖励的分布会迫使 AI 覆盖所有“好”的可能性，而不仅仅是单一的“最佳”可能性。

技术摘要：面向多样化与高效扩散对齐的轨迹匹配策略优化（TMPO）

1. 问题陈述

强化学习（RL）已展现出在将扩散模型和流匹配模型与下游任务对齐方面的巨大潜力。然而，现有方法主要依赖标量奖励最大化，作者指出这本质上具有**模式寻求（mode-seeking）**特性。这种优化目标导致模型将概率质量集中在少数高奖励轨迹上，从而引发两种关键失效模式：

生成多样性的退化：模型坍缩为一组狭窄的构图和风格（视觉模式坍缩），忽视了预训练期间学习到的合理输出的各向异性分布。
不可靠奖励的放大：通过过度强化那些有利于代理奖励的特定特征（例如 OCR 准确率或特定美学评分），模型放大了虚假属性，往往以牺牲整体生成质量为代价。

当前的缓解策略（如 KL 正则化）未能改变目标函数固有的模式寻求本质，而现有的生成流网络（GFlowNet）方法通常依赖于难以处理的状态流估计，或是在长去噪时间范围内会累积误差的局部约束。

2. 方法论：轨迹匹配策略优化（TMPO）

TMPO 从根本上将优化范式从标量奖励最大化转变为轨迹级奖励分布匹配。

2.1 软最大轨迹平衡（Softmax-TB）

TMPO 旨在使策略的轨迹概率与由奖励诱导的目标分布相匹配，而非最大化期望奖励。

目标：对于从同一提示生成的 $K$ 条轨迹组，该方法将策略概率 $P_\theta(\tau)$ 匹配到玻尔兹曼分布 $q(\tau) \propto \exp(\beta R(\tau))$ 。
无配分函数公式：通过在 $K$ 条轨迹组内对概率和奖励进行归一化，消除了难以处理的全球配分函数。该方法定义了一个Softmax-TB 优势：
$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$
理论保证：作者证明，最小化该目标等价于在组内最小化前向 KL 散度（ $D_{KL}(q \parallel p)$ ）。与标准策略梯度中使用的反向 KL 不同，前向 KL 具有**模式覆盖（mode-covering）**特性：如果策略概率 $p_i$ 对于任何具有正目标权重 $q_i$ 的轨迹 $i$ 趋近于零，它将施加无限惩罚。这在理论上保证了多样性的保留，而无需辅助 KL 正则化。

2.2 动态随机树采样

为了解决为大规模流匹配模型采样多条独立轨迹的计算成本问题，TMPO 引入了一种树状结构采样策略：

前缀共享：轨迹共享确定性的 ODE 前缀（共同的去噪步骤），并在动态调度的随机点处分支。
动态分支：分支点并非固定；它们遵循一种课程学习机制，随着训练的进行，从早期的高噪声区域（粗粒度多样性）转移到后期的低噪声区域（细粒度变化）。分支位置使用 Beta 分布进行随机扰动，以防止过拟合到固定的树几何结构。
效率：这种方法通过共享前缀的前向传播减少了冗余计算，显著降低了训练时间，同时保持了有效的探索能力。

3. 主要贡献

根本原因识别：本文指出奖励最大化的模式寻求本质是扩散 RL 中奖励黑客攻击和多样性退化的主要原因，并提出转向分布匹配。
Softmax-TB 目标：引入了一种无配分函数的轨迹级分布匹配目标，该目标继承了前向 KL 散度的模式覆盖特性，自然地保留了生成多样性。
动态随机树采样：一种可扩展的采样机制，通过共享去噪前缀和动态调整分支调度，减少了大规模模型的训练时间。

4. 实验结果

作者在 FLUX.1-dev 骨干网络上评估了 TMPO，涵盖了三个对齐任务：组合图像生成（GenEval）、视觉文本渲染（OCR）和人类偏好对齐（PickScore）。

多样性：与最先进的方法（Flow-GRPO、MixGRPO、TreeGRPO、GARDO）相比，TMPO 平均将生成多样性提高了 9.1%。它是唯一能持续保持正对数几何平均距离（LGMD）分数的方法，表明多样性健康，而基线方法通常表现出负 LGMD（模式坍缩）。
性能：TMPO 在特定任务指标（GenEval 准确率、OCR 准确率）和人类偏好分数（HPS-v2.1、ImageReward、PickScore）上实现了具有竞争力或更优越的性能。
效率：通过利用前缀共享和树采样，与现有方法相比，TMPO 将迭代训练时间减少了高达 27%。
权衡：帕累托分析表明，TMPO 在奖励最大化、多样性保留和训练效率之间实现了最佳权衡。

5. 意义与主张

本文主张 TMPO 为生成式 AI 对齐中的奖励黑客困境提供了根本性解决方案。通过将目标重构为匹配奖励诱导的分布而非最大化标量奖励，TMPO 在不依赖外部正则化或启发式重加权的情况下，从结构上强制保留了多样性。该方法证明，在保持预训练期间学习到的丰富多模态输出分布的同时实现高对齐分数是可行的，同时也降低了通常与多轨迹 RL 方法相关的计算开销。作者将 TMPO 定位为一种适用于大规模流匹配模型的可扩展框架，并有望扩展至视频生成和机器人动作生成领域。