Explorer l'espace de conception de la rétropropagation de la récompense pour le flow matching

Résumé

L'alignement des modèles de flow matching texte-image sur les préférences humaines via la rétropropagation directe de la récompense est efficace en termes d'échantillons, mais entravé par deux pathologies bien connues : les activations ne peuvent pas être stockées sur l'ensemble de la trajectoire d'échantillonnage à l'échelle des modèles modernes, et les produits de Jacobiens en chaîne à travers les étapes gonflent le gradient de la récompense lorsqu'il remonte vers les indices précoces. Les méthodes basées sur des connecteurs, telles que LeapAlign, abordent ces problèmes en remplaçant la trajectoire complète de rétropropagation par un chemin fixe court, mettant en évidence un découplage utile entre l'échantillonnage et l'optimisation. Cependant, la qualité du gradient résultant dépend de la précision avec laquelle ce chemin court approxime le déploiement complet, en particulier sur de longs intervalles. Nous proposons FlowBP, un cadre unifié de trajectoire de substitution qui traite la trajectoire de rétropropagation elle-même comme un objet de conception. FlowBP conserve un déploiement mis en cache sans gradient pour l'échantillonnage, puis construit un substitut léger de rétropropagation à partir des vitesses mises en cache et sélectivement ré-avançées. Cette vision sépare quatre choix : l'entrée du modèle de récompense, l'ensemble actif, les poids d'intégration et le couplage de pont, et retrouve les méthodes antérieures de gradient direct comme des cas particuliers. Dans ce cadre, nous instancions trois variantes : FlowBP-Sparse utilise une reconstruction d'Euler parcimonieuse, FlowBP-Bridge ajoute un couplage de pont contrôlé, et FlowBP-Lagrange augmente l'ordre de la quadrature leap. Les trois limitent la mémoire à la taille de l'ensemble actif et contraignent le chaînage du gradient à au plus un facteur Jacobien. Sur les modèles SD3.5-M, FLUX.1-dev et FLUX.2-Klein-base, en termes de préférence, de qualité et de métriques compositionnelles, les trois variantes améliorent les références de gradient direct sur la plupart des métriques.

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.