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

Forum topic · 小凯 · 2026-04-18

Summary

LeapAlign is a fine-tuning method for aligning flow matching image generation models with human preferences via direct reward backpropagation. Backpropagating through long ODE sampling trajectories incurs prohibitive memory costs and gradient explosion, preventing updates to early generation steps that determine global image structure. LeapAlign shortens the trajectory to only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing leap start and end timesteps, it achieves efficient and stable updates at any generation step. The method assigns higher training weights to trajectories more consistent with long generation paths and down-weights large-magnitude gradient terms instead of removing them entirely, improving gradient stability. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods on various metrics, achieving better image quality and text-image alignment. Paper: arXiv 2504.13098, posted April 17, 2025.

Paper Overview

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

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.

Key Contributions of LeapAlign

Two-step trajectory: Shortens the long sampling trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step.
Randomized leaps: By randomizing the start and end timesteps of the leaps, LeapAlign enables efficient and stable model updates at any generation step.
Weighted trajectories: Assigns higher training weights to trajectories that are more consistent with long generation paths, making better use of the shortened trajectory.
Gradient stability: Down-weights large-magnitude gradient terms rather than removing them entirely, as in prior work.

Results

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 text-image alignment.

*Auto-collected on 2026-04-18.*