AB-Attention: Attenzione accurata a blocchi per accelerare i transformer di diffusione video

Abstract degli autori

I transformer di diffusione video generano video di alta qualità, ma il costo quadratico dell’attenzione completa 3D domina la fase di inferenza. L’attenzione dinamica sparsa a blocchi riduce questo costo calcolando l’attenzione esatta solo sui blocchi di key/value selezionati in base alle attivazioni correnti. Tuttavia, la sua fedeltà peggiora quando la sparsità è elevata: i blocchi vengono valutati sulla base di riepiloghi grossolani o individuati tramite un costoso clustering globale, i blocchi eliminati scompaiono dalla softmax e intere tile vengono mantenute o scartate anche quando contano solo poche delle loro chiavi. Per superare questi limiti, proponiamo Accurate Block Attention (AB-Attention), un metodo che non richiede addestramento e approssima con precisione l’attenzione densa anche con un’elevata sparsità. Il suo Intra-Block Clustering Selector (IBCS) stima l’importanza dei blocchi in modo fedele e poco costoso, raggruppando le query solo all’interno di ciascun blocco di query, senza permutare i token. Il campionamento con ponderazione basata sulla probabilità di inclusione (SIP) estrae quindi blocchi in proporzione a queste stime e li ripondera in base all’inverso della probabilità di inclusione, così che un singolo passaggio stimi l’attenzione densa invece di troncarla. Infine, il kernel raccoglie blocchi di key/value molto più piccoli di una tile in tile complete per tensor core, affinando la selezione senza costi aggiuntivi per i tensor core. Su Wan2.1 e HunyuanVideo, AB-Attention salta l’89--92\% dell’attenzione, mentre i metodi di riferimento ne calcolano 1,8--3,6$\times$ tanto; eppure migliora comunque il PSNR fino a 4,4~dB rispetto al metodo di riferimento più efficace e raggiunge uno speedup end-to-end fino a 2,23$\times$. Il codice è disponibile all’indirizzo https://github.com/InfRec-Cardinal/AB-Attention.