PaperlayerView exact revision on arXiv

Source-linked AI summary

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

arXiv:2604.15311v2cs.CV

TL;DR

Direct-gradient fine-tuning of flow matching models is limited by memory costs and gradient explosion on long trajectories, especially for early generation steps. LeapAlign replaces these trajectories with randomized two-step leaps and stabilizing weights, and consistently improves Flux across image-quality and image-text-alignment metrics. The paper’s stated scope includes multi-step image generation, with non-differentiable rewards and video generation left for future work.

  • Problem

    Long-trajectory direct-gradient backpropagation incurs prohibitive memory costs and gradient explosion, limiting updates to early generation steps that determine image layout.

  • Method

    LeapAlign constructs randomized two-step leap trajectories, weights paths by similarity to the full trajectory, and discounts rather than removes large-magnitude gradient terms.

  • Results

    LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods on Flux image-generation quality and image-text alignment.

  • Takeaways & Limitations

    LeapAlign propagates reward gradients to early generation steps without excessive memory cost or sacrificing useful gradient terms.

  • Takeaways & Limitations

    The paper leaves extending LeapAlign to non-differentiable rewards and video generation as future work and focuses its discussion on multi-step models.

Abstract

from arXiv · show

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.

1 Introduction

LeapAlign addresses the memory cost, gradient explosion, and limited early-step updates that hinder direct-gradient fine-tuning of flow matching models. It uses shortened trajectories and gradient-preserving techniques, improving Flux image quality and image-text alignment.

  • Flow matching enables reward gradients to backpropagate through differentiable generation trajectories, unlike discrete LLM generation.
  • Long-trajectory backpropagation causes prohibitive memory costs and gradient explosion, so existing methods often update only late steps.This leaves early steps that largely determine image layout unchanged; stopping gradients at the model input discards substantial gradient flow.
  • LeapAlign constructs a two-step leap trajectory from a full run, enabling constant-memory reward-gradient updates at randomly selected generation steps.The method computes reward on the actual final image while backpropagating only through the shortened trajectory.
  • LeapAlign consistently outperforms GRPO-based and direct-gradient methods on Flux image-generation and image-text-alignment metrics.Reported evaluators include HPSv2.1, HPSv3, PickScore, UnifiedReward, ImageReward, and GenEval.
  • LeapAlign weights leap trajectories by path similarity and discounts large-magnitude gradient terms instead of removing them.These choices retain useful gradient information while stabilizing training.

2 Related Work

Related work includes policy-gradient, preference-optimization, and direct-gradient approaches for diffusion and flow matching models. LeapAlign differs from prior direct-gradient methods by updating early or multiple steps while preserving nested gradients.

  • Policy-gradient methods fine-tune diffusion models with PPO or REINFORCE, while DPO-based methods use preference pairs or sets.
  • Flow-matching preference methods include stochastic optimal control and forward-process reinforcement-learning approaches.
  • Direct-gradient methods exploit differentiable samplers to propagate reward gradients, with ReFL and DRaFT-LV updating only late generation steps.
  • LeapAlign constructs leap trajectories to reach early steps, unlike ReFL and DRaFT-LV, while retaining nested gradients that DRTune removes.Nested gradients capture interactions across generation timesteps.

3 Preliminaries

Flow matching learns a continuous transformation from Gaussian noise to images by estimating a velocity field. Rectified flow uses linear interpolation, enabling one-step latent predictions between arbitrary timesteps.

  • Flow matching maps Gaussian noise to real images through a learned continuous transformation.
  • The forward process interpolates between data and noise using scheduler coefficients α_t and β_t.
  • A neural network v_θ predicts the velocity field governing the interpolation.
  • Rectified flow uses α_t = 1 − t and β_t = t, yielding velocity v = x_1 − x_0.
  • One-step leap prediction estimates latent x_j from latent x_k at arbitrary timesteps, supporting two adjacent leaps along a full trajectory.

4 Proposed Approach

