Technische Samenvatting: Dual Ascent Diffusion voor Inverse Problemen

Probleemstelling

Inverse problemen betreffen het schatten van een onbekend signaal of beeld $x$ uit ruwe en beschadigde waarnemingen $y$, die worden beheerst door een voorwaartse operator $A(\cdot)$. Deze problemen komen veelvuldig voor in gebieden variërend van astrofysica en medische beeldvorming tot de bepaling van proteïnestructuren. Ze zijn doorgaans slecht gesteld, wat het gebruik van priors vereist om de oplossingsruimte te beperken.

Hoewel moderne benaderingen gebruikmaken van vooringestelde diffusiemodellen als krachtige priors, ondervinden bestaande methoden voor het oplossen van Maximum-a-Posteriori (MAP)-problemen of het bemonsteren van de posterior aanzienlijke beperkingen. Veel methoden vertrouwen op computationele benaderingen om de waarschijnlijkheid van het waarnemingsmodel te combineren met de diffusi prior. Deze benaderingen leiden vaak tot onnauwkeurige steekproeven, suboptimale reconstructiekwaliteit of een gebrek aan trouwe weergave van de waarnemingen (d.w.z. hoge residuen), wat zich kan manifesteren als "hallucinaties" in de gegenereerde oplossing. Bovendien worstelen veel state-of-the-art methoden met hoge niveaus van meetruis en lijden ze onder computationele inefficiëntie door de accumulatie van fouten of de noodzaak van aanvullende Markov Chain Monte Carlo (MCMC)-stappen.

Methodologie: Dual Ascent Diffusion (DDiff)

De auteurs stellen DDiff voor, een nieuw kader voor het oplossen van MAP-problemen met behulp van vooringestelde diffusiemodellen. In tegenstelling tot methoden die proberen op een bewezen manier te bemonsteren uit de posterior, richt DDiff zich op het afleiden van een optimalisatiestrategie die het MAP-probleem nauwkeurig en efficiënt oplost. De methode is geïnspireerd door de Alternating Direction Method of Multipliers (ADMM), maar introduceert een nieuwe dual ascent-structuur die specifiek is aangepast voor diffusiemodellen.

Kernalgoritmische Componenten

Het DDiff-kader doorloopt binnen elke stap drie primaire updates, waarbij wordt afgewisseld tussen de prymale variabelen ($x, z$) en de duale variabele ($u$):

1. x-update (Data Fideliteit): In plaats van een beperkt kleinste-kwadratenprobleem exact op te lossen, hanteert DDiff een gelijnde ADMM-benadering. Het voert een enkele gradiëntafdaalstap uit op de term voor data fideliteit $-\log p(y|x)$, waardoor de methode direct toepasbaar is op elke differentieerbare voorwaartse model $A(\cdot)$, of deze nu lineair of niet-lineair is.
   $$x \leftarrow v - \gamma \nabla_v \|y - A(v)\|_2^2, \quad \text{waarbij } v = z - u$$

2. z-update (Denoising): Deze stap vervangt de naïeve toepassing van een diffusiemodel als denoiser. De auteurs stellen vast dat standaard score-modellen $s_\theta(x, t)$ alleen nauwkeurig zijn op het diffusiemaneveld op tijdstip $t$. Om het feit aan te pakken dat de variabele $x+u$ over het algemeen niet op dit manevel ligt, construeert DDiff een recursieve variabele $x_t$ die het denoising-proces verankert. De update omvat een DDIM-stijl omgekeerde diffusiestap waarbij de "voorspelde $x_0$" de huidige prymale schatting $x$ is, en de duale variabele $u$ wordt herschaald en toegevoegd om consistentie met het signaalniveau te behouden.
   $$z \leftarrow \frac{1}{\sqrt{\bar{\alpha}_t}} (x_t + (1-\bar{\alpha}_t)s_\theta(x_t, t))$$
   waarbij $x_t$ wordt bijgewerkt via een DDIM-stap die de duale variabele $u$ integreert.

3. Duale Update: De duale variabele $u$ wordt bijgewerkt om de beperking $x=z$ af te dwingen door het residuele fout tussen de prymale variabelen op te stapelen:
   $$u \leftarrow u + x - z$$

Latente Uitbreiding

