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

Technische Samenvatting: Varianciereductie voor Niet-Log-Concaaf Samplen met Toepassingen op Inverse Problemen

Probleemstelling
Het samplen uit hoog-dimensionale, niet-log-concaaf verdelingen met ongeëvalueerde dichtheden is een fundamentele uitdaging in machine learning, met name wanneer exacte gradiënten van de potentiaalfunctie niet beschikbaar zijn. In dergelijke scenario's moeten gradiënten worden benaderd via stochastische schattingen, die vaak een hoge variantie vertonen onder een vast computationeel budget per iteratie. Hoewel technieken voor varianciereductie (bijv. SGD met momentum, STORM, PAGE) de convergentie in niet-convexe optimalisatie hebben verbeterd, blijft hun toepassing op het samplen uit niet-log-concaaf verdelingen grotendeels onontgonnen. Bovs is de huidige posterior sampling-methoden voor inverse problemen met score-gebaseerde generatieve modellen (SGM's) doorgaans gebaseerd op een full-batch regime, wat de toepasbaarheid op grote datasets die gebruikelijk zijn in 3D- en dynamische beeldvorming beperkt.

Methodologie
De auteurs stellen een verenigd framework voor dat technieken voor varianciereductie uit de niet-convexe optimalisatie aanpast aan het domein van sampling. Zij introduceren twee verenigde estimators:

1. Mean-Lipschitz Variance-Reduced Langevin Diffusion (ML-VRLD): Deze estimator verenigt SGD met momentum en de stochastische tegenhanger ervan. Het gaat uit van Lipschitz-continuïteit van de exacte gradiënten (Aanname 1) en biedt een verenigde behandeling waarbij de estimator een gewogen gemiddelde is van een mini-batch gradiënt en een momentum-gebaseerde update.
2. Sample-Lipschitz Variance-Reduced Langevin Diffusion (SL-VRLD): Deze estimator verenigt PAGE en STORM. Het gaat uit van Lipschitz-continuïteit van de stochastische gradiënten in de mean-squared zin (Aanname 4). Het bevat een correctieterm gebaseerd op het verschil tussen de huidige en voorgaande stochastische gradiënten, wat de schattingsnauwkeurigheid verbetert onder de sample-Lipschitz conditie.

Beide methoden zijn ontworpen om $\epsilon$-relatieve Fisher Informatie (FI) convergentie te bereiken met $O(1)$ stochastische gradiëntberekeningen per iteratie, in contrast met de $O(Ld\sigma^2/\epsilon)$ vereiste van standaard Stochastic Gradient Langevin Dynamics (SGLD).

Het framework wordt uitgebreid naar inverse problemen (bijv. MRI- en CT-reconstructie) door deze estimators te integreren in een geanneilde Langevin Monte Carlo (LMC) sampling proces met SGM-priors. De posterior sampling formulering maakt gebruik van een sequentie van tussenliggende doelverdelingen met annealing- en smoothing-schema's om de iteraties te leiden van een gesmoothde prior naar de ware posterior. Theoretische garanties worden vastgesteld onder aannames van begrensde SGM-schattingsfouten.

Belangrijkste Bijdragen

• Verenigde Analyse: Het artikel biedt de eerste verenigde analyse van varianciereduceerde estimators (specifiek ML-VRLD en SL-VRLD) voor het samplen uit niet-log-concaaf verdelingen.
• Verbeterde Complexiteit:
  • Zowel ML-VRLD als SL-VRLD bereiken $\epsilon$-relatieve FI-convergentie met $O(1)$ gradiëntberekeningen per iteratie, wat de per-iteratie complexiteit van SGLD aanzienlijk verbetert.
  • SL-VRLD reduceert de totale computationele complexiteit verder van $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) naar $O(L^2d^2\sigma^2/\epsilon^3)$.
• Convergentiegaranties: De auteurs stellen niet-asymptotische convergentiesnelheden vast in de gekwadrateerde total variation (TV) afstand onder een Poincaré-inequality aanname en bewijzen zwakke convergentie naar de doelverdeling voor beide algoritmen.
• Extensie naar Inverse Problemen: Het framework wordt uitgebreid naar posterior sampling met SGM-priors, waarbij dezelfde convergentiegaranties worden vastgesteld onder begrensde SGM-schattingsfouten. Dit adresseert de beperking van bestaande full-batch posterior sampling methoden.

Resultaten

• Synthetische Experimenten: Op Gaussian Mixture Models (GMM) en compressed sensing inverse problemen demonstreren de voorgestelde methoden snellere convergentie naar lagere relatieve FI-waarden vergeleken met SGLD. SGLD met een vaste batchgrootte convergeert naar een suboptimaal FI-niveau, terwijl ML-VRLD en SL-VRLD de doelnauwkeurigheid bereiken met dezelfde vaste batchgrootte. SL-VRLD bereikt consequent de doelverdeling sneller dan ML-VRLD.
• Statistische Validatie: In een posterior sampling taak met een bekende grondwaarheid produceerde SGLD samples met aanzienlijk opgeblazen variantie. ML-VRLD en SL-VRLD reduceerden deze variantie aanzienlijk, waarbij SL-VRLD de dichtstbijzijnde match bereikte met de grondwaarheid posterior variantie.
• Beeldvormingstoepassingen:
  • MRI Reconstructie: Op de fastMRI dataset met acceleratiefactoren van 8 en 12, produceerden ML-VRLD en SL-VRLD scherpere anatomische details en minder artefacten vergeleken met baselines (inclusief DPS, PnP-DM, en SGLD). SL-VRLD behaalde de hoogste PSNR- en SSIM-scores.
  • Sparse-Angle CT Reconstructie: Op de 2DeteCT dataset met 180 en 360 projectie-views presteerden de voorgestelde methoden beter dan de baselines in het behouden van fijne structurele details, wat het voordeel van varianciereductie in hoog-ruis, ondergesamplede regimes aantoont.

Betekenis
Het artikel beweert de eerste niet-asymptotische convergentiegaranties te bieden voor niet-log-concaaf sampling met varianciereduceerde estimators met $O(1)$ gradiëntberekeningen per iteratie. Door deze technieken uit te breiden naar inverse problemen met SGM-priors, maakt dit werk principiële posterior sampling mogelijk voor grootschalige beeldvormingstaken waarbij full-batch methoden computationeel onhaalbaar zijn. De empirische resultaten valideren dat varianciereductie consistent de kwaliteit van de samples en de statistische nauwkeurigheid verbetert onder vaste geheugen- en computationele beperkingen, wat een schaalbare oplossing biedt voor hoog-dimensionale Bayesiaanse inferentie.