LeapAlign shortens a full flow-matching generation trajectory into two leaps, enabling reward gradients to reach randomly selected early steps with controlled memory and gradient growth. It further stabilizes training by discounting large nested-gradient terms, weighting trajectory-consistent leaps, and evaluating a thresholded reward objective on the generated image.

  • 4.2 Leap Trajectory Construction: LeapAlign randomly selects k > j from a full trajectory and constructs two one-step leaps, from x_k to x_j and from x_j to x_0.Latent connectors align predicted states with actual states while preserving differentiability.
  • 4.2 Leap Trajectory Construction: Two-step leap trajectories keep memory cost constant while enabling stable reward-gradient backpropagation to any generation step.Random timestep selection allows updates beyond the late steps targeted by earlier direct-gradient methods.
  • 4.3 Gradient Discounting: The two-step gradient contains single-step and nested-gradient terms, with the nested term capturing interactions across generation steps.LeapAlign discounts the nested term with α instead of removing it, preserving the full gradient structure while moderating magnitude.
  • 4.4 Fine-Tuning Objective: LeapAlign uses a hinge-style objective, L_raw = max(0, λ − r(x_0)), to increase rewards beyond a threshold while limiting unstable reward maximization.Here r(·) is the reward model and λ controls the reward-maximization threshold.
  • 4.4 Fine-Tuning Objective: Rewards are computed on the generated image x_0 rather than a one-step final-image estimate, providing supervision that more directly reflects full-trajectory output quality.The method also assigns higher weights to leap trajectories whose predicted states more closely match actual states at connection points.

5 Discussions

The discussion situates LeapAlign against related direct-gradient methods and clarifies its supported reward-model scope and applicability boundary. It works with differentiable reward models, while extension to non-differentiable rewards remains future work and one- or few-step models are not its primary target.

  • Comparisons with Related Methods: LeapAlign extends one-step leap prediction into a two-leap trajectory, propagating gradients to early steps while retaining the nested gradient that captures cross-timestep dependencies.This contrasts with ReFL and DRaFT-LV, which update only a single late step, and with DRTune, which removes the nested term.
  • Reward Models: LeapAlign accommodates differentiable reward models, with effective fine-tuning demonstrated using CLIP-based HPSv2.1 and PickScore and vision-language-model-based HPSv3.Extension to non-differentiable rewards, potentially through differentiable value models, is identified as future work.
  • Applicability Scope: LeapAlign is less important for one-step or few-step image-generation models because their short trajectories already permit direct reward-gradient propagation.The discussion instead emphasizes fine-tuning methods for multi-step models, which are reported to have better image quality and alignment than one- or few-step models.

6 Experiments

Experiments evaluate LeapAlign on general preference and compositional alignment using Flux, multiple rewards, prompt sets, and baselines. LeapAlign achieves strong overall and compositional results, while analyses support two-step trajectories, gradient discounting, similarity weighting, and broad timestep selection.

  • Experimental Setup: Experiments evaluate Flux on general preference alignment with HPSv2.1, HPSv3, PickScore, UnifiedReward, ImageReward, and GenEval image-text alignment.The setup uses HPDv2 and MJHQ-30k prompts for preference alignment and GenEval prompts and tasks for compositional alignment.
  • General Preference Alignment: LeapAlign achieves the highest average scores on HPSv2.1, HPSv3, PickScore, UnifiedReward-Alignment, and UnifiedReward-IQ.Trained only with HPSv2.1, it exceeds MixGRPO on HPSv2.1 and PickScore and remains competitive on ImageReward.
  • Robustness: LeapAlign achieves the best performance across additional reward-model and prompt-set settings, confirming robustness.The experiments fine-tune Flux with PickScore on HPDv2 and HPSv3 on MJHQ-30k.
  • Compositional Alignment: 0.7420 overall GenEval score exceeds MixGRPO at 0.7232 and DRTune at 0.7101.The strongest gains occur in two objects, colors, position, and attribute binding categories.
  • Compositional Alignment: LeapAlign substantially modifies global image structure, producing compositions more faithful to text prompts than methods fine-tuning only late steps.The qualitative comparison contrasts LeapAlign with ReFL, DRaFT-LV, and the pretrained Flux model.
  • Further Analysis: α = 0.3 yields the best gradient-discounting performance; removing nested gradients lowers scores, while no discounting retains large gradients.Without the nested gradient, LeapAlign still exceeds DRTune on HPSv2.1: 0.4064 vs. 0.3882.
  • Further Analysis: Similarity weighting at both xj and x0 further improves average HPSv2.1 over single-step or unweighted variants.The analysis compares weighting based on xj, x0, both steps, and no weighting.
  • Further Analysis: Two-step leap trajectories provide the best trade-off between performance and memory usage, while three steps increase memory without improving results.One-step trajectories perform worse than two-step versions but still outperform DRTune and ReFL.

7 Conclusion

