Technische Samenvatting: SSDD: Single-Step Diffusion Decoder voor Efficiënte Image Tokenization

Probleemstelling

State-of-the-art generatieve beeldmodellen vertrouwen op tokenizers om hoogdimensionale visuele signalen te comprimeren naar compacte latente representaties. Hoewel huidige standaard tokenizers, zoals KL-regularized Variational Autoencoders (KL-VAE), effectief zijn, kampen ze met een fundamentele afweging tussen distortie (pixel-niveau gelijkenis) en distributieverschuiving (perceptuele realisme en diversiteit). Deterministische decoders minimaliseren distortie maar falen vaak in het modelleren van de complexe distributie van realistische afbeeldingen, wat leidt tot wazige outputs. Daarentegen kunnen generatieve decoders deze distributie wel modelleren, maar vereisen ze doorgaans iteratieve sampling (trage inferentie) of vertrouwen ze op adversarial losses (GANs), die moeilijk te schalen en te stabiliseren zijn.

Bestaande diffusion-gebaseerde decoders bieden een geprincipeerde aanpak voor het modelleren van de conditionele distributie van afbeeldingen gegeven latents, maar lijden onder twee belangrijke beperkingen:

1. Trage Inferentie: Ze vereisen meerdere denoising-stappen, wat een bottleneck vormt voor downstream generatieve modellen.
2. Trainingsinstabiliteit: Het bereiken van state-of-the-art resultaten vereist vaak GAN losses, die moeilijk te schalen zijn en kunnen leiden tot divergerende trainingsruns.

Methodologie

De auteurs introduceren SSDD (Single-Step Diffusion Decoder), een nieuwe architectuur die ontworpen is om deze beperkingen te overwinnen door een schaalbare hybride architectuur, GAN-vrije training en een single-step distillatiestrategie te combineren.

1. Schaalbare Hybride Architectuur

In plaats van uitsluitend te vertrouwen op convolutionele U-Nets (gebruikelijk in pixel-space diffusion) of pure Transformers (gebruikelijk in latent modeling), maakt SSDD gebruik van een U-ViT gebaseerde architectuur.

• Structuur: Het maakt gebruik van een U-Net backbone met vier niveaus van convolutionele ResNet-blocks, maar vervangt het middelste attention block door een Transformer die opereert op tokens die $8 \times 8$ pixel patches representeren.
• Conditionering: Het model conditioneert op de latente representatie $z$ door de latente grid te upsamplen en deze te concateneren met de ruisende input. Het versterkt de conditionering verder via Adaptive Group Normalization (AdaGN) en Adaptive Layer Normalization (AdaLN) binnen zowel de ResNet als de Transformer blocks.
• Schaling: De architectuur schaalt efficiënt van 13.4M (SSDD-S) naar 153.8M (SSDD-XL) parameters door het vergroten van de kanaalbreedtes en de transformer diepte.

2. GAN-Vrij Trainingsschema

SSDD elimineert de noodzaak voor adversarial training door drie specifieke loss functies te combineren:

• Flow Matching Loss: Het primaire doel gebruikt optimal transport flow matching om het velocity veld ($v = x - \epsilon$) tussen ruis en data te voorspellen, getraind met een $L_2$ loss.
• Perceptual Loss (LPIPS): Toegevoegd om generatie te sturen naar perceptueel correcte gebieden, waardoor high-level features worden uitgelijnd.
• REPA Loss: Een feature regularisatie loss die gebruik maakt van DINOv2-B features om de transformer representaties te stabiliseren en de training te versnellen.
• Multi-Stage Training: Het model wordt getraind in fasen: eerst op $128 \times 128$ crops met random resizing augmentation, gevolgd door fine-tuning op doelresoluties ($256 \times 256$).

3. Single-Step Distillatie

Om de bottleneck van de inferentiesnelheid aan te pakken, gebruikt SSDD een distillatiestrategie:

• Teacher-Student Setup: Een multi-step teacher decoder (getraind met 8 stappen) genereert hoogwaardige, diverse reconstructies. Een student decoder wordt vervolgens getraind om deze outputs in een enkele stap te reproduceren.
• Loss Preservatie: In tegen tegenover de standaard distillatie die mogelijk enkel op $L_2$ regressie vertrouwt, behoudt de SSDD distillatie de volledige training objective. De student minimaliseert de flow-matching en LPIPS losses berekend tegen de multi-step outputs van de teacher.
• Resultaat: Dit stelt het model in staat om de perceptuele en generatieve eigenschappen van de multi-step teacher over te nemen, terwijl het opereert met de throughput van een single-step model.

