SLA2: Sparse-Linear Attention with Learnable Routing and QAT

1. Il Problema

I modelli di diffusione (diffusion models), specialmente per la generazione video, richiedono un'enorme potenza di calcolo a causa della complessità quadratica dell'attenzione standard. Per accelerare questi modelli, sono state sviluppate tecniche di Attenzione Sparsa-Linear (SLA), che combinano un ramo di attenzione sparsa (per catturare le dipendenze a lungo raggio con alta precisione) e un ramo di attenzione lineare (per gestire il resto in modo efficiente).

Tuttavia, l'approccio SLA esistente presenta due limitazioni fondamentali:

1. Routing Euristico Subottimale: SLA assegna le computazioni ai rami sparsi o lineari basandosi su una semplice euristiche (selezione dei pesi di attenzione più grandi). Questa divisione non è ottimale e può portare a una distribuzione non ideale del carico computazionale.
2. Disallineamento Teorico: Esiste un errore di formulazione tra la motivazione originale della decomposizione dell'attenzione (in una parte sparsa $P_1$ e una parte a basso rango $P_2$) e l'implementazione pratica di SLA. Nello specifico, il ramo sparsa di SLA normalizza le probabilità in modo diverso rispetto alla decomposizione teorica, creando un errore di scala che il ramo lineare deve tentare di compensare, rendendo l'apprendimento più difficile e meno preciso.

2. Metodologia: SLA2

Gli autori propongono SLA2, una nuova architettura che risolve le limitazioni sopra citate attraverso tre innovazioni principali:

A. Formulazione di Attenzione Sparsa-Linear più Fedele

Invece di usare una proiezione lineare per compensare gli errori di scala, SLA2 introduce un fattore di combinazione appreso ($\alpha$).

• L'output finale è calcolato come: $O = \alpha \odot O_s + (1 - \alpha) \odot O_l$.
• Qui, $O_s$ è l'output del ramo sparsa e $O_l$ è l'output del ramo lineare.
• Il vettore $\alpha$ (appreso) bilancia dinamicamente i due rami, allineandosi perfettamente alla decomposizione teorica originale ($P \approx \alpha P_s + (1-\alpha) P_l$) ed eliminando la necessità di una proiezione di compensazione complessa.

B. Router Apprendibile (Learnable Router)

Per risolvere il problema del routing euristico, SLA2 sostituisce la selezione statica con un Router Apprendibile ($R$).

• Input: Il router prende in input le matrici Query ($Q$) e Key ($K$), ma non Value ($V$), poiché la decisione su quali token sono importanti dipende dalla similarità tra $Q$ e $K$.
• Efficienza: Per gestire sequenze lunghe, $Q$ e $K$ vengono prima sottoposti a pooling medio su blocchi di token adiacenti, riducendo la complessità.
• Meccanismo: Il router utilizza proiezioni lineari apprendibili ($proj_q, proj_k$) per generare una mappa di attenzione compressa. Durante l'addestramento, viene utilizzato un operatore SoftTop-k (differenziabile) per selezionare i token da inviare al ramo sparsa, permettendo la retropropagazione del gradiente. Durante l'inferenza, si passa a un HardTop-k per la massima efficienza.

C. Attenzione a Bassa Precisione con QAT (Quantization-Aware Training)

Per ottenere un'ulteriore accelerazione, SLA2 integra l'attenzione a bassa precisione (es. INT8 o FP8) nel ramo sparsa.

• Viene utilizzato un approccio QAT (Quantization-Aware Training): durante l'addestramento, il forward pass utilizza tensori quantizzati, mentre il backward pass rimane in precisione mista (FP16) utilizzando i valori originali.
• Questo permette al modello di adattarsi agli errori di quantizzazione, mantenendo alta la qualità della generazione video anche con pesi e attivazioni a bassa precisione.

3. Contributi Chiave

1. Analisi e Riformulazione: Gli autori hanno identificato e corretto il disallineamento teorico nell'SLA originale, proponendo una formulazione matematica più coerente che elimina la necessità di proiezioni di compensazione.
2. Routing Dinamico: Introduzione di un router apprendibile che ottimizza dinamicamente la divisione del lavoro tra i rami sparsi e lineari, superando le euristiche statiche.
3. Integrazione QAT: Dimostrazione dell'efficacia del QAT nell'ambito dell'attenzione sparsa per i modelli di diffusione, permettendo accelerazioni aggiuntive senza degradare la qualità.
4. Prestazioni Superiori: SLA2 raggiunge livelli di sparsità estremamente elevati (fino al 97%) mantenendo o migliorando la qualità rispetto all'attenzione completa.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli video diffusion di stato dell'arte (Wan2.1-1.3B e Wan2.1-14B).

• Sparsità e Velocità: SLA2 raggiunge una sparsità dell'attenzione del 97%, corrispondente a un risparmio computazionale del 96,7%.
• Speedup: Si ottiene un speedup di 18,6x nel runtime dell'attenzione rispetto all'attenzione completa (Full Attention).
• Qualità Video:
  • A una sparsità del 97%, SLA2 supera le performance di baseline con sparsità del 90% (come VMoBA e VSA) su tutte le metriche di qualità video (IQ, OC, AQ, MS, SC, VR).
  • Sorprendentemente, in alcuni casi, SLA2 supera anche l'attenzione completa (0% sparsità) dopo il fine-tuning, probabilmente grazie alla migliore qualità del dataset di addestramento utilizzato per il fine-tuning.
• Efficienza End-to-End:
  • Su Wan2.1-1.3B, la latenza end-to-end viene ridotta di 2,3x (da 97s a 7s per l'attenzione).
  • Su Wan2.1-14B, la riduzione è di 4,35x.

5. Significato e Impatto

SLA2 rappresenta un passo avanti significativo per l'efficienza dei modelli di generazione video.

• Praticità: Dimostra che è possibile raggiungere sparsità estreme (97%) senza sacrificare la qualità visiva, rendendo la generazione video ad alta risoluzione più accessibile su hardware consumer.
• Fondamentale Teorico: La correzione della formulazione matematica dell'attenzione sparsa-lineare offre una base più solida per futuri sviluppi in questo campo.
• Scalabilità: L'uso combinato di routing apprendibile e quantizzazione consapevole rende i modelli di diffusione molto più scalabili, riducendo drasticamente i costi di inferenza e l'uso di memoria VRAM.

In sintesi, SLA2 non è solo un'ottimizzazione incrementale, ma una riprogettazione metodologica che risolve limiti teorici ed euristici, offrendo una soluzione pratica e ad alte prestazioni per l'accelerazione dei modelli di diffusione video.