PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring the Design Space of Reward Backpropagation for Flow Matching

Ruoyu Wang, Boye Niu, Xiangxin Zhou, Yushi Huang, Tongliang Liu, Chi Zhang

arXiv:2606.11075v1cs.LG

TL;DR

Direct reward backpropagation for flow matching is sample-efficient but limited by activation-memory demands and multi-step Jacobian inflation. FlowBP designs a lightweight surrogate backward trajectory, and its variants consistently improve direct-gradient baselines across three backbones and multiple metrics.

  • Problem

    Direct reward backpropagation is sample-efficient, but full-trajectory activation storage is infeasible and chained Jacobians inflate gradients across steps.

  • Method

    FlowBP constructs a lightweight backward surrogate from cached and selectively re-forwarded velocities, organized around reward input, active set, integration weights, and bridge coupling.

  • Results

    Across three flow-matching backbones and preference, quality, and compositional metrics, FlowBP variants consistently improve over direct-gradient baselines.

  • Takeaways & Limitations

    The four surrogate-design axes provide actionable control over endpoint fidelity, trajectory coverage, integration accuracy, and cross-segment credit assignment.

  • Takeaways & Limitations

    Connector-based surrogates can become unreliable and destabilize training when long intervals produce large connector residuals.

Abstract

from arXiv · show

Aligning text-to-image flow matching models with human preferences via direct reward backpropagation is sample-efficient but hampered by two well-known pathologies: activations cannot be stored across the full sampling trajectory at modern model scale, and chained Jacobian products across steps inflate the reward gradient as it travels back to early indices. Connector-based methods, such as LeapAlign, address these issues by replacing the full backward trajectory with a short pinned path, highlighting a useful decoupling between sampling and optimization. However, the quality of the resulting gradient depends on how accurately this short path approximates the full rollout, especially over long intervals. We propose FlowBP, a unified surrogate-trajectory framework that treats the backward trajectory itself as the design object. FlowBP keeps a no-gradient cached rollout for sampling, then builds a lightweight backward surrogate from cached and selectively re-forwarded velocities. This view separates four choices: the reward-model input, active set, integration weights, and bridge coupling, and recovers prior direct-gradient methods as particular settings. Within this framework, we instantiate three variants: FlowBP-Sparse uses sparse Euler reconstruction, FlowBP-Bridge adds controlled bridge coupling, and FlowBP-Lagrange raises the order of leap quadrature. All three bound memory by the active-set size and limit gradient chaining to at most one Jacobian factor. Across SD3.5-M, FLUX.1-dev, and FLUX.2-Klein-base on preference, quality, and compositional metrics, the three variants improve over direct-gradient baselines on most metrics.

1 Introduction

FlowBP frames direct reward backpropagation as surrogate backward-trajectory design, addressing infeasible activation storage and exploding chained Jacobians through four controllable axes. Its Sparse, Bridge, and Lagrange variants improve preference, quality, and compositional-generation results across three flow-matching backbones.

  • Motivation: Direct reward backpropagation is sample-efficient and avoids stochastic-policy-gradient variance, but full-trajectory activation storage is infeasible and chained Jacobians inflate gradients backward.These limitations motivate replacing the full differentiable rollout with a designed surrogate trajectory.
  • Framework: FlowBP treats the backward surrogate as the design object and organizes prior direct-gradient methods around reward input, active set, integration weights, and bridge coupling.The framework retains a no-gradient cached rollout and constructs a lightweight backward surrogate from cached and selectively re-forwarded velocities.
  • Framework: FlowBP bounds memory by active-set size rather than rollout length while directly addressing the original gradient and activation-storage pathologies.Its backward surrogate is built from a cached rollout plus selectively re-forwarded velocities.
  • Variants: FlowBP-Sparse, FlowBP-Bridge, and FlowBP-Lagrange respectively explore sparse Euler reconstruction, controlled bridge coupling, and higher-order leap quadrature.Sparse removes the bridge and cross-step Jacobians, Bridge distributes an active set across coupled segments, and Lagrange replaces one-velocity hops with high-order rules.
  • Results: The three variants deliver consistent gains across three flow-matching backbones on preference and quality metrics, improve compositional generation, and achieve the strongest overall direct-gradient results.The reported evaluation spans preference, quality, and compositional-generation outcomes.

