Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

Technische Zusammenfassung: Varianzreduktion für die nicht-log-konkave Stichprobenziehung mit Anwendungen auf inverse Probleme

Problemstellung
Die Stichprobenziehung aus hochdimensionalen, nicht-log-konkaven Verteilungen mit unnormierten Dichten stellt eine fundamentale Herausforderung im maschinellen Lernen dar, insbesondere wenn exakte Gradienten der Potenzialfunktion nicht verfügbar sind. In solchen Szenarien müssen Gradienten durch stochastische Schätzungen approximiert werden, die unter einem festen Rechenbudget pro Iteration oft eine hohe Varianz aufweisen. Während Varianzreduktionstechniken (z. B. SGD mit Momentum, STORM, PAGE) die Konvergenz in der nicht-konvexen Optimierung verbessert haben, bleibt ihre Anwendung auf die Stichprobenziehung aus nicht-log-konkaven Verteilungen weitgehend unerforscht. Darüber hinaus operieren bestehende posteriore Stichprobenverfahren für inverse Probleme unter Verwendung von Score-basierten generativen Modellen (SGMs) typischerweise in einem Full-Batch-Regime, was ihre Anwendbarkeit bei großskaligen Datensätzen, wie sie in der 3D- und dynamischen Bildgebung üblich sind, einschränkt.

Methodik
Die Autoren schlagen ein vereinheitlichtes Framework vor, das Varianzreduktionstechniken aus der nicht-konvexen Optimierung auf den Bereich der Stichprobenziehung adaptiert. Sie führen zwei vereinheitlichte Schätzer ein:

1. Mean-Lipschitz Variance-Reduced Langevin Diffusion (ML-VRLD): Dieser Schätzer vereinheitlicht SGD mit Momentum und dessen stochastisches Gegenstück. Er setzt die Lipschitz-Stetigkeit der exakten Gradienten voraus (Annahme 1) und bietet eine vereinheitlichte Behandlung, bei der der Schätzer ein gewichteter Durchschnitt eines Mini-Batch-Gradienten und eines Momentum-basierten Updates ist.
2. Sample-Lipschitz Variance-Reduced Langevin Diffusion (SL-VRLD): Dieser Schätzer vereinheitlicht PAGE und STORM. Er setzt die Lipschitz-Stetigkeit der stochastischen Gradienten im Mittel-Quadrat-Sinne voraus (Annahme 4). Er beinhaltet einen Korrekturterm basierend auf der Differenz zwischen dem aktuellen und dem vergangenen stochastischen Gradienten, was die Schätzgenauigkeit unter der Sample-Lipschitz-Bedingung verbessert.

Beide Methoden sind darauf ausgelegt, eine $\epsilon$-relative Fisher-Informations-Konvergenz (FI) mit $O(1)$ stochastischen Gradientenberechnungen pro Iteration zu erreichen, im Gegensatz zum $O(Ld\sigma^2/\epsilon)$-Bedarf von Standard-Stochastic Gradient Langevin Dynamics (SGLD).

Das Framework wird auf inverse Probleme (z. B. MRT- und CT-Rekonstruktion) erweitert, indem diese Schätzer in einen Annealed Langevin Monte Carlo (LMC) Stichprobenziehungsprozess mit SGM-Priors integriert werden. Die posteriore Stichprobenformulierung nutzt eine Sequenz von intermediären Zielverteilungen mit Annealing- und Glättungsschemata, um die Iterationen von einer geglätteten Prior-Verteilung zu der wahren Posteriori-Verteilung zu führen. Theoretische Garantien werden unter Annahmen über begrenzte SGM-Schätzfehler etabliert.

Wesentliche Beiträge

• Vereinheitlichte Analyse: Die Arbeit liefert die erste vereinheitlichte Analyse von varianzreduzierten Schätzern (speziell ML-VRLD und SL-VRLD) für die Stichprobenziehung aus nicht-log-konkaven Verteilungen.
• Verbesserte Komplexität:
  • Sowohl ML-VRLD als auch SL-VRLD erreichen eine $\epsilon$-relative FI-Konvergenz mit $O(1)$ Gradientenberechnungen pro Iteration, was die pro Iteration anfallende Komplexität von SGLD signifikant verbessert.
  • SL-VRLD reduziert die gesamte Rechenkomplexität weiter von $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) auf $O(L^2d^2\sigma^2/\epsilon^3)$.
• Konvergenzgarantien: Die Autoren etablieren nicht-asymptotische Konvergenzraten in der quadratischen Total Variation (TV) Distanz unter einer Poincaré-Inequalitäts-Annahme und beweisen die schwache Konvergenz gegen die Zielverteilung für beide Algorithmen.
• Erweiterung auf inverse Probleme: Das Framework wird auf die posteriore Stichprobenziehung mit SGM-Priors erweitert, wobei dieselben Konvergenzgarantien unter begrenzten SGM-Schätzfehlern etabliert werden. Dies adresset die Einschränkung bestehender Full-Batch-Stichprobenverfahren.

Ergebnisse

• Synthetische Experimente: Bei Gaußschen Mischmodellen (GMM) und Compressed-Sensing-Inversen Problemen zeigen die vorgeschlagenen Methoden eine schnellere Konvergenz zu niedrigeren relativen FI-Werten im Vergleich zu SGLD. SGLD mit einer festen Batch-Größe konvergiert zu einem suboptimalen FI-Niveau, während ML-VRLD und SL-VRLD die Zielgenauigkeit mit derselben festen Batch-Größe erreichen. SL-VRLD erreicht die Zielverteilung konsistent schneller als ML-VRLD.
• Statistische Validierung: In einer posterioren Stichprobenaufgabe mit bekannter Grundwahrheit (Ground Truth) erzeugte SGLD Proben mit signifikant erhöhter Varianz. ML-VRLD und SL-VRLD reduzierten diese Varianz erheblich, wobei SL-VRLD die engste Übereinstimmung mit der Varianz der Ground-Truth-Posteriori erreichte.
• Bildgebende Anwendungen:
  • MRT-Rekonstruktion: Auf dem fastMRI-Datensatz mit Beschleunigungsfaktoren von 8 und 12 produzierten ML-VRLD und SL-VRLD schärfere anatomische Details und weniger Artefakte im Vergleich zu den Baselines (einschließlich DPS, PnP-DM und SGLD). SL-VRLD erzielte die höchsten PSNR- und SSIM-Werte.
  • Sparse-Angle-CT-Rekonstruktion: Auf dem 2DeteCT-Datensatz mit 180 und 360 Projektionsansichten übertrafen die vorgeschlagenen Methoden die Baselines bei der Bewahrung feiner struktureller Details, was den Nutzen der Varianzreduktion in hochrauschbehafteten, unterabgetasteten Regimen demonstriert.

Bedeutung
Das Paper beansprucht, die ersten nicht-asymptotischen Konvergenzgarantien für die nicht-log-konkave Stichprobenziehung unter Verwendung von varianzreduzierten Schätzern mit $O(1)$ Gradientenberechnungen pro Iteration bereitzustellen. Durch die Erweiterung dieser Techniken auf inverse Probleme mit SGM-Priors ermöglicht diese Arbeit eine prinzipielle posteriore Stichprobenziehung für großskalige bildgebende Aufgaben, bei denen Full-Batch-Methoden rechnerisch prohibitiv sind. Die empirischen Ergebnisse validieren, dass Varianzreduktion konsistent die Stichprobenqualität und statistische Genauigkeit unter festen Speicher- und Rechenbeschränkungen verbessert und somit eine skalierbare Lösung für die hochdimensionale Bayessche Inferenz bietet.