💻 computer science

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

LeapAlign 是一种针对流匹配模型的训练后微调方法，它通过构建随机化的两步跳跃轨迹，克服了长轨迹反向传播中的内存与梯度爆炸问题，从而实现了从奖励到早期生成步骤的高效且稳定的直接梯度更新，以获得更优的图像质量和对齐效果。

原作者： Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

发布于 2026-05-05

📖 1 分钟阅读☕ 轻松阅读

原作者： Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

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

想象一下，你正在教导一位极具天赋的艺术家（即 AI 模型）根据你的描述作画。这位艺术家从一块空白且充满噪点的画布开始，逐步细化，直到清晰的图像显现。这个过程被称为“流匹配”（flow matching）。

这篇论文要解决的问题是：我们如何教导这位艺术家画出我们确切想要的东西，尤其是整体构图（布局），而不至于让计算机内存耗尽？

以下是论文解决方案 LeapAlign 的简明解析：

1. 问题所在：“长路”过于沉重

为了教导这位艺术家，我们通常会展示一幅完成的画作，告知其好坏（即“奖励”），然后尝试将这种反馈一路传回至绘画过程的最初步骤，以修正错误。

问题所在： 如果一幅画需要 25 个步骤才能完成，那么将反馈一路传回所有 25 个步骤，就像试图从山顶向谷底大喊传递信息一样。这需要巨大的能量（计算机内存），而且信息在到达底部之前往往会失真或爆炸（梯度爆炸）。
后果： 由于很难将反馈传回起点，目前大多数方法仅修正绘画的最后几步。它们微调细节，但保持布局（狗在哪里，树在哪里）不变。如果布局错了，无论细节多么精美，整幅画都是错的。

2. 解决方案：“跳跃”捷径

作者梁展浩和杨涛发明了一种名为 LeapAlign 的方法。他们不沿着整条长路回溯以提供反馈，而是建立了一条捷径。

将绘画过程想象成从顶部（噪点）到底部（完成图像）的漫长楼梯。

旧方法： 走下每一级台阶以提供反馈。（太慢，太沉重）。
LeapAlign 方法： 你在楼梯上随机选取两个点——比如第 20 步和第 10 步。与其走完全程，你直接从第 20 步跳跃到第 10 步，然后再从第 10 步直接跳跃到完成的图像。

通过创建这种两步“跳跃轨迹”，计算机只需记住旅程中的两个步骤，而不是二十五个。这节省了海量内存，并防止信息爆炸。

3. 为何这意义重大

因为“跳跃”可以从楼梯上的任意点（随机选择）开始，系统现在可以将反馈发送到绘画过程的最开端。

结果： AI 学会了修正全局结构（布局），就像修正细节一样出色。它学会了将“红色的自行车”放在左边，将“蓝色的汽车”放在右边，而不仅仅是让自行车看起来闪闪发光。

4. 确保稳定的两个秘密配方

论文提到了两个巧妙的技巧，以确保这种跳跃平稳运行：

“音量旋钮”（梯度折扣）： 有时，当你通过跳跃发送反馈时，信号会变得太响（太强），从而破坏学习过程。以前的方法只是完全关闭信号。LeapAlign 则稍微调低音量旋钮。它保留信号（以便 AI 学习步骤之间的联系），但降低音量，以免震坏扬声器。
“信任分数”（轨迹相似度加权）： 有时，跳跃是一个看起来不像真实绘画路径的奇怪捷径。LeapAlign 会检查：“这个捷径看起来像正常的绘画路径吗？”如果是，它就给予反馈更多权重。如果捷径很奇怪，它就给予较少权重。这确保了 AI 能从最好的示例中学习。

5. 结果

作者在名为 Flux 的强大模型上测试了这种方法。

结果： LeapAlign 始终击败了其他顶级方法（如 GRPO 和 DRTune）。
证明： 它生成的图像不仅更精美，而且更好地遵循了文本指令。例如，如果你要求“一只猫坐在垫子上”，LeapAlign 确实将猫放在了垫子上，而其他方法通常会让猫漂浮在空中或放在垫子旁边。

总结： LeapAlign 是一种通过创建短小高效的反馈捷径来教导 AI 艺术家的聪明方法。这使得 AI 能够从最初就开始学习如何规划整幅画面，从而生成既美观又完全符合你要求的图像。

技术摘要：LeapAlign

问题陈述

本文解决了通过后训练将流匹配模型与人类偏好对齐的挑战。虽然直接梯度方法利用流匹配采样的可微特性，为策略梯度方法（如 GRPO）提供了一种有前景的替代方案，但在应用于长生成轨迹时面临重大障碍。具体而言，通过整个 ODE 采样步骤序列反向传播奖励梯度会产生过高的内存成本，并可能导致梯度爆炸。因此，现有的直接梯度方法通常将更新限制在靠近最终图像的单一步骤，或丢弃嵌套梯度以避免不稳定性。这一局限性阻碍了对早期生成步骤的优化，而早期步骤对于决定最终图像的全局布局和结构组成至关重要。

方法论：LeapAlign

作者提出了LeapAlign，这是一种微调方法，旨在实现高效且稳定的直接梯度传播至任意生成步骤（包括早期时间步），同时不产生过高的内存开销。核心方法论包含三个关键组件：

1. 两步跳跃轨迹构建

LeapAlign 不通过完整的生成轨迹进行反向传播，而是构建一个仅包含两个步骤的缩短跳跃轨迹，这两个步骤源自标准的完整运行轨迹。

