Riepilogo Tecnico: Campionatori e Ponti di Diffusione Discreti

Enunciato del Problema

Il lavoro affronta la sfida di campionare da una distribuzione $p(x) \propto e^{-E(x)}$ su uno spazio degli stati discreto $\mathcal{S} = \{1, \dots, C\}^d$, dove la funzione di energia $E$ è nota ma la costante di normalizzazione $Z$ è intrattabile. Sebbene i campionatori di diffusione siano stati adattati con successo per spazi continui utilizzando il controllo stocastico ottimo e l'adattamento nello spazio dei percorsi, la loro applicazione agli spazi discreti rimane poco sviluppata. I campionatori di diffusione discreti esistenti spesso non sfruttano le tecniche di apprendimento per rinforzo (RL) off-policy che sono standard negli ambienti continui, portando a una copertura subottimale delle modalità e a una qualità di campionamento inferiore, specialmente in distribuzioni altamente multimodali. Inoltre, il problema di collegare due distribuzioni arbitrarie (il ponte di Schrödinger) negli spazi discreti è stato finora limitato a contesti in cui sono disponibili campioni da entrambe le distribuzioni (da dati a dati), lasciando inesplorato il contesto "da dati a energia" (dove una distribuzione è definita solo da una densità non normalizzata).

Metodologia

Gli autori propongono un framework che integra strategie di addestramento off-policy nei campionatori di diffusione discreti e generalizza questi metodi ai ponti di Schrödinger discreti.

1. Addestramento Off-Policy per Campionatori di Diffusione Discreti

La metodologia di base consiste nell'addestrare un kernel forward $\overrightarrow{p}_\theta$ per adattarsi a una distribuzione target minimizzando una divergenza tra la distribuzione della traiettoria del processo forward e la distribuzione della traiettoria del processo backward (di rumorizzazione). Le funzioni obiettivo studiate sono divergenze del secondo momento, specificamente:

• Bilanciamento delle Traiettorie (TB): Tratta la costante di normalizzazione come un parametro apprendibile.
• Log-Varianza (LV): Minimizza la varianza del log-rapporto delle probabilità forward e backward.

Per migliorare l'esplorazione e la copertura delle modalità, gli autori introducono tre strategie off-policy adattate dal RL e dalle Reti di Flusso Generativo (GFlowNets):

• Buffer di Replay: Memorizza gli stati campionati $X_N$ e li riutilizza campionando traiettorie backward $\overleftarrow{p}^{\otimes N}$ per l'addestramento.
• Buffer Ponderato per Importanza: Assegna probabilità di campionamento agli elementi del buffer basandosi sui pesi di importanza delle traiettorie che li hanno generati, dando priorità alle regioni ad alta probabilità che il campionatore corrente potrebbe non cogliere.
• Esplorazione MCMC: Affina i campioni dal buffer di replay utilizzando passi Markov Chain Monte Carlo (MCMC) (ad esempio, Swendsen-Wang o Metropolis-Hastings) prima che vengano utilizzati per l'addestramento. Questo sfrutta direttamente la funzione di energia per guidare il campionatore verso regioni ad alta densità senza richiedere rollout del modello.

2. Ponti di Schrödinger da Dati a Energia Discreti

Il lavoro generalizza il problema del campionamento alla ricerca di un ponte di Schrödinger (SB) tra due distribuzioni arbitrarie $p_0$ e $p_N$. A differenza dei precedenti lavori su SB discreti che assumono campioni per entrambi gli estremi, questo lavoro affronta il contesto da dati a energia in cui $p_N$ è definita solo da una funzione di energia.

• Il metodo impiega una procedura di Adattamento Proporzionale Iterativo (IPF).
• Alterna l'adattamento del kernel backward (tramite massima verosimiglianza su campioni da $p_0$) e del kernel forward (tramite perdita di log-varianza su campioni dalla $p_N$ definita dall'energia).
• Le tecniche off-policy (buffer e MCMC) sono integrate nel ciclo IPF per facilitare l'esplorazione della distribuzione definita dall'energia.

3. Applicazione agli Spazi Latenti

