💻 computer science

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

LeapAlign ist eine Nachtrainierungs-Feinabstimmungsmethode für Flow-Matching-Modelle, die die Speicher- und Gradientenexplosionsprobleme der Backpropagation über lange Trajektorien durch die Konstruktion randomisierter Zwei-Schritt-Sprung-Trajektorien überwindet und damit effiziente und stabile direkte Gradientenupdates von Belohnungen zu frühen Generierungsschritten ermöglicht, um eine überlegene Bildqualität und Ausrichtung zu erreichen.

Ursprüngliche Autoren: Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

Veröffentlicht 2026-05-05

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Zhanhao Liang, Tao Yang, Jie Wu, Chengjian Feng, Liang Zheng

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie unterrichten einen sehr talentierten Künstler (das KI-Modell), Bilder basierend auf Ihren Beschreibungen zu zeichnen. Der Künstler beginnt mit einer leeren, verrauschten Leinwand und verfeinert sie schrittweise, bis ein klares Bild entsteht. Dieser Prozess wird „Flow Matching" genannt.

Das Problem, das die Arbeit angeht, lautet: Wie bringen wir diesem Künstler bei, genau das zu zeichnen, was wir wollen, insbesondere das Gesamtbild (das Layout), ohne uns mit dem Computerspeicher den Kopf zu zerbrechen?

Hier ist die einfache Aufschlüsselung der Lösung der Arbeit, LeapAlign:

1. Das Problem: Der „lange Weg" ist zu schwer

