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

📄 arXiv: 2604.15311v1 📥 PDF

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

分类: cs.CV

发布日期: 2026-04-16

备注: Accepted by CVPR 2026. Project page: https://rockeycoss.github.io/leapalign/

💡 一句话要点

LeapAlign：通过构建两步轨迹对Flow Matching模型进行任意生成步骤的后训练对齐。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Flow Matching 模型对齐 后训练 梯度反向传播 图像生成

📋 核心要点

  1. 现有Flow Matching模型微调方法在长轨迹上反向传播梯度，导致内存消耗大和梯度爆炸，难以优化早期生成步骤。
  2. LeapAlign通过构建两步跳跃轨迹，缩短反向传播路径，降低计算成本，并使奖励梯度能够有效传播到早期生成步骤。
  3. 实验结果表明，LeapAlign在图像质量和图像-文本对齐方面，优于现有的基于GRPO和直接梯度的方法。

📝 摘要（中文）

本文关注Flow Matching模型与人类偏好对齐的问题。一种有前景的方法是通过Flow Matching的可微生成过程直接反向传播奖励梯度进行微调。然而，通过长轨迹反向传播会导致巨大的内存成本和梯度爆炸。因此，直接梯度方法难以更新早期生成步骤，而这些步骤对于确定最终图像的全局结构至关重要。为了解决这个问题，我们引入了LeapAlign，一种减少计算成本并能够将奖励的直接梯度传播到早期生成步骤的微调方法。具体来说，我们通过设计两个连续的跳跃来将长轨迹缩短为仅两步，每个跳跃都跳过多个ODE采样步骤并在单个步骤中预测未来的潜在变量。通过随机化跳跃的开始和结束时间步长，LeapAlign可以在任何生成步骤实现高效且稳定的模型更新。为了更好地利用这种缩短的轨迹，我们为那些与长生成路径更一致的轨迹分配更高的训练权重。为了进一步提高梯度稳定性，我们降低了梯度项中大幅度的权重，而不是像以前的工作那样完全删除它们。在微调Flux模型时，LeapAlign在各种指标上始终优于最先进的基于GRPO和直接梯度的方法，从而实现了卓越的图像质量和图像-文本对齐。

🔬 方法详解

问题定义：论文旨在解决Flow Matching模型与人类偏好对齐时，直接通过长生成轨迹反向传播奖励梯度进行微调所面临的计算成本高昂和梯度爆炸问题。现有方法难以有效更新早期生成步骤，影响最终生成图像的全局结构。

核心思路：LeapAlign的核心思路是通过将长生成轨迹缩短为两个连续的“跳跃”步骤，从而减少反向传播的计算量和梯度爆炸的风险。每个跳跃步骤预测多个ODE采样步骤后的潜在变量，实现对长轨迹的近似。

技术框架：LeapAlign的整体框架包括以下几个关键部分：1) Flow Matching模型：作为生成图像的基础模型。2) 两步跳跃轨迹构建：通过随机化跳跃的起始和结束时间步长，构建一系列两步轨迹。3) 权重分配：根据两步轨迹与完整生成轨迹的一致性，为不同的轨迹分配不同的训练权重。4) 梯度裁剪：降低大幅度梯度项的权重，提高训练稳定性。

关键创新：LeapAlign的关键创新在于提出了两步跳跃轨迹的概念，将长轨迹的反向传播问题转化为短轨迹的反向传播问题，显著降低了计算成本和梯度爆炸的风险。此外，根据轨迹一致性分配权重和梯度裁剪策略进一步提升了训练效果。

关键设计：LeapAlign的关键设计包括：1) 跳跃步长的随机化，保证模型在各个生成阶段都能得到有效更新。2) 基于轨迹一致性的权重分配策略，使模型更关注与完整生成路径相似的轨迹。3) 梯度裁剪策略，通过降低大幅度梯度项的权重来提高训练稳定性，避免直接截断梯度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LeapAlign在Flux模型上进行了微调实验，结果表明，在图像质量和图像-文本对齐等指标上，LeapAlign始终优于最先进的基于GRPO和直接梯度的方法。这表明LeapAlign能够更有效地将人类偏好融入到生成模型中，生成更符合人类期望的图像。

🎯 应用场景

LeapAlign可应用于图像生成、文本到图像合成等领域，提升生成模型与人类偏好的对齐程度，改善生成图像的质量和相关性。该方法降低了训练成本，使得在资源有限的条件下训练高质量的生成模型成为可能，具有广泛的应用前景。

📄 摘要（原文）

This paper focuses on the alignment of flow matching models with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients through the differentiable generation process of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion. Therefore, direct-gradient methods struggle to update early generation steps, which are crucial for determining the global structure of the final image. To address this issue, we introduce LeapAlign, a fine-tuning method that reduces computational cost and enables direct gradient propagation from reward to early generation steps. Specifically, we shorten the long trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing the start and end timesteps of the leaps, LeapAlign leads to efficient and stable model updates at any generation step. To better use such shortened trajectories, we assign higher training weights to those that are more consistent with the long generation path. To further enhance gradient stability, we reduce the weights of gradient terms with large magnitude, instead of completely removing them as done in previous works. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods across various metrics, achieving superior image quality and image-text alignment.