Exploring Diffusion Models' Corruption Stage in Few-Shot Fine-tuning and Mitigating with Bayesian Neural Networks

Dit paper identificeert en analyseert een 'corruptiestadium' tijdens het few-shot fine-tunen van diffusiemodellen, waarbij beeldkwaliteit tijdelijk verslechtert door een versmalde leerverdeling, en lost dit op door Bayesiaanse neurale netwerken toe te passen om de distributie te verbreden en zo de beeldkwaliteit, diversiteit en trouw te verbeteren zonder extra inferentiekosten.

De kern

Stel je voor dat je een zeer getalenteerde kunstenaar hebt die alles kan schilderen: van landschappen tot portretten. Deze kunstenaar is getraind op miljoenen foto's en kan dus van alles maken. Dit is wat we een Diffusiemodel noemen in de wereld van kunstmatige intelligentie.

Nu wil je dat deze kunstenaar iets heel specifieks leert schilderen: bijvoorbeeld jouw eigen hond, of een heel specifieke vaas die je hebt. Maar je hebt maar één of een paar foto's van die hond of vaas. Dit noemen we "few-shot fine-tuning" (fijnafstemming met weinig voorbeelden).

De onderzoekers van dit paper ontdekten dat er iets raars gebeurt als je de kunstenaar probeert te leren met zo weinig foto's. Het proces verloopt in drie vreemde fases:

1. Het begin: De "Goede" Fase

Eerst gaat het goed! De kunstenaar kijkt naar je foto's en begint steeds beter te lijken op jouw hond. De schilderij wordt steeds trouwer aan het origineel.

2. Het mysterieuze probleem: De "Corruptie-fase"

Dan gebeurt er iets vreemds. Plotseling begint het schilderij te verpesten. Er verschijnen rare ruis, vlekken en gekke patronen op het beeld. Het lijkt alsof de kunstenaar in paniek raakt en begint te tekenen met een trillende hand. De kwaliteit daalt drastisch, terwijl je dacht dat het alleen maar beter zou worden. De onderzoekers noemen dit de "Corruptie-fase".

3. Het einde: De "Overdrijving"

Uiteindelijk stopt de ruis weer, maar dan is de kunstenaar te ver gegaan. Hij kan nu alleen nog maar exact die ene foto van je hond natekenen. Hij is vergeten hoe je een hond in een andere pose of met een andere achtergrond moet schilderen. Hij is "overgefit" (overgevoelig) geworden.

---

Waarom gebeurt dit? (De Analoge Uitleg)

De onderzoekers hebben een simpele theorie bedacht om dit uit te leggen.

Stel je voor dat de kunstenaar een groot bibliotheek heeft van alle mogelijke honden die hij ooit heeft gezien.