2 Related Work

Related work aligns diffusion and flow models with human preferences through policy-gradient, preference-optimization, and direct reward-backpropagation strategies. Direct-gradient methods exploit sampler differentiability to optimize rewards without likelihood estimation, using varied trajectory and objective designs.

  • Preference alignment for diffusion and flow models: Preference-alignment methods include policy-gradient and GRPO-style approaches that adapt reinforcement-learning objectives to diffusion and flow-model sampling trajectories.
  • Preference alignment for diffusion and flow models: Other preference-optimization methods learn from paired, step-wise, dense, or self-generated preference signals.
  • Direct reward backpropagation: Direct-gradient methods optimize reward scores by backpropagating through differentiable diffusion or flow samplers without likelihood estimation.
  • Direct reward backpropagation: Existing direct-gradient approaches vary in how they backpropagate rewards, including short late-stage denoising paths, low-variance final-step objectives, detached intermediate inputs, and composition-aware feedback.

3 Preliminaries

Diffusion and flow-matching models share a continuous-time generative formulation in which a velocity network predicts trajectory tangents between clean data and Gaussian noise. Sampling uses a discretized reverse-time Euler rollout, through which reward gradients can be backpropagated, while straight-through connectors preserve cached forward values and reroute gradients through differentiable surrogates.

  • Diffusion and Flow models: Diffusion and flow-matching models formulate the same continuous-time generative process, training vθ(x_t, t, c) to predict trajectory tangents between clean data and Gaussian noise.The trajectory is parameterized as x_t = α_t x_0 + σ_t ϵ, with clean data at t = 0 and Gaussian noise at t = 1.
  • Diffusion and Flow models: Under straight-line interpolation, the target velocity is v = ϵ − x_0, and the network is trained to predict this target.This is the specific interpolation and velocity target used in the preliminaries.
  • Sampling rollout: Sampling discretizes reverse time into an N-step rollout from initial Gaussian noise, applying Euler updates with step-dependent noise levels σ_i.The schedule satisfies σ_N = 1 > σ_{N−1} > · · · > σ_0 = 0; uniform discretization uses σ_i = i/N.
  • Sampling rollout: Because the rollout is fully differentiable, reward gradients can be backpropagated through any subset of its N steps to update θ.The trajectory terminates at the clean sample x_0.
  • Straight-through connectors: Straight-through connectors keep each latent’s forward value equal to its cached value while rerouting gradients through a differentiable surrogate estimate.The connector residual d_i = ∥x_i − x̂_i∥_2 measures discrepancy between forward and backward paths without changing the forward value.

4 Method

FlowBP replaces full-rollout backpropagation with a lightweight surrogate trajectory built from cached states and selectively re-forwarded velocities. The framework separates reward input, active velocities, integration, and bridge coupling, recovering prior methods as special cases and instantiating three new variants.

  • Challenges: Naive rollout backpropagation is impractical because storing all N activations is prohibitive and chained Jacobians can destabilize early-step gradients.Connector-based remedies add a third issue: long leaps can create large residuals and biased surrogate gradients.
  • Unified surrogate trajectory: FlowBP caches a no-autograd rollout, then builds a sparse backward graph by re-forwarding only selected velocities while evaluating reward on the original sample x0.Inactive states and velocities are detached, so memory depends on the active set rather than rollout length.
  • Design space: The framework exposes four design choices: reward-model input, active set, integration weights, and bridge coupling through nested-gradient scale α.With a bridge, α = 0 detaches segments, while larger α allows controlled cross-segment credit assignment.
  • Gradient structure: Direct active-step terms avoid full-chain differentiation, while the nested term reintroduces exactly one Jacobian factor regardless of rollout length N.Thus memory scales with |A| rather than N, and gradient amplification is bounded by a single Jacobian factor.
  • Methods: FlowBP-Sparse reconstructs x0 by Euler composition without a bridge, FlowBP-Bridge reconstructs the bridge with tunable α, and FlowBP-Lagrange raises leap quadrature order.The framework also recovers ReFL, DRaFT-LV, DRTune, and LeapAlign as special cases.

