Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

Het artikel introduceert SQDF, een nieuwe methode voor het fine-tunen van diffusiemodellen die via een herparameteriseerde policy gradient van een zachte Q-functie reward-overoptimalisatie voorkomt en zo zowel hoge doelbeloningen als de diversiteit en natuurlijkheid van gegenereerde beelden behoudt.

De kern

Stel je voor dat je een kunstenaar hebt die al jarenlang prachtige schilderijen maakt. Deze kunstenaar is getraind op duizenden bestaande foto's en weet precies hoe hij een realistische kat, een zonsondergang of een futuristische stad moet tekenen. Dit is wat we een Diffusiemodel noemen (zoals Stable Diffusion).

Het probleem is echter: deze kunstenaar volgt alleen zijn eigen smaak. Hij wil niet per se een schilderij maken dat mensen leuk vinden, of een plaatje dat perfect past bij een specifieke tekst. We willen hem "trainen" om beter te worden op basis van wat wij waarderen (bijvoorbeeld: "dit moet eruitzien als een droom" of "dit moet een hoge esthetische score hebben").

De oude manier om dit te doen, was als een strenge leraar die alleen maar schreeuwt: "Maak het mooier! Nog mooier!" De kunstenaar probeerde dan zo hard om die "mooier"-score te maximaliseren, dat hij zijn creativiteit verloor. Hij begon abstracte, gekke patronen te maken die technisch gezien een hoge score hadden, maar er niets meer op leken. Dit noemen de auteurs over-optimisatie (te veel jagen op de score, ten koste van de kwaliteit).

In dit paper introduceren ze een nieuwe methode genaamd SQDF. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Soft Q" als een Slimme Voorspeller

Stel je voor dat de kunstenaar een schilderij aan het maken is. Hij begint met een vage vlek (ruis) en werkt stap voor stap naar een scherp beeld toe.

• Het oude probleem: Om te weten of het eindresultaat goed wordt, moest de kunstenaar vaak het hele proces opnieuw doen, wat veel tijd kost en onstabiel is.
• De SQDF-oplossing: Ze gebruiken een slimme voorspeller (de Soft Q-function). In plaats van het hele schilderij af te maken, kijkt deze voorspeller naar de huidige vage vlek en zegt: "Als je nu nog één stap verder gaat, zal het eindresultaat er ongeveer zo uitzien."
• De truc: Ze gebruiken een techniek genaamd Consistency Models. Dit is alsof je een ervaren schilder hebt die, zelfs als het schilderij nog erg wazig is, al heel goed kan voorspellen hoe het eruit zal zien als het klaar is. Dit maakt de voorspelling veel betrouwbaarder dan de oude methoden.

2. De "Korting" (Discount Factor) voor de juiste prioriteiten

In het proces van het maken van een afbeelding zijn er veel stappen. De eerste stappen zijn heel wazig en hebben weinig invloed op het eindresultaat. De laatste stappen zijn cruciaal.

• De analogie: Stel je voor dat je een huis bouwt. Als je de fundering (de eerste stappen) een beetje scheef zet, is dat erg. Maar als je de laatste streepjes verf (de laatste stappen) een beetje verkeerd zet, is dat ook erg. Echter, in de oude methoden werden alle stappen even zwaar gewogen.
• De SQDF-oplossing: Ze gebruiken een korting-factor (discount factor). Dit betekent dat de kunstenaar minder zwaar leert van de vroege, wazige stappen en meer leert van de latere stappen die het beeld echt vormgeven. Het is alsof je zegt: "Wat je nu doet, is belangrijker dan wat je een uur geleden deed."

3. De "Herinneringskast" (Replay Buffer)

Vaak leren kunstenaars alleen van hun allerlaatste poging. Als ze per ongeluk een heel mooi schilderij maken, vergeten ze dat misschien de volgende dag weer.

