NanoFLUX: Distillation-Driven Compression of Large Text-to-Image Generation Models for Mobile Devices

Sintesi Tecnica: NanoFLUX

Problematica

I modelli di diffusione Text-to-Image (T2I) su larga scala, come la famiglia FLUX (che arriva fino a 17 miliardi di parametri) e Qwen-Image (20 miliardi di parametri), hanno raggiunto una qualità visiva allo stato dell'arte. Tuttavia, la loro dimensione massiccia crea una barriera significativa per l'implementazione su dispositivi (on-device), richiedendo infrastrutture GPU costose o inferenza basata su cloud. Le soluzioni esistenti per dispositivi mobili spesso si affidano a modelli più piccoli che mancano di una qualità comparabile o impiegano tecniche di compressione (pruning, quantizzazione, riduzione dei passi) che non affrontano completamente la ridondanza architettonica dei grandi Diffusion Transformer (DiT) e dei loro associati encoder testuali. La sfida consiste nel comprimere questi modelli massicci per dispositivi mobili preservando la qualità di generazione e minimizzando la latenza di inferenza.

Metodologia: NanoFLus

Gli autori introducono NanoFLUX, un modello T2I da 2,4 miliardi di parametri distillato dal docente FLUX.1-Schnell da 17 miliardi. La compressione è ottenuta attraverso una pipeline progressiva guidata dalla distillazione che mira ai due componenti più pesanti: il Diffusion Transformer (DiT) da 12 miliardi e l'encoder testuale T5-XXL da 5 miliardi.

1. Compressione del Transformer (12B $\to$ 2B)

Gli autori impiegano una strategia di compressione progressiva in quattro fasi per ridurre le dimensioni del DiT:

• Fase C1: Pruning delle Teste di Attenzione (12B $\to$ 5B):

  • Analisi: La Decomposizione ai Valori Singolari (SVD) sulle uscite delle teste di attenzione ha rivelato che non tutte le teste sono indipendenti; molte contengono informazioni ridondanti.
  • Azione: Il numero di teste di attenzione ($H$) è stato ridotto uniformemente da 24 a 16.
  • Distillazione: È stata introdotta una perdita a livello di feature per garantire che le teste rimanenti preservino collettivamente l'informazione delle teste rimosse. Lo studente è stato inizializzato con le prime 16 teste del docente.

• Fase C2: Riduzione della Dimensione della Testa (5B $\to$ 3B):

  • Analisi: La SVD è stata applicata nuovamente alle feature del modello da 5B.
  • Azione: La dimensione della testa ($d_H$) è stata ridotta da 128 a 96, mantenendo i principali componenti singolari. Ciò ha prodotto un modello da 3 miliardi di parametri.

• Fase C3: Pruning della Profondità tramite Fusione dei Blocchi (3B $\to$ 2,5B):

  • Analisi: La similarità del coseno input-output è stata calcolata per i blocchi del transformer. I blocchi Single-Stream (SS) hanno mostrato un'alta similarità (>0,85), indicando una ridondanza funzionale.
  • Azione: Invece di eliminare questi blocchi, gli autori hanno fuso una catena di 15 blocchi SS ridondanti in un singolo blocco mediando i loro parametri. Ciò ha ridotto la profondità del transformer da 38 a 24 blocchi.

• Fase C4: Normalizzazione Statica del Layer (2,5B $\to$ 1,8B):

  • Analisi: I coefficienti di Adaptive Layer Normalization (AdaLN) predetti dal modello mostravano una bassa varianza tra i campioni, suggerendo che il condizionamento dinamico dalle proiezioni CLIP fosse ridondante.
  • Azione: I coefficienti dinamici AdaLN sono stati sostituiti con medie precalcolate statiche (Static-LN). Questo ha rimosso circa 0,7 miliardi di parametri. Il modello è stato sottoposto a un breve fine-tuning con coefficienti congelati per recuperare la qualità.

2. Downsampling Progressivo dei Token

Per affrontare la complessità quadratica dell'auto-attenzione senza sacrificare il dettaglio:

