Sintesi Tecnica: SSDD: Single-Step Diffusion Decoder per una Tokenizzazione d'Immagini Efficiente

Definizione del Problema

I modelli generativi di immagini allo stato dell'arte si affidano ai tokenizer per comprimere segnali visivi ad alta dimensionalità in rappresentazioni latenti compatte. Sebbene gli attuali tokenizer standard, come i Variational Autoencoders regolarizzati KL (KL-VAE), siano efficaci, affrontano un compromesso fondamentale tra distorsione (similarità a livello di pixel) e shift della distribuzione (realismo percettivo e diversità). I decoder deterministici minimizzano la distorsione ma spesso non riescono a modellare la complessa distribuzione delle immagini realistiche, portando a output sfocati. Al contrario, i decoder generativi possono modellare questa distribuzione ma tipicamente richiedono un campionamento iterativo (inferenza lenta) o si affidano a perdite avversarie (GAN), che sono difficili da scalare e stabilizzare.

I decoder basati sulla diffusione esistenti offrono un approccio fondato per modellare la distribuzione condizionale delle immagini dati i latenti, ma soffrono di due limitazioni principali:

1. Inferenza Lenta: Richiedono molteplici step di denoising, creando un collo di bottiglia per i modelli generativi a valle.
2. Instabilità dell'Addestramento: Raggiungere risultati allo stato dell'arte spesso richiede perdite GAN, che sono difficili da scalare e possono portare a divergenze durante l'addestramento.

Metodologia

Gli autori introducono SSDD (Single-Step Diffusion Decoder), un'architettura innovativa progettata per superare queste limitazioni combinando un'architettura ibrida scalabile, un addestramento privo di GAN e una strategia di distillazione a singolo step.

1. Architettura Ibrida Scalabile

Invece di fare affidamento esclusivamente su U-Net convoluzionali (comuni nella diffusione nello spazio dei pixel) o su Transformer puri (comuni nella modellazione latente), SSDD utilizza un'architettura basata su U-ViT.

• Struttura: Impiega un backbone U-Net con quattro livelli di blocchi ResNet convoluzionali, ma sostituisce il blocco di attenzione centrale con un Transformer che opera su token rappresentanti patch di pixel $8 \times 8$.
• Condizionamento: Il modello si condiziona sulla rappresentazione latente $z$ tramite l'upsampling della griglia latente e la concatenazione con l'input rumoroso. Potenzia ulteriormente il condizionamento tramite Adaptive Group Normalization (AdaGN) e Adaptive Layer Normalization (AdaLN) all'interno sia dei blocchi ResNet che di quelli Transformer.
• Scalabilità: L'architettura scala efficientemente da 13.4M (SSDD-S) a 153.8M (SSDD-XL) di parametri aumentando la larghezza dei canali e la profondità del transformer.

2. Schema di Addestramento GAN-Free

SSDD elimina la necessità di addestramento avversario combinando tre specifiche funzioni di perdita:

• Flow Matching Loss: L'obiettivo primario utilizza il flow matching a trasporto ottimale per predire il campo di velocità ($v = x - \epsilon$) tra rumore e dati, addestrato con una perdita $L_2$.
• Perceptual Loss (LPIPS): Aggiunta per guidare la generazione verso aree percettivamente corrette, allineando le caratteristiche di alto livello.
• REPA Loss: Una perdita di regolarizzazione delle caratteristiche che utilizza le feature di DINOv2-B per stabilizzare le rappresentazioni del transformer e accelerare l'addestramento.
• Addestramento Multi-stadio: Il modello viene addestrato in fasi: prima su ritagli $128 \times 128$ con aumento di ridimensionamento casuale, seguiti dal fine-tuning alla risoluzione target ($256 \times 256$).

3. Distillazione a Singolo Step

Per affrontare il collo di bottiglia della velocità di inferenza, SSDD impiega una strategia di distillazione:

