Exploring the Design Space of Reward Backpropagation for Flow Matching

2026-06-09

AI Summary

This work addresses the challenges of high memory consumption and gradient chain explosion in direct reward backpropagation for text-to-image flow matching models. The authors propose FlowBP, a framework that treats the backward trajectory itself as a design object, decoupling sampling from optimization and unifying four key components: reward injection, activation sets, integral weighting, and bridge coupling. Through techniques including cached replay, sparse Euler reconstruction, bridge coupling control, and higher-order leapfrog integration, FlowBP effectively constrains both memory usage and gradient chain length. Experiments on SD3.5-M, FLUX.1-dev, and FLUX.2-Klein-base demonstrate that three variants of FlowBP consistently outperform direct gradient-based baselines across preference, quality, and compositional metrics.

Abstract

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.

Problem

• reward backpropagation
• flow matching
• gradient chaining
• memory constraints
• text-to-image generation

Innovation

• FlowBP
• reward backpropagation
• flow matching
• surrogate trajectory
• gradient chaining