Riepilogo Tecnico: MergeTok

Definizione del Problema

Gli attuali tokenizer visuali per la generazione di immagini sono suddivisi in due distinte famiglie, ciascuna delle quali soffre di limitazioni strutturali complementari:

1. Tokenizer basati su VAE continui: Sebbene offrano un'alta fedeltà di ricostruzione, si affidano a spazi latenti privi di un'organizzazione intrinseca. Sotto una forte compressione, i fattori semantici si intrecciano all'interno di rappresentazioni latenti dense, risultando in codici debolmente disgiunti che ostacolano il controllo semantico.
2. Tokenizer basati su VQ discreti: Sebbene consentano la generazione autoregressiva, affrontano colli di bottiglia nell'ottimizzazione. La natura non differenziabile della quantizzazione porta a aggiornamenti dei gradienti sparsi e irregolari, causando il fatto che molte voci del codebook ricevano una supervisione insufficiente. Ciò aumenta il rischio di collasso del codebook e riduce l'efficienza della rappresentazione.

Queste limitazioni creano una tensione persistente tra fedeltà di ricostruzione, struttura semantica, efficienza e stabilità dell'addestramento.

Metodologia: Framework MergeTok

MergeTok propone un tokenizer a doppio ramo unificato che ottimizza congiuntamente la tokenizzazione continua (VAE) e quella discreta (VQ) all'interno di un'unica architettura encoder-decoder. L'innovazione principale è l'uso del Token Merging (ToMe) come ponte semantico tra i due rami.

Panoramica dell'Architettura

Dato un input immagine $X$, un encoder CNN condiviso $E_c$ produce una mappa di feature $Z$. Questa mappa viene elaborata da due rami paralleli:

