LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
Summary
LeapAlign is a post-training method that improves flow matching model alignment with human preferences by reducing computational costs through two-step trajectory shortcuts while enabling stable gradient propagation to early generation steps. The method outperforms state-of-the-art approaches when fine-tuning Flux models across various image quality and text-alignment metrics.
View Cached Full Text
Cached at: 04/20/26, 08:28 AM
Paper page - LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
Source: https://huggingface.co/papers/2604.15311
Abstract
LeapAlign improves flow matching model fine-tuning by reducing computational costs and enabling stable gradient propagation through shortened trajectory steps while maintaining alignment with human preferences.
This paper focuses on the alignment of flow matching models (https://huggingface.co/papers?q=flow%20matching%20models) with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients (https://huggingface.co/papers?q=reward%20gradients) through the differentiable generation process (https://huggingface.co/papers?q=generation%20process) of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion (https://huggingface.co/papers?q=gradient%20explosion). Therefore, direct-gradient methods (https://huggingface.co/papers?q=direct-gradient%20methods) 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 (https://huggingface.co/papers?q=ODE%20sampling%20steps) 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 (https://huggingface.co/papers?q=Flux%20model), LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods (https://huggingface.co/papers?q=direct-gradient%20methods) across various metrics, achieving superior image quality and image-text alignment (https://huggingface.co/papers?q=image-text%20alignment).
View arXiv page (https://arxiv.org/abs/2604.15311)View PDF (https://arxiv.org/pdf/2604.15311)Project page (https://rockeycoss.github.io/leapalign/)Add to collection (https://huggingface.co/login?next=%2Fpapers%2F2604.15311)
Get this paper in your agent:
hf papers read 2604.15311
Don’t have the latest CLI?curl -LsSf https://hf.co/cli/install.sh | bash
Models citing this paper0
No model linking this paper
Cite arxiv.org/abs/2604.15311 in a model README.md to link it from this page.
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2604.15311 in a dataset README.md to link it from this page.
Spaces citing this paper0
No Space linking this paper
Cite arxiv.org/abs/2604.15311 in a Space README.md to link it from this page.
Collections including this paper1
Similar Articles
Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation
This paper introduces Trajectory-Shaped Discrete Flow Matching (TS-DFM), which replaces blind stochastic jumps with guided navigation to significantly improve text generation efficiency and reduce computational costs. The method achieves superior perplexity and speed compared to traditional multi-step baselines while maintaining unchanged inference costs.
Exploring the Design Space of Reward Backpropagation for Flow Matching
FlowBP introduces a unified surrogate-trajectory framework for aligning flow matching models with human preferences via reward backpropagation, reducing memory usage and gradient chaining while maintaining performance across multiple text-to-image models.
FlowLM: Few-Step Language Modeling via Diffusion-to-Flow Adaptation
FlowLM introduces a flow matching language model derived from pre-trained diffusion models via efficient fine-tuning, enabling high-quality few-step text generation that rivals 2,000-step diffusion sampling with far fewer training epochs.
Constraint-Aware Flow Matching: Decision Aligned End-to-End Training for Constrained Sampling
Proposes Constraint-Aware Flow Matching, a novel end-to-end framework that aligns the model's learning dynamics with constrained sampling procedure, mitigating distributional shift from projection corrections for high-quality constrained generation.
@HuggingPapers: Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance Naver AI eliminates unsta…
Naver AI introduces Stable-GFlowNet, a method to improve LLM red-teaming by eliminating unstable partition function estimation in Generative Flow Networks through contrastive trajectory balance.