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

The paper introduces SLA (Sparse-Linear Attention), a trainable attention method for Diffusion Transformers (DiTs) that combines sparse and linear attention to reduce computational cost.

The authors observe that attention weights can be decomposed into a small fraction of large weights with high rank and a large fraction of low-rank weights.

SLA classifies attention blocks into critical, marginal, and negligible, applying exact attention to critical blocks, linear attention to marginal blocks, and skipping negligible ones.

This approach achieves high sparsity (95%) without quality loss.

Experiments on Wan2.1-1.3B show a 20x reduction in attention computation, a 13.7x speedup in the attention kernel, and a 2.2x end-to-end speedup in video generation, with quality comparable to full attention.

SLA outperforms baselines like VSA, VMoBa, and SpargeAttn in both quality and efficiency.

The method requires only a few fine-tuning steps (e.g., 2000 steps) and is implemented as a fused GPU kernel supporting both forward and backward passes.

Ablation studies confirm the effectiveness of fusing sparse and linear attention, with softmax as the best activation function and a top-5% critical threshold providing the best trade-off.

The code is available at https://github.com/thu-ml/SLA.

The paper also includes experiments on image generation with LightningDiT, where SLA achieves better FID than full attention at 87.5% sparsity.

Additional optimizations like lookup tables and pre-aggregation are discussed in the appendix.

The work addresses limitations of existing sparse and linear attention methods, which struggle to achieve high sparsity or maintain quality in video diffusion models.

SLA's design leverages the low-rank structure of non-critical…