Riepilogo Tecnico: Oltre l'Ultimo Strato: Fusione di Rappresentazioni Multi-Strato per la Tokenizzazione Visiva

1. Enunciato del Problema

I tokenizzatori visivi all'avanguardia attuali, in particolare quelli che sfruttano Modelli Fondamentali Visivi (VFMs) preaddestrati e congelati come DINOv2 come backbones, soffrono di un critico collo di bottiglia informativo a singolo strato. I metodi esistenti estraggono esclusivamente caratteristiche dall'ultimo strato dell'encoder per definire lo spazio latente. Sebbene queste caratteristiche dell'ultimo strato siano ottimizzate per la semantica di alto livello, scartano sistematicamente i dettagli visivi di basso livello (texture, bordi, gradienti di colore) attraverso un'astrazione semantica progressiva.

Un'analisi recente indica che le informazioni di basso livello sopravvivono nell'ultimo strato solo come residui attenuati tramite connessioni residue passive, diventando sempre più lossive con ogni strato successivo. Di conseguenza, i tokenizzatori a singolo strato scartano le ricche informazioni gerarchiche distribuite tra gli strati intermedi, limitando sia la fedeltà della ricostruzione che la qualità della generazione a valle. Inoltre, rimane una questione aperta se l'aumento della "ricchezza di rappresentazione" dei tokenizzatori visivi segua una legge di scalabilità prevedibile analoga alla scalabilità della dimensione del vocabolario nei tokenizzatori NLP.

2. Metodologia: DRoRAE

Gli autori propongono DRoRAE (Depth-Routed Representation AutoEncoder), un modulo di fusione leggero progettato per aggregare caratteristiche da tutti gli strati dell'encoder mantenendo la compatibilità con decoder congelati e modelli generativi a valle.

Architettura Principale

DRoRAE inserisce un modulo di fusione addestrabile tra il backbone VFM congelato e il decoder. Il processo coinvolge tre meccanismi chiave:

1. Esperti per Strato: Ciascuno degli $L$ strati dell'encoder è associato a una rete esperta dedicata (un MLP a due strati). Questi esperti proiettano caratteristiche eterogenee degli strati su una scala comune, normalizzata dalla normalizzazione a strato del backbone per gestire le disparità di varianza.
2. Instradamento Vincolato dall'Energia: A differenza dei classici Mixture-of-Experts (MoE) che utilizzano la normalizzazione softmax, DRoRAE impiega un meccanismo di instradamento vincolato dall'energia. Una rete instradante produce pesi per token su tutti gli strati. Crucialmente, questa formulazione consente pesi negativi, permettendo all'instradante di sopprimere attivamente contributi dannosi degli strati (denoising) invece di selezionare semplicemente i vincitori. L'output è normalizzato dalla norma $\ell_2$ del vettore dei pesi per vincolare l'energia dell'output.
3. Correzione Incrementale: Invece di sostituire l'output originale dell'ultimo strato ($z_{base}$) con la rappresentazione fusa ($z_{fuse}$), DRoRAE formula l'output come una correzione incrementale:
   $$z_{final} = LN_{bb}(z_{base} + \beta \cdot (z_{fuse} - z_{base}))$$
   Qui, $\beta$ controlla la forza della fusione. Questa formulazione residua assicura che il latente fuso rimanga ancorato alla distribuzione originale, prevenendo derive arbitrarie che potrebbero compromettere la compatibilità con la generazione a valle.

Strategia di Addestramento Disaccoppiato in Tre Fasi

Per affrontare la sfida di spostare la distribuzione latente lontano da ciò che un decoder congelato si aspetta, gli autori propongono una strategia di addestramento progressiva:

• Fase 1 (Addestramento del Decoder): Addestrare il decoder sulla baseline RAE a singolo strato standard per stabilire una forte "capacità di decodifica" e definire la distribuzione target.
• Fase 2 (Addestramento del Modulo di Fusione): Congelare il backbone e il decoder della Fase 1. Addestrare solo il modulo di fusione. Il decoder congelato agisce come un vincolo distributivo implicito, costringendo il modulo di fusione ad apprendere informazioni complementari che il decoder può già invertire, prevenendo la deriva distributiva.
• Fase 3 (Fine-tuning del Decoder): Congelare il modulo di fusione ottimizzato e sbloccare il decoder. Effettuare il fine-tuning del decoder per sfruttare appieno lo spazio latente arricchito, massimizzando la qualità della ricostruzione senza danneggiare la generazione.

