SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable
  Sparse-Linear Attention
Authors: name: Jintao Zhang, name: Haoxu Wang, name: Kai Jiang, name: Shuo Yang, name: Kaiwen Zheng, name: Haocheng Xi, name: Ziteng Wang, name: Hongzhou Zhu, name: Min Zhao, name: Ion Stoica, name: Joseph E. Gonzalez, name: Jun Zhu, name: Jianfei Chen
Organization: Tsinghua University

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.