过程：在每次迭代中，从噪声 ( $x_1$ ) 到清晰图像 ( $x_0$ ) 采样一条完整轨迹。随机选择两个时间步 $k$ 和 $j$ （其中 $k > j$ ）。
跳跃：该方法利用整流流模型的一步跳跃预测属性，以单步估计未来潜在变量：
- 跳跃 1：使用 $x_k$ 处的速度场从 $x_k$ 预测潜在变量 $x_j$ ： $\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$ 。
- 跳跃 2：使用 $x_j$ 处的速度场从 $x_j$ 预测最终潜在变量 $x_0$ ： $\hat{x}_{0|j} = x_j - j v_\theta(x_j)$ 。
潜在连接器：为了在使预测状态与从完整轨迹采样的实际状态对齐的同时保持可微性，引入了使用 stop_gradient 操作的“潜在连接器”（例如 $x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$ ）。
优势：这将反向传播路径减少为两步，无论原始轨迹长度如何，内存成本保持恒定，同时允许随机选择 $(k, j)$ 以更新生成过程的任何部分。

2. 梯度折扣

为了解决梯度爆炸问题同时保留有用的学习信号，LeapAlign 引入了梯度折扣机制。

问题：通过两步反向传播引入了一个“嵌套梯度”项（ $\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$ ），该术语捕捉了跨时间步的相互作用。先前的方法（如 DRTune）通过在模型输入处停止梯度完全消除了该项，从而丢弃了有价值的依赖信息。
解决方案：LeapAlign 通过因子 $\alpha \in [0, 1]$ 对嵌套梯度项进行缩放，而不是将其移除。这是通过修改第二次跳跃中速度预测的输入来实现的： $\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$ 。
效果：这降低了大梯度项的幅度以稳定训练，同时保留了完整的梯度结构，包括嵌套梯度。

3. 轨迹相似度加权

为了确保缩短的跳跃轨迹提供可靠的监督，该方法对与真实生成路径高度相似的跳跃分配更高的权重。

机制：损失函数根据连接点（ $x_j$ 和 $x_0$ ）处预测状态（ $\hat{x}$ ）与实际状态（ $x$ ）之间的平均绝对差值的倒数进行加权。
公式： $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$ ，其中 $d$ 代表距离， $\tau$ 是最小阈值，用于防止对几乎相同的样本对过度强调。
目标：这确保了偏离真实动力学显著的跳跃对梯度更新的贡献较小，从而提高了监督的稳定性和忠实度。

4. 微调目标

该方法采用铰链式损失（hinge-style loss）以最大化奖励并防止奖励作弊： $L_{raw} = \max(0, \lambda - r(x_0))$ 。至关重要的是，与那些从跳跃预测中估计奖励的方法不同，LeapAlign 在由完整轨迹生成的实际最终图像（ $x_0$ ）上评估奖励，确保奖励模型评估真实的视觉和语义质量。

主要贡献

LeapAlign 框架：一种新颖的后训练方法，通过构建两步跳跃轨迹，实现以恒定内存成本将直接梯度反向传播至早期生成步骤。
梯度稳定性技术：引入梯度折扣来缩放（而非移除）嵌套梯度，以及轨迹相似度加权来优先处理一致的跳跃，解决了先前直接梯度方法中的不稳定性和信号丢失问题。
综合评估：证明了 LeapAlign 在多项指标上优于最先进的策略梯度（基于 GRPO）和直接梯度方法，包括人类偏好评分（HPSv2.1, HPSv3, PickScore）和组合对齐基准（GenEval）。

结果

作者使用 LeapAlign 对Flux.1-Dev模型进行了微调，并将其与 DanceGRPO、MixGRPO、ReFL、DRaFT-LV 和 DRTune 等基线进行了评估。

通用偏好对齐：LeapAlign 在域内（HPSv2.1）和域外（HPSv3, PickScore, UnifiedReward）指标上均取得了最高分。值得注意的是，它在仅使用单个奖励模型（HPSv2.1）进行训练的情况下，表现优于使用三个奖励模型的 MixGRPO。
组合对齐：在 GenEval 基准测试中，LeapAlign 显示出显著提升，特别是在“两个物体”、“颜色”、“位置”和“属性绑定”任务中。这些类别的改进凸显了该方法在优化决定全局布局的早期生成步骤方面的成功。
定性分析：视觉比较表明，LeapAlign 大幅修改了全局图像结构以更好地与文本提示对齐，而局限于后期步骤更新的方法（如 ReFL）往往保留了预训练模型的布局。
消融实验：实验证实，两步跳跃设计在性能和内存之间提供了最佳权衡。 $\alpha=0.3$ 的梯度折扣产生了最佳结果，在整个范围 $[0, 1]$ 内随机选择时间步优于固定范围，这强调了早期步骤更新的重要性。

意义与主张

本文声称 LeapAlign 成功解决了流匹配后训练中内存效率与梯度覆盖范围之间的权衡。通过在不产生过高计算成本的情况下实现对早期生成步骤的更新，它使得优化全局图像结构成为可能，这对于复杂的组合任务至关重要。作者将 LeapAlign 定位为随机策略梯度方法和不稳定直接梯度方法之外的一种稳定且有效的替代方案，实现了更优越的图像 - 文本对齐和生成质量。该工作表明，保留并缩放嵌套梯度而非丢弃它们，是提高直接梯度对齐性能的关键因素。未来的工作 noted 将把这些原则扩展到视频生成领域。