Coevolving Representations in Joint Image-Feature Diffusion

1. Il Problema

I modelli di diffusione per la sintesi di immagini hanno raggiunto risultati eccezionali, operando spesso nello spazio latente compresso di un VAE (Variational Autoencoder) o direttamente nello spazio dei pixel. Recenti approcci hanno cercato di migliorare questi modelli integrando rappresentazioni semantiche ad alto livello (estratte da encoder visivi pre-addestrati come DINOv2 o MAE) insieme alle latenti a basso livello.

Tuttavia, gli approcci esistenti presentano una limitazione fondamentale: lo spazio delle rappresentazioni semantiche è fisso e predefinito (ad esempio, ottenuto tramite una proiezione PCA statica) prima dell'inizio dell'addestramento del modello generativo. Questo spazio non si adatta all'obiettivo di generazione. Di conseguenza, le rappresentazioni semantiche potrebbero non essere ottimali per guidare il processo di sintesi dell'immagine, poiché non sono state ottimizzate congiuntamente per il compito specifico di generazione.

La domanda centrale sollevata dagli autori è: Lo spazio delle rappresentazioni che guida la diffusione dovrebbe rimanere fisso, o dovrebbe adattarsi e co-evolvere insieme al modello generativo?

2. Metodologia: CoReDi

Gli autori propongono CoReDi (Coevolving Representation Diffusion), un framework in cui la proiezione delle caratteristiche visive pre-addestrate viene appresa congiuntamente al modello di diffusione.

Componenti Chiave

Il framework prevede l'uso di un encoder visivo congelato ($VE$) che estrae caratteristiche $z_0$ dall'immagine. Invece di usare una proiezione fissa, CoReDi introduce una proiezione adattiva e apprendibile $g_\phi(\cdot)$ (implementata come un semplice strato lineare) che mappa le caratteristiche in uno spazio a dimensionalità ridotta $\tilde{z}_0$.

Per garantire una co-evoluzione stabile ed evitare soluzioni degeneri (come il collasso delle caratteristiche), gli autori identificano e implementano tre ingredienti fondamentali:

1. Stop-Gradient sul Target:
   Durante l'ottimizzazione congiunta, se sia l'input che il target della perdita di diffusione delle rappresentazioni sono apprendibili, il modello tende a minimizzare la perdita in modo banale (degenerazione). Per evitarlo, viene applicato un operatore stop-gradient sul target pulito $\tilde{z}_0$ nella funzione di perdita delle rappresentazioni. Questo permette al modello di imparare a denoisare le rappresentazioni senza che lo spazio delle rappresentazioni stesse si adatti trivialmente per ridurre la perdita.

2. Normalizzazione del Batch (Batch Normalization - BN):
   I modelli di diffusione sono sensibili alla scala degli input. Dopo la proiezione apprendibile, viene applicata una normalizzazione del batch (senza parametri affini apprendibili) per stabilizzare la scala delle caratteristiche e preservare il programma di rumore previsto. Questo agisce anche come regolarizzatore implicito contro il collasso dei campioni.

3. Regolarizzazione Esplicita contro il Collasso delle Caratteristiche:
   Anche con BN e stop-gradient, si osserva il fenomeno del "feature collapse", dove canali diversi codificano informazioni ridondanti o non variano significativamente. Per contrastare ciò, vengono proposti tre tipi di regolarizzazione:

   • Feature Variance Regularization: Penalizza i canali con varianza inferiore a una soglia, incoraggiando ogni canale a mantenere un'attività significativa.
   • Orthogonality Regularization: Penalizza la deviazione della matrice di proiezione dall'identità, incoraggiando direzioni di proiezione ortogonali e distinte.
   • Covariance Regularization: Penalizza le correlazioni tra i canali (elementi fuori diagonale della matrice di covarianza), incoraggiando la decorrelazione.

Obiettivo di Addestramento

L'obiettivo totale combina la perdita di flusso per le immagini ($\mathcal{L}_{image}$), la perdita di flusso per le rappresentazioni ($\mathcal{L}_{rep}$) e la perdita di regolarizzazione ($\mathcal{L}_{reg}$):
$$\mathcal{L}(\theta, \phi) = \mathcal{L}_{image} + \lambda_z \mathcal{L}_{rep} + \lambda_{reg} \mathcal{L}_{reg}$$
Dove $\theta$ sono i parametri del modello di diffusione e $\phi$ quelli della proiezione.

