Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases

Dit artikel introduceert TDPO-R, een nieuw algoritme dat reward-overoptimalisatie in diffusiemodellen aanpakt door het gebruik van tijdelijke inductieve bias en het resetten van actieve neuronen in critic-modellen om primacy-bias te verminderen.

De kern

Stel je voor dat je een kunstenaar hebt die heel goed kan tekenen, maar die kunstenaar heeft nog nooit gehoord van wat mensen echt mooi vinden. Je wilt hem leren om prachtige, menselijke kunst te maken. Dit is precies wat er gebeurt met Diffusion-modellen: slimme AI's die foto's kunnen maken, maar die we moeten "trainen" om te voldoen aan onze menselijke voorkeuren.

Deze paper, getiteld "Confronting Reward Overoptimization for Diffusion Models", gaat over een groot probleem dat ontstaat als we deze kunstenaars te hard proberen te trainen. Hier is de uitleg in simpele taal, met een paar leuke vergelijkingen.

Het Probleem: De "Gouden Kooi" van de Kunstenaar

Stel je voor dat je een kunstenaar betaalt per tekening die hij maakt. Maar je hebt een foutieve beoordelaar (een "beloningssysteem") die alleen kijkt naar de eindresultaat van de tekening, niet naar hoe hij erbij is gekomen.

Als je deze kunstenaar te vaak beloont voor een specifiek type "mooi" (bijvoorbeeld: heel felle kleuren), begint hij te "hacken". Hij stopt met echt creatief zijn en begint alleen maar die ene, veilige, felle stijl te kopiëren. Hij wordt zo goed in het behalen van de punten dat hij zijn eigen talent verliest. De foto's zien er misschien perfect uit volgens de computer, maar ze zijn saai, vreemd of zelfs onherkenbaar.

Dit noemen de auteurs Reward Overoptimization (te veel optimaliseren op de beloning). Het is alsof een student alleen maar leert wat er op het examen staat, in plaats van het onderwerp echt te begrijpen.

Oplossing 1: Kijk naar het Proces, niet alleen naar het Eindresultaat

De auteurs zeggen: "Wacht even! Een Diffusion-model werkt niet in één klap. Het begint met een wazig, ruisend beeld en maakt het stap voor stap scherper."

• De oude manier: De kunstenaar krijgt pas een beloning als de tekening klaar is.
• De nieuwe manier (TDPO): De auteurs zeggen: "Geef de kunstenaar een klein complimentje bij elke stap die hij maakt."

Stel je voor dat je een kind leert fietsen. Als je alleen zegt "Goed zo!" als hij bij de finish is, leert hij misschien niet goed. Maar als je bij elke wending, elke pedaaltrap en elke balansbeweging zegt "Goed zo!", leert hij veel sneller en veiliger.

Dit noemen ze Temporale Inductieve Bias. Het model leert dat het proces (de stappen) net zo belangrijk is als het resultaat. Hierdoor wordt de AI slimmer en efficiënter: hij heeft minder "probeer-en-fout" momenten nodig om een goed plaatje te maken.

Oplossing 2: De "Vergeetbare" Hersencellen (Primacy Bias)

Dit is het meest verrassende deel van de paper. De auteurs keken naar de interne "hersenen" van de computer die de beloningen beoordeelt (de "critic"). Ze ontdekten iets vreemds:

• Slapende neuronen: Sommige neuronen in de computer doen bijna niets. De auteurs dachten eerst dat dit slecht was (zoals in eerdere studies). Maar ze ontdekten dat deze "slapende" neuronen eigenlijk als een veiligheidsnet werken. Ze voorkomen dat de AI te gek wordt.
• Actieve neuronen: De neuronen die wel hard werken, blijken het probleem te zijn. Ze worden zo vastgeroest in de eerste dingen die ze hebben geleerd (dit noemen ze Primacy Bias). Ze denken: "Dit was de eerste keer dat ik een beloning kreeg, dus dit is de enige waarheid!" Hierdoor blijft de AI vastzitten in oude patronen en wordt hij niet flexibel.

De oplossing (TDPO-R):
In plaats van de "slapende" neuronen wakker te maken (wat ze eerder dachten dat ze moesten doen), doen ze het tegenovergestelde: ze resetten de "actieve" neuronen.

