Structured State-Space Regularization for Compact and Generation-Friendly Image Tokenization

1. Il Problema

I moderni modelli di visione artificiale, in particolare i modelli generativi (come i Diffusion Models) e i "world models", si basano sempre più su spazi latenti compatti invece che sullo spazio dei pixel grezzi. Un tokenizer di immagini ideale deve soddisfare due requisiti fondamentali e spesso conflittuali:

1. Compattezza e Fedeltà: Deve comprimere l'immagine in modo efficiente, preservando le informazioni essenziali per una ricostruzione fedele, evitando di sprecare capacità rappresentativa su dettagli non cruciali.
2. Compatibilità con la Generazione (Generation-Friendliness): Lo spazio latente deve essere "amico" dei modelli generativi. La ricerca recente ha dimostrato che spazi latenti con determinate proprietà strutturali (come la "diffusibilità") migliorano drasticamente la qualità della generazione. Tuttavia, spesso esiste un compromesso: un tokenizer ottimizzato per la massima fedeltà di ricostruzione tende a performare peggio nella generazione.

L'obiettivo del lavoro è sviluppare un metodo per allineare lo spazio latente a entrambi questi obiettivi, migliorando la generazione senza sacrificare significativamente la qualità della ricostruzione.

2. Metodologia: Regularizzazione dello Spazio di Stato Strutturato

Gli autori introducono una nuova tecnica di regolarizzazione che guida i tokenizer a imitare la dinamica degli stati nascosti dei Modelli di Spazio di Stato (SSM), come Mamba o HiPPO. L'idea centrale è trasferire la proprietà di "consapevolezza delle frequenze" (frequency awareness) degli SSM agli spazi latenti delle immagini.

A. Visione Generalizzata degli SSM

Il paper riformula gli SSM non solo come modelli sequenziali, ma come sistemi composti da due componenti chiave:

1. Proiezione su una Base ($c(\cdot)$): Proiettare il segnale di input su una base di funzioni (es. polinomi ortogonali o basi di Fourier).
2. Trasformazione dell'Input ($\theta(\cdot)$): Una trasformazione applicata all'input nel tempo.

Gli SSM tradizionali aggiornano lo stato nascosto basandosi sulla dinamica dei coefficienti di queste basi sotto una specifica trasformazione. Gli autori propongono di estendere questo concetto alle immagini 2D.

B. Scelta delle Componenti per le Immagini

Per adattare questo framework alle immagini, gli autori scelgono:

• Base ($c(\cdot)$): Basi ortogonali 2D, in particolare le Basi di Fourier (o polinomi ortogonali come Chebyshev, Legendre). Questo permette di rappresentare l'immagine come una combinazione lineare di componenti frequenziali.
• Trasformazione ($\theta(\cdot)$): Un processo progressivo di sfocatura Gaussiana. La sfocatura Gaussiana è matematicamente legata all'equazione del calore e rimuove progressivamente le alte frequenze, preservando le basse.

C. Derivazione della Regolarizzazione

Sfruttando le proprietà analitiche della sfocatura Gaussiana e dell'equazione del calore, gli autori derivano una regola di aggiornamento deterministica per i coefficienti della base $c(I_t)$:
$$\frac{d}{d\tau} c(I_\tau) = \mathbf{A} c(I_\tau)$$
Dove $\mathbf{A}$ è una matrice di transizione dello stato (derivata dall'operatore Laplaciano applicato alle basi). Dopo la discretizzazione, l'aggiornamento diventa:
$$c(I_t) = \bar{\mathbf{A}} c(I_{t-1})$$

La Loss di Regolarizzazione:
Invece di addestrare un SSM, si applica questa dinamica come vincolo al codificatore (encoder) $E$ di un tokenizer. Se $I_{t-1}$ e $I_t$ sono due immagini con livelli di sfocatura diversi, il modello viene regolarizzato affinché:
$$E(I_t) \approx \bar{\mathbf{A}} E(I_{t-1})$$
La funzione di perdita include una distanza tra lo stato latente previsto dalla dinamica e quello reale, più una perdita di ricostruzione per evitare il collasso del modello.

D. Implementazione Pratica

• Centratura Spaziale: Poiché la sfocatura riduce la norma dei coefficienti, gli autori introducono una funzione di centratura spaziale ($\mathcal{C}$) per stabilizzare l'addestramento e prevenire che i latenti delle immagini nitide abbiano norme eccessivamente elevate.
• Addestramento: Il tokenizer viene addestrato con una probabilità $\alpha$ a seguire questa dinamica di regolarizzazione, e con probabilità $1-\alpha$ a ricostruire l'immagine normalmente.

3. Contributi Chiave

1. Nuovo Regularizzatore Teorico: Proposta di un regolarizzatore motivato teoricamente per migliorare i tokenizer, basato sulla dinamica degli SSM.
2. Primo Framework Teorico SSM per Tokenization: Integrazione dei principi degli SSM (proiezione su base + trasformazione) nel contesto della tokenizzazione delle immagini 2D, offrendo una nuova prospettiva sulla modellazione visiva.
3. Efficacia Sperimentale: Dimostrazione che il metodo migliora le prestazioni generative su tokenizer esistenti (Flux, Cosmos) con una perdita minima nella fedeltà di ricostruzione.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-1K utilizzando tokenizer popolari come Flux (downsampling 8x) e Cosmos (downsampling 16x), valutati con un modello di diffusione latente (DiT-B/2).

• Qualità di Generazione: Il metodo proposto supera significativamente le baseline (inclusi EQVAE e FT-SE) in termini di FID generativo (gFID), Inception Score (IS) e recall. Ad esempio, su Flux, il gFID scende da 16.08 (baseline) a 13.33, con un IS che sale a 106.61.
• Qualità di Ricostruzione: La perdita di fedeltà è minima. Il PSNR e LPIPS rimangono molto vicini alle baseline, confermando che la regolarizzazione non distrugge la capacità di ricostruzione.
• Analisi Strutturale:
  • Dinamica Coerente: Le visualizzazioni mostrano che i canali latenti regolari evolvono in modo coerente (esponenziale) man mano che l'immagine viene sfocata, imitando la dinamica dei coefficienti di Fourier.
  • Separazione delle Frequenze: I canali latenti regolati codificano specifiche bande di frequenza. Svelando i canali dalle basse alle alte frequenze, si osserva una ricostruzione che passa da dettagli grossolani a dettagli fini, dimostrando una struttura spettrale ordinata assente nei tokenizer non regolari.

5. Significato e Impatto

Questo lavoro colma il divario tra la teoria degli State-Space Models (spesso usati per sequenze 1D) e la rappresentazione visiva.

• Nuova Prospettiva: Dimostra che gli SSM non sono solo architetture per l'elaborazione sequenziale, ma offrono principi di regolarizzazione potenti per strutturare spazi latenti.
• Efficienza Generativa: Fornisce un metodo per creare spazi latenti "generation-friendly" senza bisogno di riaddestrare interi modelli generativi, agendo direttamente sulla fase di tokenizzazione.
• Futuro: Apre la strada all'integrazione di principi di modellazione a stato spazio in futuri framework di apprendimento rappresentativo e generativo, suggerendo che la struttura spettrale è fondamentale per la generatività efficace.

In sintesi, il paper propone un approccio elegante che sfrutta la matematica della sfocatura e delle basi ortogonali per "insegnare" ai tokenizer a organizzare le informazioni in modo che siano sia compatte che facili da generare per i modelli di diffusione.