TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward

TDM-R1 is een nieuwe versterkingsleerparadigma dat few-step diffusiemodellen effectief optimaliseert met niet-differentieerbare beloningen door het leerproces te ontkoppelen in surrogaatbeloningslering en generatorlering, wat leidt tot state-of-the-art prestaties in diverse generatietaken.

De kern

Stel je voor dat je een kunstenaar bent die een schilderij moet maken. Normaal gesproken duurt het veel tijd en veel penseelstreken om een prachtig, gedetailleerd werk te maken. In de wereld van kunstmatige intelligentie (AI) is dit vergelijkbaar met de "stappen" die een model moet doorlopen om een afbeelding te genereren.

Het probleem:
De nieuwe, snelle AI-modellen (de "few-step" modellen) kunnen een afbeelding in slechts 4 penseelstreken maken, in plaats van de gebruikelijke 80 of 100. Dit is ongelooflijk snel en goedkoop. Maar er zit een addertje onder het gras: deze snelle kunstenaars zijn soms niet heel goed in het volgen van specifieke instructies. Als je zegt: "Teken een rode kat links van een blauwe hond", maken ze misschien een kat, maar staat hij rechts, of is hij groen.

Vroeger probeerden we deze modellen te verbeteren door ze te laten "leren" van een leraar die elke penseelstreek corrigeerde. Maar die leraar kon alleen helpen als de fouten meetbaar en wiskundig uit te rekenen waren (zoals "de kleur is 5% te rood"). Veel belangrijke dingen in de echte wereld zijn echter niet zo meetbaar: "Vindt een mens dit beeld mooi?", "Zit er precies één hond in de afbeelding?" of "Is de tekst leesbaar?". Deze dingen zijn niet-wiskundig (non-differentiable) en de oude methoden konden hier niets mee.

De oplossing: TDM-R1
De auteurs van dit paper hebben een nieuwe methode bedacht, genaamd TDM-R1. Laten we dit uitleggen met een paar creatieve vergelijkingen:

1. De "Vaste Route" in plaats van "Gokken"

Stel je voor dat je een snelle auto hebt die van punt A naar punt B moet rijden.

• De oude manier: De auto rijdt een willekeurige route (stochastisch). Omdat de route elke keer anders is, is het voor de leraar heel moeilijk om te zeggen: "Op dit specifieke moment was je te snel." De feedback is vaag en onbetrouwbaar.
• De TDM-R1 manier: De auto rijdt een vaste, voorspelbare route (deterministisch). Omdat we precies weten welke weg de auto neemt, kan de leraar op elk moment van de rit precies zeggen: "Op kilometer 5 was je te snel, en op kilometer 10 was je te traag."
  Dit maakt het mogelijk om de AI stap voor stap te corrigeren, zelfs als de eindresultaat pas aan het einde wordt beoordeeld.

2. De "Slimme Tussentijdse Oefening" (Surrogate Reward)

Het grootste probleem was: hoe geef je feedback aan een schilder die nog halverwege is, als je pas kunt zeggen of het schilderij goed is als het klaar is?

• Het oude probleem: Je wacht tot het schilderij klaar is, zegt "Dit is goed!", en probeert dan terug te rekenen welke penseelstreek daarvoor verantwoordelijk was. Dit werkt niet goed bij snelle modellen; het resultaat wordt wazig en onduidelijk.
• De TDM-R1 oplossing: Ze hebben een tussentijdse trainer bedacht. Stel je voor dat de AI een "tussenpersoon" heeft. Deze persoon kijkt naar de halve afbeelding en zegt: "Als je nu zo doorgaat, krijg je waarschijnlijk een goed resultaat."
  Deze trainer is niet de eindbeoordelaar, maar een slimme gids die de AI helpt om de goede richting op te gaan. De AI leert van deze gids, en de gids wordt zelf ook steeds slimmer door te kijken naar wat de AI produceert. Het is een soort danspartij: de AI leert de stappen, en de gids leert de muziek beter te begrijpen, zodat ze samen een perfecte dans kunnen doen.

3. Leren van "Mensen" in plaats van "Wiskunde"

Dit is het belangrijkste: TDM-R1 laat de AI leren van menselijke voorkeuren.

• Vroeger: "De afbeelding moet exact 50% rood zijn." (Wiskundig, maar saai).
• Nu met TDM-R1: "Kijk, deze afbeelding met de tekst 'Ultra-Powerful' is leuk, die andere met 'Ultra-Powerfull' (met een foutje) is niet leuk."
  De AI leert hieruit wat mensen leuk vinden, zonder dat er een wiskundige formule voor nodig is. Het is alsof je een kind leert tekenen door te zeggen: "Leuk, die hond!", in plaats van "De poten moeten 3,4 cm lang zijn."

