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

📄 arXiv: 2602.13515v1 📥 PDF

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

分类: cs.CV, cs.LG

发布日期: 2026-02-13

💡 一句话要点

SpargeAttention2：通过混合Top-k+Top-p掩码和蒸馏微调实现可训练的稀疏注意力，加速视频扩散模型。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 稀疏注意力 扩散模型 视频生成 Top-k Top-p 蒸馏 模型加速

📋 核心要点

  1. 现有无训练稀疏注意力方法虽能加速扩散模型，但稀疏性提升有限，可训练稀疏注意力面临掩码失效和微调困难。
  2. SpargeAttention2提出混合Top-k+Top-p掩码规则，结合二者优势，提升高稀疏度下的鲁棒性，并设计蒸馏微调目标。
  3. 实验表明，SpargeAttention2在视频扩散模型上实现了95%的注意力稀疏度和16.2倍的加速，且保持了生成质量。

📝 摘要（中文）

许多无训练的稀疏注意力方法在加速扩散模型方面是有效的。最近，一些研究表明，使稀疏注意力可训练可以进一步提高稀疏性，同时保持生成质量。本文研究了三个关键问题：（1）两种常见的掩码规则，即Top-k和Top-p，何时失效，以及如何避免这些失效？（2）为什么可训练的稀疏注意力能够达到比无训练方法更高的稀疏性？（3）使用扩散损失微调稀疏注意力的局限性是什么，以及如何解决它们？基于此分析，本文提出了SpargeAttention2，一种可训练的稀疏注意力方法，可以在不降低生成质量的情况下实现高稀疏性。SpargeAttention2包括（i）一种混合掩码规则，结合Top-k和Top-p，以实现更高稀疏度下的更鲁棒的掩码，（ii）一种高效的可训练稀疏注意力实现，以及（iii）一种受蒸馏启发的微调目标，以在使用稀疏注意力进行微调期间更好地保持生成质量。在视频扩散模型上的实验表明，SpargeAttention2达到了95%的注意力稀疏度和16.2倍的注意力加速，同时保持了生成质量，始终优于先前的稀疏注意力方法。

🔬 方法详解

问题定义：现有的稀疏注意力方法，特别是应用于扩散模型时，面临着两个主要问题。首先，Top-k和Top-p这两种常见的掩码规则在高稀疏度下容易失效，导致性能下降。其次，直接使用扩散损失微调稀疏注意力可能导致生成质量的损失。因此，需要一种方法能够在保持生成质量的同时，进一步提高稀疏性，并解决微调过程中的问题。

核心思路：SpargeAttention2的核心思路是结合Top-k和Top-p两种掩码规则的优点，提出一种混合掩码规则，以提高在高稀疏度下的鲁棒性。此外，采用一种受蒸馏启发的微调目标，以更好地保持生成质量。通过这种方式，可以在不牺牲生成质量的前提下，实现更高的稀疏性和更快的推理速度。

技术框架：SpargeAttention2主要包含三个组成部分：混合Top-k+Top-p掩码规则、高效的可训练稀疏注意力实现以及受蒸馏启发的微调目标。首先，使用混合掩码规则生成稀疏注意力掩码。然后，利用该掩码进行稀疏注意力计算。最后，通过受蒸馏启发的微调目标，优化稀疏注意力模型的参数。

关键创新：SpargeAttention2的关键创新在于混合Top-k+Top-p掩码规则和受蒸馏启发的微调目标。混合掩码规则结合了Top-k选择最显著的注意力权重和Top-p动态选择注意力权重的优点，从而在高稀疏度下更加鲁棒。受蒸馏启发的微调目标通过模仿原始模型的输出，更好地保持了生成质量。

关键设计：混合掩码规则的具体实现是首先使用Top-k选择一部分注意力权重，然后使用Top-p选择剩余的注意力权重。Top-k和Top-p的比例是可调节的超参数。受蒸馏启发的微调目标是最小化稀疏注意力模型的输出与原始模型的输出之间的差异。损失函数可以采用L1或L2损失。此外，还采用了高效的稀疏矩阵运算库，以加速稀疏注意力计算。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SpargeAttention2在视频扩散模型上达到了95%的注意力稀疏度和16.2倍的注意力加速，同时保持了生成质量。与之前的稀疏注意力方法相比，SpargeAttention2在稀疏性和生成质量之间取得了更好的平衡，显著提升了视频扩散模型的性能。

🎯 应用场景

SpargeAttention2在视频生成、图像生成等领域具有广泛的应用前景。通过提高扩散模型的推理速度，可以加速视频编辑、内容创作等任务。此外，该方法还可以应用于其他需要高效注意力机制的场景，例如自然语言处理和语音识别。

📄 摘要（原文）

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.