Technische Samenvatting: Inverse Problemen met Diffusiemodellen via Mode-Seeking Loss

Probleemdefinitie

Het artikel behandelt de uitdaging van het oplossen van inverse problemen met behulp van getrainde onvoorwaardelijke diffusiemodellen. Een invers probleem houdt het schatten van een originele datastructuur $x_0$ in met behulp van een gedeeltelijk gedegradeerde meting $y$, die wordt beheerst door de vergelijking $y = A(x_0) + \eta$, waarbij $A$ een bekende degradatieoperator is en $\eta$ Gaussische ruis.

Hoewel getrainde diffusiemodellen deze problemen op een zero-shot (plug-and-play) wijze kunnen oplossen, kampen bestaande benaderingen met aanzienlijke beperkingen:

1. Posterior Sampling Methoden: Deze streven ernaar om samples te trekken die proportioneel zijn aan de posterior waarschijnlijkheid $p(x_0|y)$. Echter, ze vertrouwen vaak op modelleringsbenaderingen (bijv. Gaussische benaderingen van de conditionele score) of sequentiële Monte-Carlo-methoden, wat computationeel veeleisend kan zijn en mogelijk niet strikt voldoet aan de ware posterior zonder grondwaarheid-validatie.
2. MAP Estimatie Methoden: Deze streven ernaar om de Maximum A Posteriori (MAP) schatting te vinden (de modus van de posterior). Bestaande solvers maken vaak gebruik van heuristische benaderingen, vertrouwen op variabele splitsing (bijv. HQS, ADMM), of vereisen consistentiemodellen. Deze benaderingen brengen vaak dure gradiëntberekeningen en modelleringsaannames met zich mee die mogelijk niet gelden voor alle lineaire operatoren.

De auteurs stellen dat veel beeldrestauratietaken prioriteit geven aan het genereren van een plausibel beeld dat consistent is met de metingen, in plaats van het dekken van de volledige posterior distributie. Dit komt nauwer overeen met MAP-estimatie, wat de noodzaak voor een meer principiële, benaderingsvrije MAP-strategie motiveert.

Methodologie: Variational Mode-Seeking Loss (VML)

De kernbijdrage is de Variational Mode-Seeking Loss (VML), een nieuwe objectieve functie ontworpen om gegenereerde samples te sturen naar de MAP-schatting tijdens het reverse diffusieproces.

Theoretische Fundering

De auteurs leiden VML af vanuit een variationeel perspectief door de Kullback-Leibler (KL) divergentie te minimaliseren tussen de diffusie-posterior $p(x_0|x_t)$ en de metings-posterior $p(x_0|y)$ bij elke reverse diffusie tijdstap $t$:
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

De intuïtie is dat naarmate $t \to 0$, de distributie $p(x_0|x_t)$ steeds meer gepiekt raakt rond $x_t$. Door een optimale $x_t^*$ te vinden die deze KL-divergentie bij elke stap minimaliseert, stuurt het proces de sample naar de modi van de posterior $p(x_0|y)$.

Closed-Form Afleiding voor Lineaire Operatoren

Een belangrijke theoretische prestatie is de afleiding van een closed-form expressie voor VML in het geval van lineaire inverse problemen ($A(x) = Hx$) zonder enige modelleringsbenaderingen. De volledige VML-expressie bevat:

1. Een metingsconsistentie term: $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. Prior termen: Betrokken bij de denoiser $D(x_t, t)$, de noise schedule $\sigma_t$, en de covariantie van de posterior $Cov[x_0|x_t]$.

Vereenvoudigde en Hergewogen Objectieven

De volledige VML bevat hogere-orde termen die de covariantie-matrix $Cov[x_0|x_t]$ betreffen, welke computationeel duur zijn om te berekenen en te differentiëren. De auteurs bewijzen dat deze hogere-orde termen uniform naar nul convergeren (of naar een constante) naarmate $t \to 0$. Dienovereenkomstig stellen zij een Simplified-VML ($VML_S$) voor die deze termen weglaat, waarbij alleen de metingsconsistentie en de prior termen behouden blijven.

Om de optimalisatiestabiliteit verder te verbeteren, introduceren zij een reweighted simplified-VML ($VML_S^\rho$):
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
waarbij $\rho_t$ een prior weight schedule is zodanig dat $\rho_t \to 1$ als $t \to 0$. Deze formulering maakt efficiënte gradiëntberekening mogelijk via standaard automatische differentiatie.

Algoritmen

Op basis van VML-minimalisatie stellen de auteurs het volgende voor:

