SALAD: Achieve High-Sparsity Attention via Efficient Linear Attention Tuning for Video Diffusion Transformer

1. Il Problema

I Diffusion Transformer (DiT) hanno dimostrato prestazioni eccezionali nella generazione video. Tuttavia, la loro applicazione a sequenze lunghe (video ad alta risoluzione e lunga durata) è ostacolata dalla complessità computazionale quadratica ($O(N^2)$) del meccanismo di attenzione completa (Full Attention).

Per mitigare questo costo, sono state proposte diverse attention sparse. Tuttavia, esistono due limiti principali negli approcci attuali:

• Metodi senza training (Training-free): Limitati a sparsità moderate (es. 40-60%), offrono accelerazioni modesta e spesso degradano la qualità.
• Metodi basati sul training: Possono raggiungere sparsità elevate (80-95%), ma richiedono enormi quantità di dati e risorse computazionali per il ri-addestramento (es. VMoBA richiede 182 ore GPU su dataset massicci).
• Il fallimento del LoRA: L'uso della Low-Rank Adaptation (LoRA) per affinare modelli con attenzione ultra-sparsa non riesce a recuperare completamente la qualità del modello originale. I video generati mostrano artefatti, incoerenze temporali e perdita di dettagli (es. oggetti che appaiono e scompaiono o si duplicano), poiché il LoRA non riesce a compensare efficacemente la perdita di interazioni globali tra i token causata dall'estrema sparsità.

2. Metodologia: SALAD

Gli autori propongono SALAD, un modulo di attenzione efficiente che combina un ramo di attenzione sparso (ad alta efficienza) con un ramo parallelo di attenzione lineare (a basso costo computazionale ma globale), regolato da una strategia di scalatura avanzata.

Architettura Principale

• Branch Ibrido: SALAD introduce un ramo di attenzione lineare parallelo a quello sparso. Entrambi i rami condividono le stesse proiezioni di Query (Q), Key (K) e Value (V) inizializzate dai pesi pre-addestrati.
• Condivisione dei Parametri: La maggior parte dei parametri è condivisa. Il ramo lineare aggiunge solo circa il 4.99% di parametri extra rispetto al modello pre-addestrato.
• Attenzione Lineare: Utilizza una funzione kernel (ReLU) per ridurre la complessità da $O(N^2)$ a $O(N)$, permettendo il mescolamento globale dei token (interazioni a lungo raggio) che l'attenzione sparso perde.

Strategia Chiave: Multi-level Static-Dynamic Scaling

L'analisi degli autori rivela uno sbilanciamento di rango (rank imbalance): l'output dell'attenzione lineare ha un rango molto più basso rispetto a quello sparso. Se sommati direttamente, il ramo lineare può sopprimere le strutture informative ad alto rango del ramo sparso. Per risolvere ciò, SALAD introduce una strategia di scalatura a più livelli:

1. Scalatura Statica per Livello (Layer-wise Static Scaling): Un layer di proiezione inizializzato a zero (Zero-initialized) viene inserito dopo il ramo lineare in ogni blocco Transformer. Questo permette al modello di iniziare l'adattamento partendo dal comportamento dell'attenzione sparso pura, regolando gradualmente il contributo del ramo lineare.
2. Scalatura Dinamica Consapevole di Input e Timestep (Input-Timestep-Aware Dynamic Scaling): Un modulo leggero calcola un fattore di scala globale ($s$) basato sullo stato nascosto corrente e sul timestep di denoising. Questo fattore, vincolato tra 0 e 1 tramite una funzione sigmoide, modula dinamicamente l'influenza del ramo lineare.
   • Obiettivo: Il ramo lineare agisce come un complemento controllato, iniettando informazioni globali critiche solo quando necessario, senza sovrastare l'attenzione sparso.

Post-Tuning Branch Dropping

Dopo l'adattamento, si osserva che alcuni rami lineari hanno fattori di scala molto bassi in determinati timestep. SALAD permette di eliminare dinamicamente questi rami durante l'inferenza (Branch Dropping) senza degradare la qualità, ottenendo un'ulteriore efficienza.

3. Contributi Chiave

1. Architettura SALAD: Un modulo di attenzione che integra un ramo lineare parallelo per recuperare le interazioni globali perse nell'attenzione ultra-sparsa, con un overhead parametrico minimo (~5%).
2. Strategia di Scalatura Ibrida: Una metodologia innovativa (Statica + Dinamica) per bilanciare i contributi di rami con ranghi diversi, risolvendo il problema della degradazione delle prestazioni tipico delle combinazioni naive sparse-linear.
3. Efficienza Estrema nel Fine-tuning: Il metodo richiede risorse di adattamento drasticamente ridotte:
   • Solo 2.000 campioni video (dataset Mixkit).
   • Meno di 1.600 step di training.
   • Non più di 30 ore GPU (con batch size 8).
   • Questo è un ordine di grandezza inferiore rispetto ai metodi precedenti (es. SLA richiede 20.000 video, VMoBA richiede dataset da 36M video).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Wan2.1 (1.3B e 14B parametri) con risoluzioni fino a 720p e sequenze di fino a 100k token.

• Sparsità e Velocità: SALAD raggiunge una sparsità del 90% con un speedup di inferenza da 1.52x a 2.03x rispetto all'attenzione completa.
• Qualità: A differenza dei metodi basati su LoRA che mostrano degradazione, SALAD mantiene una qualità di generazione paragonabile o superiore alla baseline Full Attention.
  • Su Wan1.3B (30k token), supera tutti i baselines (SVG2, PARO, LoRA) su metriche VBench (Consistenza del Soggetto, Sfondo, Qualità Immagine, Coerenza Testuale).
  • Su Wan14B, ottiene un speedup di 1.59x mantenendo prestazioni vicine al modello completo.
• Confronto con LoRA: SALAD supera le varianti LoRA (anche con rank più alti) utilizzando meno parametri trainabili (165M vs 189M+ per LoRA rank 256) e ottenendo risultati superiori in termini di coerenza visiva e testuale.
• Ablation Study: Dimostra che la combinazione di scalatura statica e dinamica è essenziale; l'uso di solo LoRA o solo scalatura statica non è sufficiente a recuperare le prestazioni.

5. Significato e Impatto

Il lavoro SALAD rappresenta un passo avanti significativo nell'ottimizzazione dei modelli di generazione video:

• Democratizzazione dell'Addestramento: Dimostra che è possibile adattare modelli di diffusione video ultra-sparsi con costi di calcolo e dati estremamente ridotti, rendendo accessibile l'ottimizzazione anche a gruppi di ricerca più piccoli.
• Superamento del Compromesso Qualità-Efficienza: Risolve il dilemma per cui l'alta sparsità portava inevitabilmente a una perdita di qualità, permettendo di generare video lunghi e ad alta risoluzione con costi computazionali ridotti senza sacrificare la coerenza temporale o la fedeltà al prompt.
• Generalizzabilità: L'approccio è applicabile a diverse architetture sparse (statiche e dinamiche) e a modelli di grandi dimensioni (fino a 14B parametri), offrendo una soluzione scalabile per il futuro dei video generativi.

In sintesi, SALAD offre una soluzione pratica ed efficiente per l'implementazione di attenzione ultra-sparsa nei Diffusion Transformer, bilanciando l'efficienza computazionale con la necessità di mantenere interazioni globali critiche per la qualità del video.