• Meccanismo: Un modulo di downsampling ($D$) basato su ResNet e un modulo di upsampling ($U$) sono stati inseriti nel transformer.
• Strategia Ibrida: Il modello opera su uno schema di risoluzione ibrido. I primi timestep di diffusione (che catturano la struttura globale) elaborano token sottocampionati (riducendo il numero di token di 4 volte), mentre i timestep successivi (che raffinano i dettagli) elaborano i token a risoluzione completa.
• Addestramento: Una strategia di "Progressive Token Downsampling" (PTD) è stata utilizzata per stabilizzare l'addestramento. I moduli sono stati inseriti e addestrati incrementalmente, spostando il punto di downsampling più avanti nella rete nel corso delle iterazioni.
• Risultato: 23 su 43 blocchi del transformer operano su token sottocampionati, riducendo significativamente la latenza pur mantenendo l'elaborazione ad alta risoluzione dove necessario.

3. Distillazione dell'Encoder Testuale (5B $\to$ 330M)

• Obiettivo: Distillare l'encoder T5-XXL da 5 miliardi nel modello T5-Large da 330 milioni.
• Metodo: Un processo di distillazione in due fasi:
  1. Corrispondenza degli Embedding: Un MLP viene aggiunto allo studente per far corrispondere la dimensione dell'output del docente, minimizzando l'MSE tra gli embedding dei prompt.
  2. Distillazione per Blocco: Lo studente è addestrato per corrispondere agli stati nascosti del prompt del docente in specifici layer del transformer. Fondamentalmente, il metodo sfrutta segnali visivi dai primi layer del denoiser (tramite joint attention) per guidare l'encoder testuale, mitigando l'accumulo di errori.
  3. Stabilità: I gradienti vengono staccati per i timestep ad alto rumore ($t < \hat{t}$) per evitare una supervisione instabile, concentrando la distillazione sui timestep successivi e sui primi tre layer del transformer.

Risultati Chiave

• Dimensione del Modello: Il modello finale NanoFLUX è di 2,4 miliardi di parametri (2B DiT + 330M Text Encoder), una riduzione di 7$\times$ rispetto al docente FLUX.1-Schnell da 17 miliardi.
• Performance:
  • Sui dispositivi mobili, NanoFLUX genera immagini 512 $\times$ 512 in circa 2,5 secondi.
  • I benchmark quantitativi (One-IG, DPG, GenEval, HPDv3) mostrano che NanoFLUX mantiene prestazioni comparabili al modello docente, con solo una lieve degradazione in metriche specifiche (ad esempio, HPSv3 scende da 11,45 a 10,41).
  • I risultati qualitativi indicano che il modello preserva la qualità dell'immagine e l'aderenza al prompt in modo comparabile al modello originale da 17 miliardi.
• Efficienza: La strategia di downsampling progressivo dei token bilancia con successo la latenza di inferenza e la qualità di generazione, permettendo al modello di girare efficientemente su hardware con risorse limitate.

Significato e Rivendicazioni

Il paper afferma che NanoFLUX è il primo framework che studia in modo completo la compressione di modelli su larga scala per la generazione T2I, mirando specificamente a modelli come FLUX. La sua importanza risiede in:

1. Fattibilità della Generazione On-Device: Dimostrare che la generazione T2I di alta qualità è fattibile su dispositivi mobili senza dipendere da infrastrutture cloud.
2. Approfondimento Architetturale: Fornire una via pratica per scalare verso il basso grandi modelli generativi identificando e sfruttando specifiche ridondanze architettoniche (teste di attenzione, profondità dei blocchi, dinamiche di normalizzazione) invece di fare affidamento esclusivamente su una quantizzazione generica o sulla riduzione dei passi.
3. Generalizzabilità: Offrire intuizioni e una pipeline di distillazione progressiva che possono generalizzarsi ai futuri modelli generativi su larga scala.

Gli autori posizionano questo lavoro come un passo verso l'ampliamento dell'accessibilità alla IA generativa di alta qualità per utenti con risorse computazionali limitate.