Riepilogo Tecnico: Apprendimento di Priori Autoregressivi Discreti con Flusso di Gradiente di Wasserstein

1. Enunciato del Problema

I tokenizzatori di immagini discreti sono tipicamente addestrati utilizzando una ricetta in due fasi: prima, il tokenizzatore viene addestrato per ricostruire le immagini (concentrandosi sulla verosimiglianza condizionale), e successivamente un priori autoregressivo (AR) viene addestrato sulle sequenze di token risultanti e congelate. Questa disaccoppiamento crea una discrepanza fondamentale: il tokenizzatore impara a preservare le informazioni dell'immagine ma rimane ignaro dei bias induttivi del priori AR. Di conseguenza, i token appresi possono essere eccellenti per la ricostruzione ma difficili da prevedere sequenzialmente da sinistra a destra per un modello AR.

I metodi esistenti tentano di mitigare ciò modificando il decoder (ad es. Tail Dropout, Decoder Diffusivi) per incoraggiare una struttura sequenziale. Tuttavia, questi approcci si basano ancora su una supervisione centrata sul decoder e non sfruttano direttamente la difficoltà di modellazione del priori AR per guidare l'apprendimento del tokenizzatore. Inoltre, le tecniche standard di corrispondenza a livello di distribuzione (ad es. WAE-MMD, InfoVAE) progettate per spazi latenti continui non sono applicabili direttamente ai tokenizzatori discreti rigidi a causa della mancanza di percorsi differenziabili attraverso la quantizzazione.

2. Metodologia

Quadro Teorico: Coerenza Variazionale Tripartita (TVC)

Gli autori introducono la Coerenza Variazionale Tripartita (TVC) per analizzare l'apprendimento di variabili latenti. La TVC scompone l'obiettivo generativo in tre condizioni di coerenza derivate dalla regola di Bayes:

1. Coerenza della Verosimiglianza Condizionale: Allineamento tra il decoder del modello e la posterior ancorata ai dati.
2. Coerenza del Priori: Allineamento tra il priori del modello e la posterior aggregata del codificatore.
3. Coerenza della Posterior: Allineamento tra la posterior del modello e la posterior ancorata ai dati.

Gli autori sostengono che l'addestramento standard in due fasi soddisfa la Coerenza della Verosimiglianza nella prima fase e la Coerenza del Priori nella seconda fase, ma non fornisce un segnale al tokenizzatore per garantire che la posterior aggregata sia intrinsecamente "AR-friendly". Al contrario, l'addestramento in stile ELBO impone una coerenza del priori istanza per istanza, che spesso porta al collasso della posterior.

Soluzione Proposta: WGF-DPD e wAR-Tok

Per colmare questa lacuna, il documento propone la Divergenza del Priori Distribuzionale (DPD), un termine che corrisponde alla distribuzione aggregata dei token al priori AR target senza vincolare le singole istanze. Per ottimizzare ciò per token discreti, viene impiegato il Flusso di Gradiente di Wasserstein (WGF).

Il meccanismo centrale, WGF-DPD, opera come segue:

1. Modello Proxy: Un modello AR ausiliario leggero ($q_\psi$) viene addestrato per tracciare la distribuzione aggregata corrente dei token prodotta dal tokenizzatore ($q_\phi$).
2. Aggiornamento Basato sul Punteggio: La direzione di aggiornamento per il tokenizzatore è derivata dal gradiente di Wasserstein della divergenza KL tra la posterior aggregata e il priori target. Ciò risulta in un aggiornamento del gradiente a livello di particella:
   $$\nabla_z \mathcal{L}_{DPD} \approx -\nabla_z \log p_\theta(x|z) + (\log q_\psi(z) - \log p_\theta(z))$$
   • Il primo termine è il gradiente standard di ricostruzione.
   • Il secondo termine è il punteggio di corrispondenza del priori, che rappresenta la differenza tra le log-probabilità del modello proxy e del priori AR target.
3. Ottimizzazione Efficiente: Crucialmente, questo aggiornamento richiede solo passate in avanti attraverso i modelli AR proxy e target. Non richiede la retropropagazione attraverso nessuno dei modelli AR. Il gradiente viene propagato ai parametri del codificatore tramite lo Stimatore Straight-Through (STE) attraverso il passaggio di quantizzazione non differenziabile.

