探索奖励反向传播在流匹配中的设计空间
Exploring the Design Space of Reward Backpropagation for Flow Matching
摘要
通过直接奖励反向传播的方式,将文本到图像生成模型与人类偏好进行匹配,虽然能够节省样本量,但存在两个常见问题:在现代模型规模下,激活值无法被存储在整个采样过程中;此外,不同步骤之间的雅可比乘积会使得奖励梯度在返回到早期索引时变得过大。像LeapAlign这样的基于连接器的方法通过用一条短路径替代完整的反向路径来解决这些问题,从而实现了采样与优化之间的有效分离。不过,最终得到的梯度的质量取决于这条短路径对完整采样过程的近似程度,尤其是在长时间段内。我们提出了FlowBP这一统一的替代路径框架,它将反向路径本身作为设计对象。FlowBP为采样过程保留一个无梯度的缓存结果,然后利用这些缓存的、经过选择性重新传递的速度信息来构建轻量级的反向替代路径。这种方案将四个因素分开处理:奖励-模型输入、活跃集、积分权重以及桥梁耦合机制,从而将之前的直接梯度方法转化为特定配置下的解决方案。在这个框架中,我们提出了三种变体:FlowBP-Sparse使用稀疏欧拉重建方法,FlowBP-Bridge则引入了可控的桥梁耦合机制,而FlowBP-Lagrange则提升了跃变积分的精度。这三种变体都限制了内存占用,使其不超过活跃集的大小,同时将梯度链的长度限制在一个雅可比因子以内。在SD3.5-M、FLUX.1-dev和FLUX.2-Klein-base三个数据集上,从偏好、质量和组合指标来看,这三种变体在大多数指标上都优于传统的直接梯度方法。
English 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.