4. Gedeelde Encoders

De auteurs demonstreren dat SSDD decoders getraind kunnen worden met een gedeelde encoder (specifiek de standaard KL-regularized encoder van Rombach et al., 2022). Dit maakt het mogelijk dat SSDD dient als een drop-in vervanging voor bestaande VAE decoders zonder dat de encoder opnieuw getraind hoeft te worden of een nieuwe latent space gecreëerd moet worden voor elke decoder grootte.

Belangrijkste Bijdragen

1. Eerste GAN-vrije Single-Step Diffusion Decoder: SSDD is de eerste architectuur die state-of-the-art continue tokenization bereikt zonder adversarial losses, wat trainingsstabiliteit garandeert over verschillende model schalen.
2. Hybride Architectuur: De integratie van convolutionele inductieve biases met transformer generatieve capaciteiten in een pixel-space decoder levert superieure schalingseigenschappen op vergeleken met pure U-Nets.
3. Efficiënte Distillatie: De voorgestelde distillatiemethode comprimeert succesvol multi-step diffusion gedrag naar een enkele stap met minimale kwaliteitsdegradatie, wat de throughput aanzienlijk verbetert.
4. Compatibiliteit met Gedeelde Latent Space: Het vermogen om standaard, pre-trained encoders te gebruiken, maakt SSDD direct toepasbaar op bestaande generatieve pipelines.

Experimentele Resultaten

Geëvalueerd op ImageNet-1k (256$\times$256) en COCO, vertoont SSDD significante verbeteringen ten opzichte van KL-VAE en andere diffusion decoders:

• Reconstructie Kwaliteit (f8c4 setting):
  • SSDD-S (13.4M params) verbetert de reconstructie FID (rFID) van 0.87 (KL-VAE) naar 0.46.
  • SSDD-M (48M params, vergelijkbaar met de grootte van KL-VAE) bereikt een rFID van 0.39.
  • SSDD presteert beter dan alle baselines op perceptuele metrieken (LPIPS, DreamSim) terwijl het competitieve low-level metrieken (PSPS, SSIM) behoudt, hoewel het prioriteit geeft aan perceptuele kwaliteit boven pixel-perfecte distortie.
• Throughput:
  • SSDD-S bereikt een 1.4$\times$ hogere throughput dan KL-VAE terwijl het een betere kwaliteit biedt.
  • Vergeleken met multi-step diffusion decoders (bijv. $\epsilon$-VAE), biedt SSDD tot wel 10$\times$ snellere decoding bij vergelijkbare rFID niveaus.
• Downstream Generatie:
  • Wanneer SSDD wordt gebruikt om Diffusion Transformers (DiT) te trainen, verbetert het de generatie FID (gFID) vergeleken met KL-VAE.
  • Door een hogere compressieratio (f16c4) te gebruiken, maakt SSDD 3.8$\times$ snellere sampling mogelijk voor het DiT model terwijl de generatiekwaliteit behouden blijft.
• Generalisatie: SSDD-modellen generaliserenen goed naar COCO en hogere resoluties (tot 1024$\times$1024) zonder vanaf nul te hoeven trainen, en kunnen latents decoderen van frozen encoders zoals DiTo en DC-AE.

Betekenis

Het artikel beweert dat SSDD een nieuwe standaard vestigt voor beeld tokenization door de spanning tussen snelheid, stabiliteit en generatieve fideliteit op te lossen. Door de afhankelijkheid van GANs en iteratieve sampling te verwijderen, biedt SSDD een schaalbare, drop-in vervanging voor KL-VAE die de training van kwalitatief betere en snellere generatieve modellen mogelijk maakt. Het werk demonstreert dat diffusion-gebaseerde tokenization zowel efficiënt als stabiel kan zijn, wat de weg vrijmaakt voor toekomstige grootschalige generatieve modellering zonder de computationele overhead van multi-step inferentie of de instabiliteit van adversarial training.