Het kader wordt uitgebreid tot LatentDDiff voor gebruik met Latent Diffusion Models (LDM's). Deze variant werkt op compacte latente representaties om geheugen- en rekenvereisten te verminderen. Het ondersteunt twee modi voor het afdwingen van data-consistentie:

• Latent-DC: Dwingt consistentie direct af in de latente ruimte, wat back-propagation door de decoder vereist.
• Pixel-DC: Decodeert de latente variabele naar pixelruimte, past de correctie voor data-consistentie toe en encodeert opnieuw. De auteurs merken op dat Pixel-DC consequent hogere kwaliteit reconstructies oplevert.

Belangrijkste Bijdragen

1. Dual Ascent Kader voor Diffusion Priors: Het artikel introduceert een nieuwe optimalisatiestrategie die dual ascent (geïnspireerd door ADMM) integreert met diffusiemodellen. Deze aanpak vermijdt de benaderingen die gangbaar zijn in bestaande posterior-bemonsteringsmethoden.
2. Bewijs van Convergentie naar Vast Punt: De auteurs stellen het eerste bewijs voor convergentie naar een vast punt op voor een op diffusie gebaseerde posterior-optimalisatiemethode. Ze tonen aan dat onder de aannames van een begrensd denoiser en gradiënt, de reeksen van prymale en duale variabelen convergeren naar een vast punt $(x^\star, z^\star, u^\star)$, zelfs zonder een convex prior aan te nemen.
3. Verbeterde Robuustheid en Fideliteit: De methode is ontworpen om oplossingen te produceren die de waarnemingen trouwer modelleren, wat blijkt uit residuen die dichter bij nul liggen in vergelijking met state-of-the-art methoden. Dit vermindert de kans op hallucinaties die worden geïntroduceerd door de generatieve prior.
4. Efficiëntie: DDiff bereikt snellere inferentietijden door de noodzaak van aanvullende MCMC-equilibratiestappen (vereist door methoden zoals DAPS) te elimineren en het aantal neurale functiebeoordelingen (NFE's) te verminderen dat nodig is voor resultaten van hoge kwaliteit.

Experimentele Resultaten

De auteurs hebben DDiff geëvalueerd op een breed scala aan lineaire en niet-lineaire inverse problemen, waaronder super-resolutie, inpainting, Gaussische en bewegingsontwarring, faseherstel, niet-lineaire ontwarring en beeldvorming met hoge dynamische bereik (HDR). Experimenten werden uitgevoerd op FFHQ- en ImageNet-datasets.

• Kwantitatieve Prestaties: DDiff presteerde beter dan state-of-the-art baselines (waaronder DAPS, DMPlug, DCDP, DPS, DiffPIR en DDRM) op de meeste taken. Het behaalde superieure scores voor Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM) en Learned Perceptual Image Patch Similarity (LPIPS). Opmerkelijk is dat het de laagste residuele fouten behaalde, wat wijst op betere data-consistentie.
• Robuustheid tegen Ruis: De methode toonde aanzienlijk grotere robuustheid tegen hoge niveaus van meetruis. Terwijl sommige baselines een snelle degradatie in PSNR vertoonden naarmate de ruis toenam, behield DDiff stabiele prestaties. In scenario's met hoge ruis (bijvoorbeeld faseherstel met $\sigma=0.3$) herstelde DDiff coherente globale structuren en semantics waar andere methoden beschadigde artefacten produceerden.
• Efficiëntie: In termen van tijdsefficiëntie bereikte DDiff betere perceptuele kwaliteit met minder NFE's in vergelijking met DAPS. Het vereiste ongeveer de helft van de bemonsteringstijd voor lineaire taken en een zesde van de tijd voor niet-lineaire taken om vergelijkbare of superieure resultaten te bereiken.
• Ablatiestudies: Experimenten bevestigden dat zowel de update van de duale variabele als de specifieke noisestap (Vergelijking 12 in het artikel) cruciaal zijn. Het verwijderen van de duale variabele (wat DiffPIR nabootst) of de noisestap resulteerde in aanzienlijke prestatiedalingen, met name wat betreft hoogfrequente artefacten en data-consistentie.

Betekenis en Beweringen

Het artikel beweert dat DDiff een aanzienlijke vooruitgang vertegenwoordigt in het oplossen van inverse problemen met diffusion priors door over te stappen van posterior-bemonsteringsbenaderingen naar een robuust dual-ascent optimalisatiekader. De auteurs benadrukken dat hun methode:

• Waarnemingen trouw modelleert: Door residuen te minimaliseren, vermindert het het "hallucinatie"-effect dat vaak wordt gezien in generatieve oplossers voor inverse problemen.
• Computationeel efficiënt is: Het verwijdert de bottleneck van MCMC-stappen en vermindert de computationele last van gradiëntberekeningen.
• Theoretisch onderbouwd is: Het geleverde convergentiebewijs biedt een theoretische garantie voor de stabiliteit van de methode, een kenmerk dat ontbreekt in veel eerdere op diffusie gebaseerde optimalisatiebenaderingen.

De auteurs concluderen dat DDiff een praktische oplossing is voor grootschalige, complexe inverse problemen, met name in scenario's waar meetruis prominent is, zoals reconstructie van CT-scans met lage dosis. Voor toekomstig werk wordt voorgesteld het kader uit te breiden naar hogere dimensionale data zoals 3D-volumes en video.