Improving Denoising Diffusion Models via Simultaneous Estimation of Image and Noise

Dit artikel introduceert een verbeterde denoising diffusion-methode die gebruikmaakt van een hoekgebaseerde herparameterisatie voor een stabielere ODE-oplossing en gelijktijdige schatting van het beeld en het ruis, wat resulteert in snellere generatie en hogere beeldkwaliteit.

De kern

Stel je voor dat je een schilderij moet maken, maar je begint met een doek dat volledig vol zit met statische ruis (zoals op een oud televisietoestel dat geen signaal heeft). De kunst van het "diffusiemodel" is om die ruis stap voor stap weg te werken tot er een prachtig, scherp beeld van een paard of een gezicht verschijnt.

Deze paper beschrijft een nieuwe, slimme manier om dat schilderij sneller en mooier te maken. De auteurs, Zhenkai Zhang en zijn team van de Universiteit van Melbourne, hebben twee grote verbeteringen bedacht. Laten we ze uitleggen met een paar alledaagse vergelijkingen.

1. De "Wiskundige Bocht" (Het nieuwe pad)

Het oude probleem:
Stel je voor dat je een auto moet rijden van punt A (ruis) naar punt B (het mooie plaatje). In de oude methoden (zoals DDPM) was het pad waar je overheen reed erg onhandig. Aan het begin en aan het einde van de rit waren er twee "dodenloze hoeken" of scherpe bochten waar de auto bijna vastliep. De bestuurder (de computer) moest hier heel voorzichtig zijn en langzaam rijden, wat veel tijd kostte. Wiskundig gezien waren deze punten "singulair", wat betekent dat de berekeningen daar onmogelijk of heel onstabiel werden.

De nieuwe oplossing:
De auteurs hebben het pad veranderd. In plaats van een rechte lijn met scherpe bochten, hebben ze het pad veranderd in een perfecte, soepele kwartcirkel.

• De analogie: Denk aan het verschil tussen een auto die over een hobbelig, kronkelend bergpad moet, versus een auto die over een gladde, gebogen helling van een glijbaan glijdt.
• Het resultaat: Omdat het pad nu een mooie cirkelbocht is, kunnen ze gebruikmaken van geavanceerde "rijtechnieken" (wiskundige methoden genaamd Runge-Kutta). Dit zorgt ervoor dat de auto (de computer) veel sneller en soepeler van ruis naar beeld kan gaan, zonder vast te lopen.

2. Twee ogen in plaats van één (Het simultaan schatten)

Het oude probleem:
In de oude methoden moest de computer één ding doen: raden waar de ruis zat, om die eruit te halen.

• Aan het begin: Als het beeld nog heel erg ruist, is het makkelijk om te raden waar de ruis zit.
• Aan het einde: Als het beeld bijna klaar is, is er heel weinig ruis over. Dan is het heel moeilijk om nog te raden waar die kleine beetje ruis zit, terwijl je tegelijkertijd het mooie plaatje moet zien. Het is alsof je probeert een naald te vinden in een hooiberg die bijna leeg is; je ziet de naald niet meer goed, maar je ziet het hooi wel.

De nieuwe oplossing:
De nieuwe methode laat de computer twee dingen tegelijk doen: het raden van de ruis én het raden van het eindplaatje.

• De analogie: Stel je voor dat je een verwarde foto probeert te herstellen. De oude methode was alsof je alleen naar de vlekken (ruis) keek en hoopte dat het plaatje vanzelf goed kwam. De nieuwe methode is alsof je een twee-ogen-systeem hebt: één oog kijkt naar de vlekken om ze weg te halen, en het andere oog kijkt naar het onderwerp (het paard of gezicht) om te zien hoe het eruit moet zien.
• Waarom is dit beter?
  • Aan het begin van het proces helpt het "plaatje-oog" om richting te geven (want er is nog geen duidelijk plaatje, maar het model weet al hoe een paard eruit zou moeten zien).
  • Aan het einde van het proces helpt het "ruis-oog" om de laatste kleine foutjes weg te halen.
  • Door beide tegelijk te doen, weet het model altijd precies wat het moet doen, ongeacht hoe ruizig of hoe scherp het beeld is.

