Riassunto Tecnico: PISA - Piecewise Sparse Attention per Efficient Diffusion Transformers

1. Definizione del Problema

I Diffusion Transformer (DiT) sono diventati lo standard per la generazione di immagini e video ad alta fedeltà. Tuttavia, la loro scalabilità è limitata dalla complessità quadratica ($O(L^2)$) del meccanismo di self-attention all'aumentare della lunghezza della sequenza di input ($L$) per risoluzioni più elevate o video più lunghi.

Le soluzioni esistenti presentano limitazioni significative:

• Sparse Attention ("Keep-or-Drop"): I metodi che scartano blocchi non critici di key-value (KV) accelerano il calcolo ma soffrono di un degrado delle prestazioni ad alti rapporti di sparsità a causa della perdita di contesto. Inoltre, hanno difficoltà con sequenze brevi dove l'overhead di selezione supera i benefici.
• Attention Lineare/Ibrida: Gli approcci che alterano la distribuzione dell'attenzione (ad esempio, tramite mappature di feature o combinazioni additive di rami sparsi e lineari) spesso interrompono la normalizzazione intrinseca dei modelli pre-addestrati. Ciò richiede un costoso ri-addestramento o fine-tuning, impedendo il riutilizzo diretto dei pesi pre-addestrati.
• Troncamento Hard (Hard Truncation): Scartare direttamente i blocchi porta a deviazioni nette dalla distribuzione originale dell'attenzione, causando una perdita di qualità nel contenuto generato.

La sfida principale è progettare un meccanismo che acceleri l'inferenza con una complessità sub-quadratica pur mantenendo l'accuratezza dell'attenzione completa e rimanendo compatibile con i pesi pre-addestrati senza necessità di ri-addestramento.

2. Metodologia: Piecewise Sparse Attention (PISA)

PISA introduce un nuovo paradigma "exact-or-approximate" (esatto o approssimato), andando oltre la tradizionale strategia "keep-or-drop". Tratta l'attenzione come un processo a segmenti (piecewise) che copre l'intero intervallo di attenzione con una complessità sub-quadratica.

Intuizioni Chiave

Gli autori hanno analizzato i DiT pre-addestrati e hanno scoperto due proprietà critiche:

1. Stabilità Distribuizionale: Gli score di attenzione pre-softmax dei blocchi non critici (non critici) esibiscono una distribuzione simmetrica a campana centrata intorno allo zero o a valori negativi. Questo li rende altamente suscettibili a un'accurata approssimazione tramite espansione di Taylor, piuttosto che richiedere un calcolo esatto o lo scarto.
2. Consistenza della Normalizzazione: Per mantenere la compatibilità senza addestramento, le approssimazioni devono essere integrate internamente nella numeratore e denominatore della softmax, anziché semplicemente aggiungere gli output da diversi rami (il che violerebbe le regole della somma pesata).

Algoritmo Core

PISA partiziona i blocchi KV in due set:

1. Blocchi Critici ($S_i$): Calcolati esattamente per preservare le informazioni ad alta salienza.
2. Blocchi Non Critici ($U_i$): Approssimati efficientemente utilizzando l'espansione di Taylor per blocchi.

L'output $o_t$ è derivato normalizzando l'aggregazione pesata dei valori:
$$o_t = \frac{N_t}{D_t}$$

Dove il denominatore ($D_t$) e il numeratore ($N_t$) sono composti da:

• Termini Sparsi Esatti: Calcolati per i blocchi selezionati $j \in S_i$.
• Approssimazione di Ordine Zero per Blocchi: Per i blocchi non selezionati, il contributo è approssimato usando il centroide del blocco $\bar{k}_j$ e la somma dei valori.
• Correzione Globale di Primo Ordine: Inve vest invece di calcolare un termine di primo ordine distinto per ogni blocco non selezionato (il che sarebbe limitante per la memoria), PISA impiega un'approssimazione di ordine ibrido. Utilizza una statistica globale $\bar{H}$ (la media delle matrici di primo ordine per blocchi attraverso tutti i blocchi) per approssimare il gradiente di primo ordine per l'intero set di coda.

La formulazione per il numeratore $N_t$ diventa:
$$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$$
Dove $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$.

Routing Consapevole della Covarianza (Covariance-Aware Routing)

Per minimizzare l'errore di approssimazione, PISA utilizza una Strategia Top-k Consapevole della Covarianza. Il punteggio di selezione di un blocco incorpora non solo lo score di attenzione, ma anche la norma della matrice di covarianza del blocco ($M_j = \|H_j - \bar{H}\|_2$). I blocchi con alta rilevanza semantica o alta eterogeneità sono selezionati per il calcolo esatto, mentre i blocchi più omogenei vengono delegati all'espansione di Taylor.

Implementazione Hardware

