Guiding a Diffusion Transformer with the Internal Dynamics of Itself

1. Il Problema

I modelli di diffusione (Diffusion Models), in particolare quelli basati su architetture Transformer (DiT, SiT), hanno dimostrato una capacità eccezionale nel catturare distribuzioni di dati complessi. Tuttavia, esistono due sfide principali:

• Copertura della distribuzione: I modelli tendono a faticare a generare immagini di alta qualità nelle aree a bassa probabilità della distribuzione dei dati. Durante il campionamento, strategie come la Classifier-Free Guidance (CFG) vengono utilizzate per spingere il campione verso aree ad alta probabilità, ma un coefficiente di guida eccessivo porta spesso a risultati semplificati, distorti o con una ridotta diversità.
• Limitazioni delle guide esistenti: Un approccio alternativo, noto come Autoguidance, utilizza una versione "cattiva" (degradata) del modello stesso per guidare la versione "buona". Sebbene efficace, questo metodo richiede strategie di degradazione complesse, addestramenti aggiuntivi o passaggi di campionamento extra, rendendolo poco scalabile per i generatori di immagini su larga scala.

2. Metodologia: Internal Guidance (IG)

Gli autori propongono una strategia semplice ma efficace chiamata Internal Guidance (IG), che sfrutta le dinamiche interne del Diffusion Transformer senza richiedere addestramenti aggiuntivi o passaggi extra di campionamento.

Fase di Addestramento: Supervisione Intermedia

Durante l'addestramento, viene introdotta una supervisione ausiliaria su uno strato intermedio della rete (es. l'8° strato in modelli su larga scala).

• Viene definito un output intermedio $D_i$ e un output finale $D_f$.
• La funzione di perdita totale combina la perdita finale e la perdita intermedia:
  $$L = L_{final} + \lambda L_{inter}$$
  dove $L_{inter} = \|D_i(x_t, t) - x_0\|^2$ e $L_{final} = \|D_f(x_t, t) - x_0\|^2$.
• Questo approccio semplice allevia il problema del vanishing gradient e migliora la convergenza, ottenendo performance comparabili o superiori a metodi di regolarizzazione basati su apprendimento auto-supervisionato complesso.

Fase di Campionamento: Estrapolazione delle Dinamiche Interne

Durante il campionamento, l'IG utilizza la relazione tra l'output dello strato intermedio e quello finale per guidare il processo.

• L'output intermedio $D_i$ funge da "versione debole" o "cattiva" del modello, analogamente all'Autoguidance, ma senza bisogno di addestrare un modello separato.
• La guida viene applicata tramite un'estrapolazione lineare:
  $$D_w(x; c) = D_i(x; c) + w(D_f(x; c) - D_i(x; c))$$
  dove $w$ è il coefficiente di guida.
• Questo meccanismo spinge il risultato finale lontano dalla distribuzione "povera" dello strato intermedio, migliorando la qualità mantenendo la diversità.

Compatibilità e Intervalli di Guida

• IG + CFG: L'IG è altamente compatibile con la CFG. Mentre la CFG tende a ridurre la diversità, l'IG mantiene la diversità. La loro combinazione permette di sopprimere gli outlier (campioni fuori distribuzione) senza sacrificare la varietà.
• Guidance Interval: Gli autori osservano che, a differenza della CFG che funziona meglio a bassi livelli di rumore, l'IG beneficia dell'applicazione in intervalli di rumore medio-alto (es. da 0.3 a 1.0), correggendo la traiettoria nelle fasi iniziali del processo di denoising.

3. Contributi Chiave

1. Strategia Plug-and-Play: L'IG non richiede passaggi di campionamento aggiuntivi né l'addestramento di modelli "cattivi" separati. Può essere integrata direttamente nelle architetture Transformer esistenti con un costo computazionale marginale.
2. Efficienza nell'Addestramento: La semplice introduzione di una perdita di supervisione intermedia accelera la convergenza del modello, superando o eguagliando metodi di regolarizzazione più complessi (come REPA o SRA) con meno risorse.
3. Nuova Strategia di Accelerazione: Basandosi sull'IG, gli autori propongono una nuova funzione di perdita per l'addestramento che incorpora il vettore di guida (differenza tra output finale e intermedio) direttamente nell'obiettivo di addestramento, accelerando ulteriormente la convergenza.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet 256×256 e 512×512, confrontando modelli come SiT-XL/2 e LightningDiT-XL/1.

• Performance su SiT-XL/2:
  • Con IG (senza CFG): FID = 5.31 a soli 80 epoche (superiore a SiT vanilla a 1400 epoche).
  • A 800 epoche: FID = 1.75 (senza CFG) e 1.46 (con CFG).
• Performance su LightningDiT-XL/1 (con tokenizer auto-supervisionato):
  • A 680 epoche: FID = 1.34 (senza CFG).
  • Stato dell'Arte (SOTA): Combinando IG, CFG e Guidance Interval, il modello raggiunge un FID di 1.19, stabilendo un nuovo record per la generazione di immagini su ImageNet 256×256.
• Efficienza: L'IG riduce il FID del 70% rispetto alla baseline REPA con un aumento trascurabile dei parametri e del costo computazionale (FLOPs).

5. Significato e Impatto

Il lavoro di Zhou et al. rappresenta un passo significativo nell'ottimizzazione dei modelli di diffusione su larga scala:

• Semplicità ed Efficacia: Dimostra che strategie di guida sofisticate non richiedono necessariamente architetture complesse o addestramenti multipli, ma possono emergere dalle dinamiche interne del modello stesso.
• Scalabilità: La metodologia si scala bene con l'aumento delle dimensioni del modello, offrendo miglioramenti più significativi su modelli grandi (XL).
• Versatilità: L'IG è applicabile a diverse architetture (DiT, SiT, LightningDiT) e può essere combinata con altre tecniche di guida (CFG, intervalli di rumore) per massimizzare la qualità.
• Implicazioni Future: L'idea di utilizzare le discrepanze tra strati interni per guidare l'addestramento apre la strada a nuove strategie di accelerazione della convergenza, riducendo il tempo e le risorse necessarie per addestrare modelli generativi di alta qualità.

In sintesi, l'Internal Guidance offre un metodo elegante ed efficiente per migliorare sia la qualità della generazione che l'efficienza dell'addestramento, raggiungendo lo stato dell'arte attuale con un approccio minimalista.