提出新框架FlowBP,解决文本生成模型奖励反传的内存与梯度爆炸问题。
Exploring the Design Space of Reward Backpropagation for Flow Matching

- 用缓存轨迹+轻量反向代理替代完整反向传播,降低内存占用。
- 三种变体在多个基准上优于直接梯度方法,尤其在长序列任务中更稳定。
- 适合大规模文本到图像生成模型的高效偏好对齐,可扩展性强。
通过直接奖励反传对齐文本到图像流匹配模型与人类偏好具有样本效率高优势,但面临两大难题:现代模型规模下无法存储整个采样轨迹的激活值,且跨步骤的链式雅可比乘积导致奖励梯度在回传至早期索引时被放大。基于连接器的方法(如LeapAlign)通过用短固定路径替代完整反向轨迹解决了这些问题,凸显了采样与优化之间的解耦。然而,所得梯度质量依赖于该短路径对完整轨迹的近似程度,尤其在长间隔下表现不佳。本文提出统一的代理轨迹框架FlowBP,将反向轨迹本身作为设计对象:保留无梯度的缓存采样轨迹,再从缓存和选择性重前向的速度构建轻量级反向代理。该视角分离出四个设计选择:奖励模型输入、活跃集、积分权重与桥接耦合,并可还原先前的直接梯度方法为特定设置。在此框架下,我们实现三种变体:FlowBP-Sparse采用稀疏欧拉重构,FlowBP-Bridge引入可控桥接耦合,FlowBP-Lagrange提升跳跃求积阶数。三者均将内存限制在活跃集大小内,且将梯度链长度控制在最多一个雅可比因子。在SD3.5-M、FLUX.1-dev和FLUX.2-Klein-base三个数据集上,三种变体在偏好、质量与组合性指标上多数超越直接梯度基线。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。