5 Experiments

Experiments across three flow-matching backbones show that FlowBP variants generally improve preference and quality metrics, while FlowBP-Lagrange also improves compositional generation. Ablations attribute these gains to endpoint-faithful reconstruction, accurate integration weights, and controlled surrogate coupling.

  • Experimental setup: Experiments cover SD3.5-M, FLUX.1-dev, and FLUX.2-Klein-base, comparing the base models with ReFL, DRaFT-LV, DRTune, and LeapAlign.Evaluation uses HPSv2.1, PickScore, ImageReward, UR-Align, and UR-IQ on 400 held-out HPDv2 prompts.
  • Main results: Across backbones, a FlowBP variant achieves the best score on nearly every preference and quality metric, with the strongest variant varying by backbone and metric.FlowBP-Lagrange leads most metrics on SD3.5-M and HPSv2.1 and UR-Align on FLUX.1-dev; FlowBP-Sparse leads SD3.5-M UR-IQ, while FlowBP-Bridge leads some FLUX.1-dev metrics.
  • Compositional generation: FlowBP-Lagrange improves FLUX.1-dev’s GenEval overall score from 63.25 to 69.88, achieving the best overall performance and best scores on three categories.FlowBP-Bridge is second overall and performs best on spatial position, indicating preference improvements do not degrade compositional prompt following.
  • Ablations: Endpoint-faithful reward-model inputs improve ReFL and DRTune across all metrics by reconstructing the sampled image endpoint without an endpoint-pinning residual.The cached rollout is replayed with selected active velocities re-forwarded so the forward endpoint equals the sampled image x0 by construction.
  • Ablations: FlowBP-Lagrange reduces connector error over long leaps through integrated Lagrange coefficients, producing more stable reward than Euler and equal-weight uniform connectors.The ablations also show that moderate nested coupling is beneficial, excessive coupling risks Jacobian amplification, and a compact active-step budget preserves gains while keeping the surrogate sparse.

6 Conclusion

FlowBP unifies surrogate-trajectory design for direct reward backpropagation in text-to-image flow matching models. Its variants control memory and Jacobian chaining while improving direct-gradient baselines across three backbones and multiple metric types.

  • FlowBP is a unified surrogate-trajectory framework for direct reward backpropagation in text-to-image flow matching models.
  • The framework separates reward-model input, active-set selection, integration weights, and bridge coupling, recovering prior methods as special cases.
  • FlowBP-Sparse, FlowBP-Bridge, and FlowBP-Lagrange bound memory by active velocities and avoid long Jacobian chains by construction.
  • Across three backbones and preference, quality, and compositional metrics, the variants consistently improve over direct-gradient baselines.

Appendix A. Additional Derivations

Appendix A derives each method’s unified gradient under the Section 4 notation. It formalizes cached detached rollouts and surrogate velocities parameterized by active indices, bridge coupling, and nested-gradient scaling.

  • Derivation setup: The derivations start from cached no-gradient rollouts containing detached states and velocities.These quantities follow the Section 4 notation and satisfy the cached Euler relation x_{i−1} = x_i − h_i v_i, with h_i = σ_i − σ_{i−1} > 0.
  • Surrogate parameterization: Each surrogate velocity ˜v_i is parameterized by an active set A = A_pre ∪ A_post, bridge index j, and nested-gradient scale α ∈ [0, 1].The bridge index satisfies 0 ≤ j < k, and all cached states remain constants during the second forward graph.
  • Surrogate constructions: Connector-based surrogates use a straight-through connector, whereas reconstruction-based surrogates replay cached Euler intervals with ˜v_i.Both alternatives are specified as constructions of the surrogate backward trajectory.

