Technische Samenvatting: Lokale MAP-bemonstering voor Diffusiemodellen

Probleemstelling

Diffusion Posterior Sampling (DPS) is naar voren gekomen als een krachtig raamwerk voor het oplossen van inverse problemen door te bemonsteren uit de posteriorverdeling $p(x_0 | y)$, wat een principiële onzekerheidskwantificering en multi-modale aard biedt. Echter, veel praktische toepassingen in beeldvorming en wetenschap geven prioriteit aan nauwkeurige puntenschatting boven volledige distributie-bemonstering. De Maximum A Posteriori (MAP)-schatter, gedefinieerd als $\arg \max p(x_0 | y)$, blijft een standaard reconstructiedoel.

Bestaande op optimalisatie gebaseerde methoden (bijv. Resample, DiffPIR, DCDP) proberen de reconstructieprestaties te verbeteren door te alterneren tussen denoising- en optimalisatiestappen. De theoretische connectie tussen deze iteratieve procedures, de globale MAP-oplossing en posterior-bemonstering blijft echter onduidelijk. Het is niet vastgesteld of deze methoden convergeren naar de globale MAP of of ze consistent zijn met DPS. Bovendien vereist de specifieke aard van de optimalisatiestappen in deze methoden – of ze een globaal of lokaal doel oplossen – opheldering.

Methodologie: Lokale MAP-bemonstering (LMAPS)

De auteurs introduceren Local MAP Sampling (LMAPS), een nieuw inferentieraamwerk dat iteratief lokale maximum-a-posteriori-deelproblemen oplost langs de diffusietrajectorie.

1. Theoretisch Kader

LMAPS werkt door een reeks lokale MAP-problemen op te lossen in plaats van één enkel globaal probleem. Bij elke diffusiestap $t$, gegeven een ruisbeïnvloede toestand $x_t$ en meting $y$, zoekt de methode:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
Dit wordt gevolgd door een standaard DDIM-overgangsstap:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Onderscheid met bestaande methoden:

• T.o.v. Globale MAP: Globale MAP zoekt de modus van de volledige posterior $p(x_0 | y)$. LMAPS zoekt bij elke stap de modus van de conditionele posterior $p(x_0 | x_t, y)$. De auteurs betogen dat de reeks lokale maximalisatoren niet noodzakelijkerwijs convergeert naar de globale maximalisator van $p(x_0 | y)$.
• T.o.v. DPS: DPS gebruikt de conditionele verwachting $E[x_0 | x_t, y]$ in de update. LMAPS vervangt deze verwachting door de conditionele modus. Deze zijn alleen equivalent als de conditionele posterior Gaussisch is; anders introduceert LMAPS een bias gericht op de modus, waardoor gebieden met lage dichtheid tussen modi worden vermeden.
• T.o.v. DMAP: Hoewel DMAP ook een lokaal MAP-perspectief hanteert, optimaliseert het in de ruimte van de latente diffusievariabelen ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). LMAPS stelt het probleem direct in de schone signaalruimte ($x^*_0$), wat een complementair perspectief biedt dat expliciet verbinding maakt met de meetkansen $p(y|x_0)$.

2. Praktisch Algorithmus en Benaderingen

Om het lokale MAP-probleem hanteerbaar te maken voor algemene niet-lineaire inverse problemen, stellen de auteurs specifieke benaderingen en herformuleringen voor:

• Covariantiebenadering: De auteurs benaderen de conditionele covariantie $\Sigma_{0|t}$ van de prior $p(x_0 | x_t)$ met een isotrope vorm, gemotiveerd door een Gaussische prior-aanname. Naarmate het ruisniveau $\sigma_t \to 0$, wordt de posteriorcovariantie asymptotisch isotroop. Ze benaderen $\Sigma_{0|t} \approx \frac{1}{SNR} I$, waarbij $SNR = \alpha_t^2 / \sigma_t^2$.

