PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2509.24006v2cs.LGcs.AIcs.CV

TL;DR

Long sequences make quadratic attention a bottleneck in video Diffusion Transformers, while sparse and linear attention alone have important quality or sparsity limitations. SLA combines exact sparse computation for critical weights, linear compensation for marginal weights, and skipping for negligible weights, achieving large attention and end-to-end speedups without degrading video quality.

  • Problem

    Video Diffusion Transformers require costly quadratic attention over long sequences, while existing linear and sparse methods have quality or sparsity limitations.

  • Method

    SLA is a trainable hybrid attention method that applies exact sparse attention to critical weights, linear attention to marginal weights, and skips negligible weights.

  • Results

    20× lower attention computation, 13.7× GPU-kernel speedup, and 2.2× end-to-end speedup are achieved on Wan2.1-1.3B without degrading video quality.

  • Takeaways & Limitations

    SLA enables substantially higher sparsity while preserving video-generation quality and outperforming sparse and linear baselines in efficiency and quality.

  • Takeaways & Limitations

    Linear attention in SLA requires fine-tuning because linear attention alone struggles to approximate full-attention outputs, especially in diffusion models.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Attention is the main computational bottleneck in video DiTs because its quadratic cost meets sequence lengths of 10K–100K. SLA addresses limitations of sparse and linear attention by combining their strengths through attention-weight structure.

  • 1 INTRODUCTION: Attention is the primary DiT bottleneck for video generation because its O(N^2) complexity applies to sequences typically spanning 10K–100K tokens.Other Transformer operations mostly scale linearly with sequence length.
  • 1 INTRODUCTION: Existing linear attention often severely degrades video quality, whereas sparse methods typically achieve only 40–60% sparsity below sequence lengths of 50K.Reported 80–85% sparsity generally occurs only for much longer sequences of 100K–300K.
  • 1 INTRODUCTION: SLA is motivated by attention weights consisting of a small high-rank fraction and a large extremely low-rank remainder, enabling complementary sparse and low-rank acceleration.The decomposition explains why either approach alone may be insufficient.
  • 1 INTRODUCTION: SLA partitions attention blocks into critical, marginal, and negligible categories, using exact sparse attention, linear attention, and skipping respectively.This design is intended to raise sparsity from 70% to 95% while maintaining accuracy.
  • 1 INTRODUCTION: SLA reduces attention computation by 95% without degrading video quality, while achieving 13.7× attention-kernel and 2.2× end-to-end speedups.The result holds for Wan2.1-1.3B at a moderate sequence length of 30K.

2 PRELIMINARY

Standard attention explicitly forms an N×N score and weight structure, creating quadratic cost. Sparse attention masks computation, while linear attention reorders feature-mapped products to avoid constructing those matrices.

  • 2 PRELIMINARY: Standard attention computes S = QK^T/d, P = Softmax(S), and O = PV, requiring O(N^2d) operations.The quadratic cost arises from forming attention interactions across the sequence.
  • 2 PRELIMINARY: Sparse attention applies a binary mask to attention weights and skips masked score and value multiplications.A threshold can determine which entries remain active.
  • 2 PRELIMINARY: Practical sparse FlashAttention partitions Q, K, V, scores, weights, and masks into blocks, skipping computations only for fully masked blocks.Block-level sparsity is used because element-wise sparsity is inefficient on modern GPUs.
  • 2 PRELIMINARY: Linear attention replaces standard weights with feature-map products and reorders multiplication through ϕ(K)^T V before applying ϕ(Q).This avoids explicitly constructing the N×N score and weight matrices.

3 MOTIVATION AND ANLYSIS

