Proximal Diffusion Neural Sampler

Riepilogo Tecnico: Campionatore Neurale di Diffusione Prossimale (PDNS)

Enunciato del Problema
Apprendere un campionatore neurale basato sulla diffusione per estrarre campioni da una distribuzione target non normalizzata (ad esempio, una distribuzione di Boltzmann $\pi(x) \propto e^{-\beta E(x)}$) è un compito fondamentale nella statistica computazionale, nell'inferenza bayesiana e nella fisica statistica. Questo problema è spesso formulato come un problema di Controllo Ottimo Stocastico (SOC) sulle misure di percorso, in cui una rete neurale parametrizza un processo di controllo per trasportare una semplice distribuzione di riferimento alla distribuzione target.

Tuttavia, l'addestramento di questi campionatori incontra sfide significative quando la distribuzione target è multimodale con alte barriere energetiche che separano le modalità. In tali scenari, i metodi standard basati su SOC (inclusi Cross-Entropy e le varianti Weighted Denoising Cross-Entropy) soffrono spesso di collasso delle modalità. Nelle fasi iniziali dell'addestramento, il modello cattura solo alcune modalità; la funzione di perdita globale rafforza quindi queste modalità catturate sulla base delle limitate traiettorie generate, causando la concentrazione della massa di probabilità su un sottoinsieme di modalità e ignorando le altre. Ciò porta a un'ottimizzazione instabile e a una scarsa copertura della distribuzione target.

Metodologia
Gli autori propongono il Campionatore Neurale di Diffusione Prossimale (PDNS), un quadro unificato progettato per affrontare il collasso delle modalità riformulando il processo di addestramento come una sequenza di sottoproblemi locali anziché una singola minimizzazione globale.

1. Quadro Unificato: Il lavoro unifica innanzitutto i campionatori di diffusione continui e discreti sotto un unico quadro SOC. Definisce la misura di percorso ottimale $P^*$ come la soluzione che minimizza un obiettivo di entropia relativa soggetto a una misura di percorso di riferimento $P_{ref}$ che soddisfa una condizione "senza memoria" (indipendenza tra le distribuzioni iniziale e terminale).
2. Metodo del Punto Prossimale: Per mitigare il collasso delle modalità, PDNS applica il metodo del punto prossimale allo spazio delle misure di percorso. Invece di ottimizzare direttamente l'obiettivo globale, l'algoritmo risolve iterativamente un sottoproblema locale al passo $k$:
   $$P_{\theta^*_k} = \arg\min_{P_\theta} \left( -E_{P_\theta}[r(X_T)] + KL(P_\theta \| P_{ref}) + \frac{1}{\eta_k} KL(P_\theta \| P_{\theta_{k-1}}) \right)$$
   Qui, $\eta_k$ è una dimensione del passo prossimale, e il termine $KL(P_\theta \| P_{\theta_{k-1}})$ agisce come regolarizzatore mantenendo l'iterazione corrente vicina a quella precedente.
3. Proprietà Teoriche: Gli autori dimostrano che la soluzione ottimale al $k$-esimo passo prossimale è un'interpolazione geometrica tra l'iterazione precedente e la misura di percorso ottimale. Nello specifico, la distribuzione terminale del $k$-esimo passo, $P_k$, è un'interpolazione geometrica tra la distribuzione terminale di riferimento $\nu$ e il target $\pi$. Quando $k \to \infty$, la sequenza di misure di percorso converge al target $P^*$. Questa procedura a stadi garantisce che il modello esplori gradualmente lo spazio degli stati, prevenendo un impegno prematuro su una singola modalità.
4. Istante Pratico (Proximal WDCE): Per efficienza, gli autori istanziano PDNS utilizzando un obiettivo Proximal Weighted Denoising Cross-Entropy (Proximal WDCE). Questo approccio evita l'archiviazione di traiettorie complete utilizzando pesi di importanza accumulati online.
   • Caso Continuo: L'obiettivo comporta il matching di un punteggio di denoising (o ponte) con un termine di peso derivato dalla dimensione del passo prossimale $\eta_k$ e dalla funzione di energia.
   • Caso Discreto: L'obiettivo adatta la diffusione discreta mascherata, utilizzando un simile schema di ponderazione basato sulla derivata di Radon-Nikodym delle misure di percorso.
