SLA: Voorbij sparsiteit in diffusie-transformers via fijn afstelbare Sparse-Lineaire Aandacht

Samenvatting

In Diffusion Transformer (DiT)-modellen, met name voor videogeneratie, vormt de latentie van aandacht een belangrijke bottleneck vanwege de lange sequentielengte en de kwadratische complexiteit. We ontdekken dat aandachtswaarden kunnen worden gescheiden in twee delen: een kleine fractie van grote waarden met een hoge rang en de overige waarden met een zeer lage rang. Dit suggereert natuurlijk het toepassen van versnelde sparse berekeningen op het eerste deel en versnelde low-rank berekeningen op het tweede. Op basis van deze bevinding stellen we SLA (Sparse-Linear Attention) voor, een trainbare aandachtmethode die sparse en lineaire aandacht combineert om diffusiemodellen te versnellen. SLA classificeert aandachtswaarden in kritieke, marginale en verwaarloosbare categorieën, waarbij O(N^2) aandacht wordt toegepast op kritieke waarden, O(N) aandacht op marginale waarden, en verwaarloosbare waarden worden overgeslagen. SLA combineert deze berekeningen in een enkele GPU-kernel en ondersteunt zowel voorwaartse als achterwaartse passes. Met slechts enkele fine-tuningstappen met SLA bereiken DiT-modellen een 20-voudige reductie in aandachtberekeningen, wat resulteert in een aanzienlijke versnelling zonder verlies van generatiekwaliteit. Experimenten tonen aan dat SLA de aandachtberekeningen met 95% vermindert zonder de end-to-end generatiekwaliteit te verslechteren, wat beter presteert dan baseline-methoden. Daarnaast implementeren we een efficiënte GPU-kernel voor SLA, wat een 13,7-voudige versnelling in aandachtberekeningen en een 2,2-voudige end-to-end versnelling in videogeneratie oplevert op Wan2.1-1.3B.

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: Voorbij sparsiteit in diffusie-transformers via fijn afstelbare Sparse-Lineaire Aandacht

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

Samenvatting

Support