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

1. 出发点 (Motivation)

用 reward model 给 flow matching T2I 模型对齐人类偏好,有两条路:

但 direct-gradient 有个硬约束:25 步 (或更多) 的长链 backward 会导致:

  1. 显存爆炸: 每一步的激活都要存,\(\partial x_0/\partial \theta\) 沿着 25 步链展开。
  2. 梯度爆炸: 多步乘积,数值不稳。

现有 direct-gradient 方法的"折中"都不完美:

问题是:早期步决定全局布局 (Hertz et al.; Liang et al.) — 想做组合性生成 (e.g., "a red giraffe", "two carrots") 必须改早期步。所以"无法更新早期步"是个 hard limitation。

Table 1: comparison of direct-gradient methods
Tab. 1 — Direct-gradient 方法对比。LeapAlign 是唯一四个属性 (Early Steps / Nested Gradient / Leap Trajectory / Multi-Step) 全部满足的方法。

LeapAlign 的 TL;DR 思路:"我不 backprop 整条 25 步链,我从中随机挑出 2 步,只在这 2 步上 backprop。" 通过精心的 stop_gradient 设计,前向用真实轨迹 \(x_0\) 算 reward,反向只走两步,可以 update 任何位置的步。

Performance overview
Fig. 1 — (a) 在组合性 alignment 上 LeapAlign 收敛比 DRTune 快且最终 reward 更高;(b) 6 个 evaluator (HPSv2.1, HPSv3, PickScore, UnifiedReward Align/IQ, ImageReward) 全面领先;(c) GenEval 6 个子任务,Attribute Binding 这种需要早期 layout 控制的任务上提升尤其明显 (66 vs base 45)。

2. 方法 (Method)

核心思想 (类比)

把扩散/flow matching 画图想成 25 笔素描:第 1 笔定大轮廓,第 25 笔抠细节。reward 是最后那张图的得分。你要把"反馈分数"告诉画家:

LeapAlign architecture
Fig. 2 — 总体流程。$x_1$ 是初始噪声,$x_0$ 是最终图。黑色实线是正常 25 步 ODE 采样 (在 no_grad 下跑)。蓝色实线是两个 "leap" — $\hat x_{j\mid k} = x_k - (k-j) v_\theta(x_k)$ 和 $\hat x_{0\mid j} = x_j - j v_\theta(x_j)$。橘色虚线是 "latent connector" — 把 $\hat x$ 与真实 $x$ 通过 stop_gradient 拼起来。Reward 算在真实 $x_0$ 上,梯度只通过紫色虚线 (leap trajectory) 反传。

2.1 Flow Matching 基础

Flow matching (Lipman et al. 2023; Liu et al. 2022) 学连续速度场 \(v_\theta(x_t, t)\),把噪声 \(x_1 \sim \mathcal N(0, I)\) 映回数据 \(x_0 \sim p_{\text{data}}\):

在rectified flow 下取线性 \(\alpha_t = 1-t,\ \beta_t = t\),速度场 \(v = x_1 - x_0\) 是常向量。

One-step leap prediction (论文 Eq. 3): 从时步 \(k\) 直接预测时步 \(j\) 的潜变量:

这是欧拉法 (Euler) 把 \(k\) 到 \(j\) 的所有中间步用一次速度估计跨过去 — 在长链上这是粗略近似,但在 leap trajectory 里只用来构造可微连接,精度通过 latent connector 修正。

2.2 Leap Trajectory 构造 (论文 §4.2)

每个 iteration:

  1. 从噪声 \(x_1\) 跑完整 25 步 ODE 采样得到 \(x_0\) (整条用 no_grad,但 \(x_k, x_j, x_0\) 这三个状态保留下来)。
  2. 随机抽 \(k, j \in [0, 1]\),\(k > j\)。
  3. 构造两个 one-step leap: $\(\hat x_{j\mid k} = x_k - (k - j) v_\theta(x_k) \quad (4)\)\(\)\(\hat x_{0\mid j} = x_j - j\, v_\theta(x_j) \quad (5)\)$
  4. Latent connector — 关键 trick: $\(x_j = \hat x_{j\mid k} + \text{sg}(x_j - \hat x_{j\mid k}) \quad (6)\)\(\)\(x_0 = \hat x_{0\mid j} + \text{sg}(x_0 - \hat x_{0\mid j}) \quad (7)\)\(<p>前向看:\)x_j\(的<em>值</em>仍然等于真实 ODE 采样得到的\)x_j\((因为\)\hat x_{j\mid k} + (x_j - \hat x_{j\mid k}) = x_j\()。但反向看:梯度只能流过\)\hat x_{j\mid k}\(这一项 (\)\text{sg}\(截断了\)(x_j - \hat x_{j\mid k})$ 的梯度)。这就是"forward 用真实状态、backward 只走两步"的实现机制。