5. Progettazione dello Scheduler: Il lavoro discute strategie per la selezione della dimensione del passo $\eta_k$ (o del suo equivalente $\lambda_k$). Un $\eta_k$ più piccolo (regolarizzazione più forte) preserva la copertura delle modalità ma rallenta la convergenza, mentre un $\eta_k$ più grande rischia il collasso delle modalità. Gli autori propongono sia scheduli predefiniti (diminuzione lineare di $\lambda_k$) sia una strategia adattiva che aggiusta $\eta_k$ basandosi su un vincolo di trust-region sulla divergenza KL tra il modello corrente e il prossimo target.

Contributi Chiave

1. Quadro Unificato: Il lavoro integra i campionatori basati su SOC continui e discreti in un'unica formulazione di misura di percorso.
2. Campionatore Neurale di Diffusione Prossimale (PDNS): Introduce un quadro innovativo che sfrutta iterazioni di punti prossimali nello spazio dei percorsi per stabilizzare l'apprendimento e mitigare il collasso delle modalità in contesti multimodali.
3. Proximal WDCE: Deriva istanze concrete ed efficienti dal punto di vista computazionale di PDNS per i domini continui e discreti utilizzando perdite di cross-entropy denoising pesate prossimalmente.
4. Strategie di Scheduler: Gli autori delineano strategie principiate per la selezione delle dimensioni dei passi prossimali, incluso uno scheduler adattivo che bilancia velocità di convergenza e stabilità.

Risultati Sperimentali
Gli autori valutano PDNS su una vasta gamma di benchmark continui e discreti:

• Compiti Continui:
  • Funzioni di Energia Sintetiche: PDNS supera basi solide (inclusi SMC, DDS, PIS e AS) su benchmark multimodali come il potenziale a 32 pozzi (MW54), Funnel e miscele gaussiane ad alta dimensionalità (GMM40). Raggiunge risultati all'avanguardia sulla miscela di Student's t (MoS) e su potenziali basati su particelle (sistemi di Lennard-Jones LJ-13 e LJ-55).
  • Dinamica Molecolare: Sul sistema del Dipeptide Alanina (AD), PDNS ottiene divergenze KL comparabili al metodo all'avanguardia ASBS e cattura accuratamente le strutture conformazionali nei diagrammi di Ramachandran.
• Compiti Discreti:
  • Fisica Statistica: PDNS supera significativamente LEAPS e Metropolis-Hastings (MH) nel campionamento di modelli di Ising e Potts a basse temperature (vicino o al di sotto dei punti critici), dove la separazione delle modalità è severa. Raggiunge errori inferiori nella magnetizzazione e nelle correlazioni a 2 punti.
  • Ottimizzazione Combinatoria: Applicato al problema del Maximum Cut su grafi Barabási–Albert ed Erdős–Rényi, PDNS dimostra la capacità di raggiungere qualità delle soluzioni comparabili a risolutori ground-truth (Gurobi).

Significato e Affermazioni
Il lavoro afferma che PDNS fornisce una soluzione robusta al problema del collasso delle modalità intrinseco nell'addestramento di campionatori neurali basati sulla diffusione per distribuzioni complesse e multimodali. Decomponendo l'ottimizzazione globale in una sequenza di passi prossimali, il metodo promuove un'esplorazione approfondita attraverso le modalità e stabilizza l'ottimizzazione negli spazi ad alta dimensionalità.

Gli autori notano due limitazioni principali:

1. L'indagine si concentra su una singola istanza di PDNS (proximal WDCE), lasciando spazio all'esplorazione di altri obiettivi o prior prossimali.
2. Gli esperimenti sono condotti su benchmark controllati piuttosto che su applicazioni reali su larga scala.

Il lavoro suggerisce direzioni future che includono l'esplorazione di problemi SOC con dinamiche di riferimento non senza memoria, l'estensione del quadro a spazi degli stati più generali e lo sviluppo di altre istanze dell'obiettivo prossimale.