WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

1. Il Problema: Conflitti di Traiettoria nello Spazio dei Pixel

Il lavoro affronta una sfida fondamentale nell'ambito della generazione di immagini tramite modelli di diffusione e Flow Matching (FM) operanti direttamente nello spazio dei pixel (senza l'uso di autoencoder latenti come i VAE).

• Limiti degli approcci Latenti: I modelli tradizionali operano in spazi latenti compressi per ridurre i costi computazionali, ma questo introduce un collo di bottiglia informativo che perde dettagli ad alta frequenza e genera artefatti visivi.
• Limiti degli approcci Pixel-Space: I recenti modelli che operano direttamente sui pixel (es. JiT) evitano la compressione, ma affrontano un problema di ottimizzazione critico: la mancanza di continuità semantica nel manifold dei pixel.
• Il Conflitto di Traiettoria: Poiché lo spazio dei pixel è intrinsecamente "entangled" (intrecciato) e non discriminativo, percorsi di trasporto ottimali destinati a endpoint semanticamente diversi ma visivamente simili tendono a sovrapporsi nelle regioni dense dello spazio del rumore.
  • Questo costringe la rete neurale a prevedere un campo di velocità medio per minimizzare la perdita di regressione (MSE), causando "sangue semantico" (semantic bleeding), distorsioni geometriche e una convergenza lenta.
  • Tecniche come la Classifier-Free Guidance (CFG) mitigano il problema ma non risolvono la sovrapposizione spaziale sottostante durante l'addestramento.

2. Metodologia: WiT (Waypoint Diffusion Transformers)

Gli autori propongono WiT, un paradigma generativo che risolve il conflitto di traiettoria introducendo waypoint semantici intermedi nel flusso generativo. L'idea centrale è scomporre il mapping complesso "rumore $\to$ pixel" in due fasi disaccoppiate e matematicamente più stabili.

A. Concetto di Waypoint Semantico

Invece di mappare direttamente il rumore isotropo alla distribuzione dei pixel, il flusso viene instradato attraverso un waypoint semantico discriminativo ($s_0$).

• Decomposizione della Varianza: Utilizzando la legge della varianza totale, il conflitto di traiettoria (varianza incondizionata) viene scomposto in:
  1. Una parte risolta dalla previsione del waypoint semantico (che riduce drasticamente l'ambiguità).
  2. Una varianza residua (pixel-level) che il generatore principale deve solo risolvere, essendo ora vincolato a un sottospazio affine specifico.

B. Architettura del Modello

L'architettura WiT è composta da due componenti principali:

1. Lightweight Waypoints Generator ($W_\psi$):

   • Un modello Transformer leggero (es. 21M parametri) che prende lo stato rumoroso dei pixel ($z_t$) e il timestep $t$.
   • Il suo compito è inferire il waypoint semantico pulito ($\hat{s}_0$) direttamente dal rumore.
   • I waypoint sono derivati da modelli di visione pre-addestrati (es. DINOv3) proiettati su un sottospazio a bassa dimensionalità tramite PCA (dimensione ridotta a $d=64$) per garantire separabilità e ridurre la ridondanza spaziale.
   • Viene addestrato con un obiettivo di Flow Matching nello spazio semantico.

2. Pixel Space Generator ($G_\theta$) con Just-Pixel AdaLN:

   • Un Transformer principale (es. DiT) responsabile della sintesi finale delle immagini ad alta fedeltà.
   • Riceve il waypoint $\hat{s}_0$ come condizione spaziale dinamica.
   • Meccanismo Just-Pixel AdaLN: A differenza dell'AdaLN standard che applica una modulazione globale, WiT introduce un meccanismo di modulazione spazialmente variabile. I waypoint predetti vengono proiettati e aggregati per generare parametri di scala e shift ($\gamma, \beta, \alpha$) che modulano localmente i token all'interno dei blocchi Transformer. Questo permette di guidare la generazione pixel per pixel senza alterare la sequenza di token originale.

C. Procedura di Addestramento e Inferenza

• Addestramento a due stadi: Prima si addestra il generatore di waypoint ($W_\psi$) fino a convergenza; successivamente, viene congelato e utilizzato per condizionare il generatore principale ($G_\theta$).
• Inferenza: Ad ogni passo dell'ODE, il generatore di waypoint ricalibra dinamicamente il waypoint semantico $\hat{s}_0$ dallo stato corrente rumoroso, fornendo una guida semantica continua che "raddrizza" la traiettoria verso il manifold corretto.

3. Contributi Chiave

1. Paradigma WiT: Introduzione di un nuovo framework che mitiga il conflitto di traiettoria nello spazio dei pixel ancorando i flussi generativi a manifold semantici a bassa dimensionalità.
2. Meccanismo Just-Pixel AdaLN: Una nuova tecnica di condizionamento che fornisce guida semantica spazialmente variabile, disaccoppiando la navigazione semantica dalla sintesi dei dettagli testurali.
3. Risultati SOTA: Dimostrazione che è possibile ottenere qualità generativa superiore rispetto ai modelli latenti (VAE) senza ricorrere alla compressione lossy, mantenendo la purezza dei pixel.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet 256x256.

• Performance Quantitative:

  • WiT-L/16 raggiunge un FID di 2.22 e un IS di 303.3 dopo 600 epoche, superando il baseline JiT-L/16 (FID 2.36) e battendo modelli latenti pesanti come DiT-XL/2 (FID 2.27).
  • WiT-XL/16 scende sotto la soglia di FID 2.0, raggiungendo 1.89, stabilendosi come stato dell'arte tra i modelli puramente pixel-based.
  • Efficienza: WiT accelera la convergenza dell'addestramento di 2.2x rispetto a JiT. Ad esempio, WiT-B/16 a 200 epoche supera le prestazioni di JiT-B/16 addestrato per 600 epoche.

• Analisi Qualitativa:

  • Le immagini generate mostrano una coerenza strutturale eccezionale (proporzioni corrette, prospettive rigide) grazie all'ancoraggio semantico.
  • La purezza dello spazio dei pixel preserva micro-texture ad alta frequenza (piume, ali di farfalla) spesso perse nella compressione VAE.

• Analisi del Conflitto di Traiettoria:

  • Le metriche quantitative mostrano che WiT riduce il conflitto direzionale pairwise del 1.62x rispetto a JiT, confermando che l'introduzione dei waypoint "srotola" le traiettorie sovrapposte.

5. Significato e Impatto

Il lavoro di WiT è significativo perché risolve il compromesso storico tra fedeltà dei dettagli (tipica degli spazi pixel) e stabilità dell'ottimizzazione (tipica degli spazi latenti).

• Dimostra che non è necessario ricorrere a rappresentazioni latenti lossy per addestrare modelli di diffusione scalabili.
• Fornisce una soluzione teorica e pratica al problema della sovrapposizione delle traiettorie nei campi vettoriali continui, aprendo la strada a generatori di immagini puri, ad alta risoluzione e privi di artefatti di compressione.
• L'approccio "Waypoint" potrebbe essere esteso ad altri domini di generazione dove lo spazio target è intrinsecamente entangled.

In sintesi, WiT rappresenta un avanzamento fondamentale verso la generazione di immagini fotorealistiche "pure", combinando la stabilità semantica dei modelli di visione pre-addestrati con la ricchezza informativa dello spazio dei pixel grezzi.