论文用一句紧凑的箭头表示:

实线箭头是 leap (走 \(v_\theta\)),虚线箭头是 latent connector (stop_gradient)。backward 只展开 2 步,显存 O(1),还能更新任何 \((k, j)\)。

2.3 梯度分解 + 折扣 (论文 §4.3)

把 \(\partial x_0 / \partial \theta\) 沿 leap trajectory 展开 (论文 Eq. 8,Appendix G 推导):

三项含义:

论文观察:nested gradient 是真有用,但整项保留又容易爆。所以引入gradient discounting factor \(\alpha \in [0, 1]\),把 nested 项乘 \(\alpha\) 缩小:

这一步实现非常巧妙。前向看 \(\alpha x_j + (1-\alpha) \text{sg}(x_j) = x_j\) — 没有任何数值改变。反向看:\(\partial / \partial x_j\) 这条链路只收到 \(\alpha\) 倍的梯度 (\(\text{sg}\) 那一支不传梯度)。结果是 nested gradient 乘 \(\alpha\):

消融 (Fig. 4a) 显示 \(\alpha = 0.3\) 最优,\(\alpha = 0\) (DRTune 等价) 和 \(\alpha = 1\) (整项保留) 都更差。**"不删,缩小"**是论文工程上的核心 trick。

2.4 Loss + 轨迹相似加权 (论文 §4.4, §4.5)

Hinge-style reward loss (Eq. 11):

\(r(\cdot)\) 是 reward model (默认 HPSv2.1),\(\lambda\) 是阈值 (论文设 0.55)。超过阈值后停止施压,防 reward hacking。注意:reward 算在真实 \(x_0\) 上,不是 leap 估计的 \(\hat x_{0\mid j}\) — 这跟 ReFL 不同 (ReFL 在 \(\hat x_0\) 上算 reward)。论文论证:\(\hat x_0\) 可能模糊有 artifact,reward 评估不可靠;\(x_0\) 来自完整 25 步推理,质量真实。

轨迹相似加权 (Eq. 12): Leap trajectory 是粗近似,如果两次 leap 离真实轨迹太远,梯度信号反而误导。论文定义两个 connection 点的近似误差:

加权因子 (clamp 防止过度强调几乎相同的点):

最终目标 (Eq. 13):

注意 \(w_{\text{sim}}\) 自身也走 stop_gradient — 它只是样本权重,不参与优化目标。\(\tau\) 设 0.1。leap trajectory 越接近真实轨迹,权重越高,反之 down-weight。

2.5 与代码对照

LeapAlign 的代码非常干净。核心训练循环不到 200 行,把上面所有公式逐行翻译。

(a) 随机抽两个时步 \(k, j\) — 实现 Algorithm 1 的开头。

repo/fastvideo/train_leapalign_flux.py:L106-L168 — sample_trajectory 入口

idx_bs = 2
select_indices = (
    torch.randperm(
        args.max_idx - args.min_idx,
        device="cpu",
        generator=generator,
    )[:idx_bs] + args.min_idx
).long()
# t1, t2
select_indices = torch.sort(select_indices, descending=True).values

# ... (initialize latents, set up scheduler) ...

# Indices in the descending timestep schedule where gradient is enabled
grad_on_ind_set = set(
    timesteps.size(0) - select_index.item()
    for select_index in select_indices
)

for idx, step_t in enumerate(timesteps):
    with nullcontext() if (idx in grad_on_ind_set) else torch.no_grad():
        ...

L107-L115:从 [min_idx, max_idx) 区间随机抽 2 个 unique 整数,然后降序排成 \((k, j)\)。L165-L168:把这两个索引位置加入 grad_on_ind_set,主循环里只有这两步的 forward 在 nullcontext() 下跑 (梯度开启),其他 23 步都在 no_grad 下跑。

(b) Latent connector + gradient discounting — 实现 Eq. 6, 9。论文最核心的工程实现。

repo/fastvideo/train_leapalign_flux.py:L174-L194 — latent connector @ step $j$

