LeapAlign: Ajuste Posterior de Modelos de Emparejamiento de Flujos en Cualquier Paso de Generación mediante la Construcción de Trayectorias de Dos Pasos

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

April 16, 2026

Autores: Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

cs.AI

Resumen

Este trabajo se centra en la alineación de modelos de ajuste de flujo (flow matching) con las preferencias humanas. Un enfoque prometedor es el ajuste fino mediante la retropropagación directa de los gradientes de recompensa a través del proceso de generación diferenciable del ajuste de flujo. Sin embargo, la retropropagación a través de trayectorias largas genera costos de memoria prohibitivos y explosión del gradiente. En consecuencia, los métodos de gradiente directo tienen dificultades para actualizar los pasos iniciales de la generación, los cuales son cruciales para determinar la estructura global de la imagen final. Para abordar este problema, presentamos LeapAlign, un método de ajuste fino que reduce el coste computacional y permite la propagación directa del gradiente desde la recompensa hasta los pasos iniciales de la generación. Específicamente, acortamos la trayectoria larga a solo dos pasos mediante el diseño de dos saltos consecutivos, cada uno omitiendo múltiples pasos de muestreo de EDO y prediciendo los latentes futuros en un solo paso. Al aleatorizar los pasos de tiempo de inicio y fin de los saltos, LeapAlign permite actualizaciones del modelo eficientes y estables en cualquier paso de la generación. Para aprovechar mejor estas trayectorias acortadas, asignamos pesos de entrenamiento más altos a aquellas que son más consistentes con la ruta de generación larga. Para mejorar aún más la estabilidad del gradiente, reducimos los pesos de los términos de gradiente con magnitud grande, en lugar de eliminarlos por completo como se hacía en trabajos anteriores. Al aplicar el ajuste fino al modelo Flux, LeapAlign supera consistentemente a los métodos state-of-the-art basados en GRPO y de gradiente directo en varias métricas, logrando una calidad de imagen y una alineación imagen-texto superiores.

English

This paper focuses on the alignment of flow matching models with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients through the differentiable generation process of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion. Therefore, direct-gradient methods struggle to update early generation steps, which are crucial for determining the global structure of the final image. To address this issue, we introduce LeapAlign, a fine-tuning method that reduces computational cost and enables direct gradient propagation from reward to early generation steps. Specifically, we shorten the long trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing the start and end timesteps of the leaps, LeapAlign leads to efficient and stable model updates at any generation step. To better use such shortened trajectories, we assign higher training weights to those that are more consistent with the long generation path. To further enhance gradient stability, we reduce the weights of gradient terms with large magnitude, instead of completely removing them as done in previous works. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods across various metrics, achieving superior image quality and image-text alignment.