• Normaal: Als je vraagt om "een hond", kiest hij een hond uit die grote bibliotheek.
• Few-shot (Weinig foto's): Je geeft hem maar één foto. Hij probeert zijn hele bibliotheek te vervangen door alleen die ene foto.

Het probleem is dat hij zijn bibliotheek te snel en te smal maakt. Hij denkt: "Oké, dit is de enige hond die bestaat."
Wanneer hij probeert iets te schilderen dat net iets anders is dan die ene foto (bijvoorbeeld een hond die loopt in plaats van zit), raakt hij in de war. Omdat zijn "bibliotheek" zo klein is, kan hij geen goede keuze maken. Hij probeert de ruis (de onzekerheid) te negeren, maar in plaats daarvan vermenigvuldigt hij die ruis en wordt het beeld wazig en korrelig. Dat is de corruptie.

---

De Oplossing: De "Baysean" Kunstenaar

Hoe los je dit op? De onderzoekers gebruiken een slimme truc met Bayesiaanse Neurale Netwerken (BNNs).

In plaats van dat de kunstenaar denkt: "Ik weet het zeker, dit is de enige manier om een hond te schilderen", leren ze hem: "Ik ben niet 100% zeker. Er zijn misschien wel een paar manieren om deze hond te schilderen."

• De Metafoor: Stel je voor dat je een gewone kunstenaar bent die alleen maar exacte kopieën maakt. Dan geef je hem een bril met een wazig filter (de BNN). Door dat filter moet hij niet perfect kopiëren, maar moet hij schatten en varieren.
• Het Effect: Omdat hij gedwongen wordt om te variëren en niet alleen maar één exacte kopie te maken, blijft zijn "bibliotheek" van mogelijke honden groter. Hij leert dat er ruimte is voor variatie.
• Het Resultaat: Die rare ruis en vlekken (de corruptie) verdwijnen. De kunstenaar kan nu weer prachtige, diverse foto's maken van jouw hond, zonder dat het beeld kapotgaat.

Waarom is dit geweldig?

1. Geen extra werk: Het kost geen extra tijd om het schilderij te maken (geen extra "inference costs"). Het is alsof je de bril afzet nadat je geoefend hebt; het schilderij zelf is net zo snel klaar.
2. Werkt overal: Of je nu DreamBooth, LoRA of OFT gebruikt (verschillende methoden om de kunstenaar te leren), deze truc werkt bij allemaal.
3. Beter resultaat: De foto's zien er scherper uit, lijken meer op het origineel, en je kunt er meer verschillende dingen mee doen (bijvoorbeeld "mijn hond in de sneeuw" in plaats van alleen "mijn hond").

Kort samengevat:
Wanneer je een AI probeert te leren met te weinig foto's, raakt hij in paniek en maakt hij rare ruis. Door de AI een beetje "onzekerheid" te geven (via BNNs), dwing je hem om een breder beeld te hebben. Hierdoor stopt de ruis en krijg je prachtige, diverse resultaten zonder dat de AI vergeten is hoe hij moet werken.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Exploring Diffusion Models' Corruption Stage in Few-Shot Fine-tuning and Mitigating with Bayesian Neural Networks", geschreven in het Nederlands.

Probleemstelling: Het "Corruptie"-stadium

De auteurs identificeren een onverwacht en schadelijk fenomeen tijdens het few-shot fine-tunen van Diffusiemodellen (DM's), zoals Stable Diffusion. Hoewel few-shot fine-tuning (bijv. met DreamBooth, LoRA, OFT) populair is voor het personaliseren van AI met weinig data, vertonen deze modellen een abnormaal trainingsdynamiek:

1. Initiële verbetering: De beeldkwaliteit en gelijkenis met de trainingsdata nemen aanvankelijk toe.
2. Het corruptiestadium: Vervolgens neemt de beeldkwaliteit onverwacht af. Er ontstaan ruispatronen en "corruptie" in de gegenereerde afbeeldingen, wat leidt tot een daling van de gelijkenis met de trainingsdata.
3. Overfitting: Uiteindelijk herstelt de kwaliteit zich, maar het model raakt in een staat van ernstige overfitting. Het kan dan alleen nog exacte kopieën van de trainingsafbeeldingen genereren en verliest de diversiteit en het vermogen om nieuwe variaties te creëren.

De kernoorzaak van dit probleem is een vertrouwd leerproces: bij few-shot fine-tuning leert het model een te smalle verdeling (een beperkte manifold) van de data. Het model probeert de exacte verdeling van de kleine dataset na te bootsen, wat leidt tot instabiliteit en het ontstaan van ruis voordat het volledig overfittet.

Methodologie: Bayesian Neural Networks (BNNs)

Om dit probleem op te lossen, stellen de auteurs voor om Bayesian Neural Networks (BNNs) toe te passen op de fine-tuning van diffusiemodellen. In plaats van vaste gewichten te leren, worden de parameters van het model behandeld als stochastische variabelen met een verdeling.

Kernprincipes van de methode:

• Verdeling verbreden: Door BNNs in te zetten, wordt de leerdoelstelling veranderd van het exacte kopiëren van de trainingsdata naar het leren van een bredere, robuustere verdeling. Dit voorkomt dat het model vastloopt in de smalle "corruptie"-fase.
• Variational Inference: De auteurs gebruiken variational inference om de posterior-verdeling van de parameters te benaderen. De loss-functie bestaat uit twee delen:
  1. Verwachte Diffusie-Loss ($L_{DM}$): De verwachting van de standaard diffusieloss over de verdeling van de parameters.
  2. Regularisatie ($L_r$): Een KL-divergentie-term die de variational verdeling beperkt tot de prior-verdeling (de oorspronkelijke, voorgeïnitieerde DM). Dit zorgt voor stabiliteit.
• Implementatie: De methode is compatibel met bestaande technieken (DreamBooth, LoRA, OFT). In de praktijk worden slechts een subset van de parameters (bijv. lineaire lagen of normalisatielagen) gemodelleerd als BNN's om de rekentijd te beperken.
• Inferentie: Tijdens het genereren (inference) worden de parameters vervangen door hun gemiddelde waarde ($\mu_\theta$). Dit betekent dat er geen extra inferentiekosten zijn ten opzichte van een standaard fine-tuned model.

Belangrijkste Bijdragen

1. Observatie en Analyse: De eerste identificatie en analyse van het "corruptiestadium" in few-shot fine-tuning, waarbij wordt aangetoond dat dit het gevolg is van een te beperkte leerverdeling.
2. Heuristische Modellering: De ontwikkeling van een heuristisch model (gebaseerd op Gaussische verdelingen) dat verklaart waarom corruptie optreedt (hoge variantie in de voorspelling bij beperkte data) en waarom het verdwijnt (overfitting).
3. BNN-integratie: Een innovatieve toepassing van BNNs om de leerverdeling impliciet te verbreden, waardoor corruptie wordt tegengegaan zonder extra inferentiekosten.
4. Uitgebreide Validatie: Demonstratie dat de methode werkt over verschillende modellen (SD v1.4, v1.5, v2.0), verschillende fine-tuning methoden en diverse datasets.

Resultaten

De experimenten tonen aan dat het toepassen van BNNs significante verbeteringen oplevert:

• Kwaliteitsverbetering: Er is een duidelijke stijging in beeldkwaliteit (gemeten met Clip-IQA) en beeldgelijkenis (Dino, Clip-I). De "corruptie" met ruispatronen wordt sterk verminderd.
• Diversiteit: De gegenereerde afbeeldingen behouden een hogere diversiteit (gemeten met Lpips) en overfitten minder snel dan de baseline-methoden.
• Tekst-afstemming: De consistentie met de tekst-prompten (Clip-T) verbetert, vooral bij subject-gedreven generatie.
• User Study: In een menselijke evaluatie (met 101 deelnemers) prefereerden gebruikers consequent de met BNNs gefinetuned modellen, vooral op het gebied van tekst-afstemming en algemene beeldkwaliteit, zowel in de "beste" als "gemiddelde" gevallen.
• Efficiëntie: De methode werkt goed zelfs als BNNs alleen worden toegepast op een klein percentage van de lagen (bijv. alleen normalisatielagen), wat de rekentijd en geheugengebruik tijdens training beperkt.

Significantie

Dit paper is significant omdat het een fundamenteel probleem in de huidige state-of-the-art few-shot fine-tuning van generatieve AI blootlegt en een elegante oplossing biedt.

• Theoretisch inzicht: Het verandert het perspectief op why few-shot learning faalt, niet alleen als een data-schaarste probleem, maar als een probleem van te smalle verdelingen die leiden tot instabiliteit.
• Praktische toepasbaarheid: De oplossing is plug-and-play compatibel met bestaande workflows (zoals DreamBooth en LoRA) en vereist geen extra rekenkracht tijdens het daadwerkelijke gebruik van het model.
• Toekomstige richting: Het opent de deur voor het gebruik van probabilistische methoden (BNNs) om de robustheid en generalisatie van grote generatieve modellen te verbeteren, zelfs bij zeer beperkte datasets.