Diffusion Blend: Inference-Time Multi-Preference Alignment for Diffusion Models

Het artikel introduceert Diffusion Blend, een nieuwe methode die tijdens de inferentie meerdere voorkeuren voor diffusiemodellen kan combineren door achterwaartse diffusieprocessen van fijngefineerde modellen te blenden, waardoor gebruikers dynamisch verschillende beloningen en regularisatie kunnen aanpassen zonder extra training.

De kern

Diffusion Blend: De "Mix-and-Match" voor AI-Kunst

Stel je voor dat je een supergetalenteerde schilder hebt die al duizenden prachtige schilderijen heeft gemaakt. Deze schilder (de AI) is erg goed, maar soms wil jij iets specifieks: een appel die er niet alleen echt uitziet, maar ook heerlijk smaakt, of een landschap dat er mooi uitziet maar ook precies de tekst "zonsondergang" weergeeft.

Het probleem is dat deze schilder meestal maar één opdracht tegelijk kan perfect uitvoeren. Als je hem vraagt om op te letten op de smaak, wordt de kleur misschien minder mooi. Als je hem vraagt om op de kleur te letten, wordt de tekst misschien onleesbaar.

Tot nu toe moesten mensen voor elke nieuwe wens een nieuwe schilder trainen. Dat is duur, tijdrovend en onhandig.

De Oplossing: Diffusion Blend

De auteurs van dit papier hebben een slimme truc bedacht, genaamd Diffusion Blend. In plaats van een nieuwe schilder te trainen, laten ze de bestaande schilder werken als een DJ die nummers mixt.

Hier is hoe het werkt, stap voor stap:

1. De Basis: De "Basisrecepten"

Stel je voor dat je drie verschillende recepten hebt voor een taart:

• Recept A: De taart is perfect voor de tekst (de tekst staat er duidelijk op).
• Recept B: De taart is perfect voor de smaak (zeer lekker).
• Recept C: De taart is perfect voor de presentatie (erg mooi om te zien).

Normaal gesproken zou je voor elk recept een aparte bakker moeten inhuren. Maar bij Diffusion Blend train je eerst één keer drie aparte bakkers (AI-modellen) op deze specifieke recepten.

2. De Magie: Het Muziek-Mixen (Inference-Time)

Nu komt het leuke deel. Je wilt een taart die 70% op de tekst lijkt en 30% op de smaak.

• De oude manier: Je zou een nieuwe bakker moeten trainen die precies die 70/30 verhouding kent.
• De Diffusion Blend manier: Je neemt de drie bakkers die je al hebt, en je laat ze tegelijkertijd aan het werk. Je zegt: "Bakker A, doe 70% van je werk. Bakker B, doe 30% van je werk."

De AI "mixt" de instructies van deze bakkers op het moment dat de taart wordt gemaakt (tijdens het genereren van de afbeelding). Het resultaat is een taart die precies die 70/30 verhouding heeft, zonder dat je ooit een nieuwe bakker hebt moeten inhuren.

3. De "Regelknop" voor Zekerheid (KL Regularization)

Soms wil je dat de AI heel creatief is en afwijkt van de standaard (veel vrijheid), en soms wil je dat hij heel strikt blijft bij de originele stijl (weinig vrijheid).
Diffusion Blend heeft ook een knop hiervoor. Je kunt zeggen: "Ik wil 20% creativiteit" of "Ik wil 80% creativiteit". De AI past dit direct aan door de instructies van de standaardbakker en de creatieve bakker te mengen. Het is alsof je een dimmerknop op je lamp hebt: je kunt het licht zachtjes aan of hard aan doen, zonder de lamp zelf te vervangen.

Waarom is dit zo cool?

• Snelheid: Je hoeft niet te wachten op nieuwe training. Je kunt direct wisselen van wens.
• Flexibiliteit: Je kunt elke willekeurige combinatie maken. Wil je 40% tekst, 40% smaak en 20% kleur? Geen probleem.
• Efficiëntie: De versie genaamd DB-MPA-LS is nog slimmer. In plaats van alle bakkers tegelijk te laten werken (wat zwaar is voor de computer), kiest deze versie per stap willekeurig één bakker uit, maar wel op een slimme manier die precies hetzelfde resultaat geeft. Het is alsof je een DJ hebt die snel van plaat wisselt, maar voor de luisteraar klinkt het als één perfect mix.

Samenvattend