Estensione allo Spazio dei Pixel

CoReDi non è limitato agli spazi latenti VAE. Gli autori dimostrano che il metodo può essere esteso alla diffusione nello spazio dei pixel (basandosi sull'architettura DeCo), rimuovendo il collo di bottiglia della ricostituzione del VAE e permettendo alle rappresentazioni semantiche di co-evolvere direttamente con i pixel grezzi.

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet (256x256) con diversi encoder visivi (DINOv2, MOCOv3, SigLIPv2, MAE) e architetture di modello (SiT, DeCo).

• Miglioramento della Qualità e Velocità di Convergenza:
  • Nello spazio latente, CoReDi supera i modelli basati su proiezioni fisse (ReDi, REPA). Ad esempio, con l'architettura SiT-XL/2, CoReDi raggiunge un FID di 3.3 in soli 2M di iterazioni, mentre ReDi richiede 4M iterazioni per lo stesso risultato e SiT base ne richiede 7M per un FID di 8.3.
  • Con guida senza classificatore (CFG), CoReDi ottiene un FID di 1.58 in 400 epoche, superando REPA (1.80) e ReDi (1.72) che richiedono il doppio delle epoche.
• Generalizzazione:
  I miglioramenti sono consistenti indipendentemente dall'encoder visivo utilizzato, dimostrando che l'apprendimento adattivo della proiezione è vantaggioso per diverse fonti di caratteristiche semantiche.
• Spazio dei Pixel:
  Nella diffusione pixel-based, CoReDi accelera la convergenza di un fattore 2x rispetto al baseline DeCo, raggiungendo un FID di 21.5 a 200K iterazioni contro 31.3 di DeCo.
• Analisi delle Strutture Spaziali:
  L'analisi mostra che le rappresentazioni co-evolventi sviluppano una struttura spaziale più forte e organizzata durante l'addestramento (misurata tramite metriche come LDS, CDS, RMSC) rispetto alle proiezioni PCA fisse. Questo suggerisce che l'ottimizzazione congiunta favorisce la creazione di rappresentazioni semantiche più utili per la sintesi di immagini.
• Ablazione:
  Gli esperimenti confermano che la rimozione di Stop-Gradient o Batch Normalization porta al collasso totale dell'addestramento, e che la regolarizzazione esplicita è essenziale per evitare il collasso delle caratteristiche.

4. Contributi Chiave

1. Framework CoReDi: Introduzione di un metodo in cui lo spazio delle rappresentazioni semantiche non è fisso, ma evolve congiuntamente al modello di diffusione tramite una proiezione apprendibile.
2. Stabilità dell'Addestramento: Identificazione e analisi di tre componenti critiche (Stop-Gradient, Normalizzazione del Batch, Regolarizzazione esplicita) necessarie per prevenire soluzioni degeneri e collasso delle caratteristiche.
3. Versatilità: Dimostrazione che le rappresentazioni adattive migliorano la generazione sia nello spazio latente VAE che nello spazio dei pixel, offrendo un miglioramento significativo nella qualità del campione e nella velocità di convergenza.

5. Significato e Impatto

Il lavoro CoReDi sfida il paradigma corrente secondo cui le rappresentazioni semantiche devono essere pre-calcolate e fisse. Dimostra che adattare lo spazio delle rappresentazioni all'obiettivo di generazione porta a risultati superiori.

• Efficienza: Riduce drasticamente il tempo di addestramento necessario per raggiungere lo stato dell'arte.
• Qualità: Migliora la fedeltà e la coerenza semantica delle immagini generate.
• Flessibilità: Offre una via per migliorare la generazione di immagini senza dipendere esclusivamente dalla qualità della ricostruzione del VAE, aprendo la strada a modelli di diffusione pixel-based più potenti guidati da semantiche adattive.

In sintesi, CoReDi trasforma le rappresentazioni semantiche da un "ingrediente statico" a un "componente dinamico" del processo generativo, ottimizzando direttamente la loro utilità per la sintesi di immagini.