CREPE: Controlling Diffusion with Replica Exchange

Dit artikel introduceert CREPE, een flexibele methode die replica-exchange gebruikt om diffusion-modellen tijdens de inferentie te sturen naar nieuwe constraints zonder hertraining, waarbij het ten opzichte van eerdere SMC-benaderingen sequentiële generatie, hoge diversiteit en online verfijning biedt.

De kern

Stel je voor dat je een zeer getalenteerde kunstenaar hebt die fantastische schilderijen kan maken. Deze kunstenaar is getraind op duizenden foto's van katten, auto's en landschappen. Maar nu wil je iets specifieks: "Teken een gele taxi, maar dan met een donkere achtergrond."

Als je gewoon vraagt, maakt de kunstenaar misschien een taxi, maar de achtergrond is misschien te licht, of de taxi is niet geel genoeg. De kunstenaar volgt zijn "gewoonten" (de data waarop hij is getraind) en negeert soms je specifieke wensen.

In de wereld van kunstmatige intelligentie noemen we dit diffusiemodellen. Ze werken door een beeld te maken van puur ruis (zoals statische op een oude TV) naar een helder beeld, stap voor stap.

Het probleem is: hoe krijg je deze kunstenaar om precies te doen wat jij wilt, zonder dat je hem opnieuw moet leren (retraining)?

De Oude Manier: De "SMC" Methode

Vroeger gebruikten onderzoekers een methode die lijkt op het hebben van een groot team van schilders die allemaal tegelijk werken.

1. Je geeft ze allemaal een opdracht.
2. Ze beginnen te schilderen.
3. Halverwege kijkt een manager naar de schilderijen. Die die niet goed genoeg zijn, gooit hij weg en laat hij de beste schilderijen "kopiëren" zodat er meer van die goede versies ontstaan.
4. Dit heet Sequential Monte Carlo (SMC).

Het probleem hiermee:

• Het is duur: je hebt veel rekenkracht nodig voor al die schilders tegelijk.
• Het wordt saai: als de manager te streng is, gaan alle schilders precies hetzelfde schilderij maken. Je krijgt geen diversiteit meer; iedereen tekent dezelfde saaie taxi.
• Als je halverwege zegt: "Oh, ik wil eigenlijk dat de taxi ook een dak heeft", moet je het hele team opnieuw beginnen. Je kunt niet makkelijk aan de bestaande schilderijen werken.

De Nieuwe Manier: CREPE

De auteurs van dit paper (Jiajun He en zijn team) hebben een slimme nieuwe manier bedacht, genaamd CREPE. Ze gebruiken een techniek uit de natuurkunde die "Replica Exchange" of "Parallel Tempering" heet.

Laten we dit uitleggen met een metabool:

Stel je voor dat je een groep vrienden hebt die allemaal een grote berg moeten beklimmen. De top van de berg is het perfecte schilderij dat jij wilt.

• De Oude Manier (SMC): Je stuurt 100 mensen tegelijk de berg op. Als ze vastlopen in een dal, stuur je ze terug en laat je de succesvolste mensen hun pad kopiëren. Maar vaak blijven ze allemaal in hetzelfde kleine dal hangen en missen ze de echte top.
• De Nieuwe Manier (CREPE): Je hebt niet 100 mensen die tegelijk lopen, maar je hebt één groep vrienden die op verschillende hoogtes van de berg lopen.
  • Vriend A loopt laag (dichtbij de start, veel ruis).
  • Vriend B loopt halverwege.
  • Vriend C loopt hoog (dichtbij de top, bijna klaar).

De Magie van CREPE:
In plaats van dat ze allemaal alleen lopen, mogen ze communiceren en van plek wisselen.

1. Vriend A (die laag loopt) probeert een stap omhoog te zetten.
2. Vriend C (die hoog loopt) probeert een stap omlaag te zetten.
3. Soms zeggen ze: "Hé, jouw positie hierboven is eigenlijk beter dan de mijne, wisselen we?"
4. Als de uitwisseling logisch is, wisselen ze van plek.

Waarom is dit zo cool?