Diffusion Blend is als een superkrachtige keukenmixer voor AI-kunst. Je gooit verschillende "ingrediënten" (de verschillende getrainde modellen) in de kom, draait de knop naar de verhouding die jij wilt, en de AI maakt in een handomdraai een perfect afgewerkt schilderij dat precies aan jouw wensen voldoet. Geen nieuwe training nodig, gewoon mixen en genieten!

Technische samenvatting

1. Probleemstelling

Diffusiemodellen (zoals Stable Diffusion) zijn uitstekend in het genereren van beelden, maar zijn vaak niet optimaal afgestemd op specifieke downstream doelen zoals esthetische kwaliteit, tekst-beeld consistentie of menselijke voorkeuren. Traditionele aanpakken gebruiken Reinforcement Learning (RL) om een voorgeprent model te fine-tunen op een enkele beloningsfunctie, vaak met een Kullback-Leibler (KL) regularisatie om te voorkomen dat het model te ver afwijkt van de originele verdeling (reward hacking).

De huidige beperkingen zijn echter:

• Statische afstemming: Een gefinetuned model is geoptimaliseerd voor één specifieke combinatie van beloningsfuncties en regularisatiesterkte.
• Gebrek aan flexibiliteit: Gebruikersvoorkeuren variëren per prompt, individu en context. Om een nieuwe afweging (trade-off) te maken tussen bijvoorbeeld esthetiek en tekst-accuraatheid, moet er vaak een nieuw model worden getraind, wat computatief zwaar en onpraktisch is.
• Inference-time beperkingen: Bestaande methoden voor multi-objective alignment vereisen vaak zware sampling, gradiënten of het samenvoegen van modelparameters (zoals bij "Rewarded Soup"), wat leidt tot hoge inferencekosten of suboptimale resultaten.

Het paper stelt de volgende vraag: Kan men een fine-tuning-procedure ontwerpen waarbij het model tijdens de inference (tijdens het genereren) beelden kan produceren die zijn afgestemd op elke lineaire combinatie van basisbeloningsfuncties en een willekeurige KL-regularisatiesterkte, zonder extra fine-tuning?

2. Methodologie: Diffusion Blend

De auteurs introduceren Diffusion Blend, een raamwerk dat de achterwaartse diffusieprocessen van verschillende gefinetuned modellen "mixt" om een nieuwe, controllable generatie te realiseren. De kern van de methode rust op theoretische resultaten die de relatie tussen de scorefuncties van een voorgeprent model en een met RL gefinetuned model in kaart brengen.

Theoretische Basis

De auteurs tonen aan dat de achterwaartse stochastische differentiaalvergelijking (SDE) van een model dat is afgestemd op een beloning $r$ en regularisatie $\alpha$, kan worden uitgedrukt als de SDE van het voorgeprente model plus een extra controleterm die afhangt van $r$.
Ze leiden een benadering af waarbij deze controleterm lineair kan worden benaderd door de controletermen van individuele basis-modellen te combineren. Dit maakt het mogelijk om een nieuw gedrag te synthetiseren door de drift-termen (scorefuncties) van bestaande modellen te mixen.

De Algoritmen

Het paper presenteert drie specifieke algoritmen:

1. DB-MPA (Diffusion Blend - Multi-Preference Alignment):

   • Doel: Afstemming op een lineaire combinatie van meerdere beloningsfuncties $r(w) = \sum w_i r_i$.
   • Werkwijze: Tijdens de inference wordt de drift-term $f(x_t, t)$ berekend als een gewogen som van de drift-termen van $m$ individueel gefinetuned modellen (elk getraind op een basisbeloning $r_i$).
   • Formule: $f(r(w), \alpha) \approx \sum w_i f(r_i, \alpha)$.
   • Voordeel: Het kan willekeurige voorkeuren realiseren zonder extra training.

2. DB-KLA (Diffusion Blend - KL Alignment):

   • Doel: Controle over de regularisatiesterkte (hoe ver het model mag afwijken van het origineel) zonder extra training.
   • Werkwijze: Het mixt de drift-term van het gefinetuned model (met vaste $\alpha$) met die van het voorgeprente model, gebaseerd op een door de gebruiker opgegeven factor $\lambda$.
   • Formule: $f(r, \alpha/\lambda) \approx (1-\lambda)f^{pre} + \lambda f(r, \alpha)$.

