Diffusion Controller: Framework, Algorithms and Parameterization

Het paper introduceert Diffusion Controller (DiffCon), een unificerend controle-theoretisch kader dat reverse diffusion sampling als stochastische besturing ziet, waardoor er efficiënte versterkingsleer-methoden en een lichtgewicht zijnetwerk-parameterisatie ontstaan die de prestaties van gefinetunteerde diffusiemodellen verbeteren.

De kern

Stel je voor dat je een zeer getalenteerde schilder hebt die al duizenden schilderijen heeft gemaakt. Deze schilder (de Diffusiemodel, zoals Stable Diffusion) kan prachtige beelden maken van alles wat je beschrijft: een kat in een pak, een landschap, een futuristische stad. Hij is een meester, maar hij volgt soms zijn eigen inzicht en niet altijd precies wat jij wilt.

Soms wil je dat hij iets specifieker doet: "Maak die kat niet zomaar in een pak, maar laat hem ook een sigaar roken en een ernstig gezicht trekken." Of: "Maak dit landschap zo dat het eruitziet alsof het door een oude meester is geschilderd."

Tot nu toe was het aanpassen van deze schilder een beetje als een 'puzzel'. Mensen probeerden verschillende trucjes:

1. Instructies geven tijdens het schilderen: "Nee, niet zo, probeer het anders!" (Dit heet Guidance).
2. De schilder opnieuw laten oefenen: Je gaf hem duizenden voorbeelden van wat je wilde, zodat hij het leerde. Maar dit kostte enorm veel tijd en rekenkracht, en soms verloor hij zijn oorspronkelijke talent.
3. Een kleine 'bij-schilder' toevoegen: Iemand die alleen de details aanpast zonder de hele schilderijtechniek te veranderen (zoals LoRA).

Deze nieuwe paper introduceert DiffCon (Diffusion Controller). Het is een nieuwe manier om te denken over hoe je deze kunstenaars kunt sturen. Hier is de uitleg in gewone taal:

1. Het Grote Inzicht: Het is geen "Actie", maar een "Stuur"

Stel je voor dat de schilder een auto bestuurt die van een berg afrijdt (van ruiz naar een duidelijk beeld).

• De oude manier: Mensen dachten dat ze de auto moesten sturen door op een gaspedaal te drukken (een actie toevoegen).
• De DiffCon-methode: De auteurs zeggen: "Wacht, we hoeven niet op een pedaal te drukken. We hoeven alleen maar de windrichting een beetje te veranderen."

Ze zien het proces als een stuurman die de wind (de kansverdeling) een beetje kantelt. Als de wind normaal naar het noorden waait, kan de stuurman een klein beetje de zeilen aanpassen zodat de boot net iets meer naar het oosten gaat, zonder dat de boot zelf kapotgaat. Ze noemen dit een "Lineair Oplosbaar MDP" (een wiskundige term die betekent: "we kunnen dit probleem slim oplossen zonder alles opnieuw te moeten leren").

2. De Twee Slimme Trucjes (Algoritmes)

De paper biedt twee manieren om deze stuurman te trainen:

• Truc 1: De "Beloningstabel" (Reward-Weighted Regression)
  Stel je voor dat de schilder een wedstrijd doet. Als hij een plaatje maakt dat je leuk vindt, krijg je een punt.

  • De oude methode: "Probeer het nog eens, en als het goed is, onthoud het."
  • De DiffCon-methode: "Kijk naar alle plaatjes die hij heeft gemaakt. Die plaatjes die je het meest leuk vond, laten we veel vaker oefenen. Die die je niet leuk vond, laten we negeren."
    Ze gebruiken een wiskundige formule om te beslissen hoeveel je een plaatje moet "belonen". Dit zorgt ervoor dat de schilder zich richt op wat jij wilt, zonder dat hij zijn basisvaardigheden vergeet.

• Truc 2: De "PPO" (Proximal Policy Optimization)
  Dit is iets geavanceerder. Het is alsof je de schilder een coach geeft die zegt: "Je mag je stijl veranderen om beter te worden, maar niet te veel! Blijf dicht bij je oude, goede stijl." Dit voorkomt dat de schilder ineens begint te tekenen als een kindje omdat hij probeert te hard te werken.

3. De Innovatie: De "Bij-schilder" (Parameterization)

Dit is misschien wel het coolste deel.
Stel je voor dat je de grote meester (de Backbone) wilt aanpassen.

