Explorando el espacio de diseño de la retropropagación de recompensa para Flow Matching

Resumen

Alinear modelos de flujo de texto a imagen con preferencias humanas mediante retropropagación directa de recompensas es eficiente en términos de muestras, pero se ve obstaculizado por dos patologías bien conocidas: las activaciones no pueden almacenarse a lo largo de toda la trayectoria de muestreo a la escala de los modelos modernos, y los productos jacobianos encadenados a través de los pasos inflan el gradiente de recompensa a medida que se propaga hacia los índices iniciales. Los métodos basados en conectores, como LeapAlign, abordan estos problemas reemplazando la trayectoria completa hacia atrás por un camino corto y fijo, destacando un desacoplamiento útil entre el muestreo y la optimización. Sin embargo, la calidad del gradiente resultante depende de qué tan precisamente este camino corto aproxima el despliegue completo, especialmente en intervalos largos. Proponemos FlowBP, un marco unificado de trayectorias sustitutas que trata la trayectoria hacia atrás como el objeto de diseño. FlowBP mantiene un despliegue en caché sin gradiente para el muestreo, y luego construye un sustituto ligero hacia atrás a partir de velocidades en caché y selectivamente reenviadas. Esta visión separa cuatro opciones: la entrada del modelo de recompensa, el conjunto activo, los pesos de integración y el acoplamiento puente, y recupera métodos previos de gradiente directo como configuraciones particulares. Dentro de este marco, instanciamos tres variantes: FlowBP-Sparse usa reconstrucción de Euler dispersa, FlowBP-Bridge añade acoplamiento puente controlado, y FlowBP-Lagrange eleva el orden de la cuadratura de salto. Las tres limitan la memoria por el tamaño del conjunto activo y limitan el encadenamiento de gradientes a como máximo un factor jacobiano. En SD3.5-M, FLUX.1-dev y FLUX.2-Klein-base, en métricas de preferencia, calidad y composición, las tres variantes mejoran con respecto a las líneas base de gradiente directo en la mayoría de las métricas.

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.