LeapAlign: Nachträgliches Anpassen von Flow-Matching-Modellen bei beliebigem Generierungsschritt durch Aufbau von Zwei-Schritt-Trajektorien
LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories
April 16, 2026
Autoren: Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng
cs.AI
Zusammenfassung
Dieses Papier konzentriert sich auf die Ausrichtung von Flow-Matching-Modellen an menschlichen Präferenzen. Ein vielversprechender Ansatz ist das Fine-Tuning durch direktes Backpropagieren von Belohnungsgradienten durch den differenzierbaren Generierungsprozess des Flow Matching. Das Backpropagieren über lange Trajektorien führt jedoch zu prohibitiv hohen Speicherkosten und Gradient Explosion. Daher haben Direct-Gradient-Methoden Schwierigkeiten, frühe Generierungsschritte zu aktualisieren, die für die Bestimmung der globalen Struktur des endgültigen Bildes entscheidend sind. Um dieses Problem zu lösen, stellen wir LeapAlign vor, eine Fine-Tuning-Methode, die die Rechenkosten senkt und eine direkte Gradientenpropagation von der Belohnung zu frühen Generierungsschritten ermöglicht. Konkret verkürzen wir die lange Trajektorie auf nur zwei Schritte, indem wir zwei aufeinanderfolgende "Sprünge" entwerfen, von denen jeder mehrere ODE-Sampling-Schritte überspringt und zukünftige Latents in einem einzigen Schritt vorhersagt. Durch die Randomisierung der Start- und Endzeitpunkte der Sprünge ermöglicht LeapAlign effiziente und stabile Modellaktualisierungen in jedem Generierungsschritt. Um solche verkürzten Trajektorien besser nutzen zu können, weisen wir denen höhere Trainingsgewichte zu, die konsistenter mit dem langen Generierungspfad sind. Um die Gradientenstabilität weiter zu verbessern, reduzieren wir die Gewichte von Gradiententermen mit großer Magnitude, anstatt sie vollständig zu entfernen, wie es in früheren Arbeiten der Fall war. Beim Fine-Tuning des Flux-Modells übertrifft LeapAlign durchgängig state-of-the-art GRPO-basierte und Direct-Gradient-Methoden in verschiedenen Metriken und erzielt eine überlegene Bildqualität und Bild-Text-Ausrichtung.
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.