Technische Samenvatting: DiffRGD

Probleemstelling

Diffusiemodellen zijn fundamenteel geworden voor beeldgeneratie, maar het aansturen van deze modellen voor specifieke taken (bijv. beeldrestauratie, conditionele generatie) zonder kostbare hertraining of fine-tuning blijft een uitdaging. Bestaande methoden voor sturing tijdens de inferentie (bijv. DPS, FreeDoM) sturen het bemonsteringsproces bij door gradiënten van een differentieerbaar doel direct in de latente ruimte te injecteren. Deze methoden opereren echter in de ambivalente ruimte zonder rekening te houden met de onderliggende Gaussische structuur van het diffusieproces.

Dit gebrek aan distributiebewustzijn leidt tot distributionele drift, waarbij de latente variabelen afwijken van de oorspronkelijke Gaussische manifold die door het model wordt geïnduceerd. Naarmate de bemonsteringsstappen zich ophopen, verslechtert dit "off-manifold"-fenomeen de kwaliteit van de uiteindelijke schone steekproef. Eerdere pogingen om updates te beperken, zoals MPGD (dat uitgaat van lokale lineariteit) of DSG (dat projecteert naar een hypersfeer met een vaste straal), vertrouwen ofwel op onhaalbare perfecte encoders, of falen om volledig te voldoen aan de specifieke Gaussische eigenschappen van de latente distributie, waarbij ze soms degraderen naar een uniforme distributie.

Methodologie: DiffRGD

De auteurs stellen DiffRGD voor, een plug-and-play framework voor sturing tijdens de inferentie dat expliciet de latente Gaussische structuur behoudt. De kerninzicht is dat de latente distributie op elk DDIM-bemonsteringsmomentstijdstip $t$ een isotrope Gaussische verdeling is, gecentreerd rond een voorspelde gemiddelde $\mu_t$.

1. Polaire Decompositie en Manifold Constructie

De methode maakt gebruik van een polaire decompositie-eigenschap van isotrope Gaussische distributies. Op elk tijdstip $t$ kan een latente steekproef $x_t$ worden gedecomposeerd in een statistisch onafhankelijke radiale component $r$ en een directionele component $u$:
$$x_t = \mu_t + \sigma_t r u$$
waarbij $r$ een geschaalde chi-verdeling volgt en $u$ uniform verdeeld is op een eenheidsbol.

In plaats van te optimaliseren in de ambivalente ruimte of op een bol met een vaste straal, construeert DiffRGD een sferische manifold $\mathcal{S}_{t,r}$ die wordt geïnduceerd door de specifieke latente distributie:
$$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$$
Hierbij is de straal $\sigma_t r$ gesampled uit de natuurlijke schaal van de distributie, wat ervoor zorgt dat de toegestane verzameling de Gaussische eigenschappen van het diffusieproces respecteert.

2. Riemanniaanse Gradiëntafdaling (RGD)

Het sturingsprobleem bij elk tijdstip wordt geformuleerd als een geconstrueerd optimalisatieprobleem:
$$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$$
waarbij $\mathcal{L}$ de sturingsverliesfunctie is (bijv. reconstructieverlies voor restauratie of feature-uitlijning voor conditionele generatie).

Om dit op te lossen, gebruikt DiffRGD Riemanniaanse Gradiëntafdaling op de sferische manifold $\mathcal{S}_{t,r}$:

1. Tangentruimte Projectie: De Euclidische gradiënt van de loss wordt geprojecteerd op de tangentruimte van de bol op het huidige punt, waarbij de radiale component wordt verwijderd om te garanderen dat de update-richting tangentieel is aan de manifold.
2. Retractie: Na het nemen van een stap langs de tangentruimte, brengt een retractie-operatie het bijgewerkte punt terug op de sferische manifold $\mathcal{S}_{t,r}$, waardoor de straal $\sigma_t r$ behouden blijft.

Dit proces wordt herhaald voor $K$ innerlijke iteraties (meestal $K=3$) per diffusie-tijdstip. De auteurs bewijzen dat onder aannames van $L$-gladheid, deze optimalisatie convergeert naar een stationair punt.

Belangrijkste Bijdragen

