Riepilogo Tecnico: CoReDiT

Enunciato del Problema

I Diffusion Transformers (DiT) hanno stabilito prestazioni all'avanguardia nella generazione di immagini e video modellando efficacemente le dipendenze a lungo raggio attraverso meccanismi di attenzione. Tuttavia, il loro costo computazionale scala quadraticamente con il numero di token, rendendo la sintesi ad alta risoluzione e la generazione di video multi-frame proibitivamente costose per le piattaforme mobili con risorse limitate. Sebbene una porzione significativa dei token nello spazio latente corrisponda a regioni visivamente ridondanti (ad esempio, sfondi uniformi), i metodi esistenti di pruning dei token affrontano tre sfide principali quando applicati all'inferenza di diffusione:

1. Localizzazione Efficiente: Identificare token a bassa salienza senza incorrere nel sovraccarico quadratico del calcolo dei punteggi di attenzione globale.
2. Preservazione Semantica: Mantenere la continuità visiva e prevenire artefatti (ad esempio, rotture della texture, distorsioni dei confini) quando i token vengono saltati, il che è critico per le attività generative a differenza di quelle discriminative.
3. Pianificazione Adattiva: Determinare una strategia di pruning che si adatti ai diversi livelli di ridondanza attraverso i vari blocchi del transformer e i passi di denoising.

Metodologia: CoReDiT

CoReDiT è un framework strutturato di pruning dei token progettato per accelerare i DiT pre-addestrati preservando la fedeltà dell'output. Opera attraverso tre componenti principali:

1. Selezione dei Token Basata sulla Coerenza Spaziale

Invece di fare affidamento sui punteggi di attenzione globale (che sono $O(N^2)$), CoReDiT utilizza un punteggio di Coerenza Spaziale (SC) per stimare la ridondanza locale in tempo lineare ($O(N)$).

• Meccanismo: Il reticolo dei token è suddiviso in griglie non sovrapposte. Per ogni token, il punteggio SC è calcolato come la media della similarità coseno tra l'embedding normalizzato del token e l'embedding normalizzato medio della sua griglia.
• Logica: I token con punteggi SC elevati sono considerati ridondanti perché sono ben spiegati dai loro vicini spaziali.
• Selezione: I primi $K$ token con i punteggi SC più alti sono selezionati per essere saltati ($X_S$), mentre i token rimanenti ($X_R$) procedono al livello Multi-Head Self-Attention (MHSA).

2. Ricostruzione dei Token Guidata dalla Coerenza

Per prevenire le discontinuità spaziali causate dal salto dei token, CoReDiT ricostruisce gli output di attenzione per i token saltati ($Y_S$) utilizzando gli output dei vicini mantenuti ($Y_R$).

• Approssimazione: Invece di calcolare similarità a coppie per la ricostruzione (il che aggiungerebbe sovraccarico), il metodo approssima i pesi di similarità utilizzando i punteggi SC pre-calcolati dei vicini.
• Sottogriglie Locali: Per evitare un eccessivo livellamento (over-smoothing), la ricostruzione viene eseguita all'interno di sottogriglie più piccole.
• Micro-Design:
  • Dimensioni delle Griglie Alternanti: Le dimensioni delle griglie alternano tra blocchi consecutivi per spostare gli artefatti di confine e permettere la propagazione delle informazioni attraverso i bordi delle griglie.
  • Mantenimento con Passo m (m-Stride Retention): Un pattern di passo deterministico garantisce che almeno un token sia mantenuto in ogni sottogriglia, fornendo gli ancoraggi necessari per la ricostruzione.

3. Piano di Pruning Progressivo e Adattivo ai Blocchi

Riconoscendo che la ridondanza varia tra i blocchi del transformer e i passi di denoising, CoReDiT impiega una strategia di pruning progressiva durante il fine-tuning leggero.

