探索流匹配中獎勵反向傳播的設計空間

摘要

將文本到圖像流匹配模型與人類偏好對齊，透過直接獎勵反向傳播的方法雖然樣本效率高，但卻受到兩個已知病理問題的困擾：在現代模型規模下，無法在整個採樣軌跡中儲存激活值，且跨步驟的鏈式雅可比乘積會放大回傳至早期索引的獎勵梯度。基於連接器的方法（如 LeapAlign）透過以一條短固定路徑取代完整的反向軌跡來解決這些問題，凸顯了採樣與最佳化之間有用的解耦關係。然而，所得梯度的品質取決於這條短路徑對完整展開軌跡的近似精確度，尤其是在長間隔的情況下。我們提出 FlowBP，一個統一的替代軌跡框架，將反向軌跡本身視為設計對象。FlowBP 保留一條無梯度的快取展開軌跡用於採樣，再從快取及選擇性重新前向傳播的速度中建構一個輕量級的反向替代軌跡。此觀點分離了四個選擇：獎勵模型輸入、活躍集、積分權重，以及橋耦合，並將先前的直接梯度方法歸納為特定設定。在此框架內，我們實例化三個變體：FlowBP-Sparse 使用稀疏歐拉重構，FlowBP-Bridge 加入受控橋耦合，而 FlowBP-Lagrange 則提高跳躍求積的階數。三者均將記憶體限制在活躍集大小內，並將梯度鏈限制在最多一個雅可比因子。在 SD3.5-M、FLUX.1-dev 及 FLUX.2-Klein-base 上，針對偏好、品質及組合性指標，這三個變體在多數指標上優於直接梯度基線方法。

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.