Technische Zusammenfassung: Lokale MAP-Abtastung für Diffusionsmodelle

Problemstellung

Diffusions-Posterior-Abtastung (DPS) hat sich als leistungsfähiges Framework zur Lösung inverser Probleme durch Abtastung aus der Posterior-Verteilung $p(x_0 | y)$ etabliert und bietet eine fundierte Unsicherheitsquantifizierung sowie Multi-Modalität. Viele praktische Anwendungen in der Bildgebung und den Naturwissenschaften priorisieren jedoch eine genaue Punktschätzung gegenüber einer vollständigen Verteilungsabtastung. Der Maximum-A-Posteriori (MAP)-Schätzer, definiert als $\arg \max p(x_0 | y)$, bleibt ein Standardziel für Rekonstruktionen.

Bestehende optimierungsbasierte Methoden (z. B. Resample, DiffPIR, DCDP) versuchen, die Rekonstruktionsleistung durch das Alternieren zwischen Denoising- und Optimierungsschritten zu verbessern. Der theoretische Zusammenhang zwischen diesen iterativen Verfahren, der globalen MAP-Lösung und der Posterior-Abtastung bleibt jedoch unklar. Es ist nicht geklärt, ob diese Methoden zur globalen MAP konvergieren oder ob sie mit DPS konsistent sind. Ferner bedarf die spezifische Natur der Optimierungsschritte in diesen Methoden – ob sie ein globales oder lokales Ziel lösen – einer Klärung.

Methodik: Lokale MAP-Abtastung (LMAPS)

Die Autoren stellen Local MAP Sampling (LMAPS) vor, ein neues Inferenz-Framework, das iterativ lokale Maximum-a-Posteriori-Teilprobleme entlang der Diffusionstrajektorie löst.

1. Theoretisches Framework

LMAPS funktioniert durch die Lösung einer Sequenz lokaler MAP-Probleme anstelle eines einzigen globalen Problems. Zu jedem Diffusionsschritt $t$ wird, gegeben einen verrauschten Zustand $x_t$ und eine Messung $y$, folgendes angestrebt:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
Dies wird von einem Standard-DDIM-Übergangsschritt gefolgt:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Unterschiede zu bestehenden Methoden:

• vs. Globale MAP: Die globale MAP sucht den Modus der gesamten Posterior-Verteilung $p(x_0 | y)$. LMAPS sucht zu jedem Schritt den Modus der bedingten Posterior-Verteilung $p(x_0 | x_t, y)$. Die Autoren argumentieren, dass die Sequenz lokaler Maximierer nicht notwendigerweise zum globalen Maximierer von $p(x_0 | y)$ konvergiert.
• vs. DPS: DPS verwendet den bedingten Erwartungswert $E[x_0 | x_t, y]$ im Update. LMAPS ersetzt diesen Erwartungswert durch den bedingten Modus. Diese sind nur äquivalent, wenn die bedingte Posterior-Verteilung Gaußsch ist; andernfalls führt LMAPS zu einer Modus-suchenden Verzerrung und vermeidet Regionen niedriger Dichte zwischen den Modi.
• vs. DMAP: Während DMAP ebenfalls eine lokale MAP-Perspektive nutzt, optimiert es im Raum der latenten Diffusionsvariablen ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). LMAPS formuliert das Problem direkt im Raum des sauberen Signals ($x^*_0$) und bietet eine komplementäre Sichtweise, die explizit mit der Messwahrscheinlichkeit $p(y|x_0)$ verbunden ist.

2. Praktischer Algorithmus und Approximationen

Um das lokale MAP-Problem für allgemeine nichtlineare inverse Probleme handhabbar zu machen, schlagen die Autoren spezifische Approximationen und Umformulierungen vor:

• Kovarianz-Approximation: Die Autoren approximieren die bedingte Kovarianz $\Sigma_{0|t}$ des Priors $p(x_0 | x_t)$ unter Verwendung einer isotropen Form, die durch eine Gaußsche Prior-Annahme motiviert ist. Wenn das Rauschniveau $\sigma_t \to 0$ geht, wird die Posterior-Kovarianz asymptotisch isotrop. Sie approximieren $\Sigma_{0|t} \approx \frac{1}{SNR} I$, wobei $SNR = \alpha_t^2 / \sigma_t^2$.

