CrossFlow: One-Step Generation Across Latent and Pixel Spaces

Riepilogo Tecnico: CrossFlow: Generazione in un unico step tra spazi latente e di pixel

1. Definizione del Problema

Gli attuali sistemi di generazione di immagini all'avanguardia, in particolare i modelli di diffusione latente (LDM) e i metodi di flow-matching, operano definendo il prior, il percorso di probabilità e l'obiettivo di predizione all'interno di un singolo spazio di rappresentazione. Sebbene gli LDM migliorino l'efficienza computazionale eseguendo la traiettoria generativa in uno spazio latente di un autoencoder compresso, essi si affidano a una pipeline a due stadi: denoising iterativo nello spazio latente seguito da un passaggio separato di decodifica per mappare i latenti nuovamente nei pixel.

Questa architettura introduce un disallineamento fondamentale:

• Disallineamento di Ottimizzazione: Il generatore è ottimizzato per predire spostamenti latenti, mentre la qualità finale dell'immagine dipende da un decoder addestrato separatamente.
• Spostamento della Distribuzione (Distribution Shift): Durante l'inferenza, il decoder riceve latenti generati che possono differire distribuzionalmente dai latenti puliti prodotti dall'encoder durante l'addestramento. Questo spostamento può amplificare gli errori.
• Limitazione della Supervisione: Il generatore non può essere addestrato direttamente con perdite percettive o avversarie a livello di immagine a meno che il suo output non passi attraverso il decoder separato, indebolendo la supervisione diretta nello spazio dei pixel.

Il documento pone la domanda se lo spazio del rumore e lo spazio di output debbano essere identici, proponendo un framework che preservi l'efficienza delle rappresentazioni latenti pur predendo direttamente l'immagine stessa.

2. Metodologia: CrossFlow

CrossFlow è una formulazione di flow cross-space che mappa input latenti rumorosi direttamente in immagini nello spazio dei pixel in una singola valutazione della funzione. Risolve il disallineamento tra spazi derivando obiettivi in cui la velocità nello spazio latente viene eliminata dal target supervisionato.

Formulazione Tecnica Core:

• Framework Unificato: Gli autori definiscono una funzione target unificata $Z$ per i paradigmi generativi a singolo step. I metodi standard si affidano a campi di velocità nello spazio latente, il che impedisce la generazione cross-space.
• Criteri di Design: Per abilitare la supervisione cross-space, gli autori stabiliscono due criteri:
  1. Eliminazione delle Dipendenze dalla Velocità Latente: L'obiettivo deve essere strutturalmente disaccoppiato dalle velocità nello spazio latente. L'output della rete non deve essere un vettore di spostamento latente.
  2. Condizioni al Contorno del Punto Finale: Il framework deve produrre un'immagine valida al confine di destinazione ($r=0$).
• L'Obiettivo CrossFlow: Soddisfacendo questi criteri, gli autori derivano una funzione di perdita generale (Eq. 10) dove la rete $F_\theta$ predice un'immagine nello spazio dei pixel $x_0$ piuttosto che una velocità latente.
  • La traiettoria latente (definita da $\alpha(t)$ e $\beta(t)$) detta comunque come i latenti puliti vengono corrotti per creare gli input.
  • Il target supervisionato è l'immagine pixel pulita.
  • L'obiettivo specifico utilizzato negli esperimenti (Eq. 11) coinvolge un termine di derivata temporale calcolato tramite Prodotto Jacobiano-Vettore (JVP) lungo la traiettoria latente, garantendo che il modello apprenda la corretta dinamica di flow senza predire le velocità latenti.

Dettagli di Implementazione:

• Architettura: Il modello utilizza un encoder di immagini pre-addestrato e fisso (VA-VAE) per mappare le immagini in latenti e un generatore Vision Transformer (ViT) ($F_\theta$) per mappare direttamente i latenti rumorosi in pixel.
• Pipeline di Addestramento: Un approccio in due stadi viene utilizzato per garantire la stabilità:
  1. Addestrare l'encoder $E$ utilizzando obiettivi standard di autoencoder (ricostruzione + avversario).
  2. Congelare $E$ e addestrare il generatore $F_\theta$ da zero utilizzando l'obiettivo CrossFlow.
