SpargeAttention2 : Attention Sparse Entraînable via Masquage Hybride Top-k+Top-p et Affinage par Distillation

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

February 13, 2026

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

cs.AI

papers.abstract

De nombreuses méthodes d'attention creuse non entraînables sont efficaces pour accélérer les modèles de diffusion. Récemment, plusieurs travaux suggèrent que rendre l'attention creuse entraînable peut permettre d'augmenter davantage la parcimonie tout en préservant la qualité de génération. Nous étudions trois questions clés : (1) quand est-ce que les deux règles de masquage courantes, à savoir Top-k et Top-p, échouent, et comment pouvons-nous éviter ces échecs ? (2) pourquoi l'attention creuse entraînable peut-elle atteindre une parcimonie plus élevée que les méthodes non entraînables ? (3) quelles sont les limites du réglage fin de l'attention creuse en utilisant la perte de diffusion, et comment pouvons-nous les résoudre ? Sur la base de cette analyse, nous proposons SpargeAttention2, une méthode d'attention creuse entraînable qui atteint une haute parcimonie sans dégrader la qualité de génération. SpargeAttention2 inclut (i) une règle de masquage hybride qui combine Top-k et Top-p pour un masquage plus robuste à haute parcimonie, (ii) une implémentation efficace de l'attention creuse entraînable, et (iii) un objectif de réglage fin inspiré de la distillation pour mieux préserver la qualité de génération pendant le réglage fin utilisant l'attention creuse. Les expériences sur les modèles de diffusion vidéo montrent que SpargeAttention2 atteint 95 % de parcimonie d'attention et une accélération de l'attention d'un facteur 16,2 tout en maintenant la qualité de génération, surpassant constamment les méthodes d'attention creuse antérieures.

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.