if idx == timesteps.size(0) - select_indices[1].item():
    if args.with_traj_sim_weighting:
        with torch.no_grad():
            # b
            pred_x2_weight_factor = (
                torch.abs(pred_xt2.double() - latents_xt.double())
                .mean(dim=tuple(range(1, pred_xt2.ndim)), keepdim=False)
                .clip(min=args.traj_sim_weighting_min)
            )
    # w/o grad -> w/ grad   (paper Eq. 6: x_j = x̂_{j|k} + sg(x_j - x̂_{j|k}))
    latents_xt = pred_xt2 + (latents_xt - pred_xt2).detach()
    if args.nested_grad_coe <= 0.0:
        dit_latents_xt = latents_xt.detach()
    else:
        # paper Eq. 9: pass α·x_j + (1-α)·sg(x_j) into v_θ
        dit_latents_xt = (
            args.nested_grad_coe * latents_xt +
            (1 - args.nested_grad_coe) * latents_xt.detach()
        )
else:
    dit_latents_xt = latents_xt

三个对照点:

(c) Leap step + 最终 latent connector — 实现 Eq. 4, 5, 7。

repo/fastvideo/train_leapalign_flux.py:L195-L247 — two-leap + final connector

with torch.autocast("cuda", torch.bfloat16):
    pred = transformer(
        hidden_states=dit_latents_xt, encoder_hidden_states=encoder_hidden_states,
        timestep=timestep/1000, guidance=cfg_guidance, txt_ids=text_ids,
        pooled_projections=pooled_prompt_embeds, img_ids=image_ids,
        joint_attention_kwargs=None, return_dict=False,
    )[0]

if idx == timesteps.size(0) - select_indices[0].item():
    # @ step k: one-step leap from x_k to x̂_{j|k}  (paper Eq. 4)
    t2_idx = timesteps.size(0) - select_indices[1].item()
    pred_xt2 = fm_scheduler.jump_to_step(pred, step_t, latents_xt, t2_idx)
    pred = pred.detach()

if idx == timesteps.size(0) - select_indices[1].item():
    # @ step j: standard Euler + extract x̂_{0|j}  (paper Eq. 5)
    latents_xt, latents_pred_x0 = fm_scheduler.step(
        pred, step_t, latents_xt, with_pred_x0=True,
    )
    latents_xt = latents_xt.detach()
else:
    latents_xt = fm_scheduler.step(pred, step_t, latents_xt, with_pred_x0=False)[0]

# ... loop continues ...

# Final latent connector (paper Eq. 7): x_0 = x̂_{0|j} + sg(x_0 - x̂_{0|j})
latents_pred_x0 = latents_pred_x0 + (latents_xt - latents_pred_x0).detach()
return latents_xt, latents_pred_x0, traj_sim_weight_factor

对照点:

(d) One-step leap 实现 — Scheduler

repo/fastvideo/schedulers/fm_euler_discrete_with_pred_x0.py:L25-L108 — Euler step + jump_to_step

def step(self, model_output, timestep, sample, ..., with_pred_x0=False):
    # ... validate inputs ...
    sample = sample.to(torch.float32)
    sigma = self.sigmas[self.step_index]
    sigma_next = self.sigmas[self.step_index + 1]

    prev_sample = sample + (sigma_next - sigma) * model_output           # standard Euler step

    if with_pred_x0:
        pred_x0_sample = sample + (- sigma) * model_output               # ⟵ paper Eq. 5: x̂_{0|j} = x_j - j·v_θ(x_j)
    else:
        pred_x0_sample = None
    # ...
    return (prev_sample, pred_x0_sample)

def jump_to_step(self, model_output, timestep, sample, target_idx):
    # ... validate inputs ...
    sigma = self.sigmas[self.step_index]
    sigma_next = self.sigmas[target_idx]                                 # jump straight to target sigma
    target_sample = sample + (sigma_next - sigma) * model_output         # ⟵ paper Eq. 4: x̂_{j|k} = x_k - (k-j)·v_θ(x_k)
    return target_sample

注意:jump_to_step 和普通 Euler step 的区别只是 sigma_next 的取值 — 一个步进取 sigmas[step_index+1] (下一步),一个直接取 sigmas[target_idx] (跳到 \(j\) 位置)。在 rectified flow 下 sigma 与 timestep 成线性关系,所以 sigma_next - sigma 就是 \(-(k-j)\)。

(e) Hinge reward loss + trajectory similarity weighting

repo/fastvideo/train_leapalign_flux.py:L251-L294 — get_reward_loss