• Zielformulierung: Das lokale MAP-Ziel wird als Minimierung einer gewichteten Summe aus Prior-Distanz und Datenkonsistenz formuliert:
  $$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\}$$
  Um die numerische Stabilität und Interpretierbarkeit zu verbessern, wird dies in eine skaleninvariante konvexe Kombination umformuliert:
  $$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\}$$
  wobei $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. Diese Formulierung ermöglicht eine automatische Abkühlung (Annealing): Wenn $t$ abnimmt ($\sigma_t \to 0$), verschiebt sich das Ziel von messungsgetrieben zu prior-getrieben.

• Linear vs. Nicht-linear: Für lineare inverse Probleme kann die lokale MAP-Lösung in geschlossener Form hergeleitet werden, wobei Tweedie Moment Projected Diffusion (TMPD) als Spezialfall wiederhergestellt wird. Für nichtlineare Probleme wird die Lösung mittels Gradientenabstieg ermittelt.

Hauptbeiträge

1. Theoretische Vereinheitlichung: Das Papier formuliert LMAPS, um optimierungsbasierte Methoden für inverse Probleme und Tweedie Moment Projected Diffusion (TMPD) unter einem einzigen probabilistischen Framework zu vereinen. Es klärt, dass optimierungsbasierte Methoden inhärent lokale MAP-Probleme und keine globale MAP oder vollständige Posterior-Abtastung lösen.
2. Methodische Fortschritte: Die Autoren führen eine durch eine Gauß-Verteilung motivierte Kovarianz-Approximation und eine umformulierte Zielfunktion ein, die durch automatische Abkühlung die numerische Stabilität und Interpretierbarkeit verbessert.
3. Empirische Validierung: Das Framework wird über 10 Bildwiederherstellungsaufgaben (einschließlich linearer, nichtlinearer und nicht-differenzierbarer Operatoren) und 3 wissenschaftliche inverse Probleme validiert.

Experimentelle Ergebnisse

Die Autoren evaluierten LMAPS auf den FFHQ- und ImageNet-Datensätzen für die Bildwiederherstellung sowie auf Benchmarks für wissenschaftliche inverse Probleme (Lineare Inverse Streuung, Compressed Sensing MRI und Black-Hole-Bildgebung).

• Bildwiederherstellung: LMAPS erreichte in 43 von 60 Fällen über die Metriken PSNR, SSIM und LPIPS auf FFHQ und ImageNet State-of-the-Art-Leistung. Es übertraf konsistent die eng verwandte DMAP-Baseline und zeigte Verbesserungen von über 2 dB PSNR in Aufgaben wie Bewegungsunschärfeentfernung, JPEG-Wiederherstellung und Quantisierung.
• Wissenschaftliche inverse Probleme: LMAPS erzielte konsistent den höchsten PSNR-Wert über alle wissenschaftlichen Aufgaben hinweg. Bemerkenswerterweise erzielte es im Vergleich zu früheren Methoden bei drei Aufgaben zur Linearen Inversen Streuung (LIS) Gewinne von mehr als 1,5 dB PSNR.
• Effizienz: LMAPS zeigte eine überlegene Effizienz im Vergleich zu DAPS und erzielte bessere Ergebnisse mit weniger Rechenressourcen (z. B. weniger Optimierungsschritten pro Diffusionsschritt).

Bedeutung und Behauptungen

Das Papier behauptet, dass LMAPS eine fundamentale probabilistische Interpretation für optimierungsbasierte Methoden liefert und die Lücke zwischen Posterior-Abtastung (DPS) und Punktschätzung (MAP) schließt. Durch das explizite Lösen lokaler MAP-Teilprobleme bietet die Methode eine vereinheitlichte Sichtweise, die das Verhalten bestehender optimierungsbasierter Solver erklärt.

Die Autoren betonen, dass LMAPS zwar nicht aus der vollständigen Posterior-Verteilung abtastet (und somit die Multi-Modalität nicht auf die gleiche Weise wie DPS erfasst), aber speziell für Szenarien entwickelt wurde, in denen eine genaue Punktschätzung das primäre Ziel ist. Das Framework bietet eine robuste Alternative, die die Rekonstruktionsgenauigkeit mit den strukturellen Priors von Diffusionsmodellen ausbalanciert und State-of-the-Art-Ergebnisse sowohl bei der Standard-Bildwiederherstellung als auch bei herausfordernden wissenschaftlichen Rekonstruktionsaufgaben erzielt. Die Arbeit legt nahe, dass zukünftige Forschung die Nützlichkeit der globalen MAP-Schätzung mit Diffusionspriors weiter untersuchen könnte, was eine offene Herausforderung bleibt.