Riepilogo Tecnico: Autoencoder di Rappresentazione Laminare per la Diffusione Efficiente (FlatDINO)

Definizione del Problema

I recenti progressi nella generazione di immagini hanno dimostrato che i modelli di diffusione possono operare direttamente su feature di patch apprese tramite apprendimento auto-supervisionato (SSL) (ad esempio, da DINOv2) anziché su latenti nello spazio dei pixel. Questo approccio, esemplificato da RAE (Zheng et al., 2025), evita il compromesso tra ricostruzione e generazione inerente ai VAE addestrati sui pixel, sfruttando la struttura semanticamente significativa già codificata nelle feature SSL.

Tuttavia, rimane un collo di bottiglia computazionale significativo. Gli encoder SSL standard come DINOv2 producono griglie dense bidimensionali di embedding di patch ad alta dimensionalità (ad esempio, 256 patch da 768 dimensioni per un'immagine 256×256). Questa rappresentazione è comparabile in termini di dimensioni al segnale originale, offrendo alcun vantaggio computazionale rispetto alla diffusione nello spazio dei pixel. La griglia densa contiene una significativa ridondanza spaziale poiché le patch vicine condividono un contenuto semantico sostanziale, rendendo la diffusione su queste feature inutilmente costosa in termini di FLOP e memoria.

Metodologia

Gli autori propongono FlatDINO, un autoencoder variazionale (VAE) progettato per comprimere le dense embedding di patch 2D di DINOv2 in una sequenza compatta monodimensionale di token continui.

1. Architettura

• Encoder: Un Vision Transformer (ViT-B) che prende in input le 256 patch embedding da un DINOv2 ViT-B/14 congelato (con registri) come input. L'encoder precompone $T$ token di registro apprendibili alla sequenza di patch. Dopo l'elaborazione, solo questi $T$ token di registro vengono mantenuti come rappresentazione latente, scartando la struttura spaziale 2D originale.
• Decoder: Un ViT-L che utilizza i token di registro apprendibili come query per prestare attenzione alla sequenza latente compressa e ricostruire le 256 patch embedding DINOv2 originali.
• Obiettivo di Addestramento: Il modello è addestrato utilizzando un obiettivo VAE standard con una perdita di ricostruzione MSE (Mean Squared Error) e una penalità di divergenza KL. Il peso $\beta$ della KL è normalizzato per la dimensionalità totale del latente ($T \times d$) per garantire una pressione di regolarizzazione coerente tra le diverse configurazioni.
• Generazione Latente: Un modello generativo (specificamente un DiT-XL con flow matching) è addestrato sullo spazio latente compresso di FlatDINO. Il modello impara un campo di velocità per trasportare i campioni da un prior gaussiano standard alla distribuzione dei dati lungo percorsi rettilinei.

2. Scelte Chiave di Design

• Compressione 1D: Ispirato a TiTok (Yu et al., 2024a), il metodo comprime la griglia 2D in una sequenza 1D. Gli autori hanno scoperto che 32 token forniscono un compromesso favorevole tra fedeltà di ricostruzione e lunghezza della sequenza.
• Riduzione della Dimensione delle Feature: Mentre il numero di token è ridotto da 256 a 32 (una riduzione di 8×), anche la dimensione delle feature per ogni token è compressa (ad esempio, da 768 a 128) per ridurre ulteriormente la dimensionalità totale.
• Flow Matching: Il modello generativo utilizza il flow matching invece della diffusione standard, operando direttamente sulla sequenza latente 1D con embedding posizionali appresi.

Contributi Chiave

1. FlatDINO: Il primo metodo per comprimere le feature di patch auto-supervisionate in una rappresentazione latente monodimensionale, scartando efficacemente la struttura spaziale 2D originale.
2. Guadagni di Efficienza: Dimostra che i modelli di diffusione addestrati sui latenti FlatDINO raggiungono una qualità di generazione competitiva riducendo significativamente i costi computazionali. Nello specifico, la riduzione di 8× nella lunghezza della sequenza porta a una riduzione di 8× dei FLOP per passaggio forward e fino a una riduzione di 4,5× dei FLOP per step di addestramento rispetto alla diffusione sulle feature DINOv2 non compresse.
3. Analisi della Tokenizzazione: Attraverso studi di ablazione dei token, gli autori mostrano che il modello impara campi recettivi localizzati, simili a "blob", per la maggior parte dei token, partizionando efficacemente l'immagine spazialmente. Identificano una "transizione di fase" in cui la riduzione dei token al di sotto di una certa soglia (ad esempio, da 32 a 16) costringe il modello ad adottare strategie di codifica meno efficienti (strisce orizzontali), portando a un calo netto della qualità.

Risultati

Gli esperimenti sono stati condotti su ImageNet 256×256 utilizzando un'architettura DiT-XL.

• Qualità della Generazione: Un DiT-XL addestrato su latenti FlatDINO (32×128) raggiunge un gFID di 1,85 con classifier-free guidance (CFG) e 3,34 senza CFG. Questo è competitivo con RAE (che raggiunge un gFID di 1,41 con CFG su feature non compresse) e altri modelli di diffusione latente allo stato dell'arte, nonostante la significativa compressione.
• Efficienza Computazionale:
  • Inferenza: FlatDINO richiede circa 8× meno FLOP per passaggio forward rispetto all'operazione sulle patch DINOv2 complete.
  • Addestramento: Grazie alla ridotta lunghezza della sequenza e all'inclusione dell'overhead dell'encoder, FlatDINO richiede 4,5× meno FLOP per step di addestramento per i modelli DiT-XL rispetto al baseline RAE.
• Ricostruzione: La configurazione 32×128 raggiunge un rFID di 0,77 (rispetto a 0,62 per il baseline RAE non compresso), indicando un'alta fedeltà nella ricostruzione dello spazio delle feature.
• Ablazione dei Token: Lo studio rivela che 32 token sono sufficienti per mantenere campi recettivi localizzati. Ridurre a 16 token causa uno spostamento qualitativo verso una codifica a strisce orizzontali e un calo significativo della qualità della ricostruzione.

Significato e Rivendicazioni

Il paper posiziona FlatDINO come un passo verso la sintesi di immagini efficiente e semanticamente fondata. Gli autori affermano che, sfruttando la ridondanza spaziale inerente alle griglie di patch 2D, è possibile ottenere una riduzione di 8× nel numero di token senza sacrificare i benefici di convergenza derivanti dall'operare in uno spazio latente semanticamente organizzato.

Il lavoro suggerisce che la diffusione su rappresentazioni semantiche compresse sia un approccio pratico per contesti con risorse limitate. Gli autori riconoscono che i loro risultati attuali sono preliminari e che il modello non è completamente convergente (addestrato per 600 epoche rispetto a 800 per RAE). Attribuiscono il gap rimanente nella qualità della generazione a una durata dell'addestramento insufficiente e alla necessità di ricette di diffusione specificamente adattate ai latenti semantici compressi. Si propone come lavoro futuro l'investigazione di programmi di addestramento estesi e autoencoder SSL ottimizzati per colmare questo divario.