Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models

Deze paper introduceert een nieuwe online reinforcement learning-methode voor het post-trainen van text-to-image modellen, die door het behandelen van het volledige sampleproces als één actie en het gebruik van gepaarde trajecten de variantie verlaagt, wat resulteert in snellere convergentie en hogere kwaliteit en prompt-afstemming dan bestaande technieken.

De kern

Stel je voor dat je een kunstenaar bent die een robot hebt leren schilderen. Deze robot is al een meester in het maken van prachtige landschappen en portretten (dit is het "pre-training" stadium). Maar soms wil je dat de robot iets specifieker doet: bijvoorbeeld dat de zonsondergangen nog warmer worden, of dat de tekst in de afbeelding perfect leesbaar is.

Om dit te bereiken, geef je de robot een beloningssysteem. Als hij een plaatje maakt dat je mooi vindt, krijgt hij een sterretje. Als hij iets lelijks maakt, krijgt hij geen sterretje. Dit noemen we Versterkend Leren (Reinforcement Learning).

Het probleem met de oude methoden (zoals Flow-GRPO) was dat de robot een beetje als een dronken man door de galerie liep. Hij probeerde willekeurig een paar nieuwe streken, keek of hij een sterretje kreeg, en als dat zo was, probeerde hij die beweging na te bootsen. Maar omdat hij zo willekeurig probeerde, maakte hij ook veel onnodige, rare bewegingen die niets met de beloning te maken hadden. Dit maakte het leren traag en soms kreeg de robot zelfs rare, grid-achtige vlekken op zijn schilderijen (zoals "reward hacking").

De nieuwe methode in dit papier (Finite Difference Flow Optimization) werkt als een slimme, gerichte coach. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Tweeling-Test" in plaats van Willekeur

In plaats van dat de robot willekeurig probeert, laat de coach de robot twee bijna identieke schilderijen maken, gebaseerd op precies hetzelfde beginpunt (dezelfde "ruis" of chaos).

• Schilderij A: De robot schildert een beetje zoals gewoonlijk.
• Schilderij B: De robot maakt een heel klein, willekeurig verschil in de details (bijvoorbeeld: de hoed van de figuur staat een millimeter scheef, of de kleur van de lucht is net iets anders).

2. De Vergelijking

Nu kijkt de coach naar beide schilderijen en vraagt: "Welk van deze twee is beter?"

• Als Schilderij B een hogere score krijgt (meer sterretjes), dan weet de coach precies wat het verschil was dat die betere score opleverde.
• Het verschil tussen de twee schilderijen is als een pijl die precies aangeeft: "Ga in deze richting!"

3. De Gerichte Duw

Bij de oude methode probeerde de robot duizenden willekeurige stappen en hoopte hij dat één ervan goed was. Bij deze nieuwe methode duwt de coach de robot direct in de richting van het betere schilderij.

• Het is alsof je een bal op een heuvel rolt. De oude methode duwde de bal willekeurig in alle richtingen en hoopte dat hij naar beneden rolde.
• De nieuwe methode kijkt eerst naar twee kleine proefballen, ziet welke kant naar beneden gaat, en duwt de echte bal alleen in die richting.

Waarom is dit zo goed?

• Snelheid: Omdat er geen willekeurige, nutteloze bewegingen zijn, leert de robot veel sneller. Het papier laat zien dat ze in minder dan de helft van de tijd dezelfde resultaten bereiken.
• Geen rare vlekken: Omdat de robot niet meer willekeurig rondtast, ontstaan er geen rare, grid-achtige fouten in de afbeeldingen. De kwaliteit blijft hoog.
• Precisie: De robot leert niet alleen om "mooier" te worden, maar ook om beter te luisteren naar wat je precies vraagt (bijvoorbeeld: "een kat met een hoed" wordt echt een kat met een hoed, en geen kat met een hoed die eruitziet als een hond).

De Analogie: De Dronken Toerist vs. De GPS

• De Oude Methode (Flow-GRPO): Stel je een toerist voor die een stad probeert te verkennen zonder kaart. Hij loopt een beetje links, dan rechts, dan weer links. Als hij per ongeluk een mooi uitzicht vindt, probeert hij die route te onthouden. Maar omdat hij zo willekeurig loopt, raakt hij vaak verdwaald en loopt hij in een cirkel.
• De Nieuwe Methode (Dit Papier): De toerist heeft nu een GPS. Hij laat de GPS twee routes uitrekenen die heel dicht bij elkaar liggen. De GPS ziet welke route net iets mooier is, en stuurt de toerist direct en rechtstreeks die kant op. Geen verdwalen, geen cirkels, gewoon de snelste weg naar het mooiste uitzicht.

Kortom: Dit papier introduceert een slimme manier om AI-kunstenaars te trainen. In plaats van ze blindelings te laten proberen, laat je ze twee versies maken, vergelijkt je die, en duwt je ze dan direct in de richting van de winnaar. Het resultaat: betere kunst, sneller leren, en minder rare fouten.

Technische samenvatting

1. Het Probleem

Reinforcement Learning (RL) wordt steeds vaker gebruikt voor het post-trainen van diffusion-modellen (zoals Stable Diffusion) om beelden te optimaliseren op basis van beloningssignalen (bijv. beeldkwaliteit, prompt-uitlijning). Bestaande methoden, zoals Flow-GRPO en DanceGRPO, behandelen het stochastische bemonsteringsproces als een Markov Decision Process (MDP). Hierbij wordt elke stap in het generatieproces gezien als een aparte actie.

De auteurs identificeren een fundamenteel probleem met deze aanpak:

• Hoge Variansie: De updates zijn gebaseerd op willekeurige perturbaties (stochastische stappen). Hoewel de aggregatie van deze updates de beloning verbetert, is een groot deel van de update "ruis" die niet bijdraagt aan het doel, maar de stroom (flow) in willekeurige richtingen duwt.
• Reward Hacking: Omdat de updates ruis bevatten, kunnen irrelevante dimensies vrij "driften". Dit leidt tot artefacten (zoals rasterpatronen) en een verslechtering van de algehele beeldkwaliteit na langdurig trainen, zelfs als de specifieke beloning stijgt.
• Langzame Convergentie: Door de grote hoeveelheid ruis in de updates is de signaal-ruisverhouding (signal-to-noise ratio) laag, wat resulteert in langzamere convergentie naar hoge beloningen.

2. Methodologie: Finite Difference Flow Optimization (FDFO)

De auteurs stellen een nieuwe online RL-variant voor die de variansie van modelupdates drastisch verlaagt door het hele bemonsteringsproces te behandelen als één enkele actie, in plaats van een reeks van losse acties.

Kernprincipes:

1. Paar-generatie: In plaats van een groep van trajecten te nemen, genereert de methode een paar van twee zeer vergelijkbare beelden ($x_T$ en $\hat{x}_T$) die starten vanuit dezelfde initiële ruis ($\epsilon$).
2. Stochastische Perturbatie: Tijdens het bemonsteren wordt een kleine hoeveelheid stochastische ruis toegevoegd (via een aangepaste Euler-Maruyama sampler) om variatie in de details van de beelden te creëren, terwijl de algemene lay-out behouden blijft.
3. Finite Difference Gradient: De methode berekent het verschil tussen de twee gegenereerde beelden ($\Delta x = \hat{x}_T - x_T$) en het verschil in beloning ($\Delta R = R(\hat{x}_T) - R(x_T)$).
4. Richting van Update: Het gewogen verschil $\Delta R \cdot \Delta x$ wijst gegarandeerd in de richting van het beeld met de hogere beloning. Dit wordt gebruikt als een benadering van de gradiënt.
5. Flow-Optimalisatie: De stroomvectoren (flow velocities) langs het hele generatietraject worden aangepast om zich in deze richting te buigen. Dit maakt gebruik van de "niet-rotatie" eigenschap van diffusion flows: een verandering in de tussenliggende stappen leidt tot een vergelijkbare verandering in het eindbeeld.

Technische Implementatie:

• Stochastische Sampler: De auteurs passen de EDM-stochastische sampler aan om consistentie te garanderen binnen flow-matching, waarbij ze de tijdsstappen "overschieten" en vervolgens correcte ruis toevoegen om de juiste verdeling te behouden.
• Normalisatie: Om de grootte van de updates te stabiliseren, wordt het beeldverschil $\Delta x$ genormaliseerd.
• On-policy Optimalisatie: De methode gebruikt een variant van Simple Policy Optimization (SPO), vergelijkbaar met PPO-Clip, om updates te beperken die te ver afwijken van de oorspronkelijke politiek (om overfitting op verouderde data te voorkomen).

3. Belangrijkste Bijdragen

• Nieuwe Formulering: Een verschuiving van een MDP-benadering (stap-voor-stap) naar een "single-action" benadering voor het gehele traject, wat de variansie van de updates reduceert.
• Finite Difference Flow: Het gebruik van paar-generatie en het verschil in beelden als een gradiënt-achtige signaal, zonder dat de beloningsfunctie differentieerbaar hoeft te zijn.
• Superieure Convergentie: De methode convergeert aanzienlijk sneller dan state-of-the-art (SOTA) methoden zoals Flow-GRPO.
• Vermindering van Artefacten: Door de ruis in de updates te elimineren, worden de typische "reward hacking" artefacten (zoals rasterpatronen) die bij langdurig trainen met Flow-GRPO optreden, voorkomen.

4. Resultaten

De auteurs evalueerden hun methode op Stable Diffusion 3.5 Medium met verschillende beloningsfuncties (PickScore, VLM-gebaseerde prompt-uitlijning, en een combinatie).

• Snelheid: De methode bereikt een hogere beloning veel sneller. In tests met een gecombineerde beloning was de methode 19x sneller dan Flow-GRPO in de basisconfiguratie (40 stappen) en 5x sneller in een snelle configuratie (10 stappen).
• Kwaliteit: De gegenereerde beelden tonen een hogere prompt-uitlijning en betere menselijke voorkeurscores (HPSv2) bij gelijke trainingsduur.
• Stabiliteit: Flow-GRPO begon na ongeveer 500-800 epochs artefacten (grid-achtige patronen) te vertonen die wisselden in intensiteit. De FDFO-methode vertoonde dergelijke artefacten niet, zelfs niet na 1000 epochs.
• Diversiteit: Hoewel alle RL-methoden de diversiteit iets verminderen, doet FDFO dit op een meer gecontroleerde manier zonder de stijl van de beelden willekeurig te laten "driften".

5. Betekenis en Conclusie

Dit paper biedt een krachtige alternatieve aanpak voor het post-trainen van text-to-image modellen. Door de stochastische ruis in de updates te decoupleren van de daadwerkelijke verbetering van de beloning, lost de methode het probleem van "reward hacking" en trage convergentie op.

De Finite Difference Flow Optimization kan direct worden gebruikt als een "drop-in replacement" voor bestaande RL-algoritmen in de post-training van diffusion-modellen. Het stelt onderzoekers en ontwikkelaars in staat om sneller modellen te fine-tunen met hogere kwaliteit en minder bijwerkingen, wat een belangrijke stap voorwaarts is in de richting van robuuste en betrouwbare generatieve AI-systemen. De code en getrainde modellen zijn openbaar beschikbaar gemaakt.