3. Contributi Chiave

1. Identificazione del Collo di Bottiglia: Il lavoro identifica che gli autoencoder di rappresentazione esistenti scartano informazioni visive gerarchiche affidandosi esclusivamente all'ultimo strato dell'encoder, dove i dettagli di basso livello sono meramente residui attenuati.
2. Architettura DRoRAE: Un modulo di fusione instradato per profondità (~29M parametri) che aggrega adattivamente tutti gli strati dell'encoder tramite instradamento vincolato dall'energia (che consente pesi negativi) e correzione incrementale.
3. Strategia di Addestramento: Un protocollo di addestramento disaccoppiato in tre fasi che garantisce la compatibilità con la generazione apprendendo la fusione sotto il vincolo di un decoder congelato prima di co-adattare il decoder.
4. Scoperta della Legge di Scalabilità: Il lavoro stabilisce una legge di scalabilità log-lineare ($R^2=0.86$) tra la capacità di fusione (definita dalla capacità degli esperti e dal numero di strati fusi) e la qualità della ricostruzione. Questo identifica la "ricchezza di rappresentazione" come una nuova dimensione prevedibilmente scalabile per i tokenizzatori visivi, analoga alla dimensione del vocabolario nell'NLP.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-256 (256x256) per la ricostruzione e la generazione condizionata alla classe, e su CC12M-LLaVA-Next per la sintesi testo-immagine.

• Ricostruzione: DRoRAE migliora significativamente le metriche di ricostruzione rispetto alla baseline RAE. L'rFID scende da 0.57 a 0.29, il PSNR aumenta da 18.8 a 24.32 dB, e l'LPIPS migliora da 0.256 a 0.134.
• Generazione Condizionata alla Classe: Utilizzando un generatore DiTDH-XL con AutoGuidance, DRoRAE migliora il FID di generazione (gFID) da 1.74 a 1.65 mantenendo o migliorando leggermente l'Inception Score (IS) e Precision/Recall.
• Sintesi Testo-Immagine: I guadagni si trasferiscono ai compiti testo-immagine, con DRoRAE che ottiene punteggi GenEval comparabili o leggermente migliori (0.60 contro 0.56 per RAE) attraverso le dimensioni oggetto, conteggio e spaziale.
• Studi di Ablazione:
  • Instradamento Vincolato dall'Energia: Supera l'instradamento softmax permettendo la soppressione attiva degli strati rumorosi.
  • Correzione Incrementale: Essenziale per la compatibilità con la generazione; senza di essa, la ricostruzione migliora leggermente, ma la perdita di addestramento della diffusione a valle aumenta significativamente, indicando un disallineamento distributivo.
  • Strategia di Addestramento: L'addestramento congiunto senza il vincolo della Fase 2 non riesce a raggiungere la stessa qualità di generazione, confermando la necessità dell'approccio disaccoppiato.

5. Significato e Affermazioni

Il lavoro afferma che DRoRAE dimostra che la fusione multi-strato è un aggiornamento pratico ed efficace per i tokenizzatori basati su rappresentazione, recuperando dettagli ad alta frequenza persi nell'estrazione a singolo strato.

Un significato primario di questo lavoro è la scoperta di una legge di scalabilità prevedibile per i tokenizzatori visivi. Dimostrando che l'aumento del "budget di rappresentazione" (tramite più strati fusi o capacità degli esperti più grandi) produce miglioramenti log-lineari nella qualità della ricostruzione, gli autori sostengono che la ricchezza di rappresentazione è una dimensione scalabile paragonabile alla dimensione del vocabolario nell'NLP. Ciò suggerisce che i futuri miglioramenti nella tokenizzazione visiva possano essere sistematicamente raccolti scalando la capacità di fusione, piuttosto che affidarsi esclusivamente a cambiamenti architetturali o dataset più grandi.

Il lavoro conclude che, sebbene gli esperimenti attuali utilizzino DINOv2-B, la metodologia è generalizzabile a encoder più grandi e potenzialmente alla tokenizzazione video, posizionando la fusione multi-strato come una direzione fondamentale per il potenziamento dei sistemi di generazione visiva.