JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Riepilogo Tecnico: JLT: Predizione Latente Pulita nei Trasformatori di Diffusione Latente

Enunciato del Problema

I modelli di diffusione per il denoising invertono tradizionalmente un processo di corruzione, ma l'obiettivo specifico che la rete neurale è addestrata a prevedere varia a seconda dell'architettura. Gli obiettivi comuni includono il rumore ($\epsilon$), la velocità ($v = x - \epsilon$) o i dati puliti ($x$). Sebbene questi obiettivi siano algebricamente convertibili per un tempo di corruzione fissato, ricerche recenti (ad es. JiT) suggeriscono che, in presenza di capacità di rete finita e dati finiti, la scelta dell'obiettivo di previsione altera significativamente la difficoltà del problema di regressione. Nello specifico, prevedere i dati puliti permette a un modello di concentrarsi su variazioni strutturate su varietà a bassa dimensionalità, piuttosto che ricostruire il rumore ambientale fuori dalla varietà.

La domanda centrale affrontata da questo articolo è se questo vantaggio della "predizione pulita" persista quando il modello opera in uno spazio latente compresso anziché nello spazio dei pixel grezzi. Nella diffusione latente, un Variational Autoencoder (VAE) ha già rimosso gran parte della variabilità dei pixel grezzi. Gli autori investigano se le proprietà geometriche dell'obiettivo di previsione ($x$ vs. $v$) rimangano un fattore critico anche quando lo spazio di rappresentazione è fissato e compresso.

Metodologia

Gli autori introducono JLT (Joint Latent Transformer), un quadro sperimentale controllato progettato per isolare l'effetto dell'obiettivo di previsione all'interno di uno spazio latente fisso.

1. Architettura e Configurazione:

   • Backbone: Un Transformer di scala Base (130 milioni di parametri) con 12 blocchi, dimensione nascosta 768 e 12 teste di attenzione, rispecchiando la configurazione di JiT-B/16.
   • Rappresentazione: Il modello opera su codici latenti FLUX.2 VAE congelati. Ciò garantisce che lo spazio di rappresentazione sia identico in tutte le varianti confrontate.
   • Addestramento: I modelli sono addestrati per 250K step (200 epoche) su ImageNet 256 $\times$ 256.
   • Variabili Controllate: Lo studio mantiene strettamente costanti la rappresentazione (VAE FLUX.2), la scala dell'architettura, le impostazioni dell'ottimizzatore, la dimensione del batch e i protocolli di campionamento.

2. Confronto Sperimentale:
   Gli autori confrontano due varianti principali nelle stesse impostazioni:

   • JLT (Predizione Latente Pulita): La rete è supervisionata direttamente per prevedere il vettore latente pulito $x$.
   • DiT (Predizione della Velocità): Una baseline corrispondente in cui la rete è supervisionata per prevedere il vettore velocità $v = x - \epsilon$.
   • Varianti: Gli esperimenti sono condotti su due dimensioni di patch della griglia VAE, indicate come /1 e /2 (ad es. JLT-B/1 vs. DiT-B/1).

3. Analisi Teorica:
   L'articolo fornisce un'analisi gaussiana locale per spiegare il divario empirico. Assumendo un'approssimazione lineare-gaussiana locale della distribuzione dei dati latenti ($x \sim \mathcal{N}(0, \Sigma)$) e rumore indipendente ($\epsilon \sim \mathcal{N}(0, I)$):

   • Struttura di Covarianza: La covarianza dell'obiettivo velocità è $Cov(y_v) = \Sigma + I$, mentre quella dell'obiettivo pulito è $Cov(y_x) = \Sigma$. Ciò implica che la previsione della velocità aggiunge un piano isotropo unitario a ogni direzione. Nelle direzioni in cui i dati puliti hanno bassa varianza (debolmente supportati dalla varietà), la previsione della velocità costringe il modello a imparare un obiettivo a varianza unitaria, amplificando efficacemente il rumore in quelle direzioni.
   • Ambiguità Condizionale: L'analisi mostra che la varianza residua di Bayes per la previsione della velocità è scalata di un fattore $1/(1-t)^2$ rispetto alla previsione pulita. Man mano che il valore proprio $\lambda_i$ della covarianza dei dati tende a zero, il coefficiente per la previsione pulita tende a 0 (attenuando la direzione), mentre il coefficiente della velocità tende a $-1/(1-t)$ (amplificando la direzione).

Risultati Chiave

Gli esperimenti su ImageNet 256 $\times$ 256 dimostrano un divario di prestazioni consistente e significativo a favore della predizione latente pulita:

• Prestazioni Quantitative:
  • JLT-B/1 (Pulito) ottiene un FID-50K di 2.50 (con guida senza classificatore), rispetto a 6.56 per la corrispondente DiT-B/1 (Velocità).
  • JLT-B/2 (Pulito) ottiene un FID-50K di 14.81, rispetto a 28.71 per DiT-B/2.
  • Anche l'Inception Score (IS) mostra un miglioramento sostanziale (ad es. 232.51 vs. 132.12 per le varianti /1).
• Dinamiche di Addestramento: Le curve di addestramento indicano che i modelli a latente pulito entrano nel regime a basso FID prima (intorno a 100K step) e mantengono un margine chiaro rispetto ai modelli a velocità per tutta la durata dell'addestramento.
• Robustezza: Il vantaggio si mantiene attraverso diverse granularità di tokenizzazione (/1 e /2), suggerendo che il beneficio non è un artefatto di una specifica dimensione di patch.

Significato e Affermazioni

L'articolo afferma che la scelta dell'obiettivo di previsione nella diffusione latente è una scelta di modellazione geometrica piuttosto che una mera parametrizzazione algebrica.

1. Geometria dell'Obiettivo oltre la Compressione: I risultati confutano l'idea che la sola compressione latente spieghi l'efficienza della diffusione latente. Il divario di prestazioni esiste all'interno dello stesso spazio latente fisso, dimostrando che la "geometria dell'obiettivo" (la distribuzione del segnale supervisionato) è un fattore distinto dalla "geometria della rappresentazione" (il tokenizzatore VAE).
2. Meccanismo di Miglioramento: Gli autori sostengono che la predizione pulita è superiore perché smorza le direzioni latenti a bassa varianza (che sono spesso rumore o caratteristiche della varietà debolmente supportate), whereas la predizione della velocità amplifica queste direzioni ereditando un piano di covarianza isotropo.
3. Portata delle Affermazioni: Gli autori sono modesti riguardo alle loro affermazioni teoriche. Non asseriscono che la predizione pulita sia globalmente ottimale per tutti i tokenizzatori, programmi di rumore o campionatori. Piuttosto, forniscono un meccanismo gaussiano locale che è coerente con il divario empirico misurato. Lo studio è presentato come un'ablazione controllata per isolare l'effetto dell'obiettivo, piuttosto che una proposta per una nuova architettura backbone o una sostituzione universale per tutti gli obiettivi di diffusione.

In conclusione, l'articolo stabilisce che anche negli spazi latenti compressi, prevedere direttamente lo stato latente pulito offre un vantaggio geometrico rispetto alla predizione della velocità, portando a una qualità di generazione significativamente superiore su ImageNet.