• De oude manier (White-box): Je neemt de meester mee naar school, laat hem alles uit zijn hoofd leren en hoopt dat hij het niet vergeet. Dit is zwaar en duur.
• De nieuwe manier (Gray-box / DiffCon): Je laat de meester precies doen wat hij altijd deed. Maar je plaatst een kleine, slimme assistent naast hem.
  • Deze assistent kijkt naar wat de meester aan het doen is (bijvoorbeeld: "Hij tekent nu een lijn").
  • De assistent zegt dan: "Wacht, als je die lijn hier een beetje buigt, wordt het een kat in een pak."
  • De assistent is heel klein en lichtgewicht. Hij verandert de meester niet, hij corrigeert hem alleen op het juiste moment.

Dit is belangrijk omdat je vaak geen toegang hebt tot de "geheime recepten" van de grote modellen (ze zijn gesloten). Met DiffCon kun je toch een assistent toevoegen die werkt, zelfs als je de grote meester niet mag aanraken.

4. Wat levert dit op?

De auteurs hebben dit getest op Stable Diffusion (een beroemd beeldgeneratie-model).

• Resultaat: Hun methode werkt beter dan de huidige beste methoden (zoals LoRA), zelfs met minder rekenkracht.
• Voorbeeld: Als je vraagt om "een blauwe vogel die spaghetti eet", maakt DiffCon een veel leuker en nauwkeuriger plaatje dan de standaardversie, en beter dan de andere aanpassingsmethoden.
• Efficiëntie: Omdat ze alleen een kleine assistent trainen in plaats van de hele schilder, gaat het veel sneller en kost het minder energie.

Samenvatting in één zin

DiffCon is een slimme manier om AI-kunstenaars te sturen door niet hun hele persoonlijkheid te herschrijven, maar door een kleine, slimme "stuurman" toe te voegen die de wind een beetje verandert zodat ze precies doen wat jij wilt, zonder hun talent te verliezen.

Technische samenvatting

Titel: Diffusion Controller (DiffCon): Een Unified Control-Theoretic Benadering voor Diffusiemodellen

1. Het Probleem

Hoewel diffusiemodellen (zoals Stable Diffusion) uitstekende prestaties leveren in het genereren van hoogwaardige afbeeldingen, blijft het gecontroleerd genereren (het sturen van samples om aan specifieke gebruikersintenties, beperkingen of doelen te voldoen) een uitdaging. Bestaande methoden voor controle zijn vaak een verzameling van heuristieken zonder een unified theoretisch kader:

• Inferentie-tijd mechanismen: Zoals classifier-free guidance (CFG), die vaak leiden tot kwaliteitsverlies bij sterke sturing.
• Trainings-tijd aanpassing: Zoals personalisatie via LoRA of domeinadaptatie, maar ook reinforcement learning (RL) finetuning op basis van menselijke feedback.
  Deze benaderingen worden vaak als losse "patches" gepresenteerd zonder een fundamenteel inzicht in hoe controle werkt binnen het stochastische proces van het denoisen. Er is behoefte aan een principieel kader dat deze methoden verenigt en efficiëntere, stabielere controle mogelijk maakt, zelfs wanneer de onderliggende backbone van het model niet volledig toegankelijk is (de "gray-box" setting).

2. Methodologie: Diffusion Controller (DiffCon)

De auteurs introduceren DiffCon, een raamwerk dat diffusie-sampling bekijkt als een state-only stochastisch controleprobleem binnen het kader van Linearly-Solvable Markov Decision Processes (LS-MDPs).

A. Theoretisch Kader (LS-MDP)
In tegenstelling tot standaard MDP-formuleringen die expliciete acties introduceren, ziet DiffCon de controle als het direct herwegen (reweighting) van de vooraf getrainde (passieve) reverse-time transitiekernen.

• De controle $u_t$ modificeert de overgangskern $p_{0,t}$ naar een gecontroleerde kern $P_{u,t}$ via een exponentiële tilting: $P_{u,t} \propto p_{0,t} \exp(u_t)$.
• Dit proces wordt gestuurd door een f-divergentie regularisatie (waarbij KL-divergentie een speciaal geval is). Dit creëert een optimale afweging tussen het maximaliseren van een beloning (bijv. menselijke voorkeur) en het dicht bij het vooraf getrainde model blijven om stabiliteit en kwaliteit te behouden.

B. Reinforcement Learning Finetuning (RLFT) Algoritmen
Vanuit de optimaliteitscondities van het LS-MDP kader leiden de auteurs twee praktische RL-update regels af voor het finetunen van diffusiemodellen zonder toegang tot doelsamples (alleen een beloningsmodel op het eindresultaat):