A.1 Surrogate-Fidelity Metrics

The section evaluates surrogate trajectories along reward-path fidelity and nested-coupling fidelity. These metrics distinguish forward sample quality from the gradient coupling retained across denoising steps.

  • Reward-Path Fidelity: Reward-path fidelity assesses active-step trajectory-weight coverage and endpoint-estimate accuracy against the generated sample.A large C requires both accurate endpoint reconstruction and sufficient rollout-weight coverage.
  • Nested-Coupling Fidelity: Nested-coupling fidelity measures whether the surrogate preserves an accurate gradient path across denoising steps.It is evaluated on the surrogate backward graph M, with dnest(M) counting the maximum inter-step Jacobian factors on any gradient path.
  • Nested-Coupling Fidelity: Full backpropagation has dnest(M) = N −1, detached surrogates have dnest(M) = 0, and bridged surrogates have dnest(M) = 1.Thus the metric distinguishes full, detached, and bridged gradient coupling.
  • Metric Interpretation: C describes the forward sample consumed by the reward model, whereas Fnest describes the backward coupling carried by the gradient.The two metrics separate the surrogate’s forward fidelity from its retained gradient coupling.

A.2 FlowBP-Lagrange

FlowBP-Lagrange replaces LeapAlign’s one-point leaps with integrated polynomial velocity approximations over two pinned segments. It controls gradient complexity through active supports, detached cached velocities, and attenuation of non-anchor contributions, leaving only one nested Jacobian term.

  • Integrated Lagrange quadrature: FlowBP-Lagrange replaces each LeapAlign one-point leap with an integrated polynomial approximation of the velocity field across pre- and post-bridge segments.The method retains LeapAlign’s two-segment structure and uses Lagrange quadrature supports on each segment.
  • Active-support control: The gradient-active supports are subsets of the quadrature supports, with their union bounded by the active-set budget K = |Apre ∪ Apost| ≤ M.Inactive supports still contribute detached cached velocities to the forward quadrature but are excluded from the backward graph.
  • Bridge coupling: The bridge anchor uses the bridge-case velocity, whereas other active supports use the regular active velocity evaluated at their corresponding latent and noise level.The just-pinned xj enters the bridge case when computing the bridge surrogate velocity.
  • Gradient attenuation: Non-anchor active supports are attenuated by ρi = gs, while anchor supports use ρi = 1, preserving the forward value while scaling their backward signal.For inactive supports, the surrogate velocity is detached, so they affect forward integration without contributing gradients.
  • Backward complexity: The resulting backward expression contains exactly one nested term with one Jacobian factor, scaled by α; with only segment-start anchors, each segment has a single support.This isolates the remaining chained contribution instead of propagating a full multi-step Jacobian product.

A.3 FlowBP-Sparse · A.4 FlowBP-Bridge

FlowBP-Sparse reconstructs the cached endpoint with Euler-weighted active terms and detached inactive velocities, eliminating inter-step Jacobian chaining. FlowBP-Bridge extends this trajectory with one controllable bridge whose coupling interpolates between no nested path and full one-Jacobian dependence.

  • A.3 FlowBP-Sparse: FlowBP-Sparse removes the bridge and replays the cached trajectory using Euler weights to reconstruct the endpoint.All active indices use the regular active form, while inactive indices use detached cached velocities.
  • A.3 FlowBP-Sparse: Inactive velocities telescope with active terms to match the cached x0 exactly, so FlowBP-Sparse requires no endpoint connector.Because active velocities are evaluated on detached inputs, no inter-step Jacobian appears.
  • A.3 FlowBP-Sparse: FlowBP-Sparse replaces the nested gradient path with a dense, fully decoupled collection of direct terms.
  • A.4 FlowBP-Bridge: FlowBP-Bridge adds one bridge to the full Euler-quadrature trajectory, splitting the rollout into pre- and post-segments around index j.The active sets Apre and Apost jointly define the active indices, with j serving as the bridge anchor.
  • A.4 FlowBP-Bridge: Inactive cached velocities make the reconstructed xj and x0 match the cached rollout states exactly in both FlowBP-Bridge segments.Their detached status prevents inactive velocities from contributing to the backward pass.
  • A.4 FlowBP-Bridge: At the bridge, the chain rule contributes both a direct term and a nested term through xj.This is the mechanism distinguishing bridge coupling from the fully decoupled sparse construction.
  • A.4 FlowBP-Bridge: Setting α = 0 removes the nested gradient path, whereas α = 1 preserves the full one-Jacobian bridge dependence.The resulting expression consists of direct active-step updates across both segments plus the bridge-induced nested term scaled by α.