1. Distributie-bewuste Geometrie: De auteurs identificeren een eigenschap van polaire decompositie die een door distributie geïnduceerde geometrie voor diffusie-sturing onthult. Ze stellen een Riemanniaans framework voor dat de geometrie van de latente distributie respecteert, wat de degradatie van de steekproefkwaliteit door off-manifold drift vermindert.
2. Theoretische Analyse: Het artikel biedt een analyse van distributiebehoud en bewijst de convergentie van het Riemanniaanse optimalisatiealgoritme naar een stationair punt. In de praktijk convergeert de voorgestelde geometrie-bewuste optimalisatie sneller dan ADMM-gebaseerde methoden.
3. Plug-and-Play Prestaties: Er wordt aangetoond dat DiffRGD een plug-and-play methode is die geen hertraining of fine-tuning vereist en state-of-the-art resultaten behaalt over diverse taken heen.

Experimentele Resultaten

De auteurs hebben DiffRGD geëvalueerd tegen baselines (DPS, FreeDoM, MPGD, DSG, ADMMDiff) op een enkele NVIDIA GeForce RTX 4090 GPU met behulp van de HuggingFace Diffusers framework.

Beeldrestauratie

Geëvalueerd op FFHQ 256×256 en ImageNet 256×256 voor vier taken: Inpainting, Super-Resolutie (4×), Gaussische Deblurring en Motion Deblurring.

• Kwantitatief: DiffRGD behaalde statistisch significante verbeteringen in PSNR, SSIM, LPIPS en FID ten opzichte van alle baselines in de meeste taken. Bijvoorbeeld, bij Inpainting behaalde DiffRGD een PSNR van 34,04 vergeleken met 32,38 voor de op één na beste methode (ADMMDiff).
• Robuustheid: DiffRGD behield sterke prestaties, zelfs met minder bemonsteringsstappen (100 stappen), terwijl methoden zoals DPS en ADMMDiff een merkbare degradatie vertoonden.
• Kwalitatief: Visuele resultaten vertoonden minder artefacten en een betere achtergrondconsistentie vergeleken met DPS en MPGD, met een vergelijkbare of superieure kwaliteit aan ADMMDiff maar met minder rekentijd.

Conditionele Generatie

Geëvalueerd op CelebA-HQ 256×256 voor Segmentatiekaart, Schets en FaceID sturing.

• DiffRGD presteerde consistent beter dan bestaande inferentie-tijd methoden (FreeDoM, DSG, ADMMDiff) op het gebied van uitlijningsmetrieken (mIoU, Sketch-$\ell_2$, FaceID-$\ell_2$) en generatiekwaliteit (FID, KID).
• Kwalitatieve resultaten gaven aan dat DiffRGD natuurlijkere en visueel aantrekkelijkere resultaten produceerde met een betere uitlijning met condities vergeleken met DSG, dat vaak kampte met instabiele controle.

Sturingssterkte

Experimenten op sturingssterkte ($\eta_t$) lieten zien dat DiffRGD een betere balans bereikt tussen uitlijning (mIoU) en generatiekwaliteit (FID) over een breder bereik van sterktes vergeleken met DSG, dat slecht presteerde over alle instellingen.

Betekenis en Claims

Het artikel claimt dat DiffRGD een kritiek gat in de inferentie-tijd sturing aanpakt: het falen van bestaande methoden om de stapsgewijze Gaussische latente structuur te behouden. Door sturing te formuleren als een geconstrueerde optimalisatie op een door distributie geïnduceerde sferische manifold, voorkomt DiffRGD de "off-manifold" fenomenen die de steekproefkwaliteit verslechteren.

De auteurs beweren dat geometrie-consistente sturing een levensvatbaar en superieur alternatief is voor het direct toepassen van gradiënten in de ambivalente ruimte of het gebruiken van restrictieve, niet-distributie-bewuste manifolds. De methode biedt een eenvoudige, plug-and-play oplossing die de taakmetrieken verbetert terwijl de hoge getrouwheid van de steekproefkwaliteit die inherent is aan getrainde diffusiemodellen behouden blijft, zonder de noodzaak van dure hertraining.