• De SQDF-oplossing: Ze gebruiken een replay buffer. Dit is een soort herinneringskast. Als de kunstenaar een prachtig schilderij maakt dat voldoet aan de eisen, wordt dit bewaard in de kast. Bij de volgende les haalt de kunstenaar niet alleen zijn nieuwste werk, maar ook die oude, prachtige voorbeelden uit de kast om ze opnieuw te bestuderen.
• Het voordeel: Dit zorgt ervoor dat de kunstenaar niet alleen één soort "perfect" schilderij leert maken (wat saai wordt), maar een hele variëteit aan mooie schilderijen behoudt. Het voorkomt dat hij in een hoekje vastloopt met steeds dezelfde saaie patronen.

Waarom is dit zo belangrijk?

De oude methoden waren als een student die alleen maar probeerde om een 10 te halen op een proefwerk door de antwoorden uit het hoofd te leren, maar de stof niet begreep. Uiteindelijk faalde hij op de echte toets omdat hij niets had geleerd over de logica.

SQDF is als een slimme coach die:

1. De student helpt om te voorspellen wat het eindresultaat wordt zonder het hele werk te hoeven doen.
2. Zegt: "Kijk vooral naar de details die echt tellen."
3. De student herinnert aan zijn beste oude werk, zodat hij niet vergeten hoe creatief hij eigenlijk is.

Het resultaat: De kunstenaar maakt nu niet alleen prachtige, hoge-scores schilderijen, maar ze zien er ook nog steeds natuurlijk uit en zijn divers. Hij is niet "gebroken" door de druk om perfect te zijn.

Kortom: SQDF is een nieuwe manier om AI-kunstenaars te trainen zodat ze niet alleen slimmer worden in het halen van scores, maar ook blijven doen wat ze het beste kunnen: creatief en natuurlijk kunst maken.

Technische samenvatting

Probleemstelling

Diffusiemodellen zijn uitstekend in het genereren van hoogwaardige steekproeven, maar ze moeten vaak worden afgestemd op specifieke downstream-doelstellingen (zoals esthetische kwaliteit of menselijke voorkeur). Bestaande fine-tuning-methoden voor diffusiemodellen, zoals Reinforcement Learning (RL) gebaseerde methoden (bijv. PPO/DDPO) en directe backpropagatie-methoden (bijv. DRaFT, ReFL), lijden vaak onder reward over-optimization.

Dit fenomeen treedt op wanneer een model obsessief streeft naar het maximaliseren van een beloningsfunctie, wat resulteert in:

1. Semantische ineenstorting (Semantic Collapse): De gegenereerde beelden verliezen hun betekenis en houden niet langer verband met de input-prompt.
2. Diversiteitsverlies (Diversity Collapse): De output convergeert naar een beperkt aantal patronen, waardoor de variatie in gegenereerde samples drastisch daalt.

Bestaande oplossingen die gebruikmaken van KL-divergentie regularisatie om dit te voorkomen, vereisen vaak het trainen van een onstabiele waardefunctie (Value Function) of vertrouwen op Monte Carlo-schatters met hoge variantie, wat de training inefficiënt en instabiel maakt.

Methodologie: SQDF

De auteurs stellen SQDF (Soft Q-based Diffusion Finetuning) voor, een nieuwe KL-geregulariseerde RL-methode die specifiek is ontworpen om reward over-optimization te mitigeren zonder de stabiliteit van de training te offeren.

Kernprincipes:

1. Training-vrije Soft Q-Functie: In plaats van een aparte Q-netwerk te trainen (wat instabiel is in diffusiemodellen), benadert SQDF de soft Q-functie via een één-staps posterior-middenschatting (gebaseerd op Tweedie's formule). Dit maakt gebruik van een "oracle" of proxy-model om de schone sample $\hat{x}_0$ direct te voorspellen vanuit een ruisige sample $x_{t-1}$.
2. Gereparameteriseerde Policy Gradient: Omdat de reward-functie differentieerbaar is, wordt de gradient van de reward direct gebruikt als de gradient van de Q-functie. Door de reparameterisatie-trick ($x_{t-1} = \mu_\theta(x_t, t) + \sigma_t \epsilon$) kunnen de parameters van het diffusiemodel direct worden bijgewerkt met een lage variantie, zonder backpropagatie door de volledige denoising-keten.
3. KL-Regularisatie: De methode bevat een KL-divergentie term die de getrainde policy dicht bij de oorspronkelijke, voorgeöefende distributie houdt, waardoor natuurlijke eigenschappen en diversiteit behouden blijven.

Drie Innovaties voor Stabiliteit en Efficiëntie:

• Disconteringsfactor ($\gamma$): In de vroege fasen van het denoising-proces heeft elke stap weinig invloed op het uiteindelijke resultaat. SQDF introduceert een factor $\gamma < 1$ om credits voor vroege stappen af te waarderen. Dit verbetert de "credit assignment" en vermindert de impact van onnauwkeurige schattingen in de vroege fasen.
• Integratie van Consistency Models: De standaard Tweedie's formule is onnauwkeurig bij hoge ruisniveaus (vroege stappen). SQDF gebruikt een Consistency Model om de posterior-middenschatting ($\hat{x}_0$) te verfijnen. Dit levert een betrouwbaardere schatting van de soft Q-functie op dan de traditionele methode.
• Off-Policy Replay Buffer: In tegenstelling tot eerdere on-policy methoden, gebruikt SQDF een replay buffer. Dit stelt het model in staat om zeldzame, hoog-beloonde en diverse samples opnieuw te gebruiken, wat de mode-coverage verbetert en het compromis tussen reward en diversiteit beter beheert.

Belangrijkste Resultaten

De auteurs evalueren SQDF in twee settings: text-to-image fine-tuning en online black-box optimalisatie.

1. Text-to-Image Fine-tuning:

   • Getest op Stable Diffusion 1.5 en XL met beloningen gebaseerd op LAION aesthetic scores en HPSv2 (menselijke voorkeur).
   • Resultaat: SQDF bereikt superieure target rewards terwijl het de diversiteit en semantische uitlijning behoudt. In tegenstelling tot baselines zoals DRaFT en ReFL, die snel diversiteit verliezen bij hoge rewards, blijft SQDF stabiel.
   • Vergelijking: Zelfs wanneer andere methoden worden aangepast met een extra KL-term, presteert SQDF beter op de Pareto-optimale curve (hogere rewards bij gelijke diversiteit).

2. Online Black-Box Optimalisatie:

   • In een setting met een beperkt budget voor het queryen van een "oracle" (waar de echte reward onbekend is en via een proxy wordt geschat), behaalt SQDF hoge sample-efficiëntie.
   • Het model blijft robuust en behoudt naturaliteit, terwijl methoden zoals SEIKO en PPO+KL vaak uit de distributie raken en kwaliteit verliezen.

3. Ablatie Studies:

   • Het verwijderen van de disconteringsfactor leidt tot langzamere convergentie en lagere uitlijning.
   • Het verwijderen van het Consistency Model vermindert de trainingsefficiëntie en de kwaliteit van de Q-schatting.
   • Het verwijderen van de replay buffer resulteert in een daling van de diversiteitsscores.

Bijdragen en Significantie

De belangrijkste bijdragen van dit werk zijn:

• Een nieuwe RL-formulering voor Diffusie: SQDF introduceert een methode die reward-gradients direct gebruikt via een gereparameteriseerde policy gradient, zonder de noodzaak van een onstabiel Q-netwerk.
• Oplossing voor Over-Optimalisatie: Het biedt een effectieve manier om reward-over-optimalisatie te voorkomen door een training-vrije soft Q-benadering te combineren met KL-regularisatie en een replay buffer.
• Verbeterde Schattingstechnieken: De integratie van Consistency Models voor de posterior-schatting en een disconteringsfactor voor betere credit assignment zijn cruciale technische verbeteringen die de stabiliteit van diffusie-RL aanzienlijk verhogen.
• Praktische Toepasbaarheid: De methode werkt zowel in settings met bekende gradients (text-to-image) als in black-box scenario's, wat het een veelzijdige tool maakt voor het alignen van generatieve modellen.

Conclusie:
SQDF stelt een nieuwe standaard voor het fine-tunen van diffusiemodellen. Het slaagt erin om de "Pareto-grens" te verschuiven, waardoor het mogelijk is om modellen te trainen die zowel hoge rewards behalen als de natuurlijke diversiteit en semantische consistentie van de gegenereerde content behouden. De code is open source beschikbaar gesteld.