def get_reward_loss(args, latents_pred_x0, vae, reward_model, tokenizer, caption,
                    preprocess_val, traj_sim_weight_factor):
    w, h = args.w, args.h
    with torch.autocast("cuda", dtype=torch.bfloat16):
        latents_pred_x0 = unpack_latents(latents_pred_x0, h, w, 8)
        latents_pred_x0 = (latents_pred_x0 / 0.3611) + 0.1159
        image_pred_x0 = vae.decode(latents_pred_x0, return_dict=False)[0]
        image_pred_x0 = (image_pred_x0 * 0.5 + 0.5).clamp(0, 1)

    if args.use_hpsv2:
        image_pred_x0 = preprocess_val(image_pred_x0)
        text = tokenizer(caption).to(device=image_pred_x0.device, non_blocking=True)
        with torch.amp.autocast('cuda'):
            outputs = reward_model(image_pred_x0, text)
            image_features, text_features = outputs["image_features"], outputs["text_features"]
            reward_pred_x0 = torch.einsum("bc,bc->b", image_features, text_features)
    else:
        raise NotImplementedError("Only HPSv2 is currently supported for reward computation.")

    # paper Eq. 11: L_raw = max(0, λ - r(x_0))
    reward_loss = F.relu(-reward_pred_x0 + args.loss_relu_clip) * args.loss_grad_scale
    # paper Eq. 12-13: divide by (max(d_j, τ) + max(d_0, τ))  ==  multiply by w_sim
    if traj_sim_weight_factor is not None:
        reward_loss = reward_loss / traj_sim_weight_factor
    reward_loss = torch.mean(reward_loss)
    return reward_loss, reward_pred_x0.detach().mean()

L287 实现 Eq. 11:F.relu(-r + λ) = \(\max(0, λ - r)\)。L289 实现 Eq. 12-13:reward_loss / traj_sim_weight_factor = reward_loss * w_sim (因为 w_sim = 1 / (max(d_j, τ) + max(d_0, τ)),而 traj_sim_weight_factor 就是分母)。

3. 结论 (Key Findings)

Table 2: main results
Tab. 2 — Flux 上的主结果。LeapAlign 在 in-domain HPSv2.1 0.4092 (DRTune 0.3882, MixGRPO 0.3692, base 0.3078)、out-of-domain HPSv3 / PickScore / UnifiedReward 全部第一,GenEval Overall 0.7420 (MixGRPO 0.7232, DRTune 0.7101)。Attribute Binding 子项 66.00 vs DRTune 55.50 — 这正是需要早期 layout 控制的任务。

定量收益 (vs 最强基线 DRTune):

跨 reward / 跨 prompt set 泛化 (Tab 3): 用 PickScore 训练 → 测试 PickScore;用 HPSv3 训练 → 测试 HPSv3。LeapAlign 全部第一。

Table 3: cross-reward generalization
Tab. 3 — 换 reward model / prompt set 后,LeapAlign 仍然是 best。PickScore 25.7589 vs ReFL 25.2373,HPSv3 12.5855 vs DRTune 12.0023。
Qualitative comparison
Fig. 3 — GenEval 上的定性对比。"a photo of a red giraffe" 这种需要颜色绑定的 prompt,ReFL/DRaFT-LV (只更新末步) 仍然画灰色长颈鹿;LeapAlign [0,1] 范围训练能成功绑红色。"a photo of two carrots" 也是类似:ReFL 数错,LeapAlign 数对。

消融 (Fig 4):

Ablation grid
Fig. 4 — 6 个 ablation。(a) $\alpha = 0.3$ 是 nested gradient 折扣的甜点;(b) 2-step leap 是性能/显存的最佳平衡 (3-step 没收益);(c) reward 算在 $x_0$ 上优于 $\hat x_{0\mid j}$;(d) trajectory similarity 同时考虑 $d_j$ 和 $d_0$ 优于单独;(e) 训练时步范围 $[0, 1]$ 优于 $[0, 1/2]$ — 验证"早期步重要"的 motivation;(f) 随机选 $(k, j)$ 比固定距离稍好且更简单。

额外有意思的发现 (Fig 4a):即使 \(\alpha = 0\) (砍掉 nested gradient,行为类似 DRTune),LeapAlign 仍然比 DRTune 更好 (HPSv2.1 0.4064 vs DRTune 0.3882)。说明 leap trajectory 架构本身就有收益,nested gradient discounting 是额外的 boost。

4. 实现细节 (Implementation Notes)

训练代码已发布,基于 FastVideo 框架。可重现性较高,只需 16 GPUs。

5. 批判性总结 (Critical Assessment)

5.1 优点

5.2 不足 / 疑点

5.3 适用 vs 不适用

5.4 进一步阅读

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。