• Herformulering van het Doel: Het lokale MAP-doel wordt geformuleerd als het minimaliseren van een gewogen som van de prior-afstand en de data-fideliteit:
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  Om numerieke stabiliteit en interpreteerbaarheid te verbeteren, wordt dit herformuleerd tot een schaal-invariante convexe combinatie:
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  waarbij $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. Deze formulering maakt automatische afkoeling mogelijk: naarmate $t$ afneemt ($\sigma_t \to 0$), verschuift het doel van meetgedreven naar prior-gedreven.

• Lineair versus Niet-lineair: Voor lineaire inverse problemen kan de lokale MAP-oplossing in gesloten vorm worden afgeleid, waarbij Tweedie Moment Projected Diffusion (TMPD) wordt hersteld als een speciaal geval. Voor niet-lineaire problemen wordt de oplossing verkregen via gradiëntafdalend.

Belangrijkste Bijdragen

1. Theoretische Unificatie: Het artikel formuleert LMAPS om op optimalisatie gebaseerde methoden voor inverse problemen en Tweedie Moment Projected Diffusion (TMPD) te unificeren onder één enkel probabilistisch raamwerk. Het verduidelijkt dat op optimalisatie gebaseerde methoden inherent lokale MAP-problemen oplossen in plaats van globale MAP of volledige posterior-bemonstering.
2. Methodologische Vooruitgang: De auteurs introduceren een door een Gaussische prior gemotiveerde covariantiebenadering en een herformuleerde doelfunctie die numerieke stabiliteit en interpreteerbaarheid verbetert via automatische afkoeling.
3. Empirische Validatie: Het raamwerk is gevalideerd over 10 beeldhersteltaken (inclusief lineaire, niet-lineaire en niet-differentieerbare operatoren) en 3 wetenschappelijke inverse problemen.

Experimentele Resultaten

De auteurs hebben LMAPS geëvalueerd op FFHQ- en ImageNet-datasets voor beeldherstel en op benchmarks voor wetenschappelijke inverse problemen (Lineaire Inverse Verstrooiing, Compressed Sensing MRI en Black Hole Imaging).

• Beeldherstel: LMAPS behaalde state-of-the-art prestaties in 43 van de 60 gevallen voor de metrics PSNR, SSIM en LPIPS op FFHQ en ImageNet. Het presteerde consistent beter dan de nauw verwante DMAP-baseline, met verbeteringen van meer dan 2 dB PSNR in taken zoals bewegingsontwarring, JPEG-herstel en kwantisatie.
• Wetenschappelijke Inverse Problemen: LMAPS behaalde consistent de hoogste PSNR over alle wetenschappelijke taken. Opmerkelijk is dat het winsten van meer dan 1,5 dB PSNR behaalde op drie Lineaire Inverse Verstrooiing (LIS)-taken ten opzichte van eerdere methoden.
• Efficiëntie: LMAPS toonde superieure efficiëntie in vergelijking met DAPS, waarbij betere resultaten werden behaald met minder rekenkracht (bijv. minder optimalisatiestappen per diffusiestap).

Betekenis en Claims

Het artikel beweert dat LMAPS een principiële probabilistische interpretatie biedt voor op optimalisatie gebaseerde methoden, en de kloof overbrugt tussen posterior-bemonstering (DPS) en puntenschatting (MAP). Door expliciet lokale MAP-deelproblemen op te lossen, biedt de methode een verenigd perspectief dat het gedrag van bestaande op optimalisatie gebaseerde solvers verklaart.

De auteurs benadrukken dat hoewel LMAPS niet bemonstert uit de volledige posteriorverdeling (en dus niet op dezelfde manier als DPS multi-modale aard vastlegt), het specifiek is ontworpen voor scenario's waarbij nauwkeurige puntenschatting het primaire doel is. Het raamwerk biedt een robuust alternatief dat reconstructiefideliteit afweegt tegen structurele priors van diffusiemodellen, en state-of-the-art resultaten behaalt in zowel standaard beeldherstel als uitdagende wetenschappelijke reconstructietaken. Het werk suggereert dat toekomstig onderzoek de bruikbaarheid van globale MAP-schatting met diffusiepriors verder kan verkennen, wat een open uitdaging blijft.