Исследование пространства проектирования обратного распространения вознаграждения для согласования потоков

Аннотация

Согласование моделей потокового сопоставления текст-изображение с человеческими предпочтениями с помощью прямого обратного распространения вознаграждения является эффективным с точки зрения выборки, но сдерживается двумя известными патологиями: активации не могут храниться на всей траектории выборки при современном масштабе модели, а цепные произведения Якоби по шагам раздувают градиент вознаграждения при его распространении обратно к начальным индексам. Методы на основе коннекторов, такие как 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.