MC-Sparse: analisi e riduzione del divario tra attenzione densa e sparsa nei transformer a diffusione

Abstract degli autori

L’attenzione sparsa è uno dei principali approcci per ridurre la latenza dei transformer a diffusione nei compiti di generazione di sequenze lunghe, come la generazione di video e di asset 3D ad alta risoluzione. Tuttavia, i metodi esistenti possono ridurre la qualità e la fedeltà della generazione a livelli elevati di sparsità. Attraverso confronti controllati con un oracolo, riconduciamo questo peggioramento a tre cause: i vincoli imposti dal raggruppamento dei token, la selezione imprecisa delle interazioni e i contributi all’attenzione che si perdono quando i token vengono scartati. Sulla base di questa analisi, proponiamo Meta-Cached Sparse Attention (MC-Sparse), un framework che non richiede addestramento e seleziona singoli token chiave-valore (KV), organizzando al contempo le query simili in gruppi allineati ai tile per un’esecuzione efficiente sulla GPU. MC-Sparse memorizza in cache metadati che comprendono i gruppi di query, gli indici KV selezionati in base alle probabilità esatte di attenzione e i residui tra gli output dell’attenzione densa e di quella sparsa, e li riutilizza nelle successive fasi di rimozione del rumore. Nei modelli di generazione di video e di asset 3D, MC-Sparse ottiene una maggiore fedeltà agli output dell’attenzione densa e una maggiore accelerazione della rimozione del rumore rispetto ai metodi di riferimento esistenti basati sull’attenzione sparsa, senza un peggioramento visibile della qualità. Rispetto all’attenzione densa, accelera la rimozione del rumore di $1.80\times$ su Minimax-H3-Base e di $2.32\times$ nella generazione di asset 3D, in entrambi i casi con una perdita di qualità trascurabile.