Riepilogo Tecnico: Deriva delle Caratteristiche dell'Insegnante (TFD)

Enunciato del Problema

I modelli di diffusione e di adattamento del flusso preaddestrati raggiungono una generazione di immagini ad alta fedeltà, ma soffrono di costi di inferenza elevati, richiedendo spesso decine o centinaia di passaggi in avanti per generare una singola immagine. Sebbene esistano metodi di distillazione per comprimere questi insegnanti multi-step in generatori efficienti a uno o pochi passaggi, gli approcci attuali presentano complessità significative. Metodi come la Distillazione di Adattamento della Distribuzione (DMD) e le sue varianti si basano su reti ausiliarie (ad esempio, reti di punteggio finto, discriminatori avversari) o su programmi di addestramento multi-stadio. Altri approcci, come i Modelli di Coerenza, possono soffrire di addestramento instabile o richiedere prodotti complessi tra Jacobiano e vettori. Inoltre, recenti obiettivi di "Modello in Deriva", che offrono una funzione di perdita diretta di attrazione-repulsione per la generazione a un passaggio, dipendono tipicamente da un estrattore di caratteristiche preaddestrato esterno (ad esempio, modelli auto-supervisionati o latent-MAE) per definire lo spazio di rappresentazione. Ciò introduce un sovraccarico aggiuntivo nell'addestramento e complessità architetturale.

Il problema centrale affrontato è come realizzare una distillazione della diffusione a un passaggio efficiente e di alta qualità utilizzando un singolo obiettivo senza fare affidamento su reti ausiliarie, pipeline di ottimizzazione complesse o estrattori di caratteristiche esterni.

Metodologia: Deriva delle Caratteristiche dell'Insegnante (TFD)

Gli autori propongono la Deriva delle Caratteristiche dell'Insegnante (TFD), un framework che semplifica la distillazione a un passaggio sfruttando le rappresentazioni interne del modello insegnante preaddestrato stesso.

1. Spazio di Rappresentazione Indotto dall'Insegnante

A differenza del Modello in Deriva originale, che richiede un encoder di caratteristiche esterno, la TFD utilizza gli stati nascosti intermedi dell'insegnante preaddestrato congelato come spazio di caratteristiche per calcolare la funzione di perdita di deriva.

• Meccanismo: Per un campione generato $x$ e un campione reale $x^+$, le caratteristiche vengono estratte da specifici livelli $\ell$ della rete insegnante $T$: $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$.
• Funzione di Perdita di Deriva: Il campo di deriva viene calcolato direttamente in questo spazio indotto dall'insegnante. Il generatore $f_\theta$ viene ottimizzato per minimizzare la distanza tra la sua uscita e un "target trasportato congelato" definito dal campo di attrazione-repulsione dello spazio di caratteristiche dell'insegnante. Ciò allinea l'obiettivo di distillazione con il modello che viene distillato, eliminando la necessità di una rete di rappresentazione separata.

2. Estrazione di Caratteristiche con Rumore

Una scoperta empirica chiave è che estrarre caratteristiche da input moderatamente rumorosi produce risultati migliori rispetto all'uso di campioni puliti.

• Prima di inserire un campione nell'insegnante per l'estrazione delle caratteristiche, viene perturbato con il processo di rumorizzazione in avanti dell'insegnante: $x_{tf} = x + \sigma_{tf} \xi$.
• Gli autori ipotizzano che questa perturbazione smussi la geometria delle caratteristiche dell'insegnante, sopprimendo le discrepanze ad alta frequenza e rendendo lo spazio di rappresentazione più adatto al matching a livello di distribuzione. Gli esperimenti indicano che un livello di rumore di $\sigma_{tf} = 0.1$ è ottimale.

3. Regularizzazione Margine-Ancora

