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

Sintesi Tecnica: Riduzione della Varianza per il Campionamento Non Log-Concavo con Applicazioni ai Problemi Inversi

Enunciato del Problema
Il campionamento da distribuzioni ad alta dimensionalità e non log-concave con densità non normalizzate è una sfida fondamentale nel machine learning, particolarmente quando i gradienti esatti della funzione potenziale non sono disponibili. In tali scenari, i gradienti devono essere approssimati tramite stime stocastiche, che spesso presentano un'elevata varianza sotto un budget computazionale fisso per iterazione. Sebbene le tecniche di riduzione della varianza (ad esempio, SGD con momentum, STORM, PAGE) abbiano migliorato la convergenza nell'ottimizzazione non convessa, la loro applicazione al campionamento da distribuzioni non log-concave rimane ampiamente inesplorata. Inoltre, i metodi esistenti di campionamento a posteriori per problemi inversi che utilizzano modelli generativi basati sullo score (SGM) operano tipicamente in un regime full-batch, limitandone l'applicabilità a dataset su larga scala comuni nella diagnostica per immagini 3D e dinamica.

Metodologia
Gli autori propongono un framework unificato che adatta le tecniche di riduzione della varianza dall'ottimizzazione non convessa al dominio del campionamento. Introducono due stimatori unificati:

1. Diffusione di Langevin con Riduzione della Varianza Mean-Lipschitz (ML-VRLD): Questo stimatore unifica SGD con momentum e il suo corrispettivo stocastico. Assume la continuità Lipschitz dei gradienti esatti (Assunzione 1) e fornisce un trattamento unificato in cui lo stimatore è una media ponderata tra un gradiente mini-batch e un aggiornamento basato sul momentum.
2. Diffusione di Langevin con Riduzione della Varianza Sample-Lipschitz (SL-VRLD): Questo stimatore unifica PAGE e STORM. Assume la continuità Lipschitz dei gradienti stocastici in senso di media quadratica (Assunzione 4). Incorpora un termine di correzione basato sulla differenza tra i gradienti stocastici correnti e passati, il che migliora l'accuratezza della stima sotto la condizione di sample-Lipschitz.

Entrambi i metodi sono progettati per raggiungere una convergenza di $\epsilon$-Informazione di Fisher (FI) relativa con $O(1)$ calcoli del gradiente stocastico per iterazione, contrastando con il requisito di $O(Ld\sigma^2/\epsilon)$ della standard Stochastic Gradient Langevin Dynamics (SGLD).

Il framework viene esteso ai problemi inversi (ad esempio, ricostruzione MRI e CT) integrando questi stimatori in un processo di campionamento Langevin Monte Carlo (LMC) con prior SGM. La formulazione del campionamento a posteriori utilizza una sequenza di distribuzioni target intermedie con programmi di annealing e smoothing per guidare le iterazioni da un prior liscio al vero a posteriori. Le garanzie teoriche sono stabilite sotto assunzioni di errore di stima SGM limitato.

Contributi Chiave

• Analisi Unificata: Il documento fornisce la prima analisi unificata degli stimatori a riduzione della varianza (specificamente ML-VRLD e SL-VRLD) per il campionamento da distribuzioni non log-concave.
• Complessità Migliorata:
  • Sia ML-VRLD che SL-VRLD raggiungono una convergenza di $\epsilon$-FI relativa con $O(1)$ calcoli del gradiente per iterazione, migliorando significativamente la complessità per iterazione di SGLD.
  • SL-VRLD riduce ulteriormente la complessità computazionale totale da $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) a $O(L^2d^2\sigma^2/\epsilon^3)$.
• Garanzie di Convergenza: Gli autori stabiliscono tassi di convergenza non asintotica in distanza di variazione totale (TV) al quadrato sotto un'assunzione di disuguaglianza di Poincaré e dimostrano la convergenza debole alla distribuzione target per entrambi gli algoritmi.
• Estensione ai Problemi Inversi: Il framework è esteso al campionamento a posteriori con prior SGM, stabilendo le stesse garanzie di convergenza sotto un errore di stima SGM limitato. Ciò affronta il limite dei metodi di campionamento a posteriori full-batch esistenti.

Risultati

• Esperimenti Sintetici: Sui modelli a miscela gaussiana (GMM) e sui problemi di compressed sensing, i metodi proposti dimostrano una convergenza più rapida verso valori di FI relativa inferiori rispetto a SGLD. SGLD con una dimensione del batch fissa converge a un livello di FI subottimale, mentre ML-VRLD e SL-VRLD raggiungono l'accuratezza target con la stessa dimensione di batch fissa. SL-VRLD raggiunge costantemente la distribuzione target più velocemente di ML-VRLD.
• Validazione Statistica: In un compito di campionamento a posteriori con una verità fondamentale nota, SGLD ha prodotto campioni con varianza significativamente gonfiata. ML-VRLD e SL-VRLD hanno ridotto significativamente questa varianza, con SL-VRLD che ha ottenuto l'abbinamento più vicino alla varianza del vero a posteriori.
• Applicazioni di Imaging:
  • Ricostruzione MRI: Sul dataset fastMRI con fattori di accelerazione 8 e 12, ML-VRLD e SL-VRLD hanno prodotto dettagli anatomici più nitidi e meno artefatti rispetto ai baseline (inclusi DPS, PnP-DM e SGLD). SL-VRLD ha ottenuto i punteggi PSNR e SSIM più elevati.
  • Ricostruzione CT Sparse-Angle: Sul dataset 2DeteCT con 180 e 360 viste di proiezione, i metodi proposti hanno superato i baseline nella preservazione dei dettagli strutturali fini, dimostrando il beneficio della riduzione della varianza in regimi ad alto rumore e sottocampionati.

Significatività
Il documento sostiene di fornire le prime garanzie di convergenza non asintotica per il campionamento non log-concavo utilizzando stimatori a riduzione della varianza con $O(1)$ calcoli del gradiente per iterazione. Estendendo queste tecniche ai problemi inversi con prior SGM, il lavoro abilita un campionamento a posteriori rigoroso per compiti di imaging su larga scala dove i metodi full-batch sono computazionalmente proibitivi. I risultati empirici validano che la riduzione della varianza migliora costantemente la qualità del campione e l'accuratezza statistica sotto vincoli fissi di memoria e computazione, offrendo una soluzione scalabile per l'inferenza bayesiana ad alta dimensionalità.