• Funzione di Perdita: L'obiettivo totale combina la perdita analitica CrossFlow con perdite ausiliarie nello spazio dei pixel:
  • Perdita di Ricostruzione L1: Allineamento pixel per pixel.
  • Perdita Percettiva: Somiglianza strutturale.
  • Perdita Avversaria (GAN): Fedeltà delle alte frequenze.
  • Nota: Il fattore di ponderazione analitico della funzione di perdita del flow viene omesso durante il calcolo della perdita percettiva per prevenire l'instabilità dell'addestramento.

3. Contributi Chiave

1. Framework Cross-Space: Introduzione di CrossFlow, un framework di generazione a singolo step in cui il prior di input (latente) e l'immagine di output (pixel) risiedono in spazi di rappresentazione differenti.
2. Obiettivo One-Step Privo di Velocità: Una nuova derivazione che permette a un modello di agire sia come generatore latente-pixel a singolo step che come sostituto del decoder, eliminando la necessità di un decoder separato durante l'inferenza.
3. Supervisione Diretta nello Spazio dei Pixel: La capacità di addestrare direttamente il generatore con perdite percettive e avversarie sulle immagini generate, bypassando lo stadio intermedio del decoder.
4. Validazione Empirica: Dimostrazione che un singolo modello può raggiungere punteggi FID competitivi su ImageNet-1k (256×256) con una singola valutazione della funzione (1-NFE) pur superando i metodi latenti a singolo step standard e pareggiando le baseline di diffusione latente multi-step.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-1k a risoluzione 256×256 utilizzando la generazione condizionata per classe.

• Performance Principale:

  • CrossFlow-XL (461M parametri) raggiunge un FID di 1.62 con 1-NFE.
  • È competitivo con le baseline di diffusione latente multi-step come DiT-XL/2 (2.27 FID con 250 step) e SiT-XL/2 (2.06 FID con 250 step).
  • Supera significativamente altri metodi latenti a singolo step (es. iMF-XL/2 con 1.72 FID) ed evita la necessità di un decoder separato.
  • CrossFlow-B (95M) e CrossFlow-L (315M) ottengono rispettivamente FID di 2.54 e 1.69, dimostrando scalabilità.

• Studi di Ablazione:

  • Scelta dell'Encoder: L'encoder VA-VAE (32 canali) ha prodotto i risultati migliori (2.54 FID), superando SD-VAE (4.08 FID) e FLUX.2 (3.81 FID), suggerendo che la struttura della rappresentazione latente sia critica.
  • Componenti della Perdita: Rimuovere la perdita percettiva ha causato un drastico calo del FID a 30.21, mentre rimuovere la perdita GAN ha portato il FID a 8.86. L'obiettivo di flow CrossFlow fornisce il segnale strutturale, mentre le perdite percettive e avversarie sono vitali per la fedeltà delle alte frequenze.
  • Sostituzione del Decoder: Quando utilizzato come decoder per i latenti generati da un LightningDiT-B/1, CrossFlow ha costantemente migliorato il FID rispetto al decoder VA-VAE standard durante tutto l'addestramento.

5. Significato e Rivendicazioni

Il documento afferma che CrossFlow riesce a colmare il divario tra l'efficienza delle rappresentazioni latenti e la fedeltà della supervisione diretta nello spazio dei pixel. Unificando il generatore e il decoder in un singolo modello addestrato su input latenti rumorosi e target nello spazio dei pixel, CrossFlow:

• Rimuove lo stadio esplicito del decoder durante l'inferenza.
• Espone il modello durante l'addestramento ai latenti rumorosi/imperfetti che si verificano durante la generazione, anziché solo agli output puliti dell'encoder.
• Dimostra che gli obiettivi di flow cross-space possono combinare i benefici computazionali della diffusione latente con i miglioramenti di qualità della l'addestramento avversario e percettivo diretto a livello di pixel.

Gli autori mantengono un tono modesto, osservando che sebbene la rappresentazione latente e le perdite ausiliarie siano importanti, il contributo principale sia lo spostamento architettonico e dell'obiettivo che consente la generazione cross-space in un singolo step senza richiedere moduli decoder separati.