Per affrontare il problema del collasso delle modalità (dove il generatore non riesce a coprire regioni diverse supportate dall'insegnante), gli autori introducono una leggera funzione di perdita margine-ancora.

• Ancore: Un insieme di caratteristiche ancora $\{a_i\}$ viene estratto da campioni reali o generati dall'insegnante.
• Meccanismo: La perdita incoraggia i campioni generati ad occupare regioni attorno a queste ancore. Regola adattivamente il "supporto" richiesto per ogni ancora in base alla densità locale della distribuzione delle ancore stesse.
• Obiettivo: Ciò fornisce un segnale di copertura complementare all'obiettivo locale di attrazione-repulsione di deriva, assicurando che il generatore studente copra diverse regioni supportate dall'insegnante senza introdurre nuove reti.

La funzione obiettivo finale combina la funzione di perdita di Deriva delle Caratteristiche dell'Insegnante ($L_{TFD}$) e la funzione di perdita margine-ancora ($L_{anchor}$):
$$L = L_{TFD} + \lambda_{anchor} L_{anchor}$$

Contributi Chiave

1. Framework Semplificato: La TFD propone un framework di distillazione della diffusione a un passaggio che calcola l'obiettivo di deriva direttamente nello spazio di rappresentazione nascosto dell'insegnante preaddestrato, eliminando la necessità di una rete di rappresentazione aggiuntiva.
2. Estrazione di Caratteristiche con Rumore: Il lavoro dimostra che estrarre le caratteristiche dell'insegnante da input moderatamente rumorosi migliora la geometria delle caratteristiche per la distillazione, portando a punteggi FID migliori.
3. Mitigazione del Collasso delle Modalità: Viene introdotta una leggera funzione di perdita margine-ancora per incoraggiare la diversità dei campioni e la copertura delle regioni di caratteristiche supportate dall'insegnante, affrontando una modalità di fallimento comune nella distillazione.
4. Prestazioni: Il metodo raggiunge qualità e diversità delle immagini competitive con un singolo passaggio in avanti, semplificando sostanzialmente l'intera pipeline di distillazione rispetto alle baseline più avanzate.

Risultati Sperimentali

Gli autori hanno valutato la TFD sulla generazione ImageNet-64×64 (condizionata per classe) e SDXL (da testo a immagine).

• ImageNet-64×64:

  • La TFD ha raggiunto un FID di 1.58 con un singolo passaggio in avanti dopo 200k iterazioni di addestramento.
  • Questo risultato è competitivo con DMD2 (FID 1.51) e supera significativamente altre baseline a un passaggio come DMD (2.62), EMD (2.20) e CTM (1.92).
  • Convergenza: La TFD ha raggiunto un FID $\le$ 10 in 1.0k aggiornamenti, mentre la baseline DMD2 ne ha richiesti 5.5k (5.5 volte più veloce). Ha raggiunto un FID $\le$ 3 in 12.5k aggiornamenti rispetto ai 24.0k di DMD2.

• SDXL da Testo a Immagine:

  • Sul dataset COCO 2014, la TFD ha raggiunto un FID di 18.39 con un passaggio in avanti, superando altre varianti di SDXL distillate (ad esempio, DMD2 a 19.01, SDXL-Turbo a 24.57).
  • Ha mantenuto un punteggio CLIP competitivo di 0.332, indicando una forte allineamento testo-immagine.
  • I risultati qualitativi hanno mostrato che la TFD preserva una struttura coerente, dettagli visivi e allineamento ai prompt comparabili alle baseline distillate a 4 passaggi e all'insegnante a 50 passaggi, nonostante l'uso di un solo passaggio.

Significato e Affermazioni

Il paper afferma che la TFD offre un'alternativa più semplice ed efficiente ai framework di distillazione esistenti. Utilizzando il modello insegnante stesso come estrattore di caratteristiche, la TFD evita la complessità dell'addestramento di reti ausiliarie o la progettazione di pipeline multi-stadio. L'introduzione dell'estrazione di caratteristiche con rumore e della regolarizzazione margine-ancora affronta sfide specifiche nella geometria della rappresentazione e nella copertura delle modalità, rispettivamente.

Gli autori posizionano la TFD come un metodo che raggiunge una generazione efficiente a un passaggio con qualità e diversità delle immagini competitive, dimostrando che un singolo obiettivo nello spazio di rappresentazione nativo dell'insegnante è sufficiente per una distillazione ad alta fedeltà. Il lavoro suggerisce che il concetto di "deriva" può essere realizzato efficacemente senza encoder di caratteristiche esterni, purché gli stati interni dell'insegnante siano utilizzati correttamente.

Limitazioni: Gli autori notano che il lavoro attuale si concentra su insegnanti basati su U-Net e non ha ancora esplorato architetture Diffusion Transformer (DiT). Inoltre, la selezione dei livelli di caratteristiche e dei livelli di rumore attualmente manca di guida teorica, e il metodo non è stato testato su dataset su scala più ampia o impostazioni a risoluzione più elevata oltre agli esperimenti attuali.