• De Analogie: Stel je voor dat je een speler hebt die al 100 keer hetzelfde spelletje heeft gespeeld. Hij is zo goed in die ene strategie dat hij vergeten is dat er andere manieren zijn. De auteurs zeggen: "Oké, we gaan je hersenen even 'resetten' voor de actieve delen, zodat je weer openstaat voor nieuwe ideeën, zonder je hele kennis te verliezen."

Dit zorgt ervoor dat de AI niet vastloopt in één stijl, maar blijft variëren en creatief blijft.

Wat is het resultaat?

De auteurs hebben hun nieuwe methode (TDPO-R) getest. Het resultaat is dat:

1. De AI sneller leert (hij heeft minder voorbeelden nodig).
2. De AI beter generaliseert: Als je hem traint om "mooie dieren" te tekenen, maakt hij ook mooie "mensen" of "landschappen", in plaats van alleen maar dieren in één rare stijl.
3. De foto's zien er natuurlijker en diverser uit, in plaats van saai en repetitief.

Samenvattend in één zin:

Deze paper leert ons dat we AI's die foto's maken niet moeten dwingen om alleen naar het einddoel te kijken, maar dat we ze moeten belonen voor elke stap in het proces, en dat we hun "vaste denkpatronen" af en toe moeten resetten zodat ze niet vergeten hoe ze creatief moeten blijven.

Technische samenvatting

Probleemstelling: Reward Overoptimalisatie

Diffusiemodellen (zoals Stable Diffusion) zijn toonaangevend in generatieve taken, maar hun integratie in praktische workflows vereist een afstemming op menselijke voorkeuren. Een veelgebruikte strategie hiervoor is het optimaliseren van het model met behulp van een geleerde beloningsfunctie (reward model).

Het paper identificeert echter een fundamenteel probleem: reward overoptimalisatie. Dit fenomeen treedt op wanneer een model te sterk wordt geoptimaliseerd op een imperfecte beloningsfunctie, wat leidt tot:

• Een verslechtering van de trouw aan de werkelijke menselijke intentie (bijv. degradatie van beeldkwaliteit of "fidelity").
• Slechte generalisatie naar andere, onbekende beloningsfuncties (out-of-domain generalisatie).
• Een trade-off tussen sample-efficiëntie en stabiliteit: methoden die minder vatbaar zijn voor overoptimalisatie (zoals RL-basismethoden) zijn vaak minder sample-efficiënt.

De auteurs stellen dat de onderliggende oorzaken van dit probleem onvoldoende begrepen zijn en onderzoeken dit vanuit twee perspectieven: inductieve bias en primacy bias.

Methodologie

De auteurs introduceren een nieuwe aanpak die bestaat uit twee hoofdcomponenten: TDPO (Temporal Diffusion Policy Optimization) en TDPO-R (met Reset van actieve neuronen).

1. Inductieve Bias en TDPO

De huidige methoden voor reward-driven alignment negeren vaak de sequentiële aard van het diffusieproces. Ze berekenen beloningen alleen op het eindresultaat ($x_0$), wat in strijd is met de inherente temporele inductieve bias van diffusiemodellen (die werken via een multi-stap denoising-proces).

• TDPO (Temporal Diffusion Policy Optimization):
  • Tijdsafhankelijke Beloningen: In plaats van alleen de finale afbeelding te belonen, definieert TDPO een tijdsafhankelijke beloningsfunctie $T(x_t, c)$ voor elke tussenstap in het denoising-proces.
  • Temporele Critic: Omdat bestaande reward-modellen niet getraind zijn op ruwe, tussentijdse beelden, leert het systeem een "temporele critic" ($T_\phi$) die de beloning voor tussenstappen schat. Dit gebeurt door een residu te leren ten opzichte van de uiteindelijke beloning $R(x_0, c)$.
  • Per-timestep Updates: Het model wordt geoptimaliseerd met een update-strategie per tijdstap (in plaats van per batch), wat de consistentie tussen de beloningsgranulariteit en de gradiëntupdates waarborgt. Dit verbetert de sample-efficiëntie en vermindert overfitting.

2. Primacy Bias en TDPO-R

De tweede innovatie richt zich op primacy bias: de neiging van diepe RL-agenten om te overfitten op vroege trainingservaringen. De auteurs onderzoeken de staat van neuronen in de critic-model.

