Riepilogo Tecnico: Campionamento del Posteriore Ispirato a Langevin Discreto (∆LPS)

Enunciato del Problema

Il documento affronta la sfida di risolvere problemi inversi negli spazi a stati discreti. I problemi inversi consistono nel recuperare un segnale sconosciuto $x$ da misurazioni indirette, sparse e rumorose $y$ (modellate come $y = A(x) + n$). Sebbene i modelli di diffusione continui siano diventati lo standard per risolvere tali problemi tramite campionamento del posteriore, i loro corrispettivi discreti rimangono poco esplorati.

I metodi esistenti per il campionamento del posteriore discreto soffrono di tre limitazioni principali:

1. Relassamenti Continui: Metodi come G2D2 rilassano le variabili categoriali in spazi continui per utilizzare la guida del gradiente, alterando la geometria dell'inferenza e rischiando che i token discreti proiettati non decodifichino in modo coerente.
2. Aggiornamenti Sequenziali: Metodi come SGDD si basano su aggiornamenti di tipo Gibbs, che sono sequenziali e lenti a mescolarsi negli spazi ad alta dimensionalità.
3. Specializzazione: Metodi come APS sono strettamente accoppiati a specifici processi di corruzione (ad esempio, diffusione mascherata) e paradigmi di addestramento, limitandone la generalità.

L'obiettivo è sviluppare un campionatore del posteriore che rimanga completamente discreto, aggiorni tutte le dimensioni dei token in parallelo e sia agnostico rispetto al specifico paradigma di addestramento del prior di diffusione discreto.

Metodologia: ∆LPS

Gli autori propongono ∆LPS (Campionatore del Posteriore Ispirato a Langevin Discreto), un algoritmo privo di addestramento che utilizza informazioni sul gradiente per guidare le transizioni discrete senza uscire dallo spazio degli stati discreti.

Formulazione Principale

Il problema è formulato nello spazio latente discreto $z_0 \in \{1, \dots, K\}^L$ di un modello encoder-decoder. L'obiettivo è campionare dal posteriore $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, dove $p_\theta$ è un prior di diffusione discreto preaddestrato. Il campionatore approssima la transizione del posteriore $p(z_s|z_t, y)$ inferendo prima un latente pulito $z_0$ coerente con la misurazione e il prior, quindi ri-rumorendolo al passo temporale successivo.

Proposta di Langevin Discreto

L'innovazione chiave consiste nell'adattare la dinamica di Langevin agli spazi discreti. Negli ambienti continui, un aggiornamento di Langevin sposta un campione $z_0$ nella direzione del gradiente del potenziale del log-posteriore $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Poiché $z_0$ è discreto, gli autori non aggiungono direttamente un vettore gradiente. Invece, utilizzano la distribuzione di proposta gaussiana dalla dinamica di Langevin continua per definire le probabilità di transizione sui candidati discreti validi $z'_0$:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

Espandendo l'esponente e normalizzando sull'insieme finito di token validi, la proposta si fattorizza attraverso le coordinate dei token. Ciò permette al campionatore di calcolare una distribuzione categorica per ogni posizione del token $\ell$ indipendentemente:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

Questo meccanismo abilita aggiornamenti completamente paralleli attraverso tutte le $L$ dimensioni dei token, evitando il collo di bottiglia sequenziale dei campionatori di Gibbs.

Generalizzazione e Precondizionamento

• Agnosticismo del Prior: Il metodo richiede solo probabilità predittive sugli stati categorici, rendendolo compatibile con prior di diffusione mascherata (MDLM) e diffusione a stato uniforme (Duo).
• Problemi Inversi Ciechi: Per operatori forward non differenziabili o sconosciuti, il gradiente della verosimiglianza è sostituito da un surrogato contrastivo appreso (utilizzando encoder per dati e misurazioni) per fornire un segnale di guida fluido.
• Precondizionamento Adam: Per gestire gradienti scarsamente scalati negli spazi latenti ad alta dimensionalità, il gradiente grezzo è sostituito da una direzione precondizionata con Adam. Questo riscalda adattivamente gli aggiornamenti, smorzando i grandi gradienti rumorosi preservando quelli utilmente coerenti, migliorando il mescolamento e la convergenza.

Contributi Chiave

1. Campionamento del Posteriore Completamente Discreto: ∆LPS è un campionatore privo di addestramento che mantiene ogni iterazione nello spazio degli stati discreti, evitando l'ottimizzazione di token soft e gli errori di proiezione.
2. Aggiornamenti Paralleli Efficienti: Fattorizzando la proposta attraverso le coordinate, il metodo aggiorna tutte le dimensioni dei token in parallelo, migliorando significativamente l'efficienza rispetto ai campionatori sequenziali di tipo Gibbs.
3. Design Agnostico del Prior: Il campionatore non è vincolato a un specifico processo di corruzione (mascherata vs uniforme) e funziona con qualsiasi prior di diffusione discreto preaddestrato che fornisca probabilità predittive.

Risultati Sperimentali

Gli autori valutano ∆LPS su un insieme diversificato di problemi inversi su MNIST, CIFAR-10 e FFHQ, coprendo impostazioni lineari, non lineari e cieche.

• Ripristino Immagini (FFHQ): Su compiti tra cui inpainting, deblurring gaussiano, super-risoluzione, HDR e deblurring da movimento, ∆LPS ottiene prestazioni comparabili o superiori rispetto alle baseline continue (ad esempio, DPS, PSLD) e supera le baseline discrete esistenti (SGDD, G2D2, APS) nella qualità percettiva (LPIPS).
• Compiti Discreti (MNIST/CIFAR): Il metodo dimostra robustezza attraverso vari livelli di difficoltà (Facile, Medio, Difficile). Coerentemente eguaglia o supera SGDD e G2D2 in accuratezza dei token e PSNR, anche con budget computazionali significativamente inferiori (meno valutazioni di funzioni) rispetto alla variante "completa" di SGDD.
• Problemi Inversi Ciechi: In un compito di mappatura spaziale (inferenza di piante da traiettorie di camminata dove l'operatore forward è sconosciuto), ∆LPS supera le baseline continue (CoGuide, varianti DPS) nel punteggio F1 e nell'Intersezione sull'Unione (IoU).

Significato e Affermazioni

Il documento afferma che i campionatori del posteriore completamente discreti e informati dal gradiente offrono una via scalabile e generale per risolvere problemi inversi su rappresentazioni discrete.

• Scalabilità: Il meccanismo di aggiornamento parallelo affronta le limitazioni di velocità di mescolamento dei campionatori discreti negli spazi ad alta dimensionalità.
• Generalità: Il metodo disaccoppia l'algoritmo di inferenza dal paradigma di addestramento specifico del prior, permettendo di sfruttare il crescente ecosistema di modelli di diffusione discreta (testo, codice, immagini, audio).
• Praticità: Evitando relassamenti continui, ∆LPS garantisce che i campioni finali siano token discreti validi, prevenendo il "gap di proiezione" in cui un'approssimazione continua non riesce a decodificare in un segnale discreto coerente.

Gli autori concludono che, man mano che i modelli generativi discreti si scalano, metodi di inferenza principiali come ∆LPS diventeranno sempre più critici per applicazioni nella progettazione biomolecolare, musica simbolica, completamento del codice e altri domini dove i dati sono intrinsecamente discreti.