1. VML-MAP: Een algoritme dat gradiëntafdaling gebruikt om $VML_S^\rho$ te minimaliseren bij elke reverse diffusiestap.
2. VML-MAP$_{pre}$: Een gepreconditioneerde variant voor slecht geconditioneerde lineaire operatoren. Het maakt gebruik van de Singular Value Decomposition (SVD) van de degradatiematrix $H$ om een preconditioner $P$ te construeren, wat de convergentie versnelt voor moeilijke taken zoals deblurring.
3. VML-MAP$^\tau$: Een strategie voor ruisige beeldrestauratie. Om scherpe artefacten te voorkomen die worden veroorzaakt door optimalisatieproblemen bij lage ruisniveaus, minimaliseert het algoritme VML alleen tot een diffusie-tijd drempelwaarde $\tau$, en schakelt daarna over naar standaard DDIM sampling om het uiteindelijke beeld te genereren.
4. LatentVML-MAP: Een uitbreiding naar Latent Diffusion Models (LDMs), waarbij de VML in de latente ruimte wordt benaderd.

Belangrijkste Bijdragen

1. Nieuwe Loss Functie: Introductie van VML, die theoretisch samples stuurt naar MAP-schattingen door de KL-divergentie tussen de diffusie- en metings-posteriors te minimaliseren.
2. Approximatievrije Afleiding: Voor lineaire inverse problemen is VML in closed form afgeleid zonder heuristische modelleringsbenaderingen, wat een formeel theoretisch fundament biedt voor MAP-estimatie met diffusiemodellen.
3. Convergentie Analyse: Formeel bewijs dat de vereenvoudigde en hergewogen VML-objectieven convergeren naar de negatieve log-posterior (de MAP-objectief) naarmate $t \to 0$ onder milde aannames.
4. Praktische Algoritmen: Ontwikkeling van VML-MAP en de gepreconditioneerde variant, die een empirisch effectieve oplossing bieden voor diverse beeldrestauratietaken.
5. Uitgebreide Validatie: Demonstratie van superieure prestaties ten opzichte van bestaande baselines (bijv. DDRM, $\Pi$GDM, MAPGA, DAPS) over meerdere datasets (ImageNet, FFHQ, CelebA) en taken (inpainting, super-resolutie, deblurring).

Experimentele Resultaten

De auteurs evalueerden hun methoden op beeldrestauratietaken met behulp van getrainde modellen op ImageNet256, FFHQ256 en CelebA.

• Metrieken: Het artikel legt de nadruk op perceptuele metrieken (LPIPS en FID) boven pixel-gewijze metrieken (PSNR, SSIM), met het argument dat deze laatste vaak niet in staat zijn de perceptuele kwaliteit van plausibele reconstructies in slecht gestelde inverse problemen te vatten.
• Prestaties:
  • VML-MAP presteerde consequent beter dan posterior sampling methoden (DDRM, $\Pi$GDM) en andere MAP-estimators (MAPGA, DAPS) in termen van LPIPS en FID bij inpainting-, super-resolutie- en deblurring-taken.
  • VML-MAP$_{pre}$ liet significante verbeteringen zien ten opzichte van de niet-gepreconditioneerde versie voor slecht geconditioneerde taken (bijv. deblurring), waarbij lagere LPIPS- en FID-scores werden behaald.
  • Computationele Efficiëntie: VML-MAP bereikte een betere perceptuele kwaliteit met een lagere runtime (minder neurale functie-evaluaties) vergeleken met methoden die veel reverse diffusiestappen vereisen.
• Ruisige Restauratie: De VML-MAP$^\tau$ strategie handelde metingen met ruis succesvol af, waarbij hoogwaardige reconstructies produceerde waar standaard VML-MAP moeite had met artefacten.
• Latente Modellen: LatentVML-MAP toonde effectiviteit op CelebA, hoewel de auteurs opmerken dat optimalisatie in de latente ruimte uitdagender blijft vanwege de decoder-niet-lineariteit.

Betekenis en Claims

Het artikel claimt dat VML de eerste approximatievrije modellering van MAP-estimatie is voor lineaire inverse problemen, gebruikmakend van enkel een getraind onvoorwaardelijk diffusiemodel, ondersteund door formele theoretische analyse.

De auteurs benadrukken dat hun aanpak fundamenteel verschilt van posterior sampling (die streeft naar distributie-dekking) en bestaande MAP-solvers (die vaak vertrouwen op heuristieken of consistentiemodellen). Door VML te minimaliseren, richt de methode zich effectief op de "modi" van de posterior, wat volgens de auteurs het gewenste gedrag is voor het genereren van plausibele, hoogwaardige beeldrestauraties.

Het werk is bescheiden over de beperkingen en erkent dat:

• De effectiviteit afhankelijk is van de beschikbaarheid van praktisch efficiënte optimizers voor de VML-objectief.
• Optimalisatieproblemen aanhouden voor slecht geconditioneerde of niet-lineaire operatoren, hoewel preconditioning helpt.
• Toekomstig werk nodig is om hogere-orde optimizers te ontwikkelen om optimalisatieproblemen verder te mitigeren zonder prohibitieve computationele kosten.

Het artikel concludeert dat VML-MAP een principieel, theoretisch gefundeerd en empirisch effectief kader biedt voor het oplossen van inverse problemen met diffusiemodellen, waarmee de kloof tussen theoretische MAP-estimatie en praktische beeldrestauratie wordt overbrugd.