Wat is het resultaat?
Dankzij deze methode kunnen de snelle AI-modellen (die maar 4 stappen nodig hebben) nu beter presteren dan de oude, trage modellen die 80 stappen nodig hadden.

• Ze kunnen tekst perfect in afbeeldingen schrijven.
• Ze kunnen complexe instructies volgen (bijv. "een kat links van een hond").
• Ze zien er scherper en mooier uit.

Kort samengevat:
TDM-R1 is als het geven van een GPS met live feedback aan een racewagen. In plaats van dat de wagen blindelings racet en pas aan het einde wordt beoordeeld, krijgt de wagen op elk moment van de rit een slim signaal: "Houd je aan de route, want zo komen we bij het doel dat mensen leuk vinden." Hierdoor wordt de snelle wagen (het AI-model) niet alleen sneller, maar ook slimmer en nauwkeuriger dan de oude, langzame wagens.

Technische samenvatting

1. Het Probleem

Hoewel few-step generatieve modellen (zoals gedistilleerde diffusion-modellen) revolutionair zijn voor het genereren van hoogwaardige afbeeldingen en video's met een zeer lage rekentijd (soms 50x sneller dan traditionele diffusion-modellen), blijven ze worstelen met specifieke uitdagingen zoals het nauwkeurig volgen van instructies, het renderen van complexe tekst en het correct positioneren van objecten.

