DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

1. Il Problema

I modelli di diffusione (Diffusion Models), in particolare quelli basati su architetture Transformer (DiT), hanno dimostrato prestazioni eccezionali nella generazione di immagini e video. Tuttavia, il loro meccanismo di inferenza multi-step richiede un costo computazionale enorme, rendendo difficile il loro dispiegamento in scenari reali con risorse limitate.

Le tecniche di accelerazione esistenti si basano spesso su:

• Ottimizzazione del campionatore: Riduzione del numero di passi.
• Pruning e Distillazione: Rimozione di parametri o riduzione dei passi tramite addestramento intensivo.
• Feature Caching: Riutilizzo delle caratteristiche intermedie (feature) tra passi adiacenti, sfruttando la ridondanza temporale.

Limitazioni delle metodologie attuali:

1. Sparsità manuale: I metodi di caching a livello di token (come ToCa, DuCa) richiedono allocazioni di sparsità (quanti token calcolare e quanti riutilizzare) progettate manualmente o con schedule rigidi.
2. Passi completi obbligatori: Per mantenere la qualità, molte strategie richiedono di eseguire passi di inferenza completi (senza caching) in momenti predefiniti, limitando il potenziale di accelerazione.
3. Inefficienza: Le strategie di caching attuali non ottimizzano dinamicamente la sparsità a livello di layer e di passo temporale in modo end-to-end.

2. Metodologia: DiffSparse

Gli autori propongono DiffSparse, un framework differenziabile che ottimizza l'allocazione della sparsità dei token a livello di layer per i modelli DiT. L'obiettivo è determinare automaticamente la configurazione di sparsità ottimale attraverso tutti i layer e i passi di denoising, rispettando un tasso di pruning globale target.

Il framework si compone di tre componenti principali:

A. Token Selector (Selettore di Token)

Per ogni layer e passo temporale, il selettore assegna un punteggio di importanza a ogni token. I token con punteggio più alto vengono ricalcolati, mentre gli altri riutilizzano le feature dalla cache.

• Il punteggio è una combinazione di: influenza dell'attenzione self-attention, focus dell'attenzione cross-attention, frequenza di riutilizzo della cache e una bonus spaziale per garantire una copertura uniforme.
• Viene selezionata una frazione $R$ di token per il calcolo completo, mentre il resto viene saltato.

B. Sparsity Cost Predictor (Predittore di Costo di Sparsità)

È un modulo apprendibile che stima una matrice dei costi $C \in \mathbb{R}^{(T \times L) \times |S|}$.

• Input: Passo temporale $t$, Layer $l$, e un insieme di candidati di tassi di sparsità $S$.
• Output: Il costo associato all'applicazione di una specifica configurazione di sparsità in quel layer e passo.
• Vantaggio: Le dimensioni del predittore dipendono solo dal numero di passi ($T$), layer ($L$) e candidati ($|S|$), non dalla lunghezza della sequenza di token ($N$), rendendolo scalabile ad alte risoluzioni.

C. Dynamic Programming Solver (Solver di Programmazione Dinamica)

Per soddisfare un vincolo globale di sparsità (es. ridurre il 50% dei calcoli totali), il sistema utilizza un algoritmo di programmazione dinamica.

• Minimizza il costo cumulativo previsto dalla matrice $C$ assegnando livelli di sparsità specifici a ogni layer e passo, garantendo che la somma totale rispetti il target.
• Permette di trovare la configurazione ottimale senza ricerca esaustiva.

D. Strategia di Addestramento in Due Fasi

Per eliminare la necessità di passi completi predefiniti e migliorare la convergenza:

1. Fase 1: Si addestrano separatamente le matrici dei costi per i "passi completi" (full-step) e per la sparsità a livello di layer. Si identificano i passi ottimali da mantenere completi tramite DP.
2. Fase 2: Si integra il costo dei passi completi nella matrice di sparsità dei layer (tramite un "warm-start" che riduce artificialmente il costo dei passi selezionati come completi) e si esegue un fine-tuning per redistribuire dinamicamente i calcoli, permettendo di eliminare gradualmente la dipendenza dai passi completi fissi.

Loss Function: L'ottimizzazione è guidata da una perceptual distillation loss (LPIPS), dove il modello originale (teacher) guida il modello prunato (student) per preservare la qualità percettiva dell'immagine generata.

3. Contributi Chiave

1. Ottimizzazione Differenziabile della Sparsità: DiffSparse automatizza l'allocazione della sparsità a livello di token e layer senza euristiche manuali, utilizzando un predittore di costo apprendibile e un solver DP.
2. Eliminazione dei Passi Completi Fissi: La strategia di addestramento in due fasi rimuove la necessità di mantenere passi di inferenza completi predefiniti, sbloccando il pieno potenziale di accelerazione del caching a livello di token.
3. Generalizzazione e Scalabilità: Il predittore di costo addestrato a bassa risoluzione (es. 256x256) si trasferisce efficacemente ad alte risoluzioni (es. 512x512) senza bisogno di riaddestramento, riducendo i costi computazionali di training.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli foundation (DiT-XL/2, PixArt-α, FLUX, Wan2.1) e task (immagine e video).

• Generazione Immagini (PixArt-α, 20 passi):

  • Con un pruning del 54% dei token, DiffSparse riduce il costo computazionale del 54% ottenendo un Speedup di 1.91x.
  • Qualità: Il modello ottiene un FID di 27.79, superando il modello originale (FID 28.20) e battendo nettamente gli SOTA come ToCa (FID 28.35) e TaylorSeer (FID 29.08).
  • In un setting con speedup 1.74x, DiffSparse raggiunge un FID di 26.91, migliorando la qualità rispetto al baseline.

• Generazione Class-Conditional (DiT-XL/2, ImageNet):

  • A parità di accelerazione (~2.07x), DiffSparse migliora l'FID da 3.05 (ToCa) a 2.81, dimostrando una migliore conservazione dei dettagli.

• Generazione Video (Wan2.1):

  • Su Wan2.1-1.3B, DiffSparse raggiunge il punteggio VBench più alto (43.83) con un speedup di 2.05x, superando DuCa.

• Efficienza di Training:

  • Il metodo differenziabile richiede circa 4 ore di training su 8 GPU, contro le ~16 ore richieste da metodi basati su ricerca (es. algoritmi genetici) per trovare configurazioni subottime.

5. Significato e Impatto

DiffSparse rappresenta un avanzamento significativo nell'accelerazione dei modelli di diffusione. Dimostra che è possibile ottenere accelerazioni sostanziali (fino a 2x) senza degradare, e spesso migliorando, la qualità della generazione.

La capacità di apprendere dinamicamente dove e quando risparmiare calcoli (sparsità) a livello di token e layer, eliminando la rigidità delle strategie di caching precedenti, rende i modelli DiT molto più pratici per il dispiegamento su dispositivi con risorse limitate (edge, mobile) e per applicazioni industriali che richiedono bassa latenza. Inoltre, la generalizzazione a diverse risoluzioni e architetture senza bisogno di riaddestramento massiccio ne facilita l'adozione immediata.