Attention weights are highly concentrated: a small high-rank subset carries the largest values, while most remaining weights are low-rank and can be treated differently. SLA uses this structure to combine sparse and linear computation while preserving quality at high sparsity.

  • 3.1 MOTIVATION OF SLA: Only about 8.1% of Wan2.1 attention weights exceed the average 1/N, while approximately 45% fall below 1/(100N).Skipping the smallest 45% introduces less than 3% relative L1 error.
  • 3.1 MOTIVATION OF SLA: Intermediate weights are too important to omit but too costly to compute fully, motivating critical, marginal, and negligible categories.Critical weights use sparse FlashAttention, negligible weights are skipped, and marginal weights use linear attention.
  • 3.1 MOTIVATION OF SLA: 95% sparsity in SLA preserves video quality and outperforms linear-only and 90%-sparse alternatives in Wan2.1 examples.At 95% sparsity, SLA’s complexity is nearly half that of 90% sparse attention because linear attention is nearly negligible.
  • 3.2 SEPARATING ATTENTION WEIGHTS: SPARSE FEW, LOW-RANK MANY: The full attention matrix separates into a small subset below 10% with comparable rank and a large subset above 90% with very low rank.This supports sparse acceleration for the first subset and low-rank approximation for the second.
  • 3.2 SEPARATING ATTENTION WEIGHTS: SPARSE FEW, LOW-RANK MANY: Removing top attention values leaves an extremely low-rank matrix, explaining why linear attention can replace the low-rank component after sparse masking.The proposed decomposition uses P ⊙ (1 − M) for the component targeted by linear attention.

4 SLA

SLA classifies compressed attention blocks into critical, marginal, and negligible categories, combining sparse exact attention with linear attention and skipping negligible computation. A learnable projection aligns the linear component with the overall attention output.

  • 4 SLA: SLA classifies compressed attention blocks into critical, marginal, and negligible categories, recorded in a compressed mask.The top kh% positions are critical, the bottom kl% negligible, and the remaining positions marginal.
  • 4.1 SPARSE ATTENTION IN SLA: Critical blocks use sparse FlashAttention, while negligible blocks are skipped during the forward computation.The sparse component applies blockwise online softmax only where the compressed mask marks a block as critical.
  • 4.2 LINEAR ATTENTION IN SLA: Marginal blocks are processed with linear attention rather than approximating their original softmax outputs directly.The linear component acts as a learnable compensation that improves sparse attention and requires fine-tuning of the target model.
  • 4.2 LINEAR ATTENTION IN SLA: A learnable projection transforms the linear-attention output before adding it to the sparse output, with cost O(Nd^2) versus O(N^2d) for full attention.The projection is intended to reduce distribution mismatch between softmax and linear attention.
  • 4 SLA: Figure 4 summarizes SLA by assigning different computational complexities to the three attention-weight categories and showing the compressed-weight forward algorithm.The overview links the high-level classification to the detailed forward procedure.

5 FINE-TUNING USING SLA

SLA can replace a diffusion model’s original attention and be adapted through a few fine-tuning steps, with fused forward and backward computations for its sparse and linear components. Its implementation precomputes reusable linear-attention intermediates and combines gradient contributions efficiently.

  • 5 FINE-TUNING USING SLA: Replacing original attention with SLA requires only a few fine-tuning steps on data consistent with the model’s pretraining data.The section describes both forward and backward passes and notes additional efficiency optimizations.
  • 5.1 FORWARD PASS: The forward pass precomputes h_j = ϕ(K_j)^⊤V_j and z_j = rowsum(ϕ(K_j)^⊤), reducing negligible-block work to matrix additions.The precomputed values are reused when the compressed mask marks a block for linear attention.
  • 5.1 FORWARD PASS: Algorithm 1 partitions Q, K, and V into blocks, predicts a compressed mask, and applies sparse attention or linear accumulation according to each mask value.Critical blocks use online softmax, while noncritical blocks accumulate the precomputed linear-attention intermediates.
  • 5.2 BACKWARD PASS: The backward pass computes gradients for both sparse and linear components and fuses them into a single GPU kernel.Sparse gradients follow FlashAttention-style derivations, while linear gradients are obtained through the chain rule and aggregated intermediates.
  • 5.2 BACKWARD PASS: Gradient computation propagates derivatives through sparse outputs and linear outputs to obtain gradients for Q, K, V, and feature-mapped inputs.The notation uses dOs and dOl for gradients of the sparse and linear outputs with respect to the loss.

6 EXPERIMENT

