VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations

1. Enunciazione del Problema

I modelli di generazione di immagini autoregressivi (AR) (ad esempio, LlamaGen) hanno dimostrato prestazioni competitive ma affrontano barriere significative all'adozione produttiva rispetto ai modelli di diffusione. Il collo di bottiglia principale è la scalabilità in termini di risoluzione e rapporti d'aspetto:

• Esplosione dei Token: I tokenizer 2D convenzionali (ad esempio, VQGAN) generano sequenze di token che crescono linearmente con la risoluzione dell'immagine. Per un'immagine 1024×1024, ciò si traduce in migliaia di token (ad esempio, 4.096 token).
• Costo Computazionale Quadratico: I modelli AR si basano sull'attenzione self-attention, che ha una complessità $O(T^2)$ dove $T$ è la lunghezza della sequenza. All'aumentare della risoluzione, i FLOPs di inferenza crescono quadraticamente, rendendo la generazione ad alta risoluzione computazionalmente proibitiva.
• Rigidità: La maggior parte dei modelli AR è addestrata su risoluzioni fisse (ad esempio, 256×256) e fatica a generalizzare a rapporti d'aspetto arbitrari o risoluzioni superiori senza riaddestramento o pipeline complesse di super-risoluzione.
• Dominanza della Diffusione: I modelli di diffusione gestiscono naturalmente risoluzioni arbitrarie, lasciando i modelli AR indietro in casi d'uso flessibili e di livello produttivo.

2. Metodologia

Gli autori propongono una soluzione in due parti: VibeToken (un tokenizer innovativo) e VibeToken-Gen (un generatore AR).

A. VibeToken: Tokenizzatore 1D Agnostico alla Risoluzione

VibeToken è un tokenizer di immagini basato su Transformer 1D progettato per codificare immagini di qualsiasi risoluzione in un numero fisso e controllabile dall'utente di token (ad esempio, 32–256), disaccoppiando il conteggio dei token dalla risoluzione spaziale.

Le principali innovazioni architetturali includono:

1. Embedding Posizionali a Griglia Dinamica: Invece di embedding posizionali assoluti fissi, VibeToken utilizza una griglia appresa di embedding che viene ridimensionata in modo differenziale (tramite interpolazione bilineare/bicubica) per corrispondere alla griglia delle patch di input ($T_H \times T_W$). Ciò permette al modello di gestire rapporti d'aspetto arbitrari senza degradazione della qualità.
2. Embedding delle Patch Adattivi: Ispirato da FlexiViT, il modello supporta dimensioni di patch variabili ($k$) durante l'addestramento e l'inferenza. Un operatore di ridimensionamento dei pesi mappa i pesi di proiezione di base alla dimensione della patch corrente, permettendo al modello di adattarsi a diverse risoluzioni senza riparametrizzazione.
3. Risoluzione del Decoder Adattiva: Il decoder predice i pixel a una risoluzione alta fissa (4 volte la dimensione della patch) e utilizza un livello di downsampling convoluzionale 2D apprendibile per mappare alla risoluzione target. Ciò abilita la super-risoluzione nativa senza un modulo di upscaling separato.
4. Tokenizzazione a Lunghezza Dinamica: Il modello è addestrato con un campionamento uniforme delle lunghezze latenti $L \in [32, 256]$. L'encoder produce esattamente $L$ token e il decoder ne consuma esattamente $L$, eliminando la necessità di padding o vincoli di lunghezza fissa.
5. Quantizzazione: Utilizza la Quantizzazione Multi-Vettoriale (MVQ) con 8 codebook per raggiungere una dimensione del vocabolario di 32.768, bilanciando compressione e qualità di ricostruzione.

B. VibeToken-Gen: Generatore AR Efficiente

