MacTok: Robust Continuous Tokenization for Image Generation

1. Il Problema: Collasso Posteriore nei Tokenizzatori Continui

Il lavoro affronta una sfida fondamentale nell'ambito della generazione di immagini tramite modelli latenti: il collasso posteriore (posterior collapse) nei tokenizzatori continui basati su Variational Auto-Encoders (VAE) con regolarizzazione KL (KL-VAE).

• Contesto: I tokenizzatori continui mappano le immagini in spazi latenti lisci e strutturati, ideali per modelli di generazione come Diffusion e Flow. Tuttavia, sotto una forte compressione (pochi token), la regolarizzazione KL tende a spingere la distribuzione latente verso una prior isotropa (Gaussiana standard), ignorando le informazioni specifiche dell'immagine.
• Conseguenza: L'encoder fallisce nel codificare caratteristiche informative nello spazio latente compresso. Il decoder ricostruisce l'immagine basandosi principalmente sulla prior, portando a ricostruzioni di bassa qualità e capacità generativa limitata.
• Limiti delle soluzioni attuali: I metodi precedenti tentano di risolvere il problema tramite un'attenta sintonizzazione dei pesi della KL, ma ciò richiede un'iper-parametrizzazione delicata e non risolve il problema alla radice, lasciando lo spazio latente degenere.

2. Metodologia: MacTok

Gli autori propongono MacTok (Masked Augmenting 1D Continuous Tokenizer), un framework che combina tecniche di masking (mascheramento) e allineamento delle rappresentazioni per prevenire il collasso e apprendere rappresentazioni robuste e semanticamente strutturate.

Architettura

MacTok utilizza un'architettura basata su Vision Transformer (ViT) sia per l'encoder che per il decoder.

• Input: L'immagine viene suddivisa in patch (token immagine) e concatenata con un set di token latenti apprendibili.
• Output: L'encoder produce una rappresentazione latente continua $\hat{z}$ di dimensione ridotta (es. 64 o 128 token), che viene poi decodificata per ricostruire l'immagine.

Strategie Chiave

A. Mascheramento delle Immagini (Image Masking)
Per forzare il modello a mantenere il flusso di informazioni attraverso lo spazio latente, MacTok addestra il modello a ricostruire immagini complete partendo da input parziali. Vengono utilizzate due strategie di mascheramento con probabilità uguale:

1. Mascheramento Casuale: Un sottoinsieme casuale di patch dell'immagine viene sostituito da token di mascheramento prima dell'encoding. Questo costringe i token latenti a inferire informazioni globali dal contesto visibile.
2. Mascheramento Semantico (Guidato da DINO): A differenza del mascheramento casuale, questa strategia seleziona e maschera le regioni più informative dell'immagine. Utilizzando le feature di DINOv2, viene calcolata la similarità tra il token di classificazione e le patch; le patch con la similarità più alta (le più semanticamente rilevanti) vengono mascherate. Questo costringe il modello a ricostruire le parti critiche dell'immagine basandosi su evidenze visive incomplete, trasferendo implicitamente priors semantici dallo spazio delle immagini a quello latente.

B. Allineamento delle Rappresentazioni (Global & Local Alignment)
Per garantire che lo spazio latente sia semanticamente discriminativo, MacTok allinea le sue rappresentazioni latenti con le feature pre-addestrate di DINOv2:

• Allineamento Locale: Ogni token latente viene allineato alle feature delle patch locali di DINOv2.
• Allineamento Globale: Una rappresentazione latente globale (ottenuta tramite pooling medio) viene allineata al token di classificazione di DINOv2.
  Questo approccio, più leggero e flessibile rispetto a metodi precedenti, assicura coerenza strutturale e preservazione dei dettagli su diverse lunghezze di token.

C. Obiettivi di Addestramento
La funzione di perdita combina:

• Perdita di ricostruzione (pixel-wise).
• Perdita percettiva e avversaria (GAN).
• Regolarizzazione KL.
• Perdita di Allineamento delle Rappresentazioni ($L_{RA}$): Il nuovo termine che guida l'allineamento semantico.

3. Contributi Chiave

1. Identificazione e Soluzione del Collasso: Dimostrano che il collasso posteriore è una forma di sovraregolarizzazione e che il mantenimento del flusso di informazioni tramite la ricostruzione da input incompleti (masking) stabilizza l'apprendimento latente sotto forte compressione.
2. MacTok: Introduzione di un tokenizzatore continuo 1D che integra mascheramento (casuale e semantico) e allineamento globale/locale per prevenire il collasso e preservare le caratteristiche discriminative.
3. Efficienza e Qualità: Raggiungimento di risultati all'avanguardia (SOTA) utilizzando un numero drasticamente ridotto di token (64 o 128) rispetto ai metodi esistenti che spesso ne richiedono centinaia o migliaia.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul dataset ImageNet a risoluzioni 256×256 e 512×512, utilizzando generatori come SiT-XL e LightningDiT.

• Qualità Generativa (gFID):
  • A 256×256: MacTok ottiene un gFID di 1.44 (con 128 token e CFG), superando o eguagliando i migliori metodi esistenti.
  • A 512×512: Raggiunge un gFID di 1.52 (con 128 token e CFG), stabilendo un nuovo stato dell'arte per la generazione ad alta risoluzione con pochi token.
• Efficienza: MacTok riduce l'uso dei token fino a 64 volte rispetto ai metodi basati su VQ (Vector Quantization) o KL-VAE tradizionali, mantenendo o migliorando la fedeltà.
• Ricostruzione: Ottiene un rFID competitivo (0.43 a 256×256 con 128 token), dimostrando che le rappresentazioni latenti sono ricche di informazioni semantiche.
• Analisi dello Spazio Latente: Le visualizzazioni (UMAP) mostrano che MacTok apprende uno spazio latente ben strutturato e separabile, a differenza dei modelli KL-VAE collassati che producono distribuzioni isotrope e non informative. L'analisi di linear probing conferma che una maggiore accuratezza di probing si correla con una migliore qualità generativa.

5. Significato e Impatto

MacTok rappresenta un passo significativo verso l'efficienza nella generazione di immagini. Dimostra che:

• È possibile ottenere spazi latenti continui robusti e ad alta fedeltà senza ricorrere a discretizzazioni complesse (come VQ) o a lunghe sequenze di token.
• L'uso combinato di mascheramento semantico e allineamento delle rappresentazioni è una strategia efficace per mitigare il collasso posteriore, un problema storico nei VAE.
• La riduzione drastica del numero di token (da 1024+ a 64/128) promette di accelerare notevolmente i tempi di inferenza e addestramento per i modelli generativi successivi, rendendo l'architettura scalabile e pratica per applicazioni reali.

In sintesi, MacTok risolve il compromesso tra compressione e qualità, fornendo un tokenizzatore continuo che è allo stesso tempo efficiente, stabile e capace di generare immagini ad alta fedeltà.