SpargeAttention2: ハイブリッドTop-k+Top-pマスキングと蒸留ファインチューニングによる学習可能なスパースアテンション
SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning
February 13, 2026
著者: Jintao Zhang, Kai Jiang, Chendong Xiang, Weiqi Feng, Yuezhou Hu, Haocheng Xi, Jianfei Chen, Jun Zhu
cs.AI
要旨
多くの訓練不要なスパースアテンション手法は、拡散モデルの高速化に効果的である。近年、スパースアテンションを学習可能にすることで、生成品質を維持しつつさらに高いスパース性を達成できることが複数の研究で示唆されている。本研究では以下の3つの核心的な問いを検討する：(1) Top-kおよびTop-pという2つの一般的なマスキングルールはいつ失敗するのか、またそれらの失敗をどう回避できるか？(2) 学習可能なスパースアテンションが訓練不要手法よりも高いスパース性を達成できる理由は何か？(3) 拡散損失を用いたスパースアテンションのファインチューニングにはどのような限界があり、どう解決できるか？この分析に基づき、我々はSpargeAttention2を提案する。これは生成品質を劣化させることなく高いスパース性を実現する学習可能なスパースアテンション手法である。SpargeAttention2は以下を含む：(i) 高いスパース性においてより頑健なマスキングを実現するTop-kとTop-pを組み合わせたハイブリッドマスキングルール、(ii) 効率的な学習可能スパースアテンションの実装、(iii) スパースアテンションを用いたファインチューニング時に生成品質をより良く維持するための知識蒸留に着想を得たファインチューニング目的関数。ビデオ拡散モデルを用いた実験では、SpargeAttention2が95%のアテンションスパース性と16.2倍のアテンション高速化を達成しつつ生成品質を維持し、既存のスパースアテンション手法を一貫して上回ることを示した。
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.