Riassunto Tecnico: Sticky Jump Diffusions (SJD)

Definizione del Problema

La modellazione generativa basata sulla diffusione si è estesa con successo dai dati continui alle sequenze discrete (es. testo), tuttavia gli approcci esistenti affrontano limitazioni fondamentali nell'unificare la dinamica discreta e quella continua. Esistono due famiglie dominanti:

1. Diffusione Mascherata (Masked Diffusion): Opera su spazi di token discreti, producendo token validi per costruzione. Tuttavia, tratta ogni posizione come una mappatura verso un singolo token assorbente, cancellando le informazioni graduate sulla prossimità della corruzione rispetto ai token candidati.
2. Diffusione Continua (Continuous Diffusion): Opera sugli embedding dei token utilizzando lo standard score-based learning. Il suo stato terminale risiede al di fuori del supporto discreto, richiedendo un passaggio di proiezione esterno (es. un decoder o una quantizzazione) per recuperare i token discreti, il quale non fa parte delle dinamiche apprese.
3. Diffusione Ibrida (Hybrid Diffusion): Accoppia una catena mascherata con un continuo latente. Sebbene efficace, questi modelli spesso assemblano i campionatori inversi manualmente, trattando i tassi di impegno (commitment rates), le destinazioni e i drift come scelte di design separate piuttosto che derivarli da un processo unificato.

La sfida centrale è definire un processo di Markov a tempo continuo che colleghi naturalmente questi regimi, dove le dinamiche inverse (impegno verso i token discreti e diffusione continua) siano matematicamente derivate dal processo forward tramite il bilancio del flusso, anziché essere progettate euristicamente.

Metodologia: Sticky Jump Diffusions (SJD)

Gli autori introducono le Sticky Jump Diffusions (SJD), una classe di processi di Markov a tempo continuo su $\mathbb{R}^d$ con un insieme finito o numerabile di ancore (anchors) $A$ (identificate con gli embedding dei token).

1. Dinamica Forward (Unsticking)

Nel tempo forward, i dati partono da un'ancora $a \in A$. Il processo evolve come segue:

• Fase di Sticking: La particella attende all'ancora $a$ per una durata casuale governata da un tasso di hazard $\vec{\lambda}_t$.
• Unsticking: Al verificarsi dell'evento di hazard, la particella "si stacca" (unsticks) verso un nuovo stato $Y$ campionato da un kernel $R_t(\cdot | a)$ supportato sulla regione continua $X_A = \mathbb{R}^d \setminus A$.
• Diffusione: Una volta in $X_A$, la traiettoria diffonde secondo una SDE (Equazione Differenziale Stocastica) standard: $d\vec{X}_t = \vec{b}_t(\vec{X}_t)dt + g_t d\vec{W}_t$.
• Scissione della Massa (Mass Split): La distribuzione marginale $p_t$ si scinde in una densità assolutamente continua su $X_A$ e masse atomiche sulle ancore $A$.

Innovazione Chiave (Blending Cross-Posizione): A differenza degli ibridi standard dove il kernel di unsticking dipende solo dall'ancora locale, le SJD permettono al kernel di dipendere dall'intera sequenza tramite una matrice di blending $W$. La media della distribuzione di unsticking alla posizione $i$ è una miscela ponderata degli embedding dei vicini: $\mu_i(\vec{X}_0) = (W E(\vec{X}_0))_i$. Questo codifica bias induttivi strutturali (es. località spaziale o grafi di vincolo) direttamente nel processo di corruzione.

2. Dinamica Inversa (Sticking)

La reversione temporale di una SJD produce un processo generativo con due meccanismi accoppiati:

• SDE guidata dallo Score: Nella regione continua $X_A$, il processo segue una SDE inversa guidata dalla funzione di score $\nabla \log p_t(x)$.
• Kernel di Jump Sticky: Il processo salta dalla regione continua nuovamente verso le ancore. Il tasso di hazard inverso (quando saltare) e la destinazione del salto (dove saltare) sono unicamente determinati dall'equazione di bilancio del flusso (bilancio dettagliato tra i flussi forward e reverse).
  • Fondamentalmente, il tasso di impegno e la destinazione non sono schedule progettate a mano, ma sono fissati dalla legge forward.

3. Apprendimento: Denoising Hazard Matching (DHM)

Per rendere addestinabile il processo inverso, gli autori propongono il Denoising Hazard Matching (DHM), l'analogo dell'hazard del denoising score matching.

• Classificatore Singolo: Un singolo classificatore di denoising $P_\theta(a | y, t)$ viene addestrato tramite cross-entropy standard su coppie corrotte dalla SJD $(\vec{X}_0, \vec{X}_t)$.
• Ripesatura Analitica: Il hazard inverso e le probabilità di allocazione vengono recuperati analiticamente da $P_\theta$ senza un secondo head di rete.
  • L'hazard inverso $\vec{\lambda}^\star_t(y)$ è stimato come l'aspettativa condizionale di un target noto nel forward (il target DHM) sotto la corruzione SJD.
  • La probabilità di allocazione $\pi^\star_t(a|y)$ è derivata pesando i logit del classificatore con il target DHM.
