Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

Technische Samenvatting: Cross-Space Distillatie

Probleemstelling
Moderne one-step diffusiemodellen bereiken een hoge kwaliteit van generatie door middel van distributie-gebaseerde timestep-distillatie. Deze methoden vertrouwen echter op een restrictieve Shared-Space beperking: de Teacher en de Student moeten hetzelfde latente ruimte bewonen, met identieke latente resolutie en VAE-parameterisering. Deze aanname verhindert kennisoverdracht van moderne, hoog-capacitieve Teachers (bijv. Stable Diffusion 3.5, Flux, Kolors) die opereren op hoge resoluties (bijv. 1024×1024) en geavanceerde VAE's gebruiken, naar compacte, deployment-vriendelijke Students (bijv. Stable Diffusion 1.5) die opereren op lagere resoluties (bijv. 512×512) met verschillende VAE's.

Standaard distillatie-objectieven falen in deze setting omdat de latente tensoren ($z_S$ en $z_T$) inherent incompatibel zijn vanwege:

1. Cross-Resolutie: Mismatch in ruimtelijke gridvormen (bijv. $h \times w$ vs. $H \times W$).
2. Cross-VAE: Verschillende autoencoder-parameteriseringen die resulteren in niet-uitgelijnde latente ruimtes.
3. Cross-Architectuur/Mechanisme: Verschillen tussen UNet en DiT backbones, of tussen standaard diffusie en flow matching formuleringen.

Bestaande alternatieven, zoals netwerkpruning of het bouwen van nieuwe compacte backbones vanaf nul, kampen met uitdagingen zoals harde sparsity-drempels, uitgebreide hertrainingsvereisten en het verlies van ecosysteem-compatibiliteit.

Methodologie: Cross-Space Distillatie en de Bridge
De auteurs formaliseren deze setting als Cross-Space Distillatie en introduceren Bridge ($B_\phi$), een lichtgewicht module die ontworpen is om Student-latents te mappen naar de latente ruimte van de Teacher zonder de Student-backbone aan te passen.

• Architectuur: De Bridge is gefactoriseerd in twee stadia om de mismatch efficiënt te verwerken:

  1. Spatial Prior (Frozen): De eerste $n$ decoding blocks van de Student VAE-decoder ($D_S^{(n)}$) worden bevroren gehouden. Deze blokken fungeren als een vooraf getrainde ruimtelijke upsampler, die de Student latent $z_S$ vanuit een lage resolutie naar de hoge resolutie van de Teacher ($H \times W$) uitbreidt. Dit benut de bestaande ruimtelijke kennis van de Student zonder extra trainbare parameters voor upsampling toe te voegen.
  2. Learnable Projector ($g_\phi$): Een compacte, leerbare netwerk (geïmplementeerd als SwinIR) mapt de geupsamplede features van de Student feature-ruimte naar de latente representatieruimte van de Teacher, wat een Teacher-compatibele latent $\hat{z}_T$ oplevert.

• Trainingsdoelstellingen: De Bridge wordt getraind met een synthetische dataset van gepaarde latents $(z_S, z_T)$ afgeleid van dezelfde afbeelding, gecodeerd door beide VAE's. De optimalisatie minimaliseert een gewogen som van twee losses:

  1. Latent Reconstruction ($L_{rec}$): Een $\ell_1$ loss tussen de voorspelde latent $\hat{z}_T$ en de ground-truth Teacher latent $z_T$. De auteurs merken op dat $\ell_1$ betere stabiliteit en structuurbehoud biedt dan $\ell_2$ in deze context.
  2. Attention Fidelity ($L_{attn}$): Een reverse-KL divergentie loss die de self-attention distributies van de bevroren Teacher denoiser uitlijnt wanneer deze gevoed wordt met $\hat{z}_T$ versus $z_T$. Dit zorgt ervoor dat de gemapte latent dezelfde interne structurele responsen (long-range dependencies) induceert als de ware Teacher latent, waardoor fijnmazige details behouden blijven die pixel-wise losses mogelijk missen.

