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

Sintesi Tecnica: Distillazione Cross-Space (Cross-Space Distillation)

Problema
I moderni modelli di diffusione a singolo step ottengono una generazione di alta qualità attraverso la distillazione dei timestep basata sulla distribuzione. Tuttavia, questi metodi si affidano a un vincolo restrittivo di Spazio Condiviso (Shared-Space): il Teacher e lo Student devono abitare lo stesso spazio latente, condividendo la stessa risoluzione latente e la stessa parametrizzazione VAE. Questa ipotesi impedisce il trasferimento di conoscenza da Teacher moderni ad alta capacità (es. Stable Diffusion 3.5, Flux, Kolors) che operano ad alte risoluzioni (es. 1024×1024) e utilizzano VAE avanzati, verso Student compatti e adatti alla distribuzione (es. Stable Diffusion 1.5) che operano a risoluzioni inferiori (es. 512×512) con VAE differenti.

Gli obiettivi di distillazione standard falliscono in questo scenario perché i tensori latenti ($z_S$ e $z_T$) sono intrinsecamente incompatibili a causa di:

1. Cross-Resolution: Griglie spaziali disallineate (es. $h \times w$ vs. $H \times W$).
2. Cross-VAE: Diverse parametrizzazioni dell'autoencoder che risultano in spazi latenti non allineati.
3. Cross-Architecture/Mechanism: Differenze tra backbone UNet e DiT, o tra formulazioni di diffusione standard e di flow matching.

Le alternative esistenti, come il pruning delle reti o la costruzione di nuovi backbone compatti da zero, affrontano sfide quali soglie di sparsità rigide, requisiti di riaddestramento estesi e la perdita di compatibilità con l'ecosistema.

Metodologia: Distillazione Cross-Space e il Bridge
Gli autori formalizzano questo scenario come Distillazione Cross-Space e introducono Bridge ($B_\phi$), un modulo leggero progettato per mappare i latenti dello Student nello spazio latente del Teacher senza modificare il backbone dello Student.

• Architettura: Il Bridge è fattorizzato in due stadi per gestire l'incompatibilità in modo efficiente:

  1. Prior Spaziale (Frozen): I primi $n$ blocchi di decodifica del decoder del VAE dello Student ($D_S^{(n)}$) vengono mantenuti congelati. Questi blocchi agiscono come un upsampler spaziale pre-addestrato, espandendo il latente dello Student da una bassa risoluzione alla risoluzione elevata del Teacher ($H \times W$). Ciò sfrutta la conoscenza spaziale esistente dello Student senza aggiungere parametri addestrabili per l'upsampling.
  2. Proiettore Apprendibile ($g_\phi$): Una rete compatta e apprendibile (implementata come SwinIR) mappa le feature upsampled dallo spazio delle feature dello Student allo spazio della rappresentazione latente del Teacher, producendo un latente compatibile con il Teacher $\hat{z}_T$.

• Obiettivi di Addestramento: Il Bridge viene addestrato utilizzando un dataset sintetico di latenti accoppiati $(z_S, z_T)$ derivati dalla codifica della stessa immagine da parte di entrambi i VAE. L'ottimizzazione minimizza una somma pesata di due perdite:

  1. Ricostruzione Latente ($L_{rec}$): Una perdita $\ell_1$ tra il latente predetto $\hat{z}_T$ e il ground-truth del Teacher $z_T$. Gli autori notano che la $\ell_1$ fornisce una migliore stabilità e preservazione della struttura rispetto alla $\ell_2$ in questo contesto.
  2. Fedeltà dell'Attenzione ($L_{attn}$): Una perdita di divergenza reverse-KL che allinea le distribuzioni di self-attention del denoiser del Teacher (congelato) quando alimentato con $\hat{z}_T$ rispetto a $z_T$. Ciò assicura che il latente mappato induca le stesse risposte strutturali interne (dipendenze a lungo raggio) del vero latente del Teacher, preservando i dettagli fini che le perdite pixel-wise potrebbero trascurare.

