Uni-DAD: Unified Distillation and Adaptation of Diffusion Models for Few-step Few-shot Image Generation

1. Il Problema

I modelli di diffusione (Diffusion Models - DM) hanno raggiunto lo stato dell'arte nella generazione di immagini di alta qualità e diversità. Tuttavia, presentano due limiti critici quando applicati a nuovi domini o soggetti specifici:

1. Costo Computazionale: Il processo di campionamento richiede molti passaggi di denoising (iterazioni temporali), rendendo la generazione lenta e inadatta ad applicazioni in tempo reale.
2. Adattamento a Few-Shot: Adattare un modello pre-addestrato a un nuovo dominio con pochi esempi (few-shot, es. 4-10 immagini) è complesso. Le pipeline attuali richiedono solitamente un approccio in due fasi (es. Distill-then-Adapt o Adapt-then-Distill):
   • Distill-then-Adapt: Un modello studente veloce viene adattato successivamente. Spesso porta a un "sovra-smussamento" (over-smoothing) e perdita di dettagli nei domini target.
   • Adapt-then-Distill: Il modello viene prima adattato e poi distillato. Questo rischia l'overfitting sul piccolo set di dati target e la perdita di informazioni diversificate del dominio sorgente.
   • Entrambi gli approcci aggiungono complessità progettuale e spesso degradano la qualità o la diversità delle immagini generate.

2. Metodologia: Uni-DAD

Gli autori propongono Uni-DAD, una pipeline a stadio singolo che unifica la distillazione e l'adattamento in un unico processo di training. L'obiettivo è comprimere un modello "teacher" di alta qualità (addestrato su un dominio sorgente vasto) in un modello "studente" veloce (1-4 passaggi) che sia anche adattato a un dominio target con pochi esempi.

Il metodo combina due segnali di training complementari:

A. Obiettivo di Distillazione a Doppio Dominio (Dual-Domain DMD)

Invece di allineare lo studente solo al teacher sorgente, Uni-DAD utilizza una funzione di perdita basata sulla Divergenza di Kullback-Leibler (KL) che guida lo studente verso una distribuzione comune tra:

1. Teacher Sorgente ($\epsilon_{src}$): Un modello congelato addestrato su un dataset vasto (es. FFHQ). Preserva la diversità e le conoscenze trasferibili (es. pose, sfondi, espressioni).
2. Teacher Target ($\epsilon_{trg}$): Un modello online (opzionale) addestrato sui pochi esempi del dominio target. Facilita l'adattamento a domini strutturalmente distanti.
   La perdita è una combinazione ponderata:
   $$\nabla_\theta L_{DMD}^{trg+src} = (1 - \alpha)\nabla_\theta L_{DMD}^{src} + \alpha \nabla_\theta L_{DMD}^{trg}$$
   dove $\alpha$ controlla il bilanciamento tra preservazione della sorgente e adattamento al target.

B. GAN Multi-Head

Per garantire il realismo nel dominio target e stabilizzare l'addestramento in regime few-shot, viene introdotto un GAN:

• Discriminatore Multi-Head: Utilizza i blocchi encoder del teacher "fake" (che traccia la distribuzione dello studente) per estrarre caratteristiche a diverse scale. Testa il realismo sia a livello globale che locale.
• Teacher Finto ($\epsilon_{fk}$): Un modello che viene aggiornato continuamente per imitare le uscite dello studente, fornendo segnali negativi aggiornati al discriminatore.
  Questo componente riduce l'overfitting e previene il collasso delle modalità (mode collapse), tipici problemi quando si addestra su pochi dati.

C. Processo di Training

Il training alterna l'aggiornamento di tre componenti:

1. Lo Studente (G): Minimizza la perdita DMD a doppio dominio + perdita GAN generatore.
2. Il Teacher Finto e il Discriminatore: Minimizzano la perdita MSE e la perdita GAN discriminatore.
3. Il Teacher Target (opzionale): Viene addestrato sui dati target per migliorare l'adattamento strutturale.

3. Contributi Chiave

1. Pipeline a Stadio Singolo: È il primo metodo che unisce distillazione e adattamento in un unico ciclo di training, eliminando la necessità di pipeline sequenziali complesse.
2. Obiettivi Ibridi: L'introduzione della DMD a doppio dominio e del GAN multi-head permette di mantenere la diversità del dominio sorgente mentre si affina il realismo nel dominio target con dati scarsi.
3. Checkpoint-Agnostic: Il metodo è flessibile: può utilizzare un modello target pre-adattato o un modello sorgente pre-distillato come inizializzazione senza modificare il loop di training.
4. Efficienza e Qualità: Dimostra di generare immagini di alta qualità in pochi passaggi (1-4 NFE - Neural Function Evaluations) mantenendo un'elevata diversità.

4. Risultati Sperimentali

Il metodo è stato valutato su due benchmark principali:

• FSIG (Few-Shot Image Generation): Adattamento di un modello DDPM guidato a domini come "Babies", "Sunglasses", "MetFaces" e "Cats" (10-shot).
• SDP (Subject-Driven Personalization): Adattamento di Stable Diffusion v1.5 per la generazione di soggetti specifici (es. DreamBooth) con 4-6 immagini.

Punti di forza dei risultati:

• Qualità: Uni-DAD ottiene punteggi FID (Fréchet Inception Distance) migliori o comparabili rispetto ai metodi di adattamento non distillati (come CRDI o DreamBooth standard) e supera spesso le pipeline a due stadi.
• Velocità: Riduce drasticamente i passaggi di campionamento (da 25-1000 a 1-4) mantenendo alta la qualità.
• Diversità: A differenza di altri metodi distillati che tendono a perdere diversità, Uni-DAD mantiene un alto Intra-LPIPS (diversità interna) grazie alla componente DMD sorgente e al GAN.
• Adattamento a Domini Distanti: L'uso del teacher target ($\epsilon_{trg}$) permette di adattarsi efficacemente anche a domini strutturalmente molto diversi dalla sorgente (es. da volti umani a gatti), dove i metodi puri di distillazione fallirebbero.

5. Significato e Impatto

Uni-DAD rappresenta un passo significativo verso l'uso pratico dei modelli di diffusione in scenari reali che richiedono:

• Tempo reale: Grazie alla generazione in pochi passaggi.
• Personalizzazione efficiente: Capacità di adattarsi a nuovi soggetti o stili con pochissimi dati (few-shot).
• Semplificazione: Unifica due processi complessi in uno, riducendo la complessità ingegneristica e i costi computazionali di training rispetto alle pipeline a due stadi.

Il lavoro suggerisce che distillazione e adattamento non devono essere trattati come fasi separate, aprendo la strada a generatori di immagini personalizzati, veloci e di alta qualità anche in condizioni di dati limitati e shift di dominio significativi.