Het verkennen van de ontwerpruimte van beloningsbackpropagatie voor Flow Matching

Samenvatting

Het afstemmen van tekst-naar-beeld stroommatchingmodellen op menselijke voorkeuren via directe reward-backpropagatie is sample-efficiënt, maar wordt belemmerd door twee bekende pathologieën: activeringen kunnen niet worden opgeslagen over het volledige samplingtraject bij moderne modelschaal, en aaneengeschakelde Jacobiaanproducten over stappen heen blazen de reward-gradient op terwijl deze terugreist naar vroege indices. Connector-gebaseerde methoden, zoals LeapAlign, pakken deze problemen aan door het volledige terugwaartse traject te vervangen door een kort vastgezet pad, wat een nuttige ontkoppeling tussen sampling en optimalisatie benadrukt. De kwaliteit van de resulterende gradient hangt echter af van hoe nauwkeurig dit korte pad het volledige rollout benadert, vooral over lange intervallen. We stellen FlowBP voor, een uniform surrogaattrajectraamwerk dat het terugwaartse traject zelf als ontwerpobject beschouwt. FlowBP bewaart een geen-gradient gecachte rollout voor sampling, en bouwt vervolgens een lichtgewicht terugwaarts surrogaat uit gecachte en selectief opnieuw doorgestuurde snelheden. Deze visie scheidt vier keuzes: de invoer van het rewardmodel, de actieve set, de integratiegewichten en de brugkoppeling, en herstelt eerdere directe-gradientmethoden als specifieke instellingen. Binnen dit raamwerk implementeren we drie varianten: FlowBP-Sparse gebruikt schaarse Euler-reconstructie, FlowBP-Bridge voegt gecontroleerde brugkoppeling toe, en FlowBP-Lagrange verhoogt de orde van de sprongkwadratuur. Alle drie begrenzen het geheugen door de grootte van de actieve set en beperken de gradientketening tot maximaal één Jacobiaanfactor. Over SD3.5-M, FLUX.1-dev en FLUX.2-Klein-base heen op het gebied van voorkeurs-, kwaliteits- en compositionele metrieken verbeteren de drie varianten ten opzichte van directe-gradientbaselines op de meeste metrieken.

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.