• Distillazione e Inferenza: Una volta addestrato, il Bridge viene collegato allo Student. Durante la distillazione, l'output dello Student viene passato attraverso il Bridge per generare $\hat{z}_T$, permettendo l'applicazione diretta degli obiettivi di distillazione basati sulla distribuzione (es. Variational Score Distillation, Adversarial Distillation) contro il Teacher. Fondamentalmente, il backbone dello Student rimane invariato.

Contributi Chiave

1. Formalizzazione della Distillazione Cross-Space: Il documento definisce un nuovo regime in cui la distillazione avviene tra spazi latenti di risoluzione e VAE disallineati, affrontando una lacuna critica nel trasferimento di conoscenza da modelli Teacher moderni ad alta risoluzione verso modelli compatti legacy.
2. Il Modulo Bridge: Un'interfaccia modulare e leggera che consente obiettivi di distillazione standard senza alterare l'architettura dello Student. Combina un prior spaziale congelato con un proiettore apprendibile e utilizza un obiettivo di Fedeltà dell'Attenzione per un allineamento stabile.
3. Efficacia Pratica: Il metodo permette a Student compatti (es. SD 1.5) di ereditare i prior visivi di diversi Teacher moderni (SDXL, Kolors, PixArt-Σ, Flux, SD 3.5) mantenendo l'inferenza a singolo step, la bassa latenza e l'ampia compatibilità con l'ecosistema.
4. Capacità di Model Merging: Poiché il metodo distilla diversi Teacher in Student con lo stesso backbone e spazio latente, i checkpoint risultanti possono essere facilmente uniti tramite la media dei parametri per creare un singolo modello che aggrega la conoscenza da molteplici fonti Teacher.

Risultati
Gli esperimenti dimostrano miglioramenti sostanziali nella qualità della generazione per gli Student compatti a singolo step:

• Guadagni Quantitativi: Distillare SD 3.5 Medium (Teacher) in SD 1.5 (Student) tramite Bridge migliora il punteggio HPSv3 da 5.37 a 9.42 e l'ImageReward da -0.29 a 0.62. L'unione di checkpoint provenienti da cinque diversi Teacher aumenta ulteriormente l'HPSv3 di SD 1.5 a 10.53.
• Miglioramenti Qualitativi: I confronti visivi mostrano che gli studenti abilitati da Bridge producono immagini con texture più chiare, maggiore coerenza strutturale e meno artefatti rispetto ai baseline inizializzati o ai modelli sottoposti a pruning.
• Efficienza: Il Bridge aggiunge un overhead trascurabile (circa 5M di parametri) e consente l'inferenza a singolo step. Permette inoltre l'aggiornamento della risoluzione al tempo di inferenza, dove un campione SD 1.5 da 512×512 può essere mappato a 1024×1024 tramite il Bridge e decodificato con il decoder del Teacher senza riaddestrare lo Student.
• Confronto con il Pruning: La distillazione abilitata dal Bridge supera il network pruning (es. ridurre SD 3.5 al 30% di sparsità), che tende a produrre sfocatura e artefatti.

Significatività
Il documento sostiene che la Distillazione Cross-Space sblocca il potenziale di trasferire le capacità generative di massicci Teacher moderni in backbone efficienti e pronti per la distribuzione, senza il costo di addestrare nuove architetture da zero. Trattando l'allineamento di rappresentazioni latenti eterogenee come un problema di interfaccia riutilizzabile, il lavoro abilita:

• Compatibilità con l'Ecosistema: Modelli compatti come SD 1.5 possono essere aggiornati con una qualità allo stato dell'arte pur mantenendo la compatibilità con gli strumenti e l'hardware esistenti.
• Unificazione della Conoscenza: La capacità di unire la conoscenza di molteplici Teacher architettonicamente distinti (es. Flux e SD 3.5) in un singolo modello leggero.
• Modularità: Il Bridge funge da componente plug-and-play che può essere riutilizzato sia per la distillazione che per l'aggiornamento della risoluzione al tempo di inferenza, disaccoppiando la capacità generativa dello Student dai suoi vincoli di rappresentazione latente.

Gli autori pongono questo lavoro come un passo verso un'esplorazione più ampia della distillazione attraverso spazi latenti eterogenei, applicabile oltre il text-to-image verso la generazione video e multimodale.