1. Policy Gradient met f-divergentie regularisatie: Een veralgemeende versie van PPO (Proximal Policy Optimization) die rekening houdt met de regularisatiekosten.
2. Reward-Weighted Regression (RWL): Een doelstelling waarbij de trainingssamples worden gewogen op basis van de ontvangen beloning. Voor KL-regularisatie resulteert dit in een exponentiële weging ($\exp(r/\tau)$), terwijl voor andere f-divergenties polynoom-achtige wegingen worden afgeleid. Deze methode garandeert dat de optimizer van de loss functie overeenkomt met de optimale verdeling onder de regularisatie.

C. Parameterisatie: De "Side Network" Architectuur
Een cruciale inzicht uit de LS-MDP-theorie is dat de optimale scorefunctie ($\epsilon^*$) kan worden ontbonden in een vaste, vooraf getrainde baseline ($\epsilon_0$) plus een lichte controlecorrectie. Dit leidt tot de DiffCon-parameterisatie:

• Gray-box compatibiliteit: De backbone van het model wordt bevroren. In plaats van de volledige scorefunctie te herschrijven, wordt een lichtgewicht "side network" ($s_\theta$) toegevoegd.
• Input: Het side network neemt de intermediate denoising outputs (zoals de vooraf getrainde reverse mean $\mu_0$) als input, in plaats van de ruwe noise $x_t$.
• Structuur: De output van het side network wordt gebruikt om de vooraf getrainde score te corrigeren via een gestructureerde, "gated" formule (een combinatie van een scalair gate $z$ en een vectoriële correctie $h$).
• Dit maakt DiffCon zeer efficiënt en geschikt voor scenario's waar de interne architectuur van het model verborgen is (bijv. bij commerciële API's), terwijl het toch superieure controle biedt.

3. Belangrijkste Resultaten

De auteurs evalueren DiffCon op Stable Diffusion v1.4 met drie finetuning-methoden: Supervised Finetuning (SFT), Reward-Weighted Loss (RWL) en PPO. De prestaties worden gemeten aan de hand van de HPS-v2 win rate (hoe vaak het gefinetunte model een menselijke voorkeur wint ten opzichte van het originele model).

• Superieure Prestaties: DiffCon overtreft consistent de vooraf getrainde modellen en bestaande baselines.
  • Bij SFT en RWL presteert de gray-box variant van DiffCon (met minder parameters) beter dan LoRA (een white-box methode die interne lagen aanpast).
  • Bij PPO bereiken de white-box varianten van DiffCon (DiffCon-J en DiffCon-S) win rates van boven de 93%, wat aanzienlijk hoger is dan LoRA (~90%).
• Kwaliteit-Efficiëntie Trade-off: DiffCon behoudt beter de oorspronkelijke beeldkwaliteit (gemeten via CLIP, PickScore en CLIP-Aesthetics) terwijl het de voorkeursuitlijning maximaliseert.
• Ablatie Studies: Het gebruik van de gestructureerde parameterisatie (input $\mu_0$ in plaats van $x_t$, en de gescheiden output heads) blijkt essentieel voor de prestaties. Ook de keuze van de regularisatiecoëfficiënt ($\tau$) en de wegingsfunctie heeft een grote impact.

4. Significatie en Impact

• Unificatie: Het paper biedt een eenduidig theoretisch kader dat diverse bestaande technieken (zoals guidance, LoRA, en RLFT) verenigt onder de paraplu van LS-MDPs.
• Gray-Box Revolutie: Het bewijst dat je geen volledige toegang tot de interne lagen van een groot model nodig hebt om effectieve controle te bereiken. Door slimme parameterisatie van een side network kunnen zelfs "gesloten" modellen effectief worden aangepast.
• Stabiliteit: De methode biedt een principieel mechanisme om de stabiliteit van het vooraf getrainde model te behouden tijdens het sturen naar nieuwe doelen, wat vaak een probleem is bij agressieve finetuning.
• Toekomstperspectief: De aanpak is schaalbaar naar bredere toepassingen zoals personalisatie, veiligheidsuitlijning (safety alignment) en transfer learning voor andere generatieve taken.

Kortom, DiffCon transformeert het probleem van diffusiecontrole van een verzameling heuristieken naar een wiskundig onderbouwde, efficiënte en flexibele oplossing die zowel voor open als gesloten modellen werkt.