Het grote resultaat: Sneller en Schoner

Door deze twee trucjes te combineren (de soepele cirkelbaan en het tweeoog-systeem), gebeurt er iets magisch:

1. Snelheid: Het model kan een hoogwaardig plaatje maken in veel minder stappen. Waar een ander model misschien 500 stappen nodig heeft om een duidelijk paard te zien, ziet jouw nieuwe model het al na 150 stappen. Het is alsof je een video kunt versnellen zonder dat het beeld wazig wordt.
2. Kwaliteit: De uiteindelijke plaatjes zijn scherper, realistischer en hebben minder "artefacten" (rare vlekken of vervormingen).
3. Efficiëntie: Het model heeft minder "oefentijd" nodig om te leren. Waar andere modellen miljoenen iteraties nodig hadden om goed te worden op moeilijke datasets, deed dit nieuwe model het met veel minder moeite.

Samenvattend:
De auteurs hebben de diffusiemodellen niet alleen "beter" gemaakt, maar ze hebben de hele reis van ruis naar beeld herschreven. Ze hebben een soepeler pad gebouwd en de bestuurder (de AI) uitgerust met twee ogen in plaats van één. Het resultaat? Een snellere, slimmere en mooiere manier om kunst en beelden uit het niets te creëren.

Technische samenvatting

Probleemstelling

Diffusiemodellen (zoals DDPM en DDIM) hebben zich bewezen als krachtige generatieve modellen voor het creëren van hoogwaardige afbeeldingen. Echter, deze modellen kampen met twee belangrijke beperkingen:

1. Inefficiëntie bij inferentie: Het generatieproces vereist vaak duizenden stappen om van pure ruis naar een realistische afbeelding te gaan. Vooral in de vroege fasen van het proces (wanneer het beeld nog grotendeels ruis is) is dit tijdrovend.
2. Afweging tussen ruis- en beeldschatting:
   • Ruis-gebaseerde modellen (zoals DDPM) schatten de toegevoegde ruis ($\epsilon$). Dit werkt goed in latere fasen, maar het leren van pure ruis in de vroege fasen is uitdagend.
   • Beeld-gebaseerde modellen (zoals Cold Diffusion) schatten het oorspronkelijke beeld ($x_0$) direct. Dit helpt in de vroege fasen, maar wordt onstabiel en moeilijk in de latere fasen wanneer de invoer gedomineerd wordt door ruis.
   • Bestaande modellen kiezen meestal voor één van deze twee benaderingen, wat leidt tot suboptimale prestaties of een langzame convergentie.

Methodologie

De auteurs stellen een nieuw raamwerk voor dat de voordelen van zowel ruis- als beeldschatting combineert, ondersteund door een wiskundige herschikking van het diffusieproces.

1. Reparameterisatie via een kwart-cirkelboog (Noise Scheduler)
In plaats van de traditionele lineaire parameterisatie waarbij $\sqrt{\bar{\alpha}_t}$ wordt gebruikt (wat leidt tot singulariteiten bij $t=0$ en $t=T$), stellen de auteurs een nieuwe parameterisatie voor op basis van een kwart-cirkelboog:
$$x_t = \cos(\eta_t)x_0 + \sin(\eta_t)\epsilon$$
waarbij $\eta_t = \frac{t}{T} \frac{\pi}{2}$.

• Voordelen: Deze aanpak elimineert de singulariteiten in de afgeleiden, waardoor het diffusieproces kan worden uitgedrukt als een goed gedefinieerde gewone differentiaalvergelijking (ODE). Dit maakt het gebruik van hogere-orde ODE-oplossers (zoals Runge-Kutta-methoden) mogelijk in plaats van de standaard Euler-stappen, wat de convergentiesnelheid aanzienlijk verbetert.

2. Simultane schatting van Beeld en Ruis
Het netwerk wordt getraind om zowel het oorspronkelijke beeld ($\hat{x}_0$) als de ruis ($\hat{\epsilon}$) gelijktijdig te voorspellen.

