SLA : Au-delà de la parcimonie dans les transformeurs de diffusion via une attention linéaire parcimonieuse ajustable

papers.abstract

Dans les modèles de Transformateurs de Diffusion (DiT), en particulier pour la génération de vidéos, la latence d'attention constitue un goulot d'étranglement majeur en raison de la longueur des séquences et de la complexité quadratique. Nous constatons que les poids d'attention peuvent être séparés en deux parties : une petite fraction de poids élevés de rang élevé et les poids restants de rang très faible. Cela suggère naturellement d'appliquer une accélération parcimonieuse à la première partie et une accélération de bas rang à la seconde. Sur la base de cette observation, nous proposons SLA (Sparse-Linear Attention), une méthode d'attention entraînable qui fusionne l'attention parcimonieuse et linéaire pour accélérer les modèles de diffusion. SLA classe les poids d'attention en catégories critiques, marginales et négligeables, appliquant une attention de complexité O(N²) aux poids critiques, une attention de complexité O(N) aux poids marginaux, et ignorant les poids négligeables. SLA combine ces calculs dans un seul noyau GPU et prend en charge les passes avant et arrière. Avec seulement quelques étapes de réglage fin utilisant SLA, les modèles DiT atteignent une réduction de 20x dans le calcul de l'attention, entraînant une accélération significative sans perte de qualité de génération. Les expériences montrent que SLA réduit le calcul de l'attention de 95 % sans dégrader la qualité de génération de bout en bout, surpassant les méthodes de référence. De plus, nous implémentons un noyau GPU efficace pour SLA, qui permet une accélération de 13,7x dans le calcul de l'attention et une accélération de 2,2x de bout en bout dans la génération de vidéos sur 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 : Au-delà de la parcimonie dans les transformeurs de diffusion via une attention linéaire parcimonieuse ajustable

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

papers.abstract

Support