• Diversiteit: Omdat ze op verschillende hoogtes lopen, ontdekken ze verschillende paden. Ze raken niet vast in één klein dal. Ze vinden de echte top sneller en met meer variatie.
• Online Verbetering: Stel dat je halverwege zegt: "Oh, de taxi moet ook een dak hebben!" In het oude systeem moest je alles opnieuw doen. Bij CREPE kun je gewoon een nieuwe regel toevoegen aan de "berg". De vrienden passen hun paden direct aan terwijl ze nog lopen. Je hoeft niet opnieuw te beginnen.
• Efficiëntie: Je hebt niet 100 mensen nodig die tegelijk hard werken. Je hebt een paar mensen die slim met elkaar communiceren.

Wat hebben ze bewezen?

De auteurs hebben CREPE getest op verschillende dingen:

1. Schilderijen: Ze konden specifieke instructies geven ("een gele taxi met donkere achtergrond") en kregen prachtige, diverse resultaten.
2. Moleculen: Ze hielpen bij het vinden van de beste vorm voor medicijnen (alsof je een puzzelstukje in de juiste vorm draait).
3. Tekst: Ze konden tekst genereren met een specifieke stemming (blij of verdrietig) zonder dat de tekst onleesbaar werd.

Samenvatting in één zin

CREPE is een slimme manier om een AI-kunstenaar te sturen door een team van "reizigers" op verschillende niveaus van een berg te laten lopen en met elkaar te laten wisselen, zodat ze sneller, slimmer en diverser de perfecte oplossing vinden, zonder dat je de kunstenaar opnieuw hoeft te leren.

Het is alsof je van een statische, saaie parade (SMC) overschakelt op een levendige, flexibele dans (CREPE) waarbij de dansers voortdurend van partner wisselen om de perfecte choreografie te vinden.

Technische samenvatting

Probleemstelling

Diffusiemodellen hebben de generatieve modellering revolutionair gemaakt, maar het besturen van deze modellen tijdens de inferentie (inference-time control) om aan nieuwe constraints te voldoen zonder het model opnieuw te trainen, blijft een uitdaging. Bestaande methoden vertrouwen vaak op heuristische geleiding of Sequential Monte Carlo (SMC) voor bias-correctie.

SMC-methoden hebben echter aanzienlijke beperkingen:

1. Hoge geheugeneisen: Ze vereisen het parallelle onderhoud van een groot aantal deeltjes (particles) gedurende het hele denoising-traject.
2. Lage diversiteit: SMC lijdt vaak aan "mode collapse" (het samenvoegen van deeltjes naar één modus), vooral bij een klein aantal deeltjes, wat resulteert in gebrek aan variatie in de gegenereerde samples.
3. Geen online verfijning: Eenmaal het sampling-proces voltooid is, kunnen gegenereerde samples niet verder worden verfijnd. Als de uitkomst onbevredigend is of nieuwe constraints worden toegevoegd, moeten volledig nieuwe samples worden gegenereerd.

Methodologie: CREPE

Het paper introduceert CREPE (Controlling with REPlica Exchange), een flexibel alternatief gebaseerd op Replica Exchange (ook wel bekend als Parallel Tempering of PT). CREPE biedt een computationeel dualistisch perspectief op SMC:

• Principe: In plaats van een batch deeltjes parallel langs het denoising-pad te laten evolueren (zoals bij SMC), draait CREPE ketens op verschillende denoising-stappen parallel en genereert het deeltjes sequentieel via Markov Chain Monte Carlo (MCMC).
• Het Algorithm (Algorithm 1):
  1. Anneling-pad: Er wordt een pad van verdelingen $(\pi_t)$ gedefinieerd dat interpolatie maakt tussen de doelverdeling (bij $t=0$) en een tractabele referentieverdeling (bij $t=1$, bijv. ruis).
  2. Communicatie (Swap): De kern van CREPE is de "swap"-beweging tussen buren in het anneringspad. In plaats van directe swaps, gebruikt CREPE Accelerated Parallel Tempering (APT). Dit simuleert voor- en achterwaartse voorsteltrajecten (proposal paths) tussen twee tijdstippen $t$ en $t'$. De acceptatiekans wordt berekend met behulp van de Radon-Nikodym Estimator (RNE), die de verhouding van de overgangsdynamica kwantificeert. Dit stelt het algoritme in staat om de swap-acceptatie te berekenen zonder de doelprioriteit expliciet te hoeven kennen, alleen de voorgeprogrammeerde diffusiemodellen.
  3. Lokale Exploratie: Optioneel kunnen lokale MCMC-bewegingen (zoals corrector-stappen) worden toegepast om de samples lokaal te verfijnen.
  4. Online Refinement: Een uniek kenmerk is dat CREPE "anytime" werkt. Nieuwe constraints kunnen tijdens het proces worden toegevoegd, en de bestaande samples kunnen verder worden verfijnd zonder het proces te herstarten.

