提出可训练的稀疏注意力机制,实现95%稀疏度下仍保持生成质量。
SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning
- 混合Top-k与Top-p的掩码策略提升高稀疏下的稳定性
- 在视频扩散模型上实现95%稀疏度和16.2倍加速
- 采用知识蒸馏式微调,更好保留生成质量
许多无需训练的稀疏注意力方法能有效加速扩散模型。最近研究指出,使稀疏注意力可训练可进一步提高稀疏度并保持生成质量。本文围绕三个核心问题展开:(1) Top-k与Top-p两种常见掩码规则在何时失效,如何避免?(2) 为何可训练稀疏注意力能达到比无训练方法更高的稀疏度?(3) 使用扩散损失微调稀疏注意力的局限性是什么,如何解决?基于分析,本文提出SpargeAttention2,一种可训练稀疏注意力方法,在不降低生成质量的前提下实现高稀疏度。其包含:(i) 混合Top-k与Top-p的掩码规则,增强高稀疏条件下的鲁棒性;(ii) 高效的可训练稀疏注意力实现;(iii) 基于知识蒸馏的微调目标,更好保持生成质量。在视频扩散模型上的实验表明,SpargeAttention2达到95%注意力稀疏度和16.2倍注意力速度提升,持续优于先前稀疏注意力方法。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。