• Loss-functie: De trainingsdoelfunctie combineert de fouten voor beide schattingen:
  $$\min_{\theta} \mathbb{E} [\|R_\theta(x_t, t) - x_0\| + \|\epsilon_\theta(x_t, t) - \epsilon\|]$$
• Voordeel: Dit zorgt voor een stabielere gradiëntberekening gedurende het hele proces. In de vroege fasen levert de beeldschatting waardevolle informatie, terwijl de ruischatting in de latere fasen de precisie verhoogt.

3. Sampling met Gradiënt-Updates
Het proces wordt benaderd als een iteratief optimalisatieprobleem. De auteurs gebruiken de geschatte gradiënten van het verlies om de update-stappen te sturen.

• De geschatte gradiënt $\dot{\hat{x}}_t$ wordt berekend op basis van de geschatte $\hat{x}_0$ en $\hat{\epsilon}$.
• De update-stap wordt uitgevoerd via: $x_{t-1} = x_t - \Delta t \cdot \dot{\hat{x}}_t$.
• Door gebruik te maken van hogere-orde methoden (zoals RK2 of RK4) op deze gradiënt, wordt de stabiliteit en snelheid van het generatieproces verder verhoogd.

Belangrijkste Bijdragen

1. Nieuwe Parameterisatie: Een innovatieve "noise scheduler" gebaseerd op trigonometrische functies die singulariteiten elimineert en het gebruik van geavanceerde ODE-oplossers mogelijk maakt.
2. Dual Estimation: Een architectuur die gelijktijdig leert om zowel het beeld als de ruis te schatten, waardoor de beperkingen van puur ruis- of puur beeld-gebaseerde modellen worden overwonnen.
3. Gradiënt-gebaseerde Sampling: Een methode die gradiëntinformatie gebruikt om de omgekeerde diffusiestappen nauwkeuriger en stabieler te maken, wat leidt tot snellere convergentie.

Resultaten

De modellen zijn getest op drie datasets: CIFAR-10, CelebA en LUSH (kerkgebouwen). De prestaties zijn gemeten aan de hand van FID (Frechet Inception Distance), sFID, Precision en Recall.

• Snelheid en Kwaliteit: Het voorgestelde model convergeert aanzienlijk sneller naar hoogwaardige afbeeldingen. Waar DDPM en DDIM ongeveer 400-500 stappen nodig hebben om een herkenbaar object (bijv. een paard) te vormen, doet het nieuwe model dit in ongeveer 150 stappen (ongeveer 3x sneller).
• Metingen:
  • Het model behaalt betere FID- en sFID-scores dan DDPM en DDIM, vooral bij een lager aantal stappen (tussen 50 en 200).
  • Bij 1000 stappen presteert het model vergelijkbaar of beter, afhankelijk van de dataset.
  • De precisie en recall zijn hoger, wat wijst op betere nauwkeurigheid en diversiteit.
• Efficiëntie: Het model vereist minder trainingstijd (minder iteraties) om vergelijkbare prestaties te halen als geavanceerde basismodellen. Bijvoorbeeld, op de LUSH-dataset had het model slechts 1,135,000 iteraties nodig vergeleken met 4,432,000 voor DDPM/DDIM.
• Ablatie-studies: Experimenten bevestigen dat elke component (nieuwe scheduler, trigonometrische parameterisatie, en simultane schatting) bijdraagt aan de verbetering, en dat de combinatie van alle drie de beste resultaten oplevert.

Betekenis en Impact

Dit onderzoek biedt een fundamentele verbetering in de efficiëntie en kwaliteit van diffusiemodellen. Door het probleem van de "ruis-dominantie" in vroege fasen en de "beeld-dominantie" in late fasen op te lossen via simultane schatting, en door wiskundige singulariteiten te elimineren via een nieuwe parameterisatie, maakt dit werk snellere generatie mogelijk zonder in te leveren op kwaliteit. Dit is van groot belang voor praktische toepassingen waar rekenkracht en tijd beperkt zijn, zoals real-time beeldgeneratie en data-augmentatie. De code is open-source beschikbaar gesteld, wat de reproduceerbaarheid en verdere ontwikkeling in de gemeenschap stimuleert.