Experiments evaluate SLA on Wan2.1-1.3B using video-quality and efficiency metrics, baseline comparisons, kernel measurements, and ablations. SLA preserves quality while substantially reducing attention and generation cost.

  • 6.2 EFFECTIVENESS: 19.3× efficiency gain maintains video quality comparable to Full Attention on Wan2.1-1.3B, while outperforming baselines under greater sparsity.At 95% sparsity, SLA is about 3× more efficient than its 85%-sparsity setting and still produces better video quality.
  • 6.3 EFFICIENCY: 13.7× faster forward attention and 2.2× faster end-to-end generation demonstrate SLA's practical acceleration on RTX5090.Attention latency falls from 97s to 11s, while SLA also outperforms VSA and VMoBa in forward and backward passes.
  • 6.4 ABLATION STUDY: SLA achieves the best generation quality and higher efficiency than Sparse Only and S+L, supporting fused sparse-linear attention.The ablation compares Sparse Only, Linear Only, and S+L on Wan2.1.
  • 6.4 ABLATION STUDY: Softmax generally provides better quality and efficiency than elu+1 and hedgehog in SLA's linear-attention component.This activation-function comparison is reported in Table 2.
  • 6.4 ABLATION STUDY: kh = 5% offers the best quality-efficiency trade-off, remaining close to Full Attention while saving computation versus larger kh values.It saves about half and a quarter of the computation compared with kh = 10% and kh = 20%, respectively.
  • 6.5 VISIBLE EXAMPLES: SLA produces videos comparable to Full Attention at 95% sparsity, whereas other methods show noticeable distortions below 90% sparsity.The comparison uses visible examples from Figures 5 and 7.

7 RELATED WORK

Related work improves attention efficiency mainly through sparse or linear attention, with sparse methods often masking computation at test time without training.

  • 7 RELATED WORK: Efficient attention research primarily follows sparse and linear directions as sequence lengths make quadratic attention a bottleneck.Most sparse methods accelerate inference without training by masking computation at test time.

8 CONCLUSION

The paper proposes SLA as a trainable attention method that unifies sparse and linear attention for Diffusion Transformers. It reports major reductions in attention and generation cost without degrading video quality.

  • 8 CONCLUSION: SLA assigns O(N^2) computation to critical weights, O(N) computation to marginal weights, and skips negligible computations.This importance-based allocation unifies sparse and linear attention in a trainable DiT method.
  • 8 CONCLUSION: 20× lower attention computation, 13.7× faster GPU kernels, and 2.2× faster end-to-end generation are achieved on Wan2.1-1.3B without degrading video quality.The reported gains follow a few fine-tuning steps using SLA.

A.1 MORE VISIBLE EXAMPLES

Additional visible examples compare SLA with other attention methods on Wan2.1 video generation. SLA is reported to maintain higher quality at greater sparsity levels.

  • A.1 MORE VISIBLE EXAMPLES: SLA consistently achieves higher video quality than baseline attention methods even under greater sparsity.Figure 7 provides additional Wan2.1 examples for this comparison.

A.2 EXPERIMENTS FOR IMAGE GENERATION

Image-generation experiments evaluate SLA and baselines on LightningDiT using FID for image quality and FLOPs for computational complexity. At the highest sparsity, SLA outperforms all baselines and surpasses full attention on FID, while complementary kernel optimizations improve efficiency across sparsity levels.

  • Experimental setup: The experiments use LightningDiT-1p0B/1 with 1.03B parameters, trained on ImageNet at 512 × 512 resolution.Evaluation is conducted on a pretraining task.
  • Metrics: FID measures image quality, while FLOPs measure computational complexity.
  • Results: At the highest sparsity level, SLA outperforms all baselines and even surpasses full attention on FID.The result is summarized in Table 3 and is consistent with the reported video experiments.
  • Efficiency optimizations: Lookup tables reduce mask-reading memory traffic when the critical-block mask is highly sparse.The optimization preprocesses nonzero row and column positions and accesses only the resulting lookup table.
  • Efficiency optimizations: Pre-aggregation replaces most additions with subtractions when more than 90% of mask entries are zero.This optimization precomputes row or column sums and subtracts contributions from nonzero entries.
  • Efficiency optimizations: The Method of Four Russians accelerates linear-attention computation when mask sparsity is around 50%.It precomputes subset sums for groups of consecutive blocks.
Loading 2509.24006v2…