Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?

1. Il Problema

I modelli di diffusione basati su Transformer (DiT) hanno dimostrato prestazioni eccezionali nella generazione di immagini, video e applicazioni multi-modali. Tuttavia, il loro addestramento rimane computazionalmente costoso e soffre di una convergenza lenta.
Per accelerare l'addestramento, lavori recenti come REPA hanno introdotto l'uso di feature semantiche esterne (ad esempio, estratte da encoder visivi pre-addestrati come DINO) per guidare i blocchi più superficiali del DiT. Sebbene efficace, questo approccio introduce dipendenze esterne significative: richiede il pre-addestramento di grandi modelli su dati esterni e aumenta la complessità del sistema.
La domanda centrale sollevata dagli autori è: è indispensabile la guida tramite feature esterne per accelerare l'addestramento dei DiT?
L'analisi preliminare mostra che i metodi "auto-contenuti" (che usano feature interne) esistenti (come SRA o LayerSync) hanno prestazioni inferiori rispetto a REPA perché le feature interne generate nelle fasi iniziali sono spesso rumorose, strutturalmente instabili e semanticamente poco discriminative.

2. Metodologia: Self-Transcendence

Gli autori propongono Self-Transcendence, un framework di addestramento completamente auto-guidato che supera le prestazioni dei metodi basati su feature esterne senza richiedere dati o modelli esterni.

Il metodo si basa sull'idea che le feature interne del DiT possano essere rese sufficientemente "pulite" e "discriminative" attraverso un processo a due stadi:

Requisiti delle Feature di Guida

Per essere efficaci, le feature di guida interne devono soddisfare due criteri:

1. Struttura Pulita: Aiutare i blocchi superficiali a separare il segnale dal rumore.
2. Discriminatività Semantica: Permettere ai blocchi superficiali di apprendere rappresentazioni informative.

Pipeline a Due Stadi

1. Fase 1: Guida Strutturale VAE (VAE Structure Guidance)

   • Invece di affidarsi a feature interne immature, il modello viene inizialmente guidato dalle feature latenti pulite del VAE (un componente nativo dei modelli di diffusione latente).
   • Viene introdotto un loss auxiliary ($L_{VAE-guide}$) che allinea le feature intermedie del DiT (passate attraverso un leggero MLP) con le feature latenti del VAE.
   • Questo fornisce una struttura stabile fin dall'inizio, aiutando i blocchi superficiali a distinguere il segnale dal rumore.
   • Dopo un breve periodo di "warm-up" (es. 40 epoche), il modello viene congelato e funge da "insegnante" fisso.

2. Fase 2: Allineamento delle Rappresentazioni Auto-Guidate (Self-guided Representation Alignment)

   • Per migliorare la ricchezza semantica (che il VAE da solo non possiede), il metodo estrae le feature dai blocchi più profondi del modello parzialmente addestrato.
   • Viene applicata una tecnica ispirata alla Classifier-Free Guidance (CFG), ma estesa allo spazio delle feature. Si combinano le feature condizionate ($f_c$) e incondizionate ($f_u$) dello stesso layer utilizzando una scala di guida ($\omega$):
     $$f_g = f_u + \omega \cdot (f_c - f_u)$$
   • Questa operazione ($f_g$) potenzia la discriminatività semantica delle feature interne.
   • Queste feature arricchite vengono poi utilizzate come segnali di supervisione per guidare l'addestramento di un nuovo modello DiT da zero, tramite un loss di allineamento ($L_{guide}$).

3. Contributi Chiave

• Framework Auto-Guidato: Introduzione di Self-Transcendence, il primo metodo che supera le prestazioni di REPA (basato su DINO) utilizzando esclusivamente feature interne del modello.
• Pipeline Ibrida a Due Stadi: Una strategia innovativa che combina la stabilità strutturale delle feature VAE (Fase 1) con la ricchezza semantica ottenuta tramite CFG sulle feature interne (Fase 2).
• Superamento delle Dipendenze Esterne: Dimostrazione che i DiT hanno il potenziale intrinseco per guidare il proprio addestramento, eliminando la necessità di encoder visivi esterni pre-addestrati e dei relativi costi computazionali.
• Analisi delle Feature: Identificazione che il collo di bottiglia nella convergenza dei DiT risiede nell'apprendimento di feature pulite e semanticamente ricche nei blocchi superficiali, che vengono risolti efficacemente con questo approccio.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet (256x256 e 512x512) e su task Text-to-Image, utilizzando backbones come SiT e LightningDiT.

• Qualità della Generazione e Convergenza:
  • Su SiT-XL/2 a 256x256, Self-Transcendence raggiunge un FID di 7.51 in soli 80 epoche, superando LayerSync (8.80 FID in 200 epoche) e REPA (7.90 FID in 80 epoche).
  • Su LightningDiT-XL/1, il metodo ottiene un FID di 3.55 in 64 epoche, superando REPA (4.09).
  • Nel task Text-to-Image (MS-COCO), il metodo supera REPA con un FID di 4.56 contro 4.90.
• Efficienza:
  • Il costo computazionale aggiuntivo è minimo (addestramento di un modello guida per poche epoche).
  • Il tempo totale di addestramento è paragonabile al modello vanilla e inferiore a REPA (che richiede l'uso di DINOv2).
  • L'uso della memoria GPU è inferiore rispetto a REPA, poiché non richiede il caricamento di encoder esterni pesanti durante l'intera fase di training.
• Scalabilità: I guadagni di performance aumentano all'aumentare della dimensione del modello (es. riduzione relativa del FID del 57.4% su SiT-XL/2).

5. Significato e Impatto

Questo lavoro è significativo perché ribalta la convinzione comune secondo cui l'accelerazione dell'addestramento dei DiT richieda necessariamente supervisione esterna.

• Autonomia: Dimostra che l'architettura DiT possiede un potenziale intrinseco per l'auto-miglioramento, rendendo i flussi di lavoro più semplici e accessibili.
• Efficienza dei Costi: Elimina la necessità di pre-addestrare o mantenere grandi modelli esterni (come DINO), riducendo i costi energetici e computazionali.
• Nuova Direzione: Apre la strada a future ricerche sull'auto-supervisione e sull'ottimizzazione interna delle rappresentazioni nei modelli generativi, suggerendo che la "guida" può essere appresa dinamicamente dal modello stesso piuttosto che importata.

In sintesi, Self-Transcendence offre una soluzione elegante ed efficace per accelerare l'addestramento dei Transformer di diffusione, raggiungendo e superando lo stato dell'arte senza dipendenze esterne.