Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Each language version is independently generated for its own context, not a direct translation.

🎨 De Kunst van het Fijnstomen: Hoe je een AI-kunstenaar slimmer maakt

Stel je voor dat je een zeer getalenteerde kunstenaar hebt die al duizenden schilderijen heeft gemaakt. Deze kunstenaar is een diffusiemodel (een type AI dat beelden maakt). Hij is geweldig, maar soms maakt hij kleine foutjes, of hij luistert niet goed genoeg naar wat jij precies wilt (bijvoorbeeld: "een rode auto, maar geen blauwe").

De onderzoekers van dit paper willen deze kunstenaar fijnstomen (fine-tuning). Ze willen hem leren om nog beter te doen wat jij vraagt, zonder dat hij zijn hele talent verliest.

Het probleem is echter: hoe leer je iemand iets nieuws zonder hem te verwarren of zijn natuurlijke stijl te breken? De auteurs komen met twee slimme oplossingen: P-GRAFT en Inverse Noise Correction.

1. Het Probleem: De "Te Ver" Stap

Normaal gesproken werkt zo'n AI door te beginnen met een pot vol statisch (ruis) en dit stap voor stap te "ontruisen" tot een duidelijk beeld.

Stap 1: Alles is statisch.
Stap 50: Het begint op een vaag silhouet te lijken.
Stap 100: Het is een perfect beeld.

Als je de AI wilt leren om "rode auto's" te maken, kun je kijken naar het eindresultaat. Maar als je de AI terug traint op dat eindresultaat, is het vaak te laat. De AI heeft al te veel "leerfouten" gemaakt in de tussentijd. Het is alsof je een student pas wilt leren hoe hij moet schrijven, nadat hij al een heel boek heeft geschreven met veel fouten. Het is moeilijk om die fouten later nog te corrigeren.

2. Oplossing A: P-GRAFT (De "Tussentijdse Check")

De onderzoekers bedenken een slimme truc: Wacht niet tot het einde.

Stel je voor dat je een lange reis maakt met een auto. Je wilt dat de bestuurder op de juiste weg blijft.

De oude manier: Je kijkt pas naar het eindpunt. Als je er verkeerd uitkomt, probeer je de bestuurder te straffen of te belonen. Maar hij weet niet precies waar hij de afslag miste.
De P-GRAFT manier: Je kijkt naar de bestuurder halverwege de reis (bijvoorbeeld na 25% van de weg). Op dat moment is de weg nog niet helemaal vastgelegd, maar je ziet al wel of hij de juiste richting op gaat.

De Analogie:
Stel je voor dat je een potje met modder (ruis) hebt waarin je een beeld wilt vinden.

Je laat de AI een potje modder "ontmodderen" tot het halverwege is (bijvoorbeeld 75% van de weg).
Je kijkt naar het eindresultaat dat daaruit zou kunnen komen (als je het afmaakt).
Als het eindresultaat goed is, geef je de AI een beloning voor die halve modderpot.
Je traint de AI alleen op die halve modderpotten.

Waarom werkt dit?
Het is makkelijker om een fout te corrigeren als je nog niet helemaal bij het einde bent. De "ruis" is nog niet volledig weg, maar de richting is al wel duidelijk. Dit is een perfecte balans:

Te vroeg (veel ruis): Je weet niet wat het doel is (te veel variatie).
Te laat (weinig ruis): Het is te laat om de richting te veranderen (te veel bias).
P-GRAFT pakt het moment in het midden waar het makkelijkst is om te leren.

Resultaat: De AI wordt veel beter in het volgen van instructies (zoals tekst-naar-beeld) dan eerdere methoden, en dat met minder rekenkracht.

3. Oplossing B: Inverse Noise Correction (De "Omgekeerde Toer")

Dit is voor een ander type AI-model (Flow Models), maar het idee is net zo cool.

Stel je voor dat je een machine hebt die een beeld maakt uit een willekeurig stukje ruis. Soms is de machine niet perfect ingesteld, waardoor de beelden er net niet helemaal goed uitzien.

De vraag: Kunnen we de machine niet gewoon "terugdraaien"?
De methode: De onderzoekers nemen een perfect beeld (uit de echte wereld), en laten de machine het terug draaien naar ruis.
Het inzicht: Als de machine imperfect is, dan is de "ruis" die uit die terugdraaiing komt, ook imperfect. Het is geen "normale" ruis meer, maar een "gecorrigeerde" ruis.

De Analogie:
Stel je voor dat je een bakker hebt die slechte broden bakt.

Je neemt een perfect brood (uit de supermarkt).
Je laat de bakker het brood "ongebakken" maken (terug naar deeg).
Je ziet dat het deeg er anders uitziet dan normaal (misschien te nat of te droog).
Je traint een nieuwe machine (de "Noise Corrector") om precies dat specifieke deeg te maken.
Nu laat je de originele bakker werken, maar hij begint niet met normaal deeg, maar met dat gecorrigeerde deeg.