1. Ramo VAE (Continuo):

   • Token Merging: Un encoder basato sull'attenzione dotato di moduli ToMe comprime la sequenza di token da $L$ token a $K$ token ($K < L$) fondendo i token simili. Questo processo genera una rappresentazione condensata $Z^{(vae)}_K$ e una mappa sorgente $S$ che registra l'assegnazione dei token originali ai gruppi fusi.
   • Allineamento Semantico: I token fusi vengono allineati con un modello insegnante di tipo DINO congelato (anch'esso dotato di ToMe) utilizzando una perdita di distanza del coseno ($L_{align}$). Questo regolarizza lo spazio latente continuo verso rappresentazioni semantiche e disgiunte.
   • Ricostruzione: Un decoder ibrido utilizza sia i token condensati che la mappa sorgente $S$ per "disfare il merge" (unmerge) dei gruppi e ricostruire l'immagine nello spazio dei pixel.

2. Ramo VQ (Discreto):

   • Quantizzazione Consapevole dei Gruppi: Questo ramo bypassa l'operazione di ToMe durante il passaggio in avanti (forward pass) per preservare la lunghezza completa della sequenza ($L$ token). Tuttavia, utilizza la mappa sorgente $S$ generata dal ramo VAE come prior strutturale.
   • Regolarizzazione: La mappa sorgente partiziona i $L$ token in $K$ gruppi semantici. Due perdite consapevole dei gruppi vengono applicate alle assegnazioni del codebook:
     • Diversità Intra-gruppo ($L_{div}$): Massimizza l'entropia all'interno di ogni gruppo per prevenire il collasso verso un singolo codice.
     • Esclusività Inter-gruppo ($L_{cons}$): Minimizza la sovrapposizione tra i gruppi per garantire che diversi gruppi semantici utilizzino subset di codici distinti.
   • Ricostruzione: Un decoder VQ standard ricostruisce l'immagine dai token quantizzati.

Strategia di Addestramento

• Obiettivo Unificato: Il modello è addestrato end-to-end con una perdita totale che combina regolarizzatori di ricostruzione, avversari e di merge-awareness:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Campionamento Consapevole della Granularità: Durante l'addestramento, il rapporto di merge (e quindi il numero di token mantenuti $K$) viene campionato dinamicamente da una distribuzione discreta centrata attorno a un iperparametro $K^*$. Questo espone l'encoder a diverse ampiezze di bottleneck, incoraggiando la robustezza attraverso varie granularità semantiche.
• Inferenza: All'inferenza, il tokenizer produce sempre una sequenza fissa di 256 token. I generatori a valle (sia autoregressivi che di diffusione) operano su questa sequenza fissa e sono agnostici rispetto alla struttura interna di merging.

Contributi Chiave

1. Tokenizer a Doppio Ramo Unificato: MergeTok è la prima architettura a ottimizzare congiuntamente i latenti continui e discreti sotto un unico obiettivo, utilizzando il token merging come interfaccia semantica per rafforzare la stabilità dell'addestramento VAE e i vantaggi strutturali di VQ.
2. Vincoli di Addestramento Consapevoli del Merge: Gli autori introducono due nuovi obiettivi derivati dal processo di merging:
   • Allineamento dei token fusi: Migliora la struttura semantica nello spazio latente continuo.
   • Quantizzazione consapevole dei gruppi: Stabilizza l'addestramento VQ e migliora l'utilizzo del codebook tramite diversità intra-gruppo ed esclusività inter-gruppo.
3. Campionamento del Rapporto di Merge Consapevole della Granularità: Una strategia di addestramento che campiona diverse granularità di token, consentendo al modello di apprendere una tokenizzazione unificata che è robusta attraverso diversi tassi di compressione.

Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-256 (risoluzione 256×256) con un budget fisso di 256 token.

• Fedeltà di Ricostruzione:
  • Discreto (VQ): MergeTok-BL ha raggiunto un rFID di 0.78, superando i forti baseline come GigaTok-BL (0.81) e MergeVQ-GR (1.12). Con un discriminatore DINO congelato, MergeTok-BL* ha raggiunto un rFID di 0.48.
  • Continuo (VAE): MergeTok-BL ha raggiunto un rFID di 0.47, superando marginalmente RAE (0.49) e battendo significativamente SoftVQ-L (0.61) e DC-AE-f32 (0.69).
• Qualità della Rappresentazione:
  • MergeTok-BL ha ottenuto un'accuratezza di linear probing del 78.3%, la più alta tra i tokenizer discreti elencati, superando UniTok (70.8%) e MergeVQ-GR (77.9%). Ciò indica un'organizzazione semantica superiore.
• Generazione a Valle:
  • Autoregressiva (Ramo VQ): Accoppiato con LlamaGen-XXL, MergeTok-BL ha raggiunto un gFID di 1.93 (senza CFG) e 2.13 (con CFG), superando GigaTok-BL (2.03).
  • Diffusione/Flow (Ramo VAE): Accoppiato con SiT-XL/1, MergeTok-BL ha raggiunto un gFID di 1.79 (senza CFG) e 1.29 (con CFG), eguagliando le migliori voci della tabella.

Significato e Rivendicazioni

Il paper sostiene che MergeTok dimostri che una singola architettura può dotare i tokenizer visuali sia di una robusta organizzazione semantica che di una discretezza favorevole ai generatori.

• Unificazione: Risolve il compromesso tra l'alta fedeltà di ricostruzione dei VAE e la struttura discreta richiesta per la generazione autoregressiva.
• Ponte Semantico: Sfruttando il token merging come prior strutturale, il modello trasferisce con successo le informazioni di raggruppamento semantico dal ramo continuo per stabilizzare il ramo discreto, e viceversa.
• Compatibilità: Il tokenizer produce sequenze di 256 token che sono compatibili sia con la generazione autoregressiva (es. LlamaGen) che con la generazione di diffusione/flow (es. DiT, SiT) senza richiedere modifiche ai generatori stessi.

Gli autori posizionano MergeTok non solo come uno strumento di ricostruzione, ma come un'interfaccia unificata che supporta simultaneamente la generazione autoregressiva discreta, la generazione di diffusione continua e l'apprendimento di rappresentazioni semantiche strutturate.