Belangrijkste Bijdragen

1. Formulering van PT voor Diffusie: Het paper toont aan hoe Parallel Tempering direct kan worden toegepast op voorgeprogrammeerde diffusiemodellen zonder expliciete doeldichtheden, door gebruik te maken van de RNE-relatie.
2. Afleiding van Swap-rates: De auteurs hebben specifieke swap-rates afgeleid voor diverse inferentie-taakken, waaronder tempering, reward-tilting, het debiasen van classifier-free guidance (CFG) en modelcompositie, voor zowel Gaussische als discrete (mask) diffusiemodellen.
3. Superieure Diversiteit: In tegenstelling tot SMC, behoudt CREPE van nature een hogere diversiteit in de gegenereerde samples na een "burn-in" periode, omdat het geen zware resampling van een kleine batch vereist.
4. Flexibiliteit: Het ondersteunt online aanpassing van constraints en vroege beëindiging, wat SMC niet kan.

Resultaten

De auteurs evalueren CREPE op diverse domeinen en modaliteiten:

• Boltzmann Sampling (Moleculen): Bij het temperen van moleculen (Alanine Dipeptide, Tetrapeptide, Hexapeptide) presteert CREPE beter dan SMC op de meeste metrieken (Energy TVD, Sample W2, TICA MMD). CREPE behoudt betere diversiteit en vermijdt gemiste modi, terwijl SMC neigt tot mode collapse bij kleine batchgroottes.
• Afbeeldingen (ImageNet):
  • CFG Debiasing: Bij het debiasen van Classifier-Free Guidance op ImageNet-64 en -512, presteert CREPE beter dan de SMC-baselines (FKC) wat betreft FID (Frechet Inception Distance) en diversiteit, vooral bij een groter aantal gegenereerde samples.
  • Reward-tilting: CREPE slaagt erin om diverse afbeeldingen te genereren die nauw aansluiten bij complexe prompts (bijv. "een blauwe ballon"), terwijl SMC vaak visueel vergelijkbare samples produceert.
• Model Compositie (Maze Navigation): Bij het samenvoegen van korte trajecten tot een lange, coherente route door een doolhof, bereikt CREPE succespercentages die vergelijkbaar zijn met of beter zijn dan gespecialiseerde conditional modellen, met de flexibiliteit om nieuwe constraints (zoals een tussendoelpunt) online toe te voegen.
• Discrete Diffusie (Tekst): Bij sentiment-gestuurde tekstgeneratie (CTMC) toont CREPE een aanzienlijke verbetering in perplexiteit (tot een factor 5) ten opzichte van SMC, terwijl de sentiment-accuraatheid behouden blijft.

Betekenis en Conclusie

CREPE opent een nieuwe weg voor inference-time controle van diffusiemodellen door de beperkingen van SMC (diversiteit en geheugen) te omzeilen via een MCMC-gebaseerde aanpak. Het biedt een krachtig alternatief dat:

• Efficiënter is in termen van sample-diversiteit bij beperkte rekenkracht.
• Dynamischer is door ondersteuning voor online refinements en het toevoegen van constraints tijdens het generatieproces.
• Universeel toepasbaar is op zowel continue (afbeeldingen, moleculen) als discrete (tekst) data.

Hoewel CREPE een "burn-in" periode vereist en afhankelijk is van de nauwkeurigheid van het voorgeprogrammeerde model, biedt het een robuustere en flexibeler framework voor het sturen van generatieve modellen dan de huidige state-of-the-art SMC-methoden.