• Addestramento Simulation-Free: Sotto condizioni standard (survival in forma chiusa, chiusura gaussiana della convoluzione), l'addestramento non richiede la simulazione del percorso SJD. La densità di miscela necessaria per il target viene valutata tramite una quadratura in forma chiusa, rendendo la pipeline di addestramento efficiente e priva di simulazione.

Contributi Chiave

1. La Classe di Processo SJD: Un framework unificato per le jump-diffusion dove la marginale forward si scinde in una densità continua e atomi discreti. Questo generalizza le classiche diffusion a frontiera appiccicosa (sticky-boundary diffusions) ad insiemi di ancore ad alta dimensionalità.
2. Dinamica Inversa da un Singolo Classificatore: Gli autori dimostrano che sia lo score che le intensità inverse per ogni ancora possono essere stimati da un singolo classificatore di denoising tramite DHM. Ciò elimina la necessità di separati head di hazard e si basa su una ripesatura Rao–Blackwellizzata dell'output del classificatore.
3. Visione Unificante delle Famiglie di Diffusione: Le SJD recuperano i modelli esistenti come casi limite specifici:
   • Diffusione Mascherata: Emerge quando il kernel di unsticking collassa a un singolo punto assorbente. Il "vuoto informativo" nella diffusione mascherata è spiegato strutturalmente: il kernel di unsticking è indipendente dall'ancora, quindi lo stato corrotto non trasporta evidenza del token sorgente.
   • Diffusione Continua: Emerge quando il tasso di jump è zero. La necessità di una proiezione terminale è spiegata dall'assenza di atomi nella marginale forward, che non genera jump inversi.
   • Diffusione Ibrida: Emerge quando il kernel di unsticking corrisponde al kernel di perturbazione forward. Il tasso di impegno, la destinazione e il drift seguono naturalmente il bilancio del flusso invece di essere design separati.
4. Il Jump Kernel come Spazio di Design: Variando la matrice di blending $W$, gli autori introducono un nuovo spazio di design. Il blending cross-posizione corrompe ogni posizione verso una miscela dei valori puliti dei suoi vicini, trasformando le strutture di dipendenza (spaziali, sequenziali o basate su vincoli) in un bias induttivo della corruzione stessa.

Risultati Sperimentali

Gli autori validano empiricamente l'asse del blending cross-posizione su tre dataset dove le dipendenze strutturali sono note:

• CIFAR-10 (Generazione di Immagini): Utilizzando un kernel di sfocatura dello spazio dei valori per codificare la località spaziale, la SJD con un kernel non locale ($\sigma_B = 2.5$) raggiunge un FID inferiore (14.57) rispetto all'ibrido a kernel identità (CADD, 15.88) e ad altri baseline (MDLM, Bit Diffusion) a 256 NFE.
• Text8 (Generazione di Testo): Utilizzando una convoluzione gaussiana 1D per codificare la correlazione sequenziale locale, la SJD con $\sigma_W \ge 1.0$ supera il conteggio di parole valide di CANDI e MDLM, particolarmente per lunghezze di sequenza elevate ($\ge 6$) e attraverso vari budget NFE.
• Sudoku (Soddisfacimento di Vincoli): Utilizzando un kernel di grafo di vincoli per codificare le dipendenze tra righe/colonne/box, la SJD supera significativamente l'ibrido a kernel identità (CADD). Mentre CADD mostra prestazioni bimodali (spesso collassando vicino al caso), la SJD risolve costantemente le griglie, raggiungendo il 95.65% di accuratezza e convergendo 4 volte più velocemente in termini di step di addestramento.

Significato e Rivendicazioni

Il paper sostiene che le Sticky Jump Diffusions forniscano un account a livello di processo che unifica le diffusione mascherate, continue e ibride.

• Unificazione Teorica: Dimostra che le coordinate discrete e continue nella modellazione generativa non sono entità separate da accoppiare a mano, ma un singolo stato invertito dove i tassi di impegno e le destinazioni sono fissati dal bilancio del flusso.
• Espansione dello Spazio di Design: Trattando il kernel di unsticking come uno spazio di design esplicito, le SJD permettono ai ricercatori di codificare i prior strutturali (come la località spaziale o i vincoli logici) direttamente nelle dinamiche di corruzione forward. Ciò migliora la qualità della generazione e la coerenza logica senza alterare l'architettura fondamentale del campionatore inverso.
• Efficienza Pratica: Il framework si addestra e campiona allo stesso costo degli ibridi esistenti, poiché la struttura specifica della SJD entra solo attraverso la ripesatura analitica e le quadrature in forma chiusa, non richiedendo parametri di rete aggiuntivi o overhead di simulazione.

Gli autori concludono che questa visione unificante rende più facile estendere i modelli di diffusione: nuove corruzioni, schedule e regole di impegno possono essere specificati come dinamiche forward e ottenuti in reverse gratuitamente tramite il teorema di reversione temporale.