← 返回 2026-04-17

LeapAlign：通过构建两步跳跃轨迹实现任意生成步的流匹配模型后训练 LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng 📅 2026-04-16 👍 14 2026-07-13 08:36
RLHF 人类偏好对齐 图像生成 流匹配 直接梯度法
arXiv HuggingFace PDF

两步跳跃轨迹+梯度折扣+相似度加权，让 FLUX 等流匹配模型能稳定训练任意早期步。

前置知识

流匹配模型 (Flow Matching)

流匹配是一类生成模型，通过学习连续速度场把高斯噪声平滑变换到真实图像，前向过程写作 x_t=(1-t)x_0+tx_1，rectified flow 速度场 v=x_1-x_0。

LeapAlign 微调的正是 FLUX.1-dev 和 SD3.5-M 这类 rectified flow matching 模型，要读懂其单步跳跃预测公式就必须先理解流匹配的线性插值形式和速度场含义。

人类偏好对齐与奖励模型

把人类对图像质量的偏好转成可微分的奖励 $r(x_0,c)$，再用 RLHF 思路微调文生图模型；常用奖励有 CLIP 类（HPSv2.1、PickScore）和 VLM 类（HPSv3、UnifiedReward、ImageReward）等。

LeapAlign 是一种偏好对齐算法，奖励模型的选择和评估方式（对 $x_0$ 评估而非 $\hat{x}_0$）是其大幅超过 ReFL/DRaFT-LV/DRTune 的关键设计之一。

直接梯度法 vs 策略梯度法

流匹配采样的每一步都可微，可以把奖励梯度沿 $x_t\to\theta$ 反向传播直接更新参数，这是直接梯度法（ReFL、DRaFT-LV、DRTune）；策略梯度法（GRPO 类）则把确定性 ODE 转 SDE 再用 REINFORCE/PPO 更新，引入显著方差。

全文都在对比这两类方法，LeapAlign 属于直接梯度法一族，其核心问题就是反向传播激活链过长导致的内存爆炸和梯度爆炸。

ODE 单步跳跃预测

对 rectified flow，从 $x_k$ 一步估计 $\hat{x}_{j|k}=x_k-(k-j)v_\theta(x_k,k)$，因为 $v=dx_t/dt=x_1-x_0$，线性外推即可跳到任意时间步 $j$。这是 LeapAlign 构造跳跃轨迹的数学基础。

这是 LeapAlign 能把上百步长轨迹压缩为两步跳跃轨迹 (leap trajectory) 的关键公式，没有它就无法在保持大跨度更新的同时控制反向传播开销。

研究动机

直接梯度法（如 ReFL、DRaFT-LV）的核心卖点是利用流匹配采样的可微性把奖励梯度沿长链反传到模型参数 $\theta$，但实际使用中遇到两大硬伤：其一是把奖励梯度穿过几十步 ODE 积分时，每步都要保存激活和中间速度，显存需求随步数线性增长（FLUX.1-dev 跑 25 步就直接 OOM）；其二是梯度经长链累乘后范数爆炸，参数无法稳定更新。因此现有方法只能更新靠近最后图像的少量步骤，其中 DRaFT-LV 只更新最后一步，ReFL 只能在最后 11 步里随机挑一步。而决定图像全局布局和构图恰恰是早期步（生成中间 0 到 0.5 时间步），更新不到早期步是直接梯度法的致命短板。DRTune 试图通过 stop gradient 让早期步也能训练，但又丢弃了表示跨步依赖的嵌套梯度，监督信号被严重截断，仍达不到最优。

本文的目标是本文目标明确：在不爆显存、不梯度爆炸的前提下，让奖励梯度能够稳定传到任意早期生成步，从而同时获得高质量（HPSv2.1、HPSv3、PickScore、ImageReward）和高文图对齐（GenEval）的微调模型。作者希望多个指标全面超过当前最强策略梯度基线 MixGRPO 和最强直接梯度基线 DRTune，并在反映早期布局能力的 GenEval 上拉开明显差距，从而实证证明早期步可微调对构图对齐真的有效，而非简单在某个指标上做 trick。

