让流匹配模型在任意生成步骤高效对齐人类偏好
LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
- 通过两步跳跃轨迹缩短生成路径,实现早期步骤的梯度回传
- 在Flux模型上显著提升图像质量和图文对齐效果
- 适合需要高效微调生成模型的科研与工程人员
本文聚焦于流匹配模型与人类偏好的对齐问题。直接通过可微生成过程反向传播奖励梯度是一种有前景的方法,但沿长轨迹反向传播会导致巨大的内存开销和梯度爆炸。因此,现有方法难以更新对最终图像结构至关重要的早期生成步骤。为此,我们提出LeapAlign,一种降低计算成本并支持从奖励到早期步骤直接梯度传播的微调方法。具体而言,通过设计两个连续跳跃,将长轨迹压缩为仅两步,每个跳跃跳过多个ODE采样步骤,并在单步内预测未来隐状态。通过随机化跳跃的起始和结束时间步,LeapAlign实现了任意生成步骤上的高效稳定更新。为更好利用短轨迹,我们赋予与长生成路径更一致的轨迹更高训练权重;为进一步增强梯度稳定性,我们降低大梯度项的权重,而非完全移除(如以往工作)。在微调Flux模型时,LeapAlign在多种指标上持续优于最先进的GRPO基方法和直接梯度方法,实现更优的图像质量与图文对齐。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。