• Configurazione Teacher-Student: Un decoder teacher multi-step (addestrato con 8 step) genera ricostruzioni di alta qualità e diversificate. Un decoder student viene quindi sottoposto a fine-tuning per riprodurre questi output in un singolo step.
• Preservazione della Perdita: A differenza della distillazione standard che potrebbe fare affidamento solo sulla regressione $L_2$, la distillazione di SSDD preserva l'intero obiettivo di addestramento. Lo student minimizza le perdite di flow-matching e LPIPS calcolate rispetto agli output multi-step del teacher.
• Risultato: Ciò consente al modello di ereditare le proprietà percettive e generative del teacher multi-step operando con la velocità di throughput di un modello a singolo step.

4. Encoder Condivisi

Gli autori dimostrano che i decoder SSDD possono essere addestrati con un encoder condiviso (specificamente l'encoder standard regolarizzato KL di Rombach et al., 2022). Ciò consente a SSDD di fungere da sostituzione diretta per gli esistenti decoder VAE senza richiedere il ri-addestramento dell'encoder o la creazione di un nuovo spazio latente per ogni dimensione del decoder.

Contributi Chiave

1. Primo Decoder di Diffusione a Singolo Step GAN-Free: SSDD è la prima architettura a raggiungere una tokenizzazione continua allo stato dell'arte senza perdite avversarie, garantendo la stabilità dell'addestramento attraverso diverse scale di modello.
2. Architettura Ibrida: L'integrazione dei bias induttivi convoluzionali con le capacità generative dei transformer in un decoder nello spazio dei pixel produce proprietà di scalabilità superiori rispetto ai puri U-Net.
3. Distillazione Efficiente: Il metodo di distillazione proposto comprime con successo il comportamento della diffusione multi-step in un singolo step con una degradazione minima della qualità, migliorando significativamente il throughput.
4. Compatibilità con lo Spazio Latente Condiviso: La capacità di utilizzare encoder pre-addestrati standard rende SSDD immediatamente applicabile alle pipeline generative esistenti.

Risultati Sperimentali

Valutato su ImageNet-1k (256$\times$256) e COCO, SSDD dimostra miglioramenti significativi rispetto a KL-VAE e altri decoder di diffusione:

• Qualità di Ricostruzione (impostazione f8c4):
  • SSDD-S (13.4M parametri) migliora la rFID (reconstruction FID) da 0.87 (KL-VAE) a 0.46.
  • SSDD-M (48M parametri, dimensioni simili a KL-VAE) raggiunge una rFID di 0.39.
  • SSDD supera tutti i baseline nelle metriche percettive (LPIPS, DreamSim) pur mantenendo metriche di basso livello competitive (PSNR, SSIM), sebbene dia priorità alla qualità percettiva rispetto alla distorsione perfetta del pixel.
• Throughput:
  • SSDD-S raggiunge un throughput 1.4$\times$ superiore rispetto a KL-VAE offrendo al contempo una qualità migliore.
  • Rispetto ai decoder di diffusione multi-step (es. $\epsilon$-VAE), SSDD offre fino a 10$\times$ più velocità di decoding a parità di livelli di rFID.
• Generazione a Valle:
  • Utilizzato per addestrare Diffusion Transformers (DiT), SSDD migliora la gFID (generation FID) rispetto a KL-VAE.
  • Utilizzando un rapporto di compressione più elevato (f16c4), SSDD consente un campionamento 3.8$\times$ più veloce per il modello DiT preservando la qualità di generazione.
• Generalizzazione: I modelli SSDD generalizzano bene su COCO e risoluzioni più elevate (fino a 1024$\times$1024) senza dover ri-addestrare da zero, e possono decodificare i latenti da encoder congelati come DiTo e DC-AE.

Significato

Il paper afferma che SSDD stabilisce un nuovo standard per la tokenizzazione delle immagini risolvendo la tensione tra velocità, stabilità e fedeltà generativa. Eliminando la dipendenza dalle GAN e dal campionamento iterativo, SSDD fornisce una sostituzione diretta e scalabile per KL-VAE che consente l'addestramento di modelli generativi di qualità superiore e più veloci. Il lavoro dimostra che la tokenizzazione basata sulla diffusione può essere sia efficiente che stabile, aprendo la strada a futuri modelli generativi su larga scala senza l'overhead computazionale dell'inferenza multi-step o l'instabilità dell'addestramento avversario.