Resultaat: Omdat de bakker nu met het perfecte startpunt begint, komt het eindresultaat (het brood) veel beter uit, zonder dat je de bakker zelf hoeft te vervangen of ingewikkeld te herschrijven. Het werkt zelfs zonder dat je een "beloningssysteem" (rewards) nodig hebt.

4. Wat hebben ze bewezen?

De onderzoekers hebben dit getest op verschillende gebieden:

Tekst-naar-beeld: Als je zegt "een hond in een hoed", maakt de AI nu veel sneller en beter de juiste hond in de juiste hoed dan voorheen.
Moleculen: Ze konden chemische structuren maken die stabieler zijn (belangrijk voor medicijnen).
Afbeeldingen: Ze konden de kwaliteit van gegenereerde foto's verbeteren terwijl ze minder rekenkracht (FLOPs) gebruikten.

Samenvatting in één zin

In plaats van een AI te dwingen om fouten aan het einde van het proces te herstellen, leren ze de AI om op het juiste moment (halverwege) de juiste keuzes te maken, of ze corrigeren de start van het proces zodat het eindresultaat vanzelf beter wordt.

Het is alsof je niet wacht tot de student het examen heeft gemaakt om te corrigeren, maar je hem helpt tijdens het maken van de opgaven, of je zorgt dat hij de juiste pen en papier heeft voordat hij begint.

Each language version is independently generated for its own context, not a direct translation.

Probleemstelling

Voorgeprogrammeerde generatieve modellen, zoals diffusion- en flow-modellen, vereisen vaak taakspecifieke aanpassingen om fouten te corrigeren of beter uit te lijnen met downstream-toepassingen (bijvoorbeeld prompt-gebaseerde afbeeldingengeneratie).

Beperkingen van RL-methoden: Bestaande methoden zoals Proximal Policy Optimization (PPO) met KL-geregulariseerde beloningen werken goed voor autoregressieve taalmodellen, maar zijn moeilijk toe te passen op diffusion-modellen. De marginaal waarschijnlijkheid (nodig voor exacte KL-regularisatie) is bij diffusion-modellen onberekenbaar (intractable).
Huidige oplossingen: In de praktijk wordt de KL-term vaak genegeerd (wat leidt tot instabiel training) of er wordt gebruikgemaakt van "trajectory KL" (wat suboptimale resultaten geeft en bias-problemen veroorzaakt).
Rejection Sampling: Methoden zoals RAFT (Rejection sAmpling Fine-Tuning) zijn succesvol, maar het is niet volledig duidelijk hoe ze theoretisch verhouden tot KL-regularisatie of hoe ze het beste kunnen worden toegepast op de complexe trajecten van diffusion-modellen.

Methodologie

De auteurs introduceren een theoretisch raamwerk dat rejection sampling koppelt aan KL-regularisatie en twee nieuwe methoden voorstelt: P-GRAFT en Inverse Noise Correction.

1. GRAFT (Generalized Rejection sAmpling Fine-Tuning)

De auteurs unificeren bestaande rejection sampling-strategieën (zoals RAFT en Best-of-N) onder het raamwerk Generalized Rejection Sampling (GRS).

Theoretisch inzicht: Ze bewijzen dat GRS impliciet KL-regulariseerde beloningsmaximalisatie uitvoert met een "her-vormde" beloningsfunctie ( $\hat{r}$ ), zelfs zonder dat de marginaal waarschijnlijkheid expliciet berekend hoeft te worden. Dit lost het probleem van de onberekenbare marginaal KL voor diffusion-modellen op.

2. P-GRAFT (Partial-GRAFT)

Dit is de kerninnovatie voor diffusion-modellen. In plaats van de verdeling te vormen op het eindpunt van het denoising-traject (de volledige afbeelding), vormt P-GRAFT de verdeling op een intermediair tijdstip $t$ (waarbij $0 < t < T$ ).

Werkingsprincipe: Het model genereert volledige trajecten, maar de beloning wordt toegewezen aan de deels gedenoised toestand $X_t$ op basis van de kwaliteit van de uiteindelijke afbeelding $X_0$ . Het fine-tuning gebeurt alleen voor de stappen van $T$ tot $t$ . Voor de resterende stappen ( $t$ tot $0$) wordt het originele referentiemodel gebruikt.
Bias-Variance Trade-off: De auteurs bieden een wiskundige rechtvaardiging:
- Variance: De beloning is "ruisachtiger" op vroege tijdstippen (hoge $t$ ) omdat de relatie tussen $X_t$ en de uiteindelijke kwaliteit $X_0$ minder sterk is.
- Bias: Het leren van de scorefunctie (de gradiënt van de log-dichtheid) is echter veel makkelijker op vroege tijdstippen omdat de verdeling dichter bij een eenvoudige Gaussische verdeling ligt dan bij de complexe data-verdeling.
- Conclusie: Door een "geschikt" intermediair tijdstip te kiezen, optimaliseert P-GRAFT de balans tussen de moeilijkheid van het leerprobleem (bias) en de ruis in de beloning (variance).