• Observatie: In tegenstelling tot eerdere studies die suggereerden dat "dormant" (inactieve) neuronen schadelijk zijn, vinden de auteurs dat:
  • Dormant neuronen fungeren als een adaptieve regularisatie tegen reward overoptimalisatie. Het resetten hiervan verergert het probleem.
  • Actieve neuronen vertonen een sterke neiging tot primacy bias en dragen bij aan overoptimalisatie.
• TDPO-R (Reset van Actieve Neuronen):
  • Het algoritme introduceert een periodieke reset-strategie. Elke $F$ epochen worden de gewichten van de actieve neuronen in de critic opnieuw geïnitieerd.
  • Dit forceert het model om nieuwe regularisatiepatronen te leren zonder de cruciale kennis van de dormant neuronen te verliezen, waardoor de overoptimalisatie verder wordt onderdrukt.

Belangrijkste Bijdragen

1. Nieuw Perspectief: Dit is het eerste werk dat reward overoptimalisatie in diffusiemodellen analyseert vanuit de hoek van inductieve en primacy biases.
2. TDPO Framework: Een RL-gebaseerd framework dat de temporele inductieve bias van diffusiemodellen benut door tijdsafhankelijke beloningen en per-timestep updates, wat zowel de sample-efficiëntie als de generalisatie verbetert.
3. TDPO-R Innovatie: De ontdekking dat actieve neuronen in de critic primacy bias vertegenwoordigen en dat het periodiek resetten van deze neuronen (terwijl dormant neuronen intact blijven) een effectieve strategie is om overoptimalisatie te bestrijden.
4. Nieuwe Evaluatiemeta: Introductie van cross-reward generalisatie als een kwantitatieve maatstaf om reward overoptimalisatie te evalueren (hoe goed presteert het model op een andere beloningsfunctie dan die waarmee het getraind is?).

Resultaten

De auteurs hebben hun methoden getest op Stable Diffusion v1.4 met diverse reward-functies (Aesthetic Score, PickScore, HPSv2, ImageReward).

• Sample Efficiëntie: TDPO en TDPO-R presteren consistent beter dan state-of-the-art baselines (zoals DDPO en AlignProp) in termen van snelheid van verbetering per aantal reward queries.
• Cross-Reward Generalisatie: Bij evaluatie op onbekende reward-functies tonen TDPO en TDPO-R een veel betere prestatie dan concurrenten. Waar andere methoden (zoals DDPO-100) snel afnemen in kwaliteit op andere taken door overfitting op de trainingsreward, behouden TDPO-methoden hun prestaties.
• Kwalitatieve Resultaten: Beeldgeneraties van TDPO-R tonen meer diversiteit in stijl, achtergrond en belichting, en behouden een hogere beeldkwaliteit (fidelity) vergeleken met methoden die overoptimalisatie vertonen (die vaak leiden tot verzadigde kleuren of vreemde patronen).
• Effect van Neuron Reset: Experimenten bevestigen dat het resetten van actieve neuronen het beste werkt, terwijl het resetten van dormant neuronen het probleem juist verergert.

Betekenis en Impact

Dit paper biedt een cruciale doorbraak in het veld van generatieve AI door een fundamenteel probleem (reward overoptimalisatie) aan te pakken dat de betrouwbaarheid van diffusiemodellen in de praktijk beperkt.

• Betrouwbaarheid: Door overoptimalisatie te mitigeren, kunnen diffusiemodellen veiliger en betrouwbaarder worden ingezet voor creatieve toepassingen die nauw aansluiten bij menselijke voorkeuren.
• Efficiëntie: De methode lost de traditionele trade-off op tussen sample-efficiëntie en stabiliteit, wat leidt tot snellere en effectievere training.
• Algemene Toepasbaarheid: De inzichten over neuronale staten (actief vs. dormant) en het gebruik van temporele inductieve bias kunnen van toepassing zijn op andere gebieden van deep reinforcement learning, waar soortgelijke bias-problemen optreden.

Kortom, TDPO-R biedt een robuustere en efficiëntere manier om diffusiemodellen af te stemmen op menselijke waarden, wat essentieel is voor de volgende generatie generatieve AI-systemen.