• Adattività ai Blocchi: Il budget di pruning è allocato ai blocchi con la ridondanza misurata più alta (stimata dalla somma dei punteggi SC dei token candidati).
• Decadimento Temporale: Il pruning è più aggressivo nei primi passi di denoising (dove il rumore domina e la struttura è grezza) e ridotto negli step successivi (dove i dettagli fini sono raffinati). Viene utilizzato un piano di decadimento a due fasi per bilanciare l'efficienza con la complessità della gestione dei grafi computazionali specifici per passo.
• Addestramento: Il modello viene fine-tuned utilizzando una funzione di perdita di distillazione che combina la perdita dell'output del modello originale e una perdita del blocco pruningato per recuperare le prestazioni.

Contributi Chiave

1. Stima della Ridondanza in Tempo Lineare: L'introduzione di un punteggio di coerenza spaziale che identifica i token ridondanti in tempo $O(N)$, evitando il collo di bottiglia computazionale dell'analisi dell'attenzione globale.
2. Preservazione del Reticolo Denso: Un meccanismo di ricostruzione che sintetizza gli output di attenzione saltati dai vicini, garantendo che il reticolo dei token rimanga denso per i livelli successivi e prevenendo artefatti visivi.
3. Strategia di Pruning Adattiva: Un piano progressivo e adattivo ai blocchi che alloca dinamicamente i budget di pruning basandosi sulle statistiche empiriche di ridondanza attraverso la profondità e il tempo, facilitando un'adattamento stabile.

Risultati Sperimentali

CoReDiT è stato valutato su backbones DiT all'avanguardia, inclusi PixArt-α, PixArt-Σ e MagicDrive-V2.

• Guadagni di Efficienza:
  • Ha raggiunto una riduzione fino al 55% nei FLOP di self-attention.
  • Accelerazioni nell'Inferenza: 1.33× su GPU cloud (Nvidia H100) e 1.72× su NPU mobili (Qualcomm Snapdragon 8 Elite).
  • Margine di Memoria: Ha permesso la generazione a risoluzioni fino a 1920 su dispositivi mobili, dove il modello di base incontrava errori Out-of-Memory (OOM) a 1600.
• Preservazione della Qualità:
  • Su PixArt-α-1024 con il 40% di pruning, CoReDiT ha mantenuto un FID di 28.7 (vs 27.3 del baseline) e punteggi CLIP/IS competitivi. Con la distillazione, il divario FID si è ridotto a 27.4.
  • Su PixArt-Σ-2048, un rapporto di pruning medio del 23% ha ridotto la latenza di self-attention del 22% con solo un modesto calo nelle metriche di qualità.
  • Su MagicDrive-V2 (video), una riduzione del 39% nei FLOP di self-attention ha preservato la qualità per frame (PSNR, LPIPS, SSIM) e l'allineamento condizionale (mAP, mIoU), sebbene sia stato notato un lieve calo nella coerenza temporale (FVD).
• Studi di Ablazione:
  • La selezione casuale dei token ha portato a un degrado catastrofico della qualità (FID 644.5), validando la necessità di una selezione basata sulla coerenza.
  • La disattivazione della ricostruzione ha portato a una significativa perdita di qualità (FID 30.4 vs 28.5), confermando l'importanza dello step di ricostruzione.
  • Il pruning uniforme seguito da fine-tuning ha funzionato peggio del piano adattivo progressivo proposto.

Significato e Affermazioni

Il documento afferma che CoReDiT offre una soluzione pratica per il deployment di modelli di diffusione ad alta fedeltà su dispositivi edge affrontando le sfide specifiche del pruning dei token generativi. Il suo significato risiede in:

• Scalabilità: Abilitare la sintesi ad alta risoluzione e la generazione di video su hardware mobile riducendo sia l'impronta computazionale che quella di memoria.
• Compatibilità Architettonica: Mantenere la compatibilità con le architetture DiT esistenti senza richiedere cambiamenti strutturali o moduli di routing complessi.
• Trade-off Efficienza-Qualità: Dimostrare che il salto strutturato e guidato dalla coerenza può ridurre drasticamente i costi di inferenza (fino al 55% dei FLOP di self-attention) mantenendo una qualità percettiva paragonabile al modello completo.

Gli autori notano che, sebbene l'approccio attuale si concentri sulla ridondanza spaziale, il lavoro futuro mira ad estendere l'efficienza ai componenti FFN/MLP e a incorporare obiettivi di coerenza tra frame per migliorare ulteriormente la coerenza temporale nei video.