SLA: Über Sparsity hinaus in Diffusion Transformers durch feinabstimmbare Sparse-Linear-Attention

papers.abstract

In Diffusions-Transformer (DiT)-Modellen, insbesondere bei der Videogenerierung, stellt die Latenz der Aufmerksamkeit einen wesentlichen Engpass dar, bedingt durch die lange Sequenzlänge und die quadratische Komplexität. Wir stellen fest, dass die Aufmerksamkeitsgewichte in zwei Teile separiert werden können: einen kleinen Anteil großer Gewichte mit hohem Rang und die verbleibenden Gewichte mit sehr niedrigem Rang. Dies legt nahe, für den ersten Teil eine spärliche Beschleunigung und für den zweiten Teil eine Niedrigrang-Beschleunigung anzuwenden. Basierend auf dieser Erkenntnis schlagen wir SLA (Sparse-Linear Attention) vor, eine trainierbare Aufmerksamkeitsmethode, die spärliche und lineare Aufmerksamkeit kombiniert, um Diffusionsmodelle zu beschleunigen. SLA klassifiziert die Aufmerksamkeitsgewichte in kritische, marginale und vernachlässigbare Kategorien, wobei O(N^2)-Aufmerksamkeit auf kritische Gewichte, O(N)-Aufmerksamkeit auf marginale Gewichte angewendet wird und vernachlässigbare Gewichte übersprungen werden. SLA kombiniert diese Berechnungen in einem einzigen GPU-Kernel und unterstützt sowohl Vorwärts- als auch Rückwärtspass. Mit nur wenigen Feinabstimmungsschritten unter Verwendung von SLA erreichen DiT-Modelle eine 20-fache Reduktion der Aufmerksamkeitsberechnung, was zu einer erheblichen Beschleunigung ohne Qualitätsverlust bei der Generierung führt. Experimente zeigen, dass SLA die Aufmerksamkeitsberechnung um 95 % reduziert, ohne die end-to-end-Generierungsqualität zu beeinträchtigen, und dabei Baseline-Methoden übertrifft. Zusätzlich implementieren wir einen effizienten GPU-Kernel für SLA, der eine 13,7-fache Beschleunigung der Aufmerksamkeitsberechnung und eine 2,2-fache end-to-end-Beschleunigung bei der Videogenerierung auf Wan2.1-1.3B bewirkt.

English

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.

SLA: Über Sparsity hinaus in Diffusion Transformers durch feinabstimmbare Sparse-Linear-Attention

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

papers.abstract

Support