Reinforcement Learning (RL) heeft bewezen effectief te zijn om deze problemen op te lossen bij grote taalmodellen (LLM's) en standaard diffusion-modellen. Echter, bestaande RL-methoden voor few-step modellen hebben een fundamentele beperking:

• Afhankelijkheid van differentieerbare beloningen: Bestaande algoritmen vereisen dat de beloningsfunctie (reward) differentieerbaar is, zodat de gradiënten via de generatieve output teruggepropageerd kunnen worden naar het model.
• Uitsluiting van cruciale signalen: Dit sluit de meeste real-world beloningssignalen uit, zoals menselijke voorkeuren (binair: leuk/niet leuk), discrete objecttellingen, of de correctheid van gegenereerde tekst (gecontroleerd via OCR). Deze signalen zijn per definitie niet-differentieerbaar.

Het huidige probleem is dus hoe few-step modellen effectief getraind kunnen worden met deze veelbelovende, maar niet-differentieerbare, beloningssignalen zonder de kwaliteit van de gegenereerde afbeeldingen te verstoren.

2. Methodologie: TDM-R1

De auteurs introduceren TDM-R1, een nieuw RL-paradigma gebaseerd op het Trajectory Distribution Matching (TDM) model. De kern van de methode is het ontkoppelen van het leerproces in twee fasen: het leren van een surrogaatbeloning (surrogate reward) en het optimaliseren van de generator.

De methode bestaat uit drie cruciale componenten:

A. Nauwkeurige Beloningsschatting via Deterministische Trajecten

In standaard diffusion-modellen is het moeilijk om beloningen toe te kennen aan tussentijdse stappen omdat de output van een tussentijdse stap (ruis) nog geen duidelijke afbeelding is.

• TDM-R1 maakt gebruik van deterministische trajecten (ODE-sampling) in plaats van stochastische trajecten.
• Omdat het pad van ruis naar een schone afbeelding ($x_0$) deterministisch is, kan de auteurs een onbevooroordeelde schatting maken van de beloning voor elke tussentijdse stap ($x_t$) door de verwachting te nemen over de waarschijnlijkheidsverdeling van de schone afbeelding gegeven de ruis.
• Dit vermindert de variantie van de beloningsschatting aanzienlijk en maakt het mogelijk om nauwkeurige feedback te geven voor elke stap in het generatieproces.

B. Leren van een Surrogaatbeloning (Surrogate Reward Learning)

Om niet-differentieerbare beloningen (zoals menselijke voorkeuren) bruikbaar te maken voor backpropagatie, introduceert TDM-R1 een differentieerbare surrogaatbeloning ($\tilde{r}_\phi$), die zelf een diffusion-model is.

• Groepsgebaseerde voorkeursoptimalisatie: In plaats van alleen paarsgewijze vergelijkingen, gebruikt de methode groepen van gegenereerde samples. Samples worden ingedeeld in positieve en negatieve groepen op basis van hun "advantage" (hoeveel beter of slechter ze zijn dan het gemiddelde in de groep).
• Het model leert deze surrogaatbeloning door de Bradley-Terry (BT) voorkeursverdeling te maximaliseren tussen deze groepen.
• Dynamische Referentiemodel: Om overfitting te voorkomen, wordt het referentiemodel voor de surrogaatbeloning dynamisch bijgewerkt met een Exponential Moving Average (EMA) van de parameters, in plaats van een statisch, bevroren model te gebruiken.

C. Generator Learning (Few-Step Generator Learning)

De uiteindelijke generator ($p_\theta$) wordt getraind om de surrogaatbeloning te maximaliseren, terwijl het tegelijkertijd dicht bij de oorspronkelijke verdeling van het pre-getrainde model blijft (via een reverse KL-divergentie regularisatie).

• De leerdoelstelling combineert het maximaliseren van de beloning met een regularisatieterm die zorgt dat de gegenereerde samples niet "uitwijken" naar onrealistische patronen (reward hacking).
• Dit creëert een synergetische lus: de generator produceert betere samples om de beloning te maximaliseren, en de surrogaatbeloning past zich aan om nog nauwkeurigere feedback te geven op tussentijdse stappen.

3. Belangrijkste Bijdragen

1. Eerste schaalbare RL voor niet-differentieerbare beloningen in few-step modellen: TDM-R1 is het eerste werk dat succesvol niet-differentieerbare signalen (zoals menselijke voorkeuren en OCR) toepast op few-step diffusion-modellen op grote schaal.
2. Ontkoppeling van leerprocessen: Door het scheiden van het leren van de beloning (surrogaat) en de generator, wordt het probleem van het backpropageren door niet-differentieerbare functies omzeild.
3. Deterministische Trajecten als sleutel: Het paper demonstreert dat het gebruik van deterministische trajecten (ODE) essentieel is voor het reduceren van variantie in beloningsschattingen, wat leidt tot snellere convergentie en betere resultaten dan stochastische methoden.
4. Dynamische Surrogaatbeloning: De introductie van een dynamisch bijgewerkt surrogaatmodel dat groepsvoorkeuren leert, biedt fijnmazigere en stabielere leersignalen dan bestaande methoden.

4. Resultaten

De auteurs hebben TDM-R1 uitgebreid getest op verschillende benchmarks, waaronder tekst-rendering, visuele kwaliteit en voorkeursuitlijning.

• GenEval Benchmark (Samengestelde Afbeelding): TDM-R1 (met slechts 4 NFE - Number of Function Evaluations) bereikte een score van 92%. Dit is een enorme verbetering ten opzichte van het basis-TDM-model (61%) en overtreft zelfs de dure 80-NFE basismodellen (63%) en de commerciële SOTA GPT-4o (84%).
• Visuele Tekst Rendering: Het model toonde aanzienlijke verbeteringen in het correct renderen van tekst binnen afbeeldingen, gemeten via OCR-accuraatheid.
• Out-of-Domain Kwaliteit: In tegenstelling tot sommige andere RL-methoden die de beeldkwaliteit ten koste gaan van de taakprestatie, behaalde TDM-R1 ook hogere scores op onafhankelijke kwaliteitsmetrieken (zoals Aesthetic Score, ImageReward en PickScore) dan zowel de few-step als de many-step basismodellen.
• Schalbaarheid: De methode werkt effectief op grote modellen, zoals het 6B-parameter Z-Image model. TDM-R1 overtrof hier zowel de 100-NFE variant als de few-step variant van Z-Image, opnieuw met slechts 4 NFE.
• Ablatie Studies: Vergelijkingen toonden aan dat het gebruik van deterministische trajecten en dynamische referentiemodellen cruciaal zijn; het vervangen hiervan door stochastische sampling of statische beloningen leidde tot slechtere prestaties en instabielere training.

5. Betekenis en Impact

TDM-R1 is een doorbraak in het veld van generatieve AI omdat het de "muur" doorbreekt die few-step modellen scheidt van krachtige RL-technieken.

• Efficiëntie: Het bewijst dat modellen met slechts 4 stappen (extreem snel) kunnen presteren op het niveau van of beter dan modellen die 80 of 100 stappen nodig hebben, mits ze correct getraind worden met RL.
• Praktische Toepasbaarheid: Door het gebruik van niet-differentieerbare beloningen (zoals menselijke feedback of OCR) wordt het mogelijk om modellen te trainen op real-world criteria zonder dat er dure ground-truth datasets nodig zijn.
• Toekomstperspectief: Deze methode opent de deur voor het toepassen van geavanceerde RL-strategieën (zoals die succesvol zijn bij LLM's) op snelle generatieve modellen, wat essentieel is voor de integratie van hoogwaardige generatieve AI in real-time industriële toepassingen.

Kortom, TDM-R1 combineert de snelheid van few-step generatie met de intelligentie van reinforcement learning, waardoor het mogelijk wordt om snelle, hoogwaardige en instructie-volgende generatieve modellen te bouwen die direct bruikbaar zijn in de praktijk.