LeapAlign : Alignement post-entraînement des modèles de correspondance de flux à n'importe quelle étape de génération via la construction de trajectoires en deux étapes

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

April 16, 2026

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

cs.AI

Résumé

Ce travail se concentre sur l'alignement des modèles de « flow matching » avec les préférences humaines. Une approche prometteuse consiste à effectuer un ajustement fin en rétropropagant directement les gradients de récompense à travers le processus de génération différentiable du « flow matching ». Cependant, la rétropropagation à travers de longues trajectoires entraîne des coûts mémoire prohibitifs et une explosion des gradients. Par conséquent, les méthodes à gradient direct peinent à mettre à jour les premières étapes de génération, qui sont cruciales pour déterminer la structure globale de l'image finale. Pour résoudre ce problème, nous présentons LeapAlign, une méthode d'ajustement fin qui réduit le coût computationnel et permet la propagation directe du gradient de la récompense vers les premières étapes de génération. Concrètement, nous raccourcissons la longue trajectoire en seulement deux étapes en concevant deux « sauts » consécutifs, chacun sautant plusieurs étapes d'échantillonnage d'EDO et prédisant les latents futurs en une seule étape. En randomisant les pas de temps de début et de fin des sauts, LeapAlign permet des mises à jour du modèle efficaces et stables à n'importe quelle étape de la génération. Pour mieux utiliser ces trajectoires raccourcies, nous attribuons des poids d'entraînement plus élevés à celles qui sont plus cohérentes avec le long chemin de génération. Pour améliorer davantage la stabilité du gradient, nous réduisons le poids des termes de gradient de grande magnitude, au lieu de les supprimer complètement comme le font les travaux antérieurs. Lors de l'ajustement fin du modèle Flux, LeapAlign surpasse constamment les méthodes état de l'art basées sur GRPO et à gradient direct sur diverses métriques, obtenant une qualité d'image et un alignement image-texte supérieurs.

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.