与已有工作不同的是，作者采取了一个非常独特但合理的切入角度——把上百步的 ODE 轨迹压缩成一个两步跳跃轨迹 (leap trajectory)，并保留嵌套梯度但对其进行缩放式折扣 (gradient discounting)，既不丢弃有用信号也不让范数失控。他们还引入轨迹相似度加权，让更接近真实路径的跳跃对获得更大权重，并在 $x_0$（真实最终图）而非 $\hat{x}_{0|j}$ 上评估奖励，得到更可靠的监督信号。这套组合拳与 ReFL/DRTune 都不同，是首个能让任意早期步都能在 FLUX 上稳定 fine-tune 的直接梯度方案。

核心方法

LeapAlign 的直觉非常优雅：与其沿上百步 ODE 反向传播，不如先用普通采样器跑一条完整轨迹，从中随机挑两个时间步 $k>j$，用流匹配的一次跳跃公式分别从 $x_k$ 跳到 $x_j$ 再跳到 $x_0$，把这条两步轨迹当成反向传播路径。奖励计算在真实最终图 $x_0$ 上，因此监督信号忠实；为了让反向传播不爆梯度，对跨越两步的嵌套梯度乘以折扣系数 $\alpha\in[0,1]$（默认 0.3）来缩放其幅度而不是直接截断；再用一个轨迹相似度权重 $w_{sim}=1/[\max(d_j,\tau)+\max(d_0,\tau)]$ 来强调更贴近真实路径的跳跃对。整体上，LeapAlign 把直接梯度法改造成小成本、稳定、可训练早期步的版本。

与已有方法的本质区别在三点：(1) ReFL/DRaFT-LV 只更新最后一步或最后几步，DRTune 通过 stop gradient 让早期步可训练但丢弃嵌套梯度；LeapAlign 显式构造两步跳跃轨迹，使反向传播只需穿过两个模型调用，且嵌套梯度以折扣方式被保留下来；(2) 评估奖励时 LeapAlign 用真实采样得到的 $x_0$，而 ReFL/DRTune 用一步跳跃近似 $\hat{x}_0$，后者包含显著噪声，使奖励模型评估失真；(3) LeapAlign 引入轨迹相似度加权动态调整每条跳跃轨迹的权重，进一步去除监督噪声。三者结合，使早期步终于能在 FLUX 上稳定拿到奖励信号。

方法步骤详情

流程每次迭代四步：(1) 用 ODE 采样器以 25 步、CFG=3.5、$720\times 720$ 在 FLUX.1-dev 上生成一张图，得轨迹 $\{x_T,\dots,x_0\}$；(2) 随机采两个时间步 $k>j$，从 $x_k$ 用 $\hat{x}_{j|k}=x_k-(k-j)v_\theta(x_k,k)$ 跳到 $x_j$，用 latent connector 对齐真实再跳到 $x_0$，奖励在真实 $x_0$ 上评估；(3) 在 $x_j$ 引入梯度折扣 $\hat{x}_{0|j}=x_j-j\,v_\theta(\alpha x_j+(1-\alpha)\text{sg}(x_j))$，让嵌套梯度被 $\alpha=0.3$ 缩放而非移除；(4) hinge loss $\max(0,\lambda-r(x_0))$（$\lambda=0.55$）乘相似度权重 $w_{sim}$，AdamW（lr=$1\times 10^{-5}$）批量 64 更新 FLUX DiT。$k,j$ 覆盖 $[0,1]$ 全程，16 卡 300 iter 收敛。

技术新颖性

技术上三个新颖点环环相扣：(1) 跳跃轨迹的概念，首次显式把用一两个大步替代长链反向传播的思想形式化，并证明其对 rectified flow 成立；(2) 梯度折扣机制，揭示了 DRTune 的 stop-gradient 等价于 $\alpha=0$ 的特例，而 $\alpha=0.3$ 是一个更好的折中（附录 Fig. 6 显示 $\alpha=1$ 时梯度范数过大反而损害性能）；(3) 轨迹相似度加权，把跳跃是否靠谱的度量显式注入损失，缓解了 $\hat{x}\neq x$ 时的奖励噪声。三者缺一就会退化成 ReFL/DRTune/单纯 hinge loss，因此它们组成一个连贯设计：跳跃轨迹给可行性，梯度折扣给稳定性，相似度加权给精度。

Overview of LeapAlign
Figure 2: Overview of LeapAlign
Further analysis of design components in LeapAlign
Figure 4: Further analysis of design components in LeapAlign
Analysis of the nested gradient
Figure 6: Analysis of the nested gradient

实验结果