PISA è implementato tramite un kernel fuso che intercala il calcolo esatto e la scansione approssimativa all'interno del processo di softmax online (stile FlashAttention). La pipeline prevede:

1. Pre-calcolo: Calcolo dei centroidi dei blocchi, delle somme dei valori e della statistica globale $\bar{H}$.
2. Fase 1 (Esatta): Caricamento dei blocchi selezionati in SRAM per l'attenzione esatta.
3. Fase 2 (Approssimazione): Scansione dei vettori aggregati dei blocchi per i blocchi non selezionati utilizzando l'espansione di ordine zero.
4. Fase 3 (Correzione): Iniezione del termine di correzione di primo ordine globale.

Questo design assicura che il metodo sia efficiente dal punto di vista della memoria ed evita i colli di bottiglia associati al caricamento di matrici distinte per ogni blocco.

3. Contributi Chiave

1. Nuova Attention Sparse a Segmenti (Piecewise Sparse Attention): Un meccanismo senza addestramento che consente la copertura dell'intero intervallo di attenzione con complessità sub-quadratica, unificando il calcolo sparso esatto e l'approssimazione per blocchi.
2. Schema di Approssimazione di Ordine Ibrido: Una strategia che combina l'espansione di ordine zero per blocchi con una correzione di primo ordine globale. Questo aumenta significativamente l'accuratezza con un costo computazionale trascurabile rispetto alla standard sparse attention.
3. Routing Consapevole della Covarianza: Una derivazione dall'analisi dell'errore che minimizza la divergenza dell'approssimazione selezionando i blocchi in base sia agli score di attenzione che alle norme di covarianza del blocco.
4. Prestazioni allo Stato dell'Arte: Dimostrata un'efficienza e una qualità superiori in diversi compiti (generazione di video e immagini) senza richiedere il ri-addestramento.

4. Risultati Sperimentali

Gli autori hanno valutato PISA su modelli di generazione video (Wan2.1-1.3B/14B, Hunyuan-Video) e modelli di generazione di immagini (SD 3.5, FLUX.1).

• Generazione Video:

  • Wan2.1-14B: Ha ottenuto un accelerazione di 1.91× rispetto all'attenzione densa (1468s $\to$ 818s) con l'87.5% di sparsità, mantenendo o migliorando leggermente le metriche di qualità (VBench, PSNR, SSIM) rispetto all'attenzione densa.
  • Hunyuan-Video: Ha ottenuto un'accelerazione di 2.57× (1651s $\to$ 641s) con una costante preservazione della qualità.
  • Confronto: Ha superato i metodi sparsi esistenti come SpargeAttn e SVG2 sia in velocità che in qualità, in particolare nella preservazione dei dettagli fini e della coerenza temporale.

• Generazione Immagini:

  • FLUX.1-dev: Ha ottenuto un'accelerazione di 1.2× (8.32s $\to$ 6.87s) con l'85% di sparsità. Notevolmente, PISA ha ottenuto punteggi più alti nelle metriche di preferenza umana (ImageReward, HPSv2) e nelle metriche di similarità (PSNR, SSIM) rispetto sia all'attenzione densa che a SpargeAttn.
  • SD 3.5: Ha dimostrato accelerazioni significative (ad esempio, 1.4× su SD 3.5-M) mantenendo punteggi FID comparabili o migliori rispetto all'attenzione densa.

• Efficienza del Kernel:

  • PISA ha costantemente superato FlashAttention-2/3 e SpargeAttn in termini di latenza attraverso varie lunghezze di sequenza (da 4K a 32K) e densità.
  • A differenza di SpargeAttn, che degrada su sequenze brevi, PISA ha mantenuto il vantaggio di velocità anche a 4K token.

5. Significato e Rivendicazioni

Il paper sostiene che PISA rappresenti un cambio di paradigma dal "keep-or-drop" al "exact-or-approximate" sparse attention. Il suo significato primario risiede nel:

• Colmare il divario Velocità-Qualità: Colma efficacemente il divolo tra velocità di inferenza e qualità di generazione, ottenendo un'approssimazione quasi priva di perdite con un overhead trascurabile.
• Compatibilità Senza Addestramento: Integrando le approssimazioni direttamente nel framework di normalizzazione della softmax, PISA consente il riutilizzo diretto dei pesi pre-addestrati, eliminando la necessità del costoso ri-addestramento o fine-tuning richiesto dai metodi di attenzione lineare o ibrida.
• Scalabilità: Fornisce una soluzione unificata sia per la generazione di immagini ad alta risoluzione che per la generazione di video di lunga durata, affrontando il collo di bottiglia della complessità quadratica senza sacrificare la fedeltà del contenuto generato.

Gli autori sottolineano che il loro metodo è il primo a integrare nativamente l'espansione di Taylor per blocchi e la sparse attention in un framework unificato che preserva la distribuzione intrinseca dei modelli pre-addestrati.