SpargeAttention2: Atención Dispersa Entrenable mediante Enmascaramiento Híbrido Top-k+Top-p y Ajuste Fino por Destilación

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

Resumen

Muchos métodos de atención dispersa sin entrenamiento son efectivos para acelerar los modelos de difusión. Recientemente, varios trabajos sugieren que hacer que la atención dispersa sea entrenable puede aumentar aún más la dispersión preservando la calidad de la generación. Estudiamos tres preguntas clave: (1) ¿cuándo fallan las dos reglas de enmascaramiento comunes, es decir, Top-k y Top-p, y cómo podemos evitar estas fallas? (2) ¿por qué la atención dispersa entrenable puede alcanzar una dispersión mayor que los métodos sin entrenamiento? (3) ¿cuáles son las limitaciones de afinar la atención dispersa utilizando la pérdida de difusión, y cómo podemos abordarlas? Basándonos en este análisis, proponemos SpargeAttention2, un método de atención dispersa entrenable que logra una alta dispersión sin degradar la calidad de la generación. SpargeAttention2 incluye (i) una regla de enmascaramiento híbrida que combina Top-k y Top-p para un enmascaramiento más robusto con alta dispersión, (ii) una implementación eficiente de atención dispersa entrenable, y (iii) un objetivo de ajuste fino inspirado en la destilación para preservar mejor la calidad de la generación durante el ajuste fino utilizando atención dispersa. Los experimentos en modelos de difusión de video muestran que SpargeAttention2 alcanza un 95% de dispersión de atención y una aceleración de la atención de 16.2x manteniendo la calidad de la generación, superando consistentemente a los métodos de atención dispersa anteriores.

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.