VibeToken-Gen è un modello AR condizionato per classe (basato sull'architettura LlamaGen) che opera sulle brevi sequenze di token agnostiche alla risoluzione prodotte da VibeToken.

• Budget Computazionale Costante: Poiché la lunghezza della sequenza di token $L$ è fissa (ad esempio, 64 token) indipendentemente dalla risoluzione di output, i FLOPs di inferenza rimangono costanti.
• Condizionamento: Il modello si condiziona sulla risoluzione target $(H, W)$ e sul rapporto d'aspetto insieme all'etichetta di classe per prevenire artefatti di distorsione geometrica.
• Addestramento: Addestrato su ImageNet-1k con risoluzioni miste (da 256×256 a 512×512) e rapporti d'aspetto, utilizzando un budget di token comparabile ai baseline a risoluzione fissa.

3. Contributi Chiave

1. Tokenizzazione Agnostica alla Risoluzione: Il primo tokenizer Transformer 1D che supporta nativamente risoluzioni e rapporti d'aspetto arbitrari mantenendo un budget di token breve e dinamico (32–256 token).
2. Efficienza Computazionale: Disaccoppia il costo di inferenza AR dalla risoluzione dell'immagine. VibeToken-Gen mantiene 179 GFLOPs costanti per qualsiasi risoluzione, mentre baseline come LlamaGen scalano a ~11 TFLOPs per immagini 1024×1024 (un guadagno di efficienza di 63,4×).
3. Super-Risoluzione Nativa: L'architettura supporta la generazione di immagini ad alta risoluzione da input a bassa risoluzione (o target arbitrari) senza moduli di upscaling separati.
4. Prestazioni all'Avanguardia: Raggiunge qualità di ricostruzione (rFID) e generazione (gFID) competitive rispetto sia ai tokenizer 2D che ai modelli di diffusione, nonostante utilizzi significativamente meno token.

4. Risultati Sperimentali

Il paper valuta VibeToken e VibeToken-Gen sui dataset ImageNet-1k e FFHQ su varie risoluzioni (da 256×256 a 1024×1024) e rapporti d'aspetto.

• Qualità di Ricostruzione:
  • VibeToken-LL raggiunge 0,40 rFID a 256×256 e 2,40 rFID a 1024×1024.
  • Supera altri tokenizer 1D (ad esempio, TiTok, One-D-Piece) nella generalizzazione della risoluzione e corrisponde o supera i tokenizer 2D (ad esempio, IBQ, MAGViT) in efficienza di compressione.
• Qualità di Generazione:
  • VibeToken-Gen (XXL) sintetizza immagini 1024×1024 utilizzando solo 64 token con un gFID di 3,94.
  • Confronto:
    • Vs. Diffusione (NiT): VibeToken-Gen è 2,35× più veloce (0,46s vs 1,08s) e raggiunge una qualità migliore (3,94 gFID vs 5,87 gFID) a 1024×1024.
    • Vs. Baseline AR (LlamaGen): LlamaGen richiede ~11 TFLOPs per 1024×1024, mentre VibeToken-Gen utilizza 179 GFLOPs costanti.
• Efficienza:
  • Il conteggio dei token è ridotto da ~4.096 (nei tokenizer 2D) a 64 token per immagini 1024×1024.
  • La velocità di inferenza è sub-secondo (0,46s) per la generazione ad alta risoluzione, rispetto a 32+ secondi per i modelli AR baseline a risoluzioni simili.

5. Significato

Questo lavoro colma il divario tra modelli autoregressivi e modelli di diffusione in termini di flessibilità ed efficienza.

• Fattibilità Produttiva: Consentendo ai modelli AR di gestire risoluzioni arbitrarie con calcolo costante, VibeToken rende la generazione AR praticabile per applicazioni reali dove le dimensioni di input/output variano (ad esempio, app mobili, formati di visualizzazione diversi).
• Scalabilità: Dimostra che i modelli AR possono scalare ad alte risoluzioni senza l'esplosione computazionale quadratica, sfidando l'idea che la diffusione sia l'unica strada verso l'IA generativa ad alta risoluzione.
• Direzione Futura: L'approccio apre la porta a modelli multimodali unificati (testo, immagine, video) che possono operare sotto un budget computazionale fisso e controllabile dall'utente, indipendentemente dalla complessità spaziale del contenuto generato.

In sintesi, VibeToken risolve il problema del "scaling della risoluzione" nella generazione AR spostando il carico di adattamento della risoluzione dal modello generativo a un tokenizer flessibile e agnostico alla risoluzione, risultando in un sistema più veloce, economico e flessibile rispetto alle attuali alternative all'avanguardia.