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

TL;DR

LeapAlign is a direct-gradient (non-RL) post-training recipe for rectified flow-matching image models that propagates reward gradients to any generation step — including the early-step layout decisions that prior direct-gradient methods (ReFL, DRaFT-LV) can only update at the very end of the trajectory. The mechanism is a two-step “leap trajectory”: at each iteration, randomly pick two timesteps t > s along the full sampling chain and build a shortened two-step path noise → (latent at t) → (final image) using one-step leap predictions, evaluate reward on the actual end image, and backprop only through the two leaps — keeping memory constant while reaching early steps. Two stabilizers tame the resulting gradient: a discounting factor α that scales down (rather than zeros out, as DRTune does) the cross-step nested gradient term, and a trajectory-similarity weighting that down-weights leap trajectories that drift too far from the real ODE path. Fine-tuning FLUX.1-dev with HPSv2.1 as the reward, LeapAlign beats DanceGRPO, MixGRPO, ReFL, DRaFT-LV, and DRTune across HPSv2.1 / HPSv3 / PickScore / UnifiedReward / ImageReward and the GenEval compositional benchmark (overall 0.7420 vs MixGRPO 0.7232, DRTune 0.7101).

Key claims

Flow-matching sampling is continuous and differentiable, so reward gradients can flow back through the generation trajectory — a fundamentally different optimization surface than the discrete-token policy-gradient setting GRPO was designed for; LeapAlign argues this differentiability should be exploited rather than discarded by SDE-reformulation tricks [§1, §2 Related Work].
Direct backpropagation through the full sampling chain is foreclosed by (a) prohibitive activation memory across the long trajectory and (b) gradient explosion; prior direct-gradient methods (ReFL, DRaFT-LV) work around this by updating only a single late step per iteration, leaving early steps — which largely determine global image layout — untrained [§1, §4.3].
A two-step “leap trajectory” t > s with one-step leap prediction (Eq. 3, derived in Appendix 14) lets the model update any timestep via constant-memory backprop, and randomizing (t, s) across iterations covers the full trajectory [§4.2, Eqs. 4–7].
Backpropagation through two flow-matching steps produces a gradient with two single-step terms and one nested gradient that captures cross-step interactions; DRTune drops the nested term by stop-gradient on model inputs, LeapAlign keeps it but multiplies it by a discount α [§4.3, Eq. 8 vs Eq. 10].
A trajectory-similarity weighting term down-weights leap trajectories whose intermediate predictions diverge from the real ODE path, on the grounds that gradient signals from off-trajectory leaps are misleading [§4.5, Eq. 12–13].
Reward is evaluated on the actual end image x₀ rather than the one-step leap estimate x̂₀ used by ReFL/DRTune — the leap estimate may contain noise that biases the reward model [§4.4].
On FLUX.1-dev fine-tuned with HPSv2.1, LeapAlign beats every filed baseline (DanceGRPO, MixGRPO, ReFL, DRaFT-LV, DRTune) on HPSv2.1 (0.4092), HPSv3 (15.77), PickScore (23.71), UnifiedReward-Align/IQ (3.50/3.72), and ImageReward (1.51); MixGRPO is the only baseline close on ImageReward, and it was trained with three rewards [Table 2].
On the GenEval compositional benchmark, LeapAlign reaches overall 0.7420 vs MixGRPO 0.7232 (best policy-gradient baseline) and DRTune 0.7101 (best direct-gradient baseline); largest gains are on the layout-sensitive categories — Two Objects 96.46, Position 30.25, Attribute Binding 66.00 — consistent with the claim that the wins come specifically from early-step updates [Table 2].
Restricting the (t, s) range to late timesteps [0.5, 1.0] (i.e. matching DRTune-style late-only updates) degrades GenEval, while the full [0, 1] range wins — direct evidence that early-step training is the source of the compositional gains [§6.3, Fig. 4(e)].
Ablation: removing the nested gradient (α=0) underperforms α=0.5; setting α=1 (no discounting) destabilizes training; the discount mechanism is doing real work [§6.3, Fig. 4(a)]. Two-step leap trajectories beat both one-step (worse quality) and three-step (no gain, more memory) [Fig. 4(b)].
Generalizes across reward models (HPSv2.1, PickScore, HPSv3 all work) and across flow-matching backbones (FLUX.1-dev + SD3.5-M in Appendix 10) [Table 3, §6.2].

Method

LeapAlign post-trains a rectified flow-matching model (default FLUX.1-dev) against a differentiable reward. At each iteration:

Generate a full trajectory. Run standard ODE sampling from Gaussian noise to image x₀ with 25 steps, CFG scale 3.5, and record the latent at every step.
Sample a leap (t, s). Pick two timesteps t > s uniformly from the full range [0, 1]. Use the one-step leap prediction (a closed-form expression in rectified flow that estimates the latent at any timestep s from any timestep t using the model’s velocity prediction at t) to build x̂_s from x_1=noise via the velocity at t=1, then x̂_0 from x̂_s via the velocity at s. Latent connectors (Eq. 6–7) replace x̂_s with the actual x_s plus a stop-gradient correction so the predicted and real latents stay aligned while preserving differentiability through the two flow steps.
Backprop through the two leaps only. This gives gradient ∂L/∂θ = single-step₁ + single-step₂ + α · nested, where the nested term cross-couples the two leap predictions via the chain rule on the model input of the second leap. DRTune zeros the nested term by stop-grad; LeapAlign keeps it scaled by α=0.5 (Eq. 9–10), which the paper argues preserves the cross-step learning signal at controlled magnitude.
Evaluate reward on x₀ (the real end image, not x̂₀). Use a hinge-style objective ℓ = max(0, threshold − r(x₀)) to prevent reward hacking past a saturation point [§4.4, Eq. 11].
Trajectory-similarity weighting. Multiply the loss by w = clamp((dist(x̂_s, x_s) + dist(x̂_0, x_0)) / 2, min) → 1/w (Eq. 12), down-weighting leaps that drift far from the real ODE path.

Training: AdamW lr 1e-5, 16 GPUs, 300 iterations, batch 32, EMA 0.9999. Hinge threshold tuned per reward model. The recipe is reward-agnostic — any differentiable reward (CLIP-based HPSv2.1, VLM-based HPSv3, PickScore) plugs in unchanged.

Results

General preference alignment on FLUX (HPDv2 prompts, HPSv2.1 reward). HPSv2.1 0.3078 → 0.4092 (LeapAlign) vs 0.3852 ReFL, 0.3859 DRaFT-LV, 0.3882 DRTune, 0.3451 DanceGRPO, 0.3692 MixGRPO [Table 2]. Out-of-domain HPSv3 15.77 (LeapAlign) vs 15.56 DRTune (next best) and PickScore 23.71 vs 23.64 DRaFT-LV.
GenEval compositional benchmark. Overall 0.7420 (LeapAlign) vs 0.7232 MixGRPO, 0.7101 DRTune. Category breakdown shows LeapAlign winning particularly on layout-driven categories: Two Objects 96.46 (vs 93.69 MixGRPO/DRTune), Position 30.25 (vs 27.50 DRTune, 24.25 MixGRPO), Attribute Binding 66.00 (vs 57.00 ReFL, 56.25 MixGRPO) [Table 2]. Single-object generation saturates near 99% for all methods.
Cross-reward generality. Fine-tuning with PickScore on HPDv2 and HPSv3 on MJHQ-30k both work, beating ReFL/DRaFT-LV/DRTune on the matched test sets [Table 3].
Ablations [§6.3, Fig. 4]. (a) Nested-gradient discounting α=0.5 best; α=0 (DRTune-equivalent) and α=1 (no discount) both worse — but LeapAlign at α=0 still beats DRTune (0.4064 vs 0.3882 HPSv2.1), so the leap-trajectory design itself is a separate contribution. (b) Two leaps optimal; one leap worse, three leaps no better and more memory. (e) Full-range [0, 1] timestep sampling beats late-only [0.5, 1.0] on GenEval — confirming early-step training drives the compositional wins. (f) Random (t, s) selection marginally better than fixed gap |t−s|.
Qualitative. Fig. 3 shows ReFL/DRaFT-LV outputs hewing to the pretrained model’s layout (those methods only update late steps), while LeapAlign substantially restructures global composition to match the prompt.

Why it’s interesting

LeapAlign is a third move in the live “how to post-train a flow-matching image model” debate this wiki has been tracking, and it’s the direct-gradient faction’s strongest entry to date. The contrast with two filed neighbors is clean:

DanceGRPO: Unleashing GRPO on Visual Generation is the canonical GRPO-on-flow recipe — reformulate sampling as an SDE, take group-relative advantages, drop the KL term. It works, but as Liang Zheng’s tweet emphasizes, it imports policy-gradient stochasticity that flow matching’s differentiable sampler doesn’t actually need. LeapAlign beats DanceGRPO on every metric in Table 2.
Finite Difference Flow Optimization (FDFO): RL Post-Training for Flow-Based Image Generators is the other direct-gradient alternative also filed this month — it argues that backpropagating through the sampling chain is poorly conditioned (composing per-step Jacobians distorts the gradient’s frequency content) and that copying an image-space gradient uniformly across all steps works better. LeapAlign takes the opposite bet: backprop through two steps is fine if you use leap predictions plus α-discounted nested gradients, and you actually need the cross-step coupling for early-step training. The two methods now define a real open question — is the right credit-assignment unit (a) one image-space gradient copied to all steps (FDFO), (b) full backprop through two leaps with nested-gradient discounting (LeapAlign), or (c) an SDE-reformulated MDP with group-relative advantages (DanceGRPO/MixGRPO)? All three claim wins on FLUX-class models against different baselines; a controlled head-to-head doesn’t yet exist.

This also sharpens Reasoning RL‘s scope. The page’s TL;DR currently frames GRPO + shaped rewards as the default post-training recipe across modalities; LeapAlign is an explicit pushback in the image-diffusion sub-domain — gradient descent, not RL, when the sampler is differentiable. Worth tracking whether the same argument extends to video diffusion (where Human detectors are surprisingly powerful reward models has GRPO on Wan 2.1 winning 73% over base) and to audio diffusion (PrismAudio: Decomposed Chain-of-Thoughts and Multi-dimensional Rewards for Video-to-Audio Generation uses GRPO).