LeapAlign: Allineamento Post-Addestramento di Modelli di Flow Matching a Qualsiasi Passo di Generazione mediante la Costruzione di Traiettorie a Due Passi

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

April 16, 2026

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

cs.AI

Abstract

Questo lavoro si concentra sull'allineamento dei modelli di flusso con le preferenze umane. Un approccio promettente consiste nella messa a punto mediante la retropropagazione diretta dei gradienti di ricompensa attraverso il processo di generazione differenziabile del flusso. Tuttavia, la retropropagazione attraverso traiettorie lunghe comporta costi di memoria proibitivi ed esplosione del gradiente. Di conseguenza, i metodi a gradiente diretto faticano ad aggiornare le fasi iniziali della generazione, cruciali per determinare la struttura globale dell'immagine finale. Per affrontare questo problema, introduciamo LeapAlign, un metodo di messa a punto che riduce il costo computazionale e permette la propagazione diretta del gradiente dalla ricompensa alle fasi iniziali della generazione. Nello specifico, accorciamo la lunga traiettoria in soli due passi progettando due salti consecutivi, ciascuno dei quali salta multipli passi di campionamento ODE e predice i latent futuri in un singolo passo. Randomizzando i tempi di inizio e fine dei salti, LeapAlign consente aggiornamenti del modello efficienti e stabili in qualsiasi fase della generazione. Per utilizzare al meglio tali traiettorie abbreviate, assegniamo pesi di addestramento più elevati a quelle più coerenti con il percorso di generazione lungo. Per migliorare ulteriormente la stabilità del gradiente, riduciamo i pesi dei termini di gradiente con magnitudo elevata, invece di rimuoverli completamente come fatto in lavori precedenti. Durante la messa a punto del modello Flux, LeapAlign supera costantemente i metodi all'avanguardia basati su GRPO e a gradiente diretto in varie metriche, raggiungendo una qualità dell'immagine e un allineamento immagine-testo superiori.

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.