在 FLUX.1-dev 上以 HPSv2.1 为奖励，LeapAlign 在 8 项指标上刷新基线：HPSv2.1 0.4092（MixGRPO 0.3692 / DRTune 0.3882）、PickScore 23.7137、HPSv3 15.7678、UR-Align 3.4984、UR-IQ 3.7244，IR 1.5104 仅次于联合训多奖励的 MixGRPO 1.6155。GenEval 整体 0.7420，明显高于 MixGRPO 0.7232 和 DRTune 0.7101，其中 Colors 80.59、Position 30.25、AttrB 66.00 提升最大，证明早期步可训练对构图对齐至关重要。奖励曲线（图 1a）显示 LeapAlign 在微调早期就拉开 DRTune，定性（图 3、图 7）证实它能改对黄自行车+红摩托等组合。在 SD3.5-M 上 HPSv2.1 0.3915 同样领先（附录表 4），PickScore/HPSv3 奖励下稳定拿第一（0.4092/25.7589/12.5855），证明对模型和奖励都通用。

Comparison of direct-gradient methods
Table 1: Comparison of direct-gradient methods
Comparing different post-training methods
Table 2: Comparing different post-training methods
Comparison of post-training methods using various rewards and prompt sets
Table 3: Comparison of post-training methods using various rewards and prompt sets
Comparison of post-training methods on Stable Diffusion 3.5 Medium
Table 4: Comparison of post-training methods on Stable Diffusion 3.5 Medium
Impact of the loss threshold λ
Table 5: Impact of the loss threshold λ
Qualitative comparison on GenEval benchmark
Figure 3: Qualitative comparison on GenEval benchmark
Comparison of GenEval score improvement during fine-tuning
Figure 5: Comparison of GenEval score improvement during fine-tuning
Additional qualitative comparisons on the GenEval benchmark
Figure 7: Additional qualitative comparisons on the GenEval benchmark
Qualitative results of Flux fine-tuned with LeapAlign using HPSv3 as the reward model
Figure 8: Qualitative results of Flux fine-tuned with LeapAlign using HPSv3 as the reward model
Qualitative results of Flux fine-tuned with LeapAlign using HPSv3 as the reward model (cont.)
Figure 9: Qualitative results of Flux fine-tuned with LeapAlign using HPSv3 as the reward model (cont.)
查看结构化数据
任务指标本文基线提升
Human Preference Alignment on HPDv2 (FLUX.1-dev) HPSv2.1 (in-domain) 0.4092 MixGRPO 0.3692 / DRTune 0.3882 +0.0210 vs MixGRPO, +0.0210 vs DRTune
Human Preference Alignment on HPDv2 (FLUX.1-dev) PickScore (out-of-domain) 23.7137 MixGRPO 23.5184 / DRTune 23.5185 +0.1953 vs MixGRPO
Human Preference Alignment on HPDv2 (FLUX.1-dev) HPSv3 (out-of-domain) 15.7678 MixGRPO 14.7530 / DRTune 15.5606 +1.0148 vs MixGRPO
Human Preference Alignment on HPDv2 (FLUX.1-dev) UnifiedReward-IQ 3.7244 MixGRPO 3.6241 +0.1003
Compositional Alignment (FLUX.1-dev) GenEval Overall 0.7420 MixGRPO 0.7232 / DRTune 0.7101 +0.0188 vs MixGRPO
Compositional Alignment (FLUX.1-dev) GenEval Position 30.25 MixGRPO 24.25 / DRTune 27.50 +6.00 vs MixGRPO
Compositional Alignment (FLUX.1-dev) GenEval Attribute Binding 66.00 MixGRPO 56.25 / DRTune 55.50 +9.75 vs MixGRPO
Robustness across rewards (PickScore/HPSv3) PickScore / HPSv3 25.7589 / 12.5855 DRTune 25.1021 / 12.0023 +0.6568 / +0.5832
Generalization to SD3.5-M (HPSv2.1 reward) HPSv2.1 0.3915 ReFL 0.3833 / DRTune 0.3828 +0.0082 vs DRTune

局限与改进

