DASH: Dual-Branch Score Distillation for Guidance-Calibrated Compact Diffusion Models

Sintesi Tecnica: DASH: Dual-Branch Score Distillation per la Guida Calibrata di Modelli di Diffusione Compatti

Problema
I modelli di diffusione condizionati alla classe che utilizzano la Classifier-Free Guidance (CFG) si affidano a una singola rete per produrre sia le predizioni condizionali ($\epsilon_y$) che quelle incondizionali ($\epsilon_\emptyset$). Il meccanismo di guida amplifica il divario $\Delta = \epsilon_y - \epsilon_\emptyset$ per guidare la generazione. Tuttavia, la compressione di questi modelli tramite la riduzione dei parametri introduce un fallimento strutturale critico negli obiettivi di distillazione esistenti.

La distillazione standard supervisiona solo la predizione composta guidata ($\tilde{\epsilon}_S \approx \tilde{\epsilon}_T$) o il ramo condizionale da solo. Ciò crea un problema di ottimizzazione sottodeterminato: molteplici decomposizioni dei rami dello studente possono soddisfare la funzione di perdita, incluse soluzioni degenerate in cui il divario di guida collassa ($\Delta_S \to 0$). In questo stato, le predizioni condizionali e incondizionali convergono, rendendo inefficace il peso di guida $w$ durante l'inferenza, anche se la perdita di addestramento composita è minimizzata. Inoltre, i metodi esistenti non riescono a trasferire il curriculum di importanza per timestep (importance weighting) appreso dal docente, costringendo lo studente a riapprendere questo schema entro budget di distillazione limitati, il che spesso corrompe la traiettoria ottimale di addestramento.

Metodologia: Framework DASH
Il documento propone DASH (Dual-Branch Score Distillation), un framework progettato per risolvere l'sottodeterminazione e preservare la calibrazione della guida attraverso due meccanismi primari:

1. Supervisione Indipendente a Doppio Ramo (Dual-Branch Independent Supervision):
   Inveve di far corrispondere l'output composito, DASH supervisiona indipendentemente entrambi i rami dello score utilizzando termini di perdita distinti:

   • Perdita di Imitazione ($L_{im}$): fa corrispondere il ramo condizionale ($\epsilon^y_S \approx \epsilon^y_T$).
   • Perdita Incondizionale ($L_{un}$): fa corrispondere il ramo incondizionale ($\epsilon^\emptyset_S \approx \epsilon^\emptyset_T$).
   • Perdita di Ancoraggio ($L_{an}$): un regolarizzatore leggero che ancora il ramo condizionale al rumore ground-truth ($\epsilon$) per prevenire la deriva durante le prime fasi dell'addestramento.

   Imponendo vincoli su entrambi i rami simultaneamente, il framework specifica univocamente gli output target per ogni campione, eliminando lo spazio delle soluzioni ammissibili che permette il collasso del divario di guida. Gli autori definiscono un rapporto di divario di guida $\rho = E[\|\Delta_S\|] / E[\|\Delta_T\|]$ per monitorare la calibrazione, dove $\rho \approx 1$ indica una corretta preservazione.

2. Trasferimento TIRT (Timestep-Importance Rebalanced Training Transfer):
   Il modello docente utilizza TIRT, un'estensione del weighting Min-SNR che apprende fattori di scala per timestep ($\lambda_t$) per concentrare l'addestramento sulle regioni di media traiettoria più difficili dal punto di vista del contenuto. DASH introduce il TIRT Transfer, che copia direttamente il curriculum di importanza per timestep convergente del docente ($\hat{\omega}_T$) nello studente come prior congelato. Ciò impedisce allo studente di corrompere il curriculum con gradienti di distillazione che premiano la facilità di corrispondenza della predizione piuttosto che la difficoltà del contenuto.

Contributi Chiave

• Risoluzione del Sottodeterminazione Strutturale: Il documento identifica e risolve il collasso degenerato del divario di guida nei modelli di diffusione compressi nei parametri. Dimostra che la supervisione incondizionale non è solo benefica ma essenziale, rappresentando oltre il 60% del guadagno totale di distillazione.
• Trasferimento del Curriculum: L'introduzione del TIRT Transfer, che inizializza lo studente con il peso temporale appreso dal docente come un prior congelato, eliminando la necessità per lo studente di riapprendere il curriculum.
• Framework a Doppio Ramo: Una pipeline di distillazione unificata che vincola indipendentemente i rami condizionale e incondizionale, validata come empiricamente essenziale per mantenere la fedeltà della guida sotto compressione.

Risultati Sperimentali
Gli esperimenti sono stati condotti su CIFAR-10 e CIFAR-100 utilizzando un modello docente (35.8M parametri) e uno studente compresso (6.1M parametri, compressione di 5.9×).

• Performance: Lo studente DASH raggiunge un FID di 8.87 (CIFAR-10) e 10.47 (CIFAR-100) con campionamento DDIM a 50 step. Questo è entro 4 punti FID dal docente (FID 5.47 / 6.80) e supera significativamente l'addestramento da zero (FID 20.47 / 26.14) e le baseline a singolo vincolo.
• Calibrazione della Guida: DASH mantiene un rapporto di divario di guida $\rho$ di 0.91 (CIFAR-10) e 0.89 (CIFAR-100), con un coseno direzionale $\cos(\Delta)$ vicino a 0.94. Al contrario, le baseline di distillazione composita soffrono di un parziale collasso ($\rho \approx 0.68$), e la supervisione solo condizionale porta a un quasi totale collasso ($\rho \approx 0.09$), degradando le performance al di sotto dell'addestramento da zero.
• Studi di Ablazione:
  • Rimuovere $L_{un}$ (perdita incondizionale) causa un calo di 7.77 punti FID su CIFAR-10 e fa collassare $\rho$ a 0.61, confermando la sua natura di componente dominante.
  • Rimuovere il TIRT Transfer (sostituendolo con un Min-SNR fisso) comporta un calo di 2.60 punti FID, validando la necessità di trasferire il curriculum appreso.
  • Il rapporto di performance cross-dataset (1.18× per DASH rispetto a 1.51× per le baseline) conferma una robusta generalizzazione attraverso diverse complessità di etichetta.

Significato e Rivendicazioni
Il documento afferma che DASH stabilisce un nuovo standard per la diffusione a parametri efficienti affrontando un difetto strutturale nei processi di distillazione esistenti. Gli autori asseriscono che:

1. La supervisione incondizionale è non negoziabile: Non può essere sostituita da obiettivi compositi o condizionali; senza di essa, la CFG diventa inefficace sotto compressione.
2. Il trasferimento del curriculum è distinto dalla distillazione delle feature: Trasferire lo schema di addestramento (TIRT) come un prior congelato è un contributo originale che impedisce la corruzione del weighting adattivo ai dati appreso dal docente.
3. Efficienza senza Compromessi: Il framework consente una significativa riduzione dei parametri (5.9×) preservando il meccanismo di guida calibrato, permettendo una generazione di alta qualità in contesti con risorse limitate senza aumentare i passi di inferenza.

Il lavoro si posiziona all'intersezione tra compressione dei parametri e calibrazione a livello di ramo, offrendo una soluzione che è ortogonale ai metodi di riduzione dei passi e composibile con essi.