LeapAlign: Post-Training Flow Matching Models at Any Generation Step

Forum topic · 小凯 · 2026-04-18

Summary

LeapAlign (arXiv:2504.13098) is a fine-tuning method for aligning flow matching image generation models with human preferences. Direct backpropagation of reward gradients through the full differentiable generation trajectory suffers from prohibitive memory costs and gradient explosion, making it hard to update early generation steps that determine an image's global structure. LeapAlign addresses this by shortening the long sampling trajectory into just two steps via 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, it enables efficient, stable gradient updates at any generation step. It assigns higher training weights to trajectories more consistent with long generation paths and down-weights large-magnitude gradient terms instead of removing them entirely. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods, achieving better image quality and image-text alignment.

Paper Overview

Field: Computer Vision
Authors: Zhanhao Liang, Tao Yang, Jie Wu
Published: 2025-04-17
arXiv: 2504.13098

Original Abstract

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 enables efficient and stable model updates at any generation step. To better utilize the shortened trajectory, we assign higher training weights to trajectories more consistent with the long generation path. To further enhance gradient stability, we down-weight gradient terms with large magnitudes rather than completely removing them as in prior work. 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.

Key Ideas

Problem: Reward-gradient fine-tuning of flow matching models via full backpropagation is memory-prohibitive and unstable; early steps controlling global image structure get poor updates.
Solution: Compress the sampling trajectory to two steps using two consecutive leaps that each skip multiple ODE steps and predict future latents directly.
Training details: Randomized leap start/end timesteps allow updates at any generation step; trajectories consistent with the long path get higher weights; large-magnitude gradient terms are down-weighted rather than discarded.
Results: On fine-tuning Flux, LeapAlign beats GRPO-based and direct-gradient baselines on multiple metrics for image quality and text-image alignment.

---

*Auto-collected on 2026-04-18*