3. DB-MPA-LS (Diffusion Blend - Multi-Preference Alignment - LoRA Sampling):

   • Doel: Het verminderen van de inferencekosten van DB-MPA.
   • Probleem: DB-MPA vereist dat bij elke denoising-stap alle $m$ modellen worden uitgevoerd, wat de tijd $m$-voudig vergroot.
   • Oplossing: In plaats van alle modellen te mixen, wordt er bij elke stap stochastisch één LoRA-adapter (van de gefinetuned modellen) geselecteerd met een kans evenredig aan de gewenste gewichten $w_i$.
   • Theoretisch bewijs: De auteurs bewijzen dat deze stochastische selectie leidt tot dezelfde marginale verdeling als de deterministische mix, waardoor de inferencekosten terugvallen naar dat van één enkel model (geen lineaire schaalvergroting).

3. Belangrijkste Bijdragen

• Theoretische Afleiding: Een nieuwe theoretische benadering die toont hoe de achterwaartse diffusie van een multi-objective model kan worden benaderd door een lineaire combinatie van de drift-termen van single-objective modellen.
• Inference-Time Flexibiliteit: Een framework dat gebruikers toelaat om tijdens het genereren de afweging tussen verschillende doelen (bijv. esthetiek vs. tekst-accuraatheid) en de regularisatiesterkte dynamisch aan te passen zonder het model opnieuw te hoeven trainen.
• Efficiëntie: De introductie van DB-MPA-LS, die de inferencekosten van multi-reward alignment reduceert tot die van een standaard diffusiemodel, een significant voordeel ten opzichte van bestaande methoden.
• Uitgebreide Validatie: Experimenten tonen aan dat de methode dicht in de buurt komt van de theoretische bovengrens (Oracle MORL) en beter presteert dan bestaande baselines zoals Rewarded Soup, CoDe en Reward Gradient Guidance.

4. Resultaten

De auteurs hebben hun methoden geëvalueerd met Stable Diffusion v1.5 (en uitgebreid naar SDXL) op diverse datasets (DrawBench, GenEval) met meerdere beloningsmodellen (ImageReward, VILA, PickScore, JPEG-compressibility).

• Pareto-optimaliteit: DB-MPA bereikt een Pareto-front die zeer dicht ligt bij die van MORL (Multi-Objective RL), wat betekent dat het de beste mogelijke trade-offs haalt tussen conflicterende doelen.
• Prestatie vs. Baselines: DB-MPA en DB-MPA-LS overtreffen consistent methoden zoals Rewarded Soup (RS), CoDe en RGG. Bijvoorbeeld, bij een gewicht $w=0.5$ voor twee beloningen, presteerde DB-MPA aanzienlijk beter dan RS en CoDe.
• Schaalbaarheid: De prestaties van DB-MPA-LS blijven stabiel naarmate het aantal beloningsfuncties toeneemt (van 2 naar 4), terwijl de prestaties van Rewarded Soup sterk dalen bij meer doelen.
• Inference Snelheid: DB-MPA-LS heeft een inference-tijd die vergelijkbaar is met het standaard Stable Diffusion-model, terwijl DB-MPA (zonder LS) ongeveer $m$ keer zo lang duurt (waarbij $m$ het aantal rewards is). Bestaande methoden zoals CoDe en RGG hebben veel hogere kosten door multi-particle sampling of gradiëntberekeningen.
• Visuele Kwaliteit: Visuele vergelijkingen tonen aan dat DB-MPA realistischere en semantisch accuratere beelden genereert dan baselines, zelfs bij conflicterende beloningen (bijv. esthetiek vs. JPEG-compressibiliteit).

5. Betekenis en Impact

Dit werk is significant omdat het een fundamentele beperking van huidige RL-gebaseerde afstemming van diffusiemodels oplost: de rigiditeit van statische fine-tuning.

• Gebruikersvriendelijkheid: Het stelt ontwikkelaars en eindgebruikers in staat om real-time de "toon" van de generatie aan te passen (bijv. meer artistiek vs. meer tekst-accuraat) zonder zware computercapaciteit voor hertraining.
• Efficiëntie: Door de inferencekosten te minimaliseren via DB-MPA-LS, wordt multi-preference alignment praktisch toepasbaar in real-time toepassingen.
• Toekomstperspectief: Het biedt een schaalbare route naar gepersonaliseerde generatieve AI, waarbij één set gefinetuned modellen kan dienen als basis voor een oneindig aantal gebruikersvoorkeuren.

Kortom, Diffusion Blend transformeert multi-objective alignment van een dure, statische trainingstak naar een flexibele, inference-time operatie, met behoud van hoge kwaliteit en theoretische onderbouwing.