Appendix B. Additional Experimental Details · Appendix C. Algorithmic Templates for Reward-Gradient Updates

Appendix B details the training, sampling, reconstruction, and evaluation procedures for FlowBP variants, while Appendix C formalizes posterior-mean and cached-rollout surrogate update graphs. Together, the appendices specify how active velocities, surrogate paths, and reward targets are selected and differentiated.

  • Appendix B. Additional Experimental Details: Dirichlet sampling with parameters (2.5, 6.0, 2.0) partitions the reverse-index range into three segments to stabilize leap-index selection.Uniformly sampling two indices can increase reward-gradient variance when endpoints are too close or near the noisy trajectory end.
  • Appendix B. Additional Experimental Details: All runs use HPSv2.1 as the training reward with 25-step online rollouts at 512 × 512 resolution.
  • Appendix B. Additional Experimental Details: The strict two-segment convention uses 0 < j < k ≤ N, while FlowBP-Bridge uses only split j and fixes the noisy-side endpoint at N.For SD3.5-M and FLUX.2-Klein-base, krev is capped at Kmax = 10 to avoid unstable noisy-end placements.
  • Appendix B. Additional Experimental Details: All FlowBP variants cache a no-gradient rollout, re-forward selected active velocities, and reconstruct surrogate paths using Euler, bridge-coupled Euler, or Lagrange supports.Active steps for FlowBP-Sparse and FlowBP-Bridge use a late-biased distribution, while shared gradient rescaling applies to active velocities.
  • Appendix B. Additional Experimental Details: Evaluation reports average HPDv2 scores and gains over the base model on a 400-prompt test split with 50-step sampling, alongside GenEval task accuracies and overall score.Figure 8 tracks PickScore, ImageReward, and HPSv2.1 over training across FLUX.1-dev, FLUX.2-Klein-base, and SD3.5-M.
  • C.2 Cached-Rollout Surrogate Updates: Cached-rollout surrogate updates run a full trajectory without gradients, expose selected velocities to autograd, and update parameters through active calls only.LeapAlign and FlowBP variants differ in index choices, active sets, weights, connector or reconstruction rules, and surrogate endpoint construction.
  • C.1 Posterior-Mean Reward Updates: Posterior-mean updates unify ReFL, DRTune, and DRaFT-LV through different stop indices, active velocity sets, and numbers of clean reward targets.DRaFT-LV constructs additional low-variance targets by re-noising the shared posterior-mean target.

Appendix D. Additional Qualitative Results

The appendix presents additional qualitative examples generated from HPDv2 test prompts for SD3.5-M, FLUX.1-dev, and FLUX.2-Klein-base, shown in Figures 9–11.

  • SD3.5-M: Additional qualitative results for SD3.5-M use prompts from the HPDv2 test split.These examples are presented in Figure 9.
  • FLUX.1-dev: Additional qualitative results for FLUX.1-dev use prompts from the HPDv2 test split.These examples are presented in Figure 10.
  • FLUX.2-Klein-base: Additional qualitative results for FLUX.2-Klein-base use prompts from the HPDv2 test split.These examples are presented in Figure 11.
Loading 2606.11075v1…