LeapAlign: Nabewerking van Flow Matching-modellen op Willekeurige Generatiestappen door Bouw van Twee-Staps Trajectorieën
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
Samenvatting
Dit artikel richt zich op de afstemming van flow matching-modellen op menselijke voorkeuren. Een veelbelovende aanpak is finetunen door beloningsgradiënten direct terug te propageren door het differentieerbare generatieproces van flow matching. Het terugpropageren door lange trajecten resulteert echter in onhoudbare geheugenkosten en gradiëntexplosie. Daarom hebben directe-gradiëntmethoden moeite met het updaten van vroege generatiestappen, die cruciaal zijn voor het bepalen van de globale structuur van de uiteindelijke afbeelding. Om dit probleem aan te pakken, introduceren wij LeapAlign, een finetunemethode die de rekenkosten verlaagt en directe gradiëntpropagatie van beloning naar vroege generatiestappen mogelijk maakt. Concreet verkorten we het lange traject tot slechts twee stappen door twee opeenvolgende sprongen te ontwerpen, waarbij elke sprong meerdere ODE-bemonsteringsstappen overslaat en toekomstige latente representaties in één stap voorspelt. Door de start- en eindtijdstappen van de sprongen te randomiseren, leidt LeapAlign tot efficiënte en stabiele modelupdates bij elke generatiestap. Om dergelijke verkorte trajecten beter te benutten, kennen we hogere trainingsgewichten toe aan trajecten die consistenter zijn met het lange generatiepad. Om de gradiëntstabiliteit verder te verbeteren, verlagen we de gewichten van gradiënttermen met een grote omvang, in plaats van ze volledig te verwijderen zoals in eerdere werken. Bij het finetunen van het Flux-model presteert LeapAlign consistent beter dan state-of-the-art GRPO-gebaseerde en directe-gradiëntmethoden op verschillende metrieken, en bereikt het superieure beeldkwaliteit en beeld-tekstafstemming.
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.