LeapAlign:通过构建两步轨迹在任意生成步对流匹配模型进行后训练
LeapAlign: Post-training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
LeapAlign 通过将长 ODE 采样轨迹压缩为两步跳跃,使奖励梯度能直接回传到早期生成步,从而以低显存代价对齐流匹配模型;该技术可复用于任何需要端到端优化生成过程(如扩散/流模型)的视觉任务,例如文本到图像生成中的偏好微调。
原文依据 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.
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论本文关注流匹配模型与人类偏好的对齐问题。一种有前景的方法是通过流匹配可微生成过程直接反向传播奖励信号来进行微调。然而,通过长轨迹反向传播会导致高昂的内存开销和梯度爆炸。因此,直接梯度方法难以更新早期生成步,而这些步对于决定最终图像的全局结构至关重要。为解决该问题,我们提出 LeapAlign,一种降低计算成本并使奖励梯度能直接传播到早期生成步的微调方法。具体而言,我们通过设计两个连续的跳跃将长轨迹缩短为仅两步,每个跳跃跳过多个 ODE 采样步并在单步中预测未来潜变量。通过随机化跳跃的起始和结束时间步,LeapAlign 能够在任意生成步实现高效且稳定的模型更新。为更好地利用这种缩短的轨迹,我们为与长生成路径更一致的轨迹分配更高的训练权重。为进一步增强梯度稳定性,我们降低大幅度梯度项的权重,而非像先前工作那样完全移除它们。在微调 Flux 模型时,LeapAlign 在多项指标上持续优于最优的基于 GRPO 和直接梯度的方法,实现了更优的图像质量和图像-文本对齐。
英文原文摘要 1502 字符
This paper focuses on the alignment of flow-matching models with human preference. A promising way is fine-tuning by directly backpropagating reward signals 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.