SpargeAttention2: Atenção Esparsa Treinável por meio de Mascaramento Híbrido Top-k+Top-p e Fine-Tuning por Destilação

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

February 13, 2026

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

cs.AI

Resumo

Muitos métodos de atenção esparsa sem treinamento são eficazes para acelerar modelos de difusão. Recentemente, vários trabalhos sugerem que tornar a atenção esparsa treinável pode aumentar ainda mais a esparsidade, preservando a qualidade da geração. Estudamos três questões-chave: (1) quando as duas regras de mascaramento comuns, Top-k e Top-p, falham, e como podemos evitar essas falhas? (2) por que a atenção esparsa treinável pode atingir maior esparsidade do que os métodos sem treinamento? (3) quais são as limitações do ajuste fino da atenção esparsa usando a perda de difusão, e como podemos abordá-las? Com base nesta análise, propomos o SpargeAttention2, um método de atenção esparsa treinável que alcança alta esparsidade sem degradar a qualidade da geração. O SpargeAttention2 inclui (i) uma regra de mascaramento híbrida que combina Top-k e Top-p para um mascaramento mais robusto em alta esparsidade, (ii) uma implementação eficiente de atenção esparsa treinável, e (iii) um objetivo de ajuste fino inspirado em destilação para preservar melhor a qualidade da geração durante o ajuste fino usando atenção esparsa. Experimentos em modelos de difusão de vídeo mostram que o SpargeAttention2 atinge 95% de esparsidade de atenção e uma aceleração de atenção de 16,2x, mantendo a qualidade da geração, superando consistentemente os métodos anteriores de atenção esparsa.

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.