SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention

📄 arXiv: 2509.24006v2 📥 PDF

作者: Jintao Zhang, Haoxu Wang, Kai Jiang, Shuo Yang, Kaiwen Zheng, Haocheng Xi, Ziteng Wang, Hongzhou Zhu, Min Zhao, Ion Stoica, Joseph E. Gonzalez, Jun Zhu, Jianfei Chen

分类: cs.LG, cs.AI, cs.CV

发布日期: 2025-09-28 (更新: 2025-11-19)

🔗 代码/项目: GITHUB

💡 一句话要点

提出SLA：一种可微调的稀疏线性注意力机制，加速Diffusion Transformer模型。

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

关键词: Diffusion Transformer 稀疏注意力 线性注意力 视频生成 模型加速 可训练注意力 GPU内核优化

📋 核心要点

  1. Diffusion Transformer模型在视频生成中面临长序列和二次复杂度带来的注意力计算瓶颈。
  2. SLA通过将注意力权重分类，对关键权重使用稀疏注意力，对边缘权重使用线性注意力，忽略可忽略权重。
  3. 实验表明，SLA在不损失生成质量的前提下，显著减少了注意力计算量，并实现了端到端加速。

📝 摘要（中文）

本文提出了一种名为SLA（Sparse-Linear Attention）的可训练注意力方法，旨在融合稀疏注意力和线性注意力，从而加速Diffusion Transformer (DiT) 模型，尤其是在视频生成任务中。研究发现，注意力权重可以分解为两部分：一小部分具有高秩的大权重和剩余的具有极低秩的权重。SLA将注意力权重分为关键、边缘和可忽略三类，分别应用O(N^2)复杂度、O(N)复杂度和跳过操作。SLA将这些计算融合到单个GPU内核中，并支持前向和后向传播。通过少量微调，DiT模型使用SLA可实现20倍的注意力计算量减少，从而在不损失生成质量的情况下显著加速。实验表明，SLA在不降低端到端生成质量的情况下，将注意力计算量减少了95%，优于基线方法。此外，SLA的高效GPU内核在Wan2.1-1.3B数据集上实现了13.7倍的注意力计算加速和2.2倍的端到端视频生成加速。代码已开源。

🔬 方法详解

问题定义：Diffusion Transformer (DiT) 模型在处理长序列（如视频生成）时，注意力机制的计算复杂度为O(N^2)，成为性能瓶颈。现有的稀疏注意力方法虽然可以降低计算复杂度，但可能导致生成质量下降。

核心思路：论文的核心思路是将注意力权重分解为重要和不重要的部分，并对不同部分采用不同的加速策略。具体来说，将注意力权重分为高秩的重要权重和低秩的不重要权重，分别使用稀疏注意力和线性注意力进行加速。

技术框架：SLA方法包含以下几个主要步骤：1) 将注意力权重分为“关键”、“边缘”和“可忽略”三类。2) 对“关键”权重应用标准的O(N^2)注意力计算。3) 对“边缘”权重应用线性注意力，将复杂度降低到O(N)。4) 直接跳过“可忽略”的权重。5) 将所有计算合并到一个高效的GPU内核中，以实现加速。

关键创新：SLA的关键创新在于提出了一种可训练的注意力机制，能够自适应地学习哪些注意力权重是重要的，哪些是不重要的。通过这种方式，SLA可以在保证生成质量的前提下，最大程度地减少注意力计算量。与传统的稀疏注意力方法相比，SLA不需要手动设计稀疏模式，而是通过学习自动确定。

关键设计：SLA的关键设计包括：1) 使用一个可学习的分类器来将注意力权重分为“关键”、“边缘”和“可忽略”三类。2) 设计了一个高效的GPU内核，能够将稀疏注意力和线性注意力计算融合在一起。3) 使用少量微调来使DiT模型适应SLA。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SLA在DiT模型上实现了显著的加速效果。实验结果表明，SLA可以将注意力计算量减少95%，同时保持与原始DiT模型相当的生成质量。在Wan2.1-1.3B数据集上，SLA实现了13.7倍的注意力计算加速和2.2倍的端到端视频生成加速，明显优于其他基线方法。这些结果表明SLA是一种高效且有效的注意力加速方法。

🎯 应用场景

SLA具有广泛的应用前景，尤其是在需要处理长序列数据的生成任务中，例如高清视频生成、3D内容生成等。通过降低注意力计算复杂度，SLA可以显著减少计算资源消耗，并提高生成速度，从而使得更大规模、更高质量的生成模型成为可能。此外，SLA还可以应用于其他Transformer模型，以加速其在各种任务中的推理速度。

📄 摘要（原文）

In Diffusion Transformer (DiT) models, particularly for video generation, attention latency is a major bottleneck due to the long sequence length and the quadratic complexity. We find that attention weights can be separated into two parts: a small fraction of large weights with high rank and the remaining weights with very low rank. This naturally suggests applying sparse acceleration to the first part and low-rank acceleration to the second. Based on this finding, we propose SLA (Sparse-Linear Attention), a trainable attention method that fuses sparse and linear attention to accelerate diffusion models. SLA classifies attention weights into critical, marginal, and negligible categories, applying O(N^2) attention to critical weights, O(N) attention to marginal weights, and skipping negligible ones. SLA combines these computations into a single GPU kernel and supports both forward and backward passes. With only a few fine-tuning steps using SLA, DiT models achieve a 20x reduction in attention computation, resulting in significant acceleration without loss of generation quality. Experiments show that SLA reduces attention computation by 95% without degrading end-to-end generation quality, outperforming baseline methods. In addition, we implement an efficient GPU kernel for SLA, which yields a 13.7x speedup in attention computation and a 2.2x end-to-end speedup in video generation on Wan2.1-1.3B. The code is available at https://github.com/thu-ml/SLA.