• Distillatie en Inferentie: Zod eens getraind, wordt de Bridge aan de Student gekoppeld. Tijdens de distillatie wordt de output van de Student door de Bridge geleid om $\hat{z}_T$ te genereren, waardoor standaard distributie-gebaseerde objectieven (bijv. Variational Score Distillation, Adversarial Distillation) direct tegen de Teacher toegepast kunnen worden. Cruciaal is dat de Student backbone ongewijzigd blijft.

Belangrijkste Bijdragen

1. Formalisering van Cross-Space Distillatie: Het artikel definieert een nieuw regime waarbij distillatie plaatsvindt over mismatchende latente resoluties en VAE-ruimtes, waarmee een kritiek gat wordt gedicht in de overdracht van kennis van moderne high-resolution modellen naar compacte legacy-modellen.
2. De Bridge Module: Een lichtgewicht, modulaire interface die standaard distillatie-objectieven mogelijk maakt zonder de Student architectuur te wijzigen. Het combineert een bevroren spatial prior met een leerbare projector en maakt gebruik van een Attention Fidelity objectief voor stabiele uitlijning.
3. Praktische Effectiviteit: De methode stelt compacte Students (bijv. SD 1.5) in staat om de visuele priors van diverse, moderne Teachers (SDXL, Kolors, PixArt-Σ, Flux, SD 3.5) over te nemen, terwijl ze één-stap inferentie, lage latentie en brede ecosysteem-compatibiliteit behouden.
4. Model Merging Capaciteit: Omdat de methode diverse Teachers distilleert naar Students met dezelfde backbone en latente ruimte, kunnen de resulterende checkpoints eenvoudig worden samengevoegd via parameter averaging om een enkel model te creëren dat de kennis van meerdere Teacher-bronnen aggregeert.

Resultaten
Experimenten tonen substantiële verbeteringen in generatiekwaliteit voor compacte one-step Students:

• Kwantitatieve Winsten: Het distilleren van SD 3.5 Medium (Teacher) naar SD 1.5 (Student) via Bridge verbetert de HPSv3 score van 5.37 naar 9.42 en ImageReward van -0.29 naar 0.62. Het mergen van checkpoints van vijf verschillende Teachers verhoogt de SD 1.5 HPSv3 verder naar 10.53.
• Kwalitatieve Verbeteringen: Visuele vergelijkingen laten zien dat Bridge-enabled students afbeeldingen produceren met duidelijkere texturen, sterkere structurele consistentie en minder artefacten vergeleken met geïnitialiseerde baselines of gepruide modellen.
• Efficiëntie: De Bridge voegt verwaarloosbare overhead toe (ongeveer 5M parameters) en maakt één-stap inferentie mogelijk. Het maakt ook inference-time resolution upgrades mogelijk, waarbij een 512×512 Student sample via de Bridge naar 1024×1024 kan worden gemapt en gedecodeerd met de Teacher decoder zonder de Student te hertrainen.
• Vergelijking met Pruning: Bridge-enabled distillatie presteert beter dan netwerkpruning (bijv. het prunen van SD 3.5 naar 30% sparsity), wat de neiging heeft om onscherpte en artefacten te produceren.

Betekenis
Het artikel stelt dat Cross-Space Distillatie het potentieel ontsluit om de generatieve capaciteiten van massieve, moderne Teachers over te dragen naar efficiënte, deploybare backbones zonder de kosten van het trainen van nieuwe architecturen vanaf nul. Door de uitlijning van heterogene latente representaties te behandelen als een herbruikbaar interfaceprobleem, maakt dit werk het mogelijk om:

• Ecosysteem-Compatibiliteit: Compacte modellen zoals SD 1.5 kunnen worden geüpgraded met state-of-the-art kwaliteit terwijl ze compatibel blijven met bestaande tools en hardware.
• Kennis-Unificatie: De mogelijkheid om de kennis van meerdere, architecturaal verschillende Teachers (bijv. Flux en SD 3.5) samen te voegen in een enkele, lichtgewicht Student model.
• Modulariteit: De Bridge dient als een plug-and-play component die zowel voor distillatie als voor resolutie-verbetering tijdens inferentie kan worden hergebruikt, waardoor de generatiecapaciteit van de Student wordt ontkoppeld van zijn latente representatie-beperkingen.

De auteurs positioneren dit werk als een stap naar een bredere exploratie van distillatie over heterogene latente ruimtes, toepasbaar buiten text-to-image naar video en multimodale generatie.