Um den Künstler zu unterrichten, zeigen wir ihm normalerweise eine fertige Zeichnung, sagen ihm, wie gut sie ist (eine „Belohnung"), und versuchen dann, dieses Feedback den gesamten Weg zurück zum allerersten Schritt des Zeichnungsprozesses zu senden, um Fehler zu korrigieren.

Das Problem: Wenn die Zeichnung 25 Schritte benötigt, um fertig zu werden, ist es, als würde man versuchen, eine Nachricht vom Gipfel eines Berges hinunter in das Tal zu schreien. Es erfordert eine enorme Menge an Energie (Computerspeicher), und die Nachricht wird oft verzerrt oder explodiert (Gradientenexplosion), bis sie den Boden erreicht.
Die Folge: Da es zu schwierig ist, Feedback an den Anfang zu senden, korrigieren die meisten aktuellen Methoden nur die letzten paar Schritte der Zeichnung. Sie justieren die Details, lassen aber das Layout (wo der Hund ist, wo der Baum ist) unverändert. Wenn das Layout falsch ist, ist das Bild falsch, egal wie hübsch die Details sind.

2. Die Lösung: Der „Sprung"-Abkürzungsweg

Die Autoren, Zhanhao Liang und Tao Yang, haben eine Methode namens LeapAlign entwickelt. Anstatt den ganzen langen Weg rückwärts zu gehen, um Feedback zu geben, bauen sie eine Abkürzung.

Stellen Sie sich den Zeichnungsprozess als eine lange Treppe vom oberen Ende (Rauschen) zum unteren Ende (fertiges Bild) vor.

Alter Weg: Gehen Sie jeden einzelnen Schritt hinunter, um Feedback zu geben. (Zu langsam, zu schwer).
LeapAlign-Weg: Sie wählen zwei zufällige Stellen auf der Treppe aus – sagen wir, Schritt 20 und Schritt 10. Anstatt den ganzen Weg zu gehen, springen Sie von Schritt 20 direkt zu Schritt 10 und dann von Schritt 10 direkt zum fertigen Bild.

Indem diese zweistufige „Sprung-Trajektorie" erstellt wird, muss der Computer nur zwei Schritte der Reise im Gedächtnis behalten anstatt fünfundzwanzig. Dies spart enorme Mengen an Speicher und verhindert, dass die Nachricht explodiert.

3. Warum das eine große Sache ist

Da der „Sprung" von jedem Punkt auf der Treppe starten kann (zufällig gewählt), kann das System nun Feedback an den aller Anfang des Zeichnungsprozesses senden.

Das Ergebnis: Die KI lernt, die globale Struktur (das Layout) genauso gut zu korrigieren wie die Details. Sie lernt, das „rote Fahrrad" links und das „blaue Auto" rechts zu platzieren, anstatt nur dafür zu sorgen, dass das Fahrrad glänzend aussieht.

4. Zwei geheime Zutaten für Stabilität

Die Arbeit erwähnt zwei clevere Tricks, um diesen Sprung reibungslos funktionieren zu lassen:

Der „Lautstärkeregler" (Gradientenabsenkung): Manchmal wird das Signal, wenn Sie Feedback durch einen Sprung senden, zu laut (zu stark) und zerstört den Lernprozess. Frühere Methoden schalteten das Signal einfach komplett aus. LeapAlign dreht den Lautstärkeregler leicht herunter. Es hält das Signal aufrecht (damit die KI die Verbindung zwischen den Schritten lernt), dämpft die Lautstärke jedoch, damit die Lautsprecher nicht durchbrennen.
Der „Vertrauens-Score" (Trajektorien-Ähnlichkeits-Gewichtung): Manchmal ist ein Sprung eine seltsame Abkürzung, die nicht wie ein echter Zeichnungsweg aussieht. LeapAlign prüft: „Sieht diese Abkürzung wie ein normaler Zeichnungsweg aus?" Wenn ja, gibt es dem Feedback mehr Gewicht. Wenn die Abkürzung seltsam ist, gibt es ihr weniger Gewicht. Dies stellt sicher, dass die KI aus den besten Beispielen lernt.

5. Die Ergebnisse

Die Autoren testeten dies an einem leistungsstarken Modell namens Flux.

Das Ergebnis: LeapAlign schlug konsistent andere Top-Methoden (wie GRPO und DRTune).
Der Beweis: Es generierte Bilder, die nicht nur hübscher waren, sondern auch den Textanweisungen viel besser folgten. Wenn Sie zum Beispiel nach „eine Katze, die auf einer Matte sitzt" fragten, platzierte LeapAlign die Katze tatsächlich auf der Matte, während andere Methoden die Katze oft schwebend in der Luft oder neben der Matte darstellten.

Zusammenfassend: LeapAlign ist eine intelligente Art, KI-Künstler zu unterrichten, indem kurze, effiziente Abkürzungen für Feedback erstellt werden. Dies ermöglicht es der KI, zu lernen, das gesamte Bild von Anfang an zu planen, was zu Bildern führt, die sowohl schön als auch genau das sind, was Sie verlangt haben.

Technische Zusammenfassung: LeapAlign

Problemstellung

Die Arbeit adressiert die Herausforderung, Flow-Matching-Modelle durch Nachtraining mit menschlichen Präferenzen in Einklang zu bringen. Während direkte Gradientenmethoden eine vielversprechende Alternative zu Policy-Gradient-Ansätzen (wie GRPO) darstellen, indem sie die Differenzierbarkeit des Flow-Matching-Samplings nutzen, stoßen sie bei der Anwendung auf lange Generierungstrajektorien auf erhebliche Hürden. Insbesondere verursacht die Rückwärtspropagierung von Belohnungsgradienten durch die gesamte Sequenz von ODE-Sampling-Schritten prohibitive Speicherkosten und birgt das Risiko einer Gradientenexplosion. Folglich beschränken bestehende direkte Gradientenmethoden Aktualisierungen typischerweise auf einen einzigen Zeitschritt in der Nähe des finalen Bildes oder verwerfen verschachtelte Gradienten, um Instabilität zu vermeiden. Diese Einschränkung verhindert die Optimierung früher Generierungsschritte, die entscheidend für die Bestimmung des globalen Layouts und der strukturellen Komposition des finalen Bildes sind.

Methodik: LeapAlign

Die Autoren schlagen LeapAlign vor, eine Feinabstimmungsmethode, die eine effiziente und stabile direkte Gradientenpropagierung zu jedem Generierungsschritt, einschließlich früher Zeitschritte, ermöglicht, ohne übermäßige Speicherkosten zu verursachen. Die Kernmethodik umfasst drei Hauptkomponenten:

1. Zwei-Schritt-Leap-Trajektorien-Konstruktion

Anstatt durch die gesamte Generierungstrajektorie zurückzupropagieren, konstruiert LeapAlign eine verkürzte Leap-Trajektorie, die nur aus zwei Schritten besteht, die von einer standardmäßigen Volllauf-Trajektorie abgeleitet werden.

Prozess: Während jeder Iteration wird eine vollständige Trajektorie von Rauschen ( $x_1$ ) zu einem sauberen Bild ( $x_0$ ) gesampelt. Zwei Zeitschritte, $k$ und $j$ (wobei $k > j$ ), werden zufällig ausgewählt.
Sprünge: Die Methode nutzt die Eigenschaft der Ein-Schritt-Leap-Vorhersage von rektifizierten Flow-Modellen, um zukünftige Latente in einem einzigen Schritt zu schätzen:
- Sprung 1: Vorhersage des Latenten $x_j$ aus $x_k$ unter Verwendung des Geschwindigkeitsfelds bei $x_k$ : $\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$ .
- Sprung 2: Vorhersage des finalen Latenten $x_0$ aus $x_j$ unter Verwendung des Geschwindigkeitsfelds bei $x_j$ : $\hat{x}_{0|j} = x_j - j v_\theta(x_j)$ .
Latente Konnektoren: Um die Differenzierbarkeit beizubehalten, während vorhergesagte Zustände mit tatsächlichen, aus der Volltrajektorie gesampelten Zuständen abgeglichen werden, werden „latente Konnektoren" unter Verwendung von stop_gradient-Operationen eingeführt (z. B. $x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$ ).
Vorteil: Dies reduziert den Rückwärtspropagierungspfad auf zwei Schritte, hält die Speicherkosten unabhängig von der ursprünglichen Trajektorienlänge konstant und ermöglicht gleichzeitig die zufällige Auswahl von $(k, j)$ , um jeden Teil des Generierungsprozesses zu aktualisieren.

2. Gradienten-Discounting

Um Gradientenexplosion zu adressieren und gleichzeitig nützliche Lernsignale zu erhalten, führt LeapAlign einen Gradienten-Discounting-Mechanismus ein.

Das Problem: Die Rückwärtspropagierung durch zwei Schritte führt zu einem „verschachtelten Gradienten"-Term ( $\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$ ), der Interaktionen über Zeitschritte hinweg erfasst. Bisherige Methoden wie DRTune eliminieren diesen Term vollständig, indem sie Gradienten am Modelleingang stoppen, was wertvolle Abhängigkeitsinformationen verwirft.
Die Lösung: LeapAlign skaliert den verschachtelten Gradiententerm mit einem Faktor $\alpha \in [0, 1]$ , anstatt ihn zu entfernen. Dies wird erreicht, indem die Eingabe für die Geschwindigkeitsvorhersage im zweiten Sprung modifiziert wird: $\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$ .
Wirkung: Dies reduziert die Magnitude großer Gradiententerme, um das Training zu stabilisieren, behält jedoch die vollständige Gradientenstruktur, einschließlich des verschachtelten Gradienten, bei.

3. Trajektorien-Ähnlichkeits-Gewichtung

Um sicherzustellen, dass die verkürzten Leap-Trajektorien eine zuverlässige Überwachung bieten, weist die Methode Leap-Trajektorien, die dem wahren Generierungspfad stark ähneln, höhere Gewichte zu.

Mechanismus: Der Verlust wird mit dem Kehrwert des durchschnittlichen absoluten Unterschieds zwischen den vorhergesagten Zuständen ( $\hat{x}$ ) und den tatsächlichen Zuständen ( $x$ ) an den Verbindungspunkten ( $x_j$ und $x_0$ ) gewichtet.
Formel: $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$ , wobei $d$ die Distanz darstellt und $\tau$ eine minimale Schwelle ist, um eine übermäßige Betonung nahezu identischer Paare zu verhindern.
Ziel: Dies stellt sicher, dass Sprünge, die erheblich von der wahren Dynamik abweichen, weniger zum Gradientenupdate beitragen, was die Stabilität und Zuverlässigkeit der Überwachung verbessert.

4. Feinabstimmungs-Ziel

Die Methode verwendet einen Hinge-Style-Verlust, um Belohnungen zu maximieren und gleichzeitig Reward-Hacking zu verhindern: $L_{raw} = \max(0, \lambda - r(x_0))$ . Im Gegensatz zu Methoden, die Belohnungen aus Leap-Vorhersagen schätzen, bewertet LeapAlign die Belohnung am tatsächlichen finalen Bild ( $x_0$ ), das von der Volltrajektorie generiert wird. Dies stellt sicher, dass das Belohnungsmodell die wahre visuelle und semantische Qualität beurteilt.

Hauptbeiträge

LeapAlign-Rahmenwerk: Eine neuartige Nachtrainingsmethode, die Zwei-Schritt-Leap-Trajektorien konstruiert, um eine direkte Gradientenrückwärtspropagierung zu frühen Generierungsschritten mit konstanten Speicherkosten zu ermöglichen.
Techniken zur Gradientenstabilität: Die Einführung von Gradienten-Discounting, um verschachtelte Gradienten zu skalieren (anstatt sie zu entfernen), sowie Trajektorien-Ähnlichkeits-Gewichtung, um konsistente Sprünge zu priorisieren. Dies adressiert Instabilitäts- und Signalverlustprobleme bei früheren direkten Gradientenmethoden.
Umfassende Evaluierung: Nachweis, dass LeapAlign über mehrere Metriken hinweg, einschließlich menschlicher Präferenzwerte (HPSv2.1, HPSv3, PickScore) und Benchmarks zur kompositorischen Ausrichtung (GenEval), State-of-the-Art-Policy-Gradient- (GRPO-basiert) und direkte Gradientenmethoden übertrifft.

Ergebnisse

Die Autoren haben das Flux.1-Dev-Modell mit LeapAlign feinabgestimmt und gegen Baselines wie DanceGRPO, MixGRPO, ReFL, DRaFT-LV und DRTune evaluiert.

Allgemeine Präferenz-Ausrichtung: LeapAlign erzielte die höchsten Werte bei in-domain (HPSv2.1) und out-of-domain (HPSv3, PickScore, UnifiedReward) Metriken. Bemerkenswerterweise übertraf es MixGRPO (welches drei Belohnungsmodelle verwendet), obwohl es nur mit einem einzigen Belohnungsmodell (HPSv2.1) trainiert wurde.
Kompositorische Ausrichtung: Auf dem GenEval-Benchmark zeigte LeapAlign signifikante Verbesserungen, insbesondere bei den Aufgaben „zwei Objekte", „Farben", „Position" und „Attributbindung". Die Verbesserung in diesen Kategorien unterstreicht den Erfolg der Methode bei der Optimierung früher Generierungsschritte, die das globale Layout bestimmen.
Qualitative Analyse: Visuelle Vergleiche zeigten, dass LeapAlign globale Bildstrukturen erheblich modifiziert, um besser mit Textprompts übereinzustimmen, während Methoden, die auf späte-Schritt-Aktualisierungen beschränkt sind (wie ReFL), oft das Layout des vortrainierten Modells beibehielten.
Ablationsstudien: Experimente bestätigten, dass das Zwei-Schritt-Leap-Design den besten Kompromiss zwischen Leistung und Speicher bietet. Gradienten-Discounting mit $\alpha=0.3$ ergab optimale Ergebnisse, und die zufällige Auswahl von Zeitschritten über den gesamten Bereich $[0, 1]$ war überlegenen festen Bereichen, was die Bedeutung von frühen-Schritt-Aktualisierungen unterstreicht.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass LeapAlign erfolgreich den Kompromiss zwischen Speichereffizienz und Gradientenabdeckung beim Nachtraining von Flow-Matching-Modellen auflöst. Durch die Ermöglichung von Aktualisierungen früher Generierungsschritte ohne prohibitive Rechenkosten ermöglicht es die Optimierung der globalen Bildstruktur, was für komplexe kompositorische Aufgaben entscheidend ist. Die Autoren positionieren LeapAlign als stabile und effektive Alternative sowohl zu stochastischen Policy-Gradient-Methoden als auch zu instabilen direkten Gradientenansätzen und erzielen eine überlegene Bild-Text-Ausrichtung sowie Generierungsqualität. Die Arbeit legt nahe, dass das Beibehalten und Skalieren verschachtelter Gradienten, anstatt sie zu verwerfen, ein Schlüsselfaktor für die Verbesserung der Leistung direkter Gradienten-Ausrichtung ist. Als zukünftige Arbeit wird die Erweiterung dieser Prinzipien auf die Videogenerierung notiert.