作者自己在文中承认两点局限：(1) LeapAlign 需要可微奖励，目前仅验证了 HPSv2.1、PickScore、HPSv3 等基于 CLIP/VLM 的奖励，对 ImageReward 这类更复杂的奖励或非可微的人类标注奖励只能借助可微 value model 支持（Future Work）。(2) 对单步/少步生成模型（如 consistency model、SDXL-Turbo）反而意义不大，因为这些模型轨迹本来就很短，奖励梯度本就能直接传播；LeapAlign 主要服务于多步高质量流匹配模型。除此之外可观察到的限制包括：训练算力大（16 GPU × 300 iter，文中未明示具体型号）、$\tau=0.1$、$\alpha=0.3$、$\lambda=0.55$ 等超参对不同模型要重新调，奖励曲线显示前期增益明显但后期饱和（附录 Fig. 5），说明仍有进一步压缩训练开销的空间；定性结果也显示对 two carrots、bench left of bear 这类含细小物体或强空间关系的 prompt 仍偶有失败。

独立分析的弱点

独立审视可指出四个潜在弱点及改进方向。第一，$\alpha=0.3$、$\tau=0.1$、$\lambda=0.55$ 三个超参对性能很敏感（$\alpha\in\{0,0.3,1\}$ 间差距明显，$\lambda$ 倒 U 形），可改进为基于梯度范数或轨迹相似度自适应调整，省去网格搜索。第二，仅支持可微奖励，无法直接用于 ImageReward 这类 VQA 类或人类标注奖励，可借鉴 Adjoint Matching 用可微 value model 包装任意奖励。第三，随机选 $k,j$ 是无信息采样，效率偏低，可引入基于 critic 的优先级采样，让提升潜力大的早期步被多采。第四，300 iter × 16 GPU 的预算对学术复现不友好，可结合 LoRA 或轨迹蒸馏压缩成本；定性结果对双对象、颜色绑定偶尔失败，说明单 hinge+HPSv2.1 监督粒度有限，可叠加 Layout-aware 损失或 latent 几何约束。

未来方向

作者明确指出未来有两个方向：(a) 把 LeapAlign 推广到视频生成（同时面临长轨迹和早期步决定的全局时序问题，可与 Sora、MovieGen 等同源方法结合）；(b) 通过可微 value model 让非可微奖励也能用 LeapAlign。基于成果还天然可延伸：(c) 将跳跃轨迹与 GRPO 思路结合，让 on-policy 优势估计与可微梯度同行；(d) 把两步跳跃推广到 $N$ 步连续跳跃，配合自适应 $\alpha$ 调度，可能进一步逼近 full-trajectory 反传但成本可控；(e) 用 LeapAlign 探索文生视频、文生 3D、文生音频等统一框架，因为这些模态同样需要更新早期步以决定全局结构；(f) 进一步研究嵌套梯度的统计特性，把 $\alpha$ 替换成学习得到的标量甚至更一般的修正矩阵。

复现评估

复现难度中等偏上，作者开源度有限但给出充分细节：算法伪代码在附录 Algorithm 1，主超参与训练配置（lr=$1\times 10^{-5}$ 或 $8\times 10^{-6}$, batch=64, wd=$1\times 10^{-4}$, EMA 0.995, $\alpha=0.3$, $\tau=0.1$, $\lambda=0.55$, 25 采样步, CFG 3.5, $720\times 720$, 16 GPU × 300 iter）全在正文中。基线 DRaFT-LV、DRTune、ReFL 的官方实现都缺失，作者只能基于论文伪代码自行复现，这一点对公平对比的可重复性是隐患。奖励模型 HPSv2.1/PickScore/HPSv3 均可从 Hugging Face 下载，FLUX.1-dev 也是开源权重，但混合精度、显存分配等工程细节并未披露，需要 16 张 H100/A100 量级 GPU 才能复现主表。一般研究者建议先在 SD3.5-M 小规模上重跑奖励模型=PickScore 的子实验（附录 B）以验证流程，再考虑 FLUX 全量。

论文图表

Performance overview of LeapAlign
Figure 1: Performance overview of LeapAlign

(a) 显示在组合对齐任务上 LeapAlign（粉红）相对 DRTune（绿色）的 ImageReward 增长曲线更快、更高；(b) FLUX.1-dev 基线、FLUX.1-Dev+DRTune、FLUX.1-Dev+LeapAlign 三组在 HPSv2.1、PickScore、UnifiedReward-ImageQuality、ImageReward 上的多维度雷达图，LeapAlign 各项外推；(c) GenEval 各子任务雷达图，LeapAlign 在 Position、Counting、Attribute Binding 等布局相关项上明显领先 MixGRPO 和 DRTune。

一张图同时给出奖励曲线、通用偏好分、组合对齐分，是论文最直观展示全面超过现有方法的视觉证据，对读者判断贡献大小非常关键。