Semantic-Aware Prefix Learning for Token-Efficient Image Generation

1. Il Problema

Gli attuali modelli di generazione di immagini latente (come VQ-GAN, VAE, e modelli autoregressivi) dipendono da tokenizzatori visivi per mappare immagini ad alta dimensionalità in rappresentazioni latenti compatte. Tuttavia, esistono due limitazioni fondamentali:

• Obiettivi di addestramento dominati dalla ricostruzione: La maggior parte dei tokenizzatori esistenti viene addestrata con l'obiettivo primario di minimizzare l'errore di ricostruzione pixel-per-pixel. Questo porta a rappresentazioni latenti che catturano bene i dettagli visivi a basso livello ma sono debolmente allineate con la semantica di alto livello (concetti, struttura globale).
• Segnali semantici come regolarizzazione ausiliaria: Gli approcci recenti tentano di migliorare l'allineamento semantico trattando i segnali semantici (es. etichette di classe) come vincoli di regolarizzazione secondari. Il paper sostiene che questo non è sufficiente: la semantica non deve essere solo "allineata", ma deve diventare un componente funzionale e indispensabile del processo di apprendimento della rappresentazione. Senza questo, lo spazio latente non è ottimale per la generazione, specialmente sotto budget di token ridotti.

2. Metodologia Proposta

Il lavoro introduce due componenti principali: SMAP (il tokenizzatore) e CARD (il generatore).

A. SMAP: Tokenizzatore Consapevole della Semantica (Semantic-Aware Prefix Tokenizer)

SMAP è un tokenizzatore 1D basato su query che forza l'informazione semantica a diventare un "prefisso" stabile e necessario per la ricostruzione.

• Iniezione Semantica Esplicita: A differenza dei metodi precedenti, SMAP inietta le condizioni semantiche (embedding di classe) direttamente nella sequenza di token sia nell'encoder che nel decoder. Le condizioni semantiche agiscono come elementi intermedi che interagiscono con i patch visivi e le query latenti tramite l'attenzione self-attention.
• Strategia di "Tail Token Dropping" (Rimozione dei Token di Coda): Questa è l'innovazione chiave. Durante l'addestramento, il tokenizzatore applica una strategia di truncation progressiva:
  • Ad ogni iterazione, viene campionata una lunghezza di prefisso $k$ (da 0 a $K$).
  • Vengono mantenuti solo i primi $k$ token latenti; i token di coda ($k+1$ a $K$) vengono rimossi o mascherati.
  • Caso Estremo ($k=0$): Il decoder deve ricostruire l'immagine utilizzando solo gli embedding semantici di classe e i token maschera, senza alcun token latente.
  • Effetto: Questa perturbazione forza il modello a caricare la responsabilità della struttura globale e dell'identità della categoria sui token di prefisso (inclusi i segnali semantici). I token latenti successivi vengono quindi obbligati a rifinire i dettagli specifici dell'istanza (texture, dettagli locali).
• Risultato: Si ottiene una sequenza di token ordinata per informazione, dove il prefisso contiene la semantica e la struttura globale, e la coda contiene i dettagli.

B. CARD: Generatore Ibrido Causale Autoregressivo-Diffusion

Per sfruttare lo spazio latente ordinato di SMAP, viene proposto CARD:

• Architettura Ibrida: Combina un modulo autoregressivo causale (basato su Transformer) che modella le dipendenze strutturali tra i token latenti (catturando la struttura globale) con un modulo di rifinitura continua basato su Diffusion/Flow Matching (MLP) che denoisa le variabili latenti per migliorare la fedeltà.
• Condivisione delle Condizioni Semantiche: CARD riutilizza direttamente gli embedding semantici appresi durante l'addestramento di SMAP come input di condizione per la generazione. Questo elimina la necessità di un codificatore di classe separato e garantisce una coerenza semantica perfetta tra tokenizzazione e generazione.

3. Contributi Chiave

1. Identificazione del Limite: Dimostrare che l'allineamento semantico "lasso" non è sufficiente; la semantica deve essere resa funzionalmente necessaria per la ricostruzione.
2. SMAP: Un tokenizzatore che implementa l'apprendimento di prefissi consapevoli della semantica tramite iniezione esplicita e truncation progressiva, creando rappresentazioni latenti ordinate e robuste.
3. CARD: Un generatore ibrido che sfrutta la struttura ordinata di SMAP, dimostrando che un tokenizzazione semanticamente fondata migliora le prestazioni di generazione anche con budget di token molto ridotti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-1K (256x256).

• Qualità di Ricostruzione: SMAP supera costantemente i tokenizzatori baselines (come TiTok e TA-TiTok) sia in setting discreti (VQ) che continui (KL, SoftVQ).
  • Con un budget di 128 token, SMAP (KL) raggiunge un rFID di 0.75 (vs 1.02 del baseline) e un IS di 308.9.
  • La strategia di tail token dropping porta a miglioramenti significativi, specialmente con budget di token limitati (es. 32 o 64 token).
• Generazione (Downstream):
  • Il modello SMAP(KL) + CARD (1.1B parametri) raggiunge un gFID di 1.85 (con CFG) e un IS di 325.1 utilizzando solo 128 token.
  • Questo risultato è competitivo con modelli autoregressivi e diffusion-based molto più grandi che utilizzano centinaia o migliaia di token (es. LlamaGen con 576 token o LDM-4 con 4096 token).
• Analisi Ablativa:
  • La rimozione dei token di coda dimostra che il modello riesce a ricostruire la struttura di base e l'identità della categoria usando solo il prefisso semantico.
  • La condivisione degli embedding semantici tra tokenizzatore e generatore migliora ulteriormente l'allineamento tra condizione e contenuto generato.

5. Significato e Impatto

Questo lavoro cambia il paradigma di progettazione dei tokenizzatori per la generazione di immagini:

• Efficienza: Dimostra che è possibile ottenere generazioni di alta qualità con un numero drasticamente ridotto di token (128 vs 256-4096), riducendo i costi computazionali per la fase di generazione.
• Ridefinizione del Ruolo della Semantica: Sposta la semantica da un ruolo di "regolarizzatore" a quello di "scheletro strutturale" obbligatorio. Questo rende lo spazio latente più facile da modellare per i generatori successivi.
• Generalità: L'approccio funziona sia per tokenizzazione discreta che continua, offrendo una soluzione unificata per migliorare sia la ricostruzione che la generazione.

In sintesi, SMAP e CARD dimostrano che un'apprendimento della rappresentazione "consapevole della semantica" e "ordinato" è la chiave per sbloccare l'efficienza e la qualità nella generazione di immagini latente.