Gli autori applicano questi campionatori al campionamento "esternalizzato" del posterior. Addestrano campionatori di diffusione nello spazio latente discreto di un VQ-VAE (Vector Quantized Variational Autoencoder) pre-addestrato. L'obiettivo è campionare dal posterior $p(z|y) \propto p_{latent}(z)p(y|f(z))$ senza osservare campioni condizionati non distorti, condizionando efficacemente il modello generativo imparando a campionare nel suo spazio latente.

Contributi Chiave

1. Algoritmi Off-Policy per Campionatori Discreti: L'introduzione di tecniche RL off-policy (buffer di replay, ponderazione per importanza ed esplorazione MCMC) nell'addestramento dei campionatori di diffusione discreti. Il lavoro dimostra che queste tecniche migliorano significativamente le prestazioni rispetto alle baseline on-policy.
2. Ponti di Schrödinger da Dati a Energia: La prima generalizzazione dei campionatori di diffusione discreti al problema del ponte di Schrödinger da dati a energia, consentendo la costruzione di mappe di trasporto tra una distribuzione basata su campioni e una distribuzione definita da energia.
3. Campionamento del Posterior negli Spazi Latenti: La dimostrazione di campionatori di diffusione discreti per l'inferenza del posterior negli spazi latenti discreti dei modelli generativi di immagini, mostrando la capacità di condizionare su etichette di classe (ad esempio, cifre dispari vs pari) senza riaddestrare il modello generativo.

Risultati Sperimentali

Gli autori valutano i loro metodi su benchmark sintetici e applicazioni reali:

• Modelli di Ising e Potts: Su reticoli $16 \times 16$, i metodi off-policy (in particolare TB con Buffer + MCMC) superano costantemente i metodi on-policy e la baseline Masked Diffusion Neural Sampler (MDNS). Raggiungono distanze di Sinkhorn inferiori e errori di magnetizzazione/correlazione ridotti, specialmente a basse temperature (alto $\beta$) dove il collasso delle modalità è prevalente.
• Densità Sintetiche Discretizzate: Su distribuzioni multimodali complesse (40GMM e ManyWell) discretizzate tramite codici Gray, i metodi off-policy scoprono con successo tutte le modalità, mentre i metodi on-policy e MDNS spesso falliscono nel coprire l'intera distribuzione.
• Ponti di Schrödinger: Nei compiti di ponte tra distribuzioni semplici e multimodali, i metodi off-policy scalano efficacemente verso target complessi dove i campionatori on-policy collassano.
• Campionamento Esternalizzato: Le visualizzazioni di cifre MNIST generate dallo spazio latente mostrano che sia i campionatori on-policy che quelli off-policy possono generare con successo cifre delle classi desiderate (ad esempio, dispari vs pari), validando l'utilità dell'approccio per la generazione condizionata.

Significato e Affermazioni

Il lavoro afferma che l'integrazione di tecniche di addestramento off-policy rappresenta un passo critico in avanti per i campionatori di diffusione discreti, affrontando il problema del collasso delle modalità e migliorando l'efficienza del campionamento in spazi discreti ad alta dimensionalità e multimodali. Colmando il divario tra diffusione discreta e RL off-policy, gli autori dimostrano che questi campionatori possono gestire problemi complessi di fisica statistica (Ising/Potts) e compiti di ottimizzazione combinatoria in modo più efficace rispetto ai metodi precedenti.

Inoltre, il lavoro estende la portata teorica dei ponti di Schrödinger al contesto da dati a energia nei domini discreti, offrendo un nuovo strumento per il trasporto tra distribuzioni in cui una è definita da una funzione di energia. Infine, l'applicazione agli spazi latenti dimostra una via pratica per condizionare modelli generativi pre-addestrati senza richiedere l'accesso all'intera distribuzione dei dati o il riaddestramento del modello generativo stesso.

Gli autori riconoscono le limitazioni, inclusa la necessità di fissare un kernel di rumorizzazione e il costo computazionale dello stoccaggio delle traiettorie, ma suggeriscono che l'addestramento a scale temporali più grossolane (come supportato dalle estensioni della teoria a tempo continuo) potrebbe mitigare alcuni costi. Identificano direzioni future come l'apprendimento di funzioni di energia dai dati e la scalabilità di questi metodi verso modelli generativi di immagini più grandi.