Il tokenizzatore risultante, denominato wAR-Tok, viene addestrato congiuntamente con l'obiettivo di ricostruzione e il segnale WGF-DPD.

3. Contributi Chiave

• Quadro TVC: Una scomposizione teorica dell'apprendimento di variabili latenti che identifica esplicitamente la "Coerenza del Priori" come obiettivo distinto, spiegando perché l'addestramento in due fasi lascia il tokenizzatore ignaro della difficoltà di modellazione del priori.
• Algoritmo WGF-DPD: Un metodo di ottimizzazione innovativo per tokenizzatori discreti che allinea la distribuzione aggregata dei token a un priori AR utilizzando il flusso di gradiente di Wasserstein. Evita l'intrattabilità della retropropagazione attraverso quantizzatori discreti o l'instabilità della corrispondenza avversariale.
• wAR-Tok: Un'implementazione pratica che integra il segnale WGF-DPD nell'addestramento del tokenizzatore, richiedendo solo passate in avanti attraverso modelli AR ausiliari per generare segnali di allineamento lato codificatore.

4. Risultati Sperimentali

Il metodo è stato valutato su CIFAR-10 e ImageNet 256×256, confrontandolo con baseline tra cui Tokenizzazione 1D, Tail Dropout, Decoder Diffusivi e Tied Timesteps.

• CIFAR-10:
  • A parità di qualità di ricostruzione (rFID $\approx$ 5.0), wAR-Tok ha raggiunto un FID di generazione (gFID) di 14.50 e una perdita di generazione AR (gLoss) di 0.42, superando significativamente la baseline più forte (ImageFolder/One-D-Piece: gFID 17.60, gLoss 0.78).
  • A parità di iterazioni di addestramento, wAR-Tok ha raggiunto un gFID di 11.30 e un gLoss di 0.34, dimostrando una convergenza più rapida e una qualità superiore.
• ImageNet 256×256:
  • Utilizzando la stessa architettura di codificatore-decoder e priori AR della baseline, wAR-Tok ha ridotto il gFID da 6.32 a 5.42 e la perdita AR di valutazione da 7.89 a 7.07, con un compromesso minore nella qualità di ricostruzione (rFID 2.05 $\to$ 2.35).
• Dinamiche di Addestramento: Le curve di addestramento AR mostrano che wAR-Tok mantiene una perdita di entropia crociata inferiore e una maggiore accuratezza nella previsione dei token durante tutto l'addestramento rispetto alle baseline, indicando che i token appresi sono intrinsecamente più facili da modellare per il priori AR.

5. Significato e Affermazioni

Il documento afferma che wAR-Tok risolve efficacemente la discrepanza tra tokenizzatori discreti e priori autoregressivi introducendo un segnale di corrispondenza del priori a livello di distribuzione durante l'addestramento del tokenizzatore.

• Efficienza: Il metodo è computazionalmente efficiente, aggiungendo solo due passate in avanti (proxy e target) per passo di addestramento ed evitando la retropropagazione costosa attraverso i modelli AR.
• Generalità: Sebbene focalizzato sui priori AR, il quadro TVC e l'approccio WGF-DPD sono presentati come indipendenti dal priori, applicabili a qualsiasi priori discreto con log-probabilità trattabili.
• Modestia: Gli autori dichiarano esplicitamente di non rivendicare prestazioni assolute State-of-the-Art (SOTA), poiché hanno deliberatamente escluso tecniche ausiliarie comuni nei tokenizzatori recenti (ad es. allineamento DINO, distillazione, post-addestramento con decoder diffusivi) per isolare l'effetto di WGF-DPD. I risultati dimostrano l'efficacia dell'obiettivo di allineamento proposto in un contesto controllato.

Il lavoro suggerisce che l'ottimizzazione esplicita per la "prevedibilità" dei token da parte del modello generativo a valle, piuttosto che solo per la loro ricostruibilità, porta a sistemi generativi discreti più efficienti e di qualità superiore.