End-to-End Autoregressive Image Generation with 1D Semantic Tokenizer

1. Enunciato del Problema

La generazione di immagini autoregressiva (AR) si è affermata come una valida alternativa ai modelli di diffusione, basandosi su token visivi discreti per comprimere le immagini in rappresentazioni latenti. Tuttavia, gli approcci esistenti affrontano due sfide principali:

• Disallineamento Spaziale: La maggior parte dei tokenizzatori per immagini utilizza strutture a griglia 2D (preservando il layout spaziale). Ciò induce dipendenze bidirezionali tra i token, in conflitto con la fattorizzazione unidirezionale richiesta dalla modellazione AR standard in ordine raster.
• Paradigmi di Addestramento Subottimali: I metodi attuali impiegano tipicamente una strategia di addestramento in due fasi: prima si addestra un tokenizzatore per la ricostruzione (poi congelato), quindi si addestra un modello AR sui token risultanti. Ciò impedisce al tokenizzatore di adattarsi alle esigenze specifiche del compito generativo.
• Divario nella Funzione di Perdita: La perdita di Predizione del Prossimo Token (NTP) utilizzata nei modelli AR opera nello spazio dei token discreti e non correla direttamente con la qualità finale della generazione nello spazio dei pixel. Ottimizzare esclusivamente per la NTP può portare al "collasso dello spazio latente", dove il tokenizzatore utilizza solo pochi token per minimizzare la perdita di previsione, degradando la fedeltà dell'immagine.

2. Metodologia: EOSTok

Gli autori propongono EOSTok (End-to-end One-dimensional Semantic Tokenizer), un framework che ottimizza congiuntamente un tokenizzatore visivo 1D e un modello generativo AR in una pipeline a fase singola.

A. Tokenizzatore Vision Transformer 1D

• Architettura: Basata su una struttura di autoencoder ViT (simile a TiTok).
• Processo: I patch 2D dell'immagine vengono appiattiti e concatenati con token di query apprendibili. Un encoder ViT causale elabora questa sequenza, producendo in output solo i token latenti 1D (scartando gli embedding dei patch spaziali).
• Quantizzazione: Utilizza IBQ (Improved Bottleneck Quantization) per mappare i vettori latenti continui su un dizionario discreto, abilitando la modellazione AR.
• Decoder: Un decoder ViT 1D ricostruisce le immagini concatenando i token quantizzati con token di maschera apprendibili.

B. Addestramento Congiunto End-to-End

Invece di congelare il tokenizzatore, l'encoder, il decoder e il modello AR vengono addestrati simultaneamente. La funzione di perdita totale combina:

1. Perdita di Ricostruzione ($L_{recon}$): Perdite VQ-VAE standard (L1/L2, Perceptual, GAN) per garantire la fedeltà dell'immagine.
2. Perdita di Predizione del Prossimo Token ($L_{NTP}$): Perdita di entropia incrociata standard per il modello AR.
3. Perdita di Ricostruzione della Predizione Autoregressiva (APR) ($L_{APR}$):
   • L'Innovazione: Per colmare il divario tra la previsione dei token e la qualità dei pixel, i token previsti dal modello AR (in modalità teacher-forcing) vengono decodificati nello spazio dei pixel dal decoder del tokenizzatore.
   • Meccanismo: La perdita confronta questi pixel decodificati con l'immagine ground-truth.
   • Effetto: Ciò fornisce una supervisione generativa end-to-end al tokenizzatore, impedendogli di "hackerare" la perdita NTP collassando lo spazio latente (utilizzando troppi pochi token).

C. Allineamento della Rappresentazione Semantica (VFM)

Per iniettare informazioni semantiche globali senza imporre priorità spaziali 2D (che danneggerebbero le prestazioni AR), gli autori propongono un Allineamento Implicito:

• Allineamento Diretto (Rifiutato): Allineare direttamente i token latenti 1D alle caratteristiche VFM 2D impone una struttura 2D allo spazio 1D, degradando le prestazioni.
• Sostituzione Diretta (Rifiutata): Sostituire i patch dell'immagine con caratteristiche VFM come input.
• Allineamento Implicito (Proposto): Allinea gli embedding dei patch nascosti dell'encoder 1D (prima dell'estrazione dei token di query) con le caratteristiche VFM (ad esempio, da DINOv2). Questo distilla la conoscenza semantica nello spazio latente sequenziale senza imporre vincoli spaziali 2D espliciti.
• Allineamento del Decoder: Anche il decoder è allineato con le caratteristiche VFM per trattare la ricostruzione come un compito di generazione condizionata, favorendo la convergenza.

3. Contributi Chiave

1. Framework di Addestramento End-to-End: Una pipeline innovativa a fase singola che ottimizza congiuntamente ricostruzione e generazione, permettendo al tokenizzatore di apprendere rappresentazioni specificamente adattate alla modellazione AR.
2. Perdita APR: Una funzione di perdita critica che decodifica le previsioni AR nello spazio dei pixel, assicurando che il tokenizzatore apprenda uno spazio latente sia ricostruttivo che prevedibile generativamente, prevenendo il collasso latente.
3. Allineamento Semantico Implicito: Una strategia per sfruttare i Modelli Fondamentali Visivi (VFM) per tokenizzatori 1D che migliora la qualità della generazione senza reintrodurre i bias spaziali 2D che ostacolano la modellazione AR.
4. Prestazioni all'Avanguardia: Raggiungimento di risultati SOTA sulla generazione ImageNet 256×256 senza guida del classificatore.

4. Risultati Sperimentali

Il modello è stato valutato sul benchmark ImageNet-1K 256×256.

• Qualità della Generazione:
  • EOSTok-H (644M parametri): Raggiunge un gFID di 1.48 (senza guida) e un IS di 239.5. Questo è un nuovo stato dell'arte per i modelli AR senza guida, superando i modelli basati su diffusione come REPA-XL/2 (gFID 1.42 con guida) e altre baseline AR.
  • EOSTok-L (312M parametri): Raggiunge un gFID di 1.74 senza guida.
• Qualità della Ricostruzione:
  • EOSTok-L raggiunge un rFID di 0.73, dimostrando capacità di ricostruzione ad alta fedeltà.
• Studi di Ablazione:
  • Perdita APR: Senza APR, l'utilizzo del dizionario crolla drasticamente (da ~99% a ~51%) e il gFID peggiora significativamente (8.01 contro 3.32).
  • Allineamento Semantico: L'allineamento implicito migliora il gFID da 12.27 (baseline) a 3.32. L'allineamento diretto peggiora le prestazioni, confermando che le priorità 2D sono dannose per la modellazione AR 1D.
  • Ordinamento della Sequenza: I modelli addestrati con supervisione end-to-end mostrano prestazioni significativamente migliori sulle sequenze di token originali rispetto a sequenze invertite o casuali, dimostrando che lo spazio latente è ottimizzato per la previsione causale.
  • Scalabilità: Le prestazioni migliorano costantemente con la dimensione del modello (da S a H) e la dimensione del dizionario, con modelli più grandi che mitigano il compromesso ricostruzione-generazione.

5. Significato

Questo paper sposta fondamentalmente il paradigma della generazione di immagini autoregressiva:

• Da Due Fasi a Una Fase: Dimostra che l'addestramento congiunto del tokenizzatore e del generatore produce risultati superiori rispetto all'approccio tradizionale con tokenizzatore congelato.
• Fattibilità della Tokenizzazione 1D: Dimostra che i tokenizzatori 1D, se opportunamente addestrati con guida semantica e feedback end-to-end, possono superare i tokenizzatori basati su griglia 2D per la generazione AR eliminando i conflitti di dipendenza spaziale.
• Colmare il Divario della Perdita: L'introduzione della perdita APR risolve il problema critico per cui l'ottimizzazione a livello di token non si traduce in qualità a livello di pixel, un collo di bottiglia comune nella modellazione generativa discreta.
• Scalabilità: Il metodo scala efficacemente, raggiungendo risultati SOTA con un modello da 644M parametri, suggerendo una strada promettente per scalare i modelli visivi AR a dimensioni maggiori.