Diffusione Video Più Veloce con Attenzione Sparse Addestrabile

Faster Video Diffusion with Trainable Sparse Attention

May 19, 2025

Autori: Peiyuan Zhang, Haofeng Huang, Yongqi Chen, Will Lin, Zhengzhong Liu, Ion Stoica, Eric P. Xing, Hao Zhang

cs.AI

Abstract

Il ridimensionamento dei transformer per la diffusione video (DiTs) è limitato dalla loro attenzione quadratica 3D, nonostante la maggior parte della massa di attenzione si concentri su un piccolo sottoinsieme di posizioni. Trasformiamo questa osservazione in VSA, un'attenzione sparsa addestrabile ed efficiente dal punto di vista hardware che sostituisce l'attenzione completa sia durante l'addestramento che durante l'inferenza. In VSA, una fase iniziale leggera raggruppa i token in tessere e identifica i token critici ad alto peso; una fase fine calcola l'attenzione a livello di token solo all'interno di quelle tessere, sottoponendosi a un layout di calcolo a blocchi per garantire efficienza hardware. Ciò porta a un singolo kernel differenziabile che si addestra end-to-end, non richiede profilazione post-hoc e mantiene l'85\% dell'MFU di FlashAttention3. Eseguiamo un'ampia serie di studi di ablazione e esperimenti di legge di ridimensionamento preaddestrando DiTs da 60M a 1.4B parametri. VSA raggiunge un punto di Pareto che riduce i FLOPS di addestramento di 2.53 volte senza alcuna perdita nella perdita di diffusione. Il retrofit del modello open-source Wan-2.1 accelera il tempo di attenzione di 6 volte e riduce il tempo di generazione end-to-end da 31s a 18s con una qualità comparabile. Questi risultati stabiliscono l'attenzione sparsa addestrabile come un'alternativa pratica all'attenzione completa e un abilitatore chiave per un ulteriore ridimensionamento dei modelli di diffusione video.

English

Scaling video diffusion transformers (DiTs) is limited by their quadratic 3D attention, even though most of the attention mass concentrates on a small subset of positions. We turn this observation into VSA, a trainable, hardware-efficient sparse attention that replaces full attention at both training and inference. In VSA, a lightweight coarse stage pools tokens into tiles and identifies high-weight critical tokens; a fine stage computes token-level attention only inside those tiles subjecting to block computing layout to ensure hard efficiency. This leads to a single differentiable kernel that trains end-to-end, requires no post-hoc profiling, and sustains 85\% of FlashAttention3 MFU. We perform a large sweep of ablation studies and scaling-law experiments by pretraining DiTs from 60M to 1.4B parameters. VSA reaches a Pareto point that cuts training FLOPS by 2.53times with no drop in diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention time by 6times and lowers end-to-end generation time from 31s to 18s with comparable quality. These results establish trainable sparse attention as a practical alternative to full attention and a key enabler for further scaling of video diffusion models.