The conclusion presents LeapAlign as an efficient and stable method for propagating reward gradients to early generation steps. It attributes performance gains to two-step trajectories, gradient down-scaling, and trajectory-similarity weighting.

  • Conclusion: LeapAlign constructs two-step leap trajectories for efficient and stable reward-gradient backpropagation.The method targets early generation steps without excessive memory cost or sacrificing useful gradient terms.
  • Conclusion: Down-scaling large-magnitude gradient terms and up-weighting trajectories similar to original trajectories are useful design choices.The conclusion identifies both mechanisms as contributing components of LeapAlign.
  • Conclusion: LeapAlign consistently improves general image preference and image-text alignment across a wide range of metrics.The paper notes future implementation and improvement of LeapAlign for video generation.

A Visualization of GenEval Score Improvement During Fine-Tuning for Direct-Gradient Methods

The figure compares GenEval improvement curves for direct-gradient methods during fine-tuning. LeapAlign improves more rapidly and reaches a higher final GenEval score than the compared baselines.

  • Score Improvement: LeapAlign shows both more rapid GenEval improvement and a higher final score than DRTune, DRaFT-LV, and ReFL.The figure evaluates score improvement throughout fine-tuning.
  • Score Improvement: Methods updating early generation steps improve more strongly than methods that do not, underscoring early-step fine-tuning for compositional alignment.LeapAlign achieves the greatest improvement across the entire fine-tuning process.

B Additional Results on Stable Diffusion 3.5 Medium

Additional experiments test LeapAlign on Stable Diffusion 3.5 Medium under the general preference alignment setup. LeapAlign again performs best across all evaluators compared with other direct-gradient methods.

  • Experimental Setup: Experiments fine-tune and evaluate Stable Diffusion 3.5 Medium at 512 × 512 for 200 iterations.Other settings follow the general preference alignment experiment using HPSv2.1.
  • Results: LeapAlign achieves the best performance across all evaluators compared with other direct-gradient methods.The results indicate generalization to another flow matching model while maintaining strong improvements.

C Additional Analysis

Additional analyses show that moderate nested-gradient discounting and an intermediate loss threshold improve the trade-off between optimization and performance. Qualitative comparisons further indicate stronger global-structure and preference alignment improvements with LeapAlign.

  • Analysis of the nested gradient: α = 0.3 reduces gradient magnitude and improves performance over removing the nested gradient, whereas α = 1 increases gradient norms and degrades performance.The experiment uses HPSv2.1 and reports average test-set HPSv2.1 alongside average gradient norm.
  • Impact of the loss threshold λ: λ = 0.55 achieves the best overall performance among tested loss thresholds.Too-small thresholds under-optimize the model, while too-large thresholds make optimization overly aggressive and hurt out-of-domain generalization and reward scores.

E Additional Implementation and Training Details

The implementation derives one-step leap predictions for arbitrary timesteps and trains LeapAlign through randomized two-step trajectories with discounted gradients and consistency-based weighting. Additional details specify scheduler assumptions, gradient handling, hyperparameters, and qualitative evaluation settings.

  • Training details: Training uses reward-dependent gradient discounting and learning rates: α is 0.3 for HPSv2.1, 0.1 for PickScore or HPSv3, and HPSv3 uses learning rate 8e−6.The corresponding loss thresholds are 0.55, 0.4, and 13.5, respectively.
  • Baseline hyperparameters: Baseline configurations use recommended settings, including two re-noising steps for DRaFT-LV, two training timesteps for DRTune, and early-stop sampling from the last 11 of 25 steps for DRTune and ReFL.ReFL fine-tuning omits the pre-training loss in the reported configuration.
  • One-step leap prediction: One-step leap prediction extends a state at timestep k to an arbitrary earlier timestep j using scheduler derivatives and the velocity field.The practical prediction uses the pretrained flow matching model vθ(xk, k) to approximate the velocity.
  • Leap trajectory gradients: LeapAlign randomly selects timesteps k > j and propagates gradients through the resulting two-step leap trajectory while treating full-trajectory rollout states as detached constants.Gradient discounting modifies backpropagation while preserving the forward-pass velocity evaluation.
  • Qualitative evaluation: Additional GenEval comparisons evaluate pretrained Flux, ReFL, DRaFT-LV, DRTune, and LeapAlign, while further figures show HPSv3-based LeapAlign outputs.The reported qualitative results span diverse styles, themes, and scenarios.
  • Algorithm 1: The implementation samples LeapAlign timesteps tk and tj with tk > tj, applies the leap prediction at tj, and evaluates the reward-gradient update at the predicted final image.Algorithm 1 also includes consistency-based weighting through differences between full-trajectory and leap-predicted states.
Loading 2604.15311v2…