Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling

Sintesi Tecnica: Tensor Train Diffusion

Definizione del Problema
Il documento affronta la sfida di campionare da densità di probabilità complesse e ad alta dimensione $p_{\text{target}} = \rho_{\text{target}}/Z$, dove la densità non normalizzata $\rho_{\text{target}}$ è valutabile puntualmente, ma la costante di normalizzazione $Z$ è intrattabile e non sono disponibili campioni da $p_{\text{target}}$. Questo è un problema centrale nelle scienze computazionali, inclusi la dinamica molecolare e la statistica bayesiana.

Sebbene i modelli di diffusione offrano un framework potente per tali compiti, invertendo un processo di addizione di rumore tramite equazioni differenziali stocastiche (SDE), essi richiedono tipicamente l'apprendimento di una funzione di score da campioni tramite denoising score matching. In questo specifico contesto, lo score deve essere derivato direttamente dalla densità non normalizzata senza accesso ai campioni. La dinamica inversa è governata da un'equazione alle derivate parziali di Hamilton-Jacobi-Bellman (HJB). I metodi esistenti per risolvere questa PDE ad alta dimensione, come le Physics-Informed Neural Networks (PINN) o le tecniche basate su traiettorie, soffrono di lunghi tempi di addestramento, alta sensibilità agli iperparametri e della necessità di valutazioni onerose di $\rho_{\text{target}}$ e delle sue derivate di ordine superiore.

Metodologia
Gli autori propongono Tensor Train Diffusion (TTD), un nuovo solver per l'equazione HJB sottostante che sfrutta le rappresentazioni Functional Tensor Train (FTT).

1. Formulazione Matematica:

   • Il problema è inquadrato come la ricerca del controllo ottimale (drift) per una SDE di campionamento che inverte un processo di addizione di rumore (noising) in avanti.
   • La log-densità del processo in avanti, $V = -\log \vec{p}_Y$, soddisfa una specifica PDE HJB.
   • Utilizzando la connessione tra le equazioni HJB e le Backward Stochastic Differential Equations (BSDE), gli autori riformulano la soluzione della PDE come una sequenza di problemi di regressione. Ciò consente alla soluzione di essere computata tramite uno schema iterativo a ritroso nel tempo.

2. Approssimazione Functional Tensor Train (FTT):

   • Invece di utilizzare reti neurali, la funzione valore $V$ e il suo gradiente sono approssimati tramite FTT. Questo formato sfrutta le strutture a basso rango latenti in funzioni ad alta dimensione, rappresentandole come un prodotto di core tensoriali a basso rango.
   • Il metodo utilizza Extended Tensor Trains (xTT), che combinano il formato TT con specifiche funzioni di base univariante (ad esempio, polinomi di Legendre, B-spline, modi di Fourier).
   • L'ottimizzazione dei coefficienti FTT è eseguita tramite Alternating Least Squares (ALS), un approccio basato sulla regressione che evita le instabilità del gradiente stocastico (SGD).

3. Caratteristiche Algoritmiche:

   • Raffinamento Iterativo: L'algoritmo impiega un ciclo esterno in cui la politica di controllo viene aggiornata iterativamente. Partendo da una stima iniziale (ad esempio, dinamica di Langevin), il metodo affina le traiettorie di campionamento per allineare meglio la distribuzione di addestramento alla distribuzione target, scoprendo efficacemente nuovi modi.
   • Adattività: Il framework include strategie adattive per:
     • Selezione del Rango (Rank Selection): Regolazione dinamica del rango TT basata sul decadimento dei valori singolari per bilanciare accuratezza e costo computazionale.
     • Grado delle Basi (Basis Degrees): Adattamento del grado polinomiale delle funzioni di base in base alla regolarità della soluzione ai diversi istanti temporali.
     • Regolarizzazione: Tuning automatico dell'entità della regolarizzazione per stabilizzare l'apprendimento delle derivate.
   • Gestione del Dominio: Il metodo gestisce domini di approssimazione mobili ed estende la politica appresa alle regioni esterne al dominio di addestramento utilizzando estensioni lineari affini per garantire la stabilità durante il campionamento.

Contributi Chiave

• Nuovo Solver: Introduzione di TTD, un solver per le equazioni HJB che integra il campionamento basato sulla diffusione, la teoria HJB e le BSDE per ottenere un favorevole compromesso tra costo computazionale e accuratezza.
• Robustezza e Flessibilità: Dimostrazione della flessibilità del metodo incorporando diverse funzioni di base (Legendre, B-spline, Fourier) e migliorando la robustezza rispetto ai precedenti solver di PDE basati su tensor-train, particolarmente nella gestione delle traiettorie al di fuori del dominio primario di addestramento.
• Performance: Dimostrazione che TTD può campionare da distribuzioni complesse e multimodali di varie dimensioni, superando i campionatori basati su diffusione esistenti sia in velocità che in accuratezza.

Risultati Sperimentali
Gli autori valutano TTD su diversi benchmark impegnativi:

• Problemi Multi-well ad Alta Dimensione: In dimensioni $d=10$ e $d=50$ con molteplici modi, TTD ottiene un campionamento ad alta fedeltà. Il metodo evita il collasso dei modi (mode collapse), un problema comune nel campionamento basato sulla diffusione. Il tempo di esecuzione scala linearmente con il numero di passi temporali ed è significativamente più veloce delle alternative basate su reti neurali (es. DIS).
• Target Gaussian Anisotropi: Gli esperimenti confermano la capacità dell'algoritmo di ridurre adattivamente i ranghi man mano che la distribuzione target converge verso una normale standard, validando il meccanismo di adattività del rango.
• Teoria dei Campi Scalari $\phi^4$: Il metodo campiona con successo dal modello di Ginzburg-Landau discretizzato, un sistema con esponenzialmente molti modi metastabili e forti correlazioni tra vicini, dove i metodi MCMC standard tipicamente falliscono a causa del collasso dei modi.
• Confronto: TTD supera costantemente i baseline allo stato dell'arte (inclusi PIS, DDS, DIS e varie varianti MCMC) in termini di Dimensione del Campione Effettivo (ESS), divergenza della log-varianza e accuratezza delle costanti di normalizzazione stimate, ottenendo spesso risultati migliori in meno tempo.

Significatività e Rivendicazioni
Il documento afferma che TTD fornisce un'alternativa veloce, accurata e stabile alle tecniche basate su reti neurali per risolvere equazioni HJB ad alta dimensione in contesti di campionamento. Sfruttando le strutture tensoriali a basso rango, il metodo elude la "maledizione della dimensionalità" associata ai metodi tradizionali basati su mesh e agli alti costi di addestramento e alla sensibilità agli iperparametri degli algoritmi di apprendimento basati su SGD. Gli autori posizionano TTD come una base robusta per risolvere le equazioni HJB, richiedendo meno valutazioni della densità target e offrendo una maggiore efficienza per problemi di campionamento ad alta dimensione. Riconoscono i limiti riguardanti la garanzia di identificare tutti i modi in distribuzioni estremamente complesse e i limiti di espressività delle FTT per certe densità target, ma enfatizzano la performance superiore del metodo nei regimi testati.