SpargeAttention2: Trainbare Sparse Attention via Hybride Top-k+Top-p Maskering en Distillatie Fine-Tuning

SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning

February 13, 2026

Auteurs: Jintao Zhang, Kai Jiang, Chendong Xiang, Weiqi Feng, Yuezhou Hu, Haocheng Xi, Jianfei Chen, Jun Zhu

cs.AI

Samenvatting

Veel trainingsvrije sparse attention-methoden zijn effectief voor het versnellen van diffusiemodellen. Recentelijk suggereren verschillende werken dat het trainbaar maken van sparse attention de sparsiteit verder kan vergroten zonder de generatiekwaliteit aan te tasten. Wij bestuderen drie kernvragen: (1) wanneer falen de twee gebruikelijke maskeringsregels, namelijk Top-k en Top-p, en hoe kunnen we deze fouten vermijden? (2) waarom kan trainbare sparse attention een hogere sparsiteit bereiken dan trainingsvrije methoden? (3) wat zijn de beperkingen van het fine-tunen van sparse attention met behulp van het diffusieverlies, en hoe kunnen we deze aanpakken? Op basis van deze analyse stellen we SpargeAttention2 voor, een trainbare sparse attention-methode die een hoge sparsiteit bereikt zonder de generatiekwaliteit te verslechteren. SpargeAttention2 omvat (i) een hybride maskeringsregel die Top-k en Top-p combineert voor robuustere maskering bij hoge sparsiteit, (ii) een efficiënte implementatie van trainbare sparse attention, en (iii) een op distillatie geïnspireerd fine-tuningdoel om de generatiekwaliteit beter te behouden tijdens het fine-tunen met sparse attention. Experimenten met videodiffusiemodellen tonen aan dat SpargeAttention2 een attentiesparsiteit van 95% en een attentieversnelling van 16,2x bereikt met behoud van generatiekwaliteit, waarbij het consistent beter presteert dan eerdere sparse attention-methoden.

English

Many training-free sparse attention methods are effective for accelerating diffusion models. Recently, several works suggest that making sparse attention trainable can further increase sparsity while preserving generation quality. We study three key questions: (1) when do the two common masking rules, i.e., Top-k and Top-p, fail, and how can we avoid these failures? (2) why can trainable sparse attention reach higher sparsity than training-free methods? (3) what are the limitations of fine-tuning sparse attention using the diffusion loss, and how can we address them? Based on this analysis, we propose SpargeAttention2, a trainable sparse attention method that achieves high sparsity without degrading generation quality. SpargeAttention2 includes (i) a hybrid masking rule that combines Top-k and Top-p for more robust masking at high sparsity, (ii) an efficient trainable sparse attention implementation, and (iii) a distillation-inspired fine-tuning objective to better preserve generation quality during fine-tuning using sparse attention. Experiments on video diffusion models show that SpargeAttention2 reaches 95% attention sparsity and a 16.2x attention speedup while maintaining generation quality, consistently outperforming prior sparse attention methods.