Erkundung des Designraums der Reward-Backpropagation für Flow Matching

Zusammenfassung

Die Ausrichtung von Text-zu-Bild-Flow-Matching-Modellen an menschlichen Präferenzen mittels direkter Belohnungs-Rückpropagation ist stichprobeneffizient, wird jedoch durch zwei bekannte Pathologien behindert: Aktivierungen können in modernen Modellskalen nicht über die gesamte Sampling-Trajektorie gespeichert werden, und verkettete Jacobi-Produkte über mehrere Schritte blähen den Belohnungsgradienten auf, während er zu frühen Indizes zurückwandert. Connector-basierte Methoden wie LeapAlign begegnen diesen Problemen, indem sie die vollständige Rückwärtstrajektorie durch einen kurzen fixierten Pfad ersetzen und so eine nützliche Entkopplung von Sampling und Optimierung aufzeigen. Allerdings hängt die Qualität des resultierenden Gradienten davon ab, wie genau dieser kurze Pfad den vollständigen Rollout approximiert, insbesondere über lange Intervalle. Wir schlagen FlowBP vor, ein einheitliches Surrogat-Trajektorien-Framework, das die Rückwärtstrajektorie selbst als Designobjekt behandelt. FlowBP behält einen gradientenfreien zwischengespeicherten Rollout für das Sampling bei und baut dann ein leichtgewichtiges Rückwärtssurrogat aus zwischengespeicherten und selektiv erneut weitergeleiteten Geschwindigkeiten auf. Diese Sichtweise trennt vier Entscheidungen: die Eingabe des Belohnungsmodells, die aktive Menge, die Integrationsgewichte und die Bridge-Kopplung, und fasst frühere direkte Gradientenmethoden als spezielle Fälle. Innerhalb dieses Frameworks instantiieren wir drei Varianten: FlowBP-Sparse verwendet eine dünnbesetzte Euler-Rekonstruktion, FlowBP-Bridge fügt eine kontrollierte Bridge-Kopplung hinzu und FlowBP-Lagrange erhöht die Ordnung der Sprungquadratur. Alle drei begrenzen den Speicher durch die Größe der aktiven Menge und beschränken die Gradientenverkettung auf höchstens einen Jacobi-Faktor. Über SD3.5-M, FLUX.1-dev und FLUX.2-Klein-base hinweg erzielen die drei Varianten bei Präferenz-, Qualitäts- und Kompositionsmetriken Verbesserungen gegenüber direkten Gradienten-Baselines bei den meisten Metriken.

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.