3. Inverse Noise Correction (voor Flow Modellen)

Voor rectified flow-modellen (die deterministische ODE's gebruiken) stellen de auteurs een methode voor om fouten in voorgeprogrammeerde modellen te corrigeren zonder expliciete beloningen.

Principe: Aangezien het eindresultaat volledig bepaald wordt door de initiële ruis, kunnen fouten in de gegenereerde verdeling worden gecorrigeerd door de verdeling van de initiële ruis aan te passen.
Implementatie:
1. Gebruik het voorgeprogrammeerde model om data naar "ruis" te draaien (via omgekeerde Euler-integratie).
2. Deze gegenereerde "inverse ruis" vormt een nieuwe doelverdeling.
3. Train een klein "Noise Corrector"-model om standaard Gaussische ruis om te zetten in deze gecorrigeerde inverse ruisverdeling.
4. Tijdens inferentie wordt eerst de gecorrigeerde ruis gegenereerd en vervolgens de afbeelding gegenereerd met het oorspronkelijke model.

Belangrijkste Bijdragen

Unificatie van GRAFT: Een theoretisch raamwerk dat rejection sampling koppelt aan KL-regulariseerde optimalisatie, waardoor marginal KL-constraints mogelijk worden voor diffusion-modellen.
P-GRAFT: Een nieuwe strategie die fine-tuning beperkt tot intermediaire denoising-stappen. Dit leidt tot betere prestaties door de bias-variance trade-off te benutten.
Inverse Noise Correction: Een parameter-efficiënte adapter-methode om de kwaliteit van flow-modellen te verbeteren zonder extra beloningsmodellen, door de initiële ruisverdeling te corrigeren.
Empirische Validatie: Uitgebreide experimenten op tekst-naar-beeld, layout-generatie, molecuulgeneratie en onvoorwaardelijke beeldgeneratie.

Resultaten

De methoden zijn getest op diverse benchmarks:

Tekst-naar-beeld (Stable Diffusion v2):
- P-GRAFT presteert significant beter dan policy-gradient methoden (zoals DDPO) en de basis SDv2 op benchmarks zoals GenAI-Bench, T2ICompBench++ en GenEval.
- Op GenAI-Bench boekte P-GRAFT een 8,81% relatieve verbetering ten opzichte van het basismodel in termen van VQAScore.
- P-GRAFT generaliseert goed naar onbekende prompts.
- Ablatiestudies tonen aan dat een intermediair tijdstip van $0.25N$ (waarbij $N$ het totaal aantal stappen is) vaak de beste balans biedt.
Layout en Molecuulgeneratie:
- Op PubLayNet (layout) en QM9 (moleculen) verbeterde P-GRAFT de uitlijning en stabiliteit.
- Voor moleculen werd mode-collapse voorkomen door een specifieke de-duplicatie-strategie binnen GRAFT toe te passen.
Inverse Noise Correction:
- Op CelebA-HQ en LSUN-Church verbeterde deze methode de FID (Fréchet Inception Distance) aanzienlijk.
- Efficiëntie: Het systeem bereikte betere kwaliteit met minder FLOPs per afbeelding. Bijvoorbeeld, het combineren van 100 stappen van de Noise Corrector en 100 stappen van het basismodel presteerde beter dan 1000 stappen van alleen het basismodel.

Betekenis en Impact

Dit paper biedt een fundamentele doorbraak in het fine-tunen van generatieve modellen:

Theoretische Strengh: Het lost het probleem van de onberekenbare marginaal KL bij diffusion-modellen op door een brug te slaan tussen rejection sampling en RL.
Efficiëntie: P-GRAFT is computatie-efficiënter dan traditionele RL-methoden omdat het geen backpropagatie vereist over het volledige denoising-traject tijdens het trainen, en het gebruik van LoRA maakt het goedkoop.
Generalisatie: De methoden werken niet alleen voor diffusion-modellen, maar ook voor discrete-continue modellen en flow-modellen.
Nieuwe Richting: De "Inverse Noise Correction" opent een nieuwe weg voor het verbeteren van generatieve modellen zonder de noodzaak van dure beloningsmodellen, wat vooral waardevol is voor domeinen waar beloningen moeilijk te definiëren zijn.

Samenvattend stelt dit werk dat het vormgeven van verdelingen op intermediaire tijdstippen een krachtigere en efficiëntere strategie is voor het fine-tunen van generatieve modellen dan het focussen op het eindresultaat of het gebruik van traditionele policy-gradient methoden.