SLA: Más allá de la dispersión en transformadores de difusión mediante atención lineal dispersa ajustable

Resumen

En los modelos de Transformador de Difusión (DiT), particularmente para la generación de video, la latencia de atención es un cuello de botella importante debido a la longitud de las secuencias y la complejidad cuadrática. Descubrimos que los pesos de atención pueden separarse en dos partes: una pequeña fracción de pesos grandes con rango alto y los pesos restantes con rango muy bajo. Esto sugiere naturalmente aplicar aceleración dispersa a la primera parte y aceleración de bajo rango a la segunda. Basándonos en este hallazgo, proponemos SLA (Atención Lineal Dispersa), un método de atención entrenable que fusiona atención dispersa y lineal para acelerar los modelos de difusión. SLA clasifica los pesos de atención en categorías críticos, marginales y despreciables, aplicando atención O(N^2) a los pesos críticos, atención O(N) a los pesos marginales y omitiendo los despreciables. SLA combina estos cálculos en un único núcleo de GPU y soporta tanto pasos hacia adelante como hacia atrás. Con solo unos pocos pasos de ajuste fino utilizando SLA, los modelos DiT logran una reducción de 20x en el cálculo de atención, lo que resulta en una aceleración significativa sin pérdida de calidad en la generación. Los experimentos muestran que SLA reduce el cálculo de atención en un 95% sin degradar la calidad de generación de extremo a extremo, superando a los métodos de referencia. Además, implementamos un núcleo de GPU eficiente para SLA, que produce una aceleración de 13.7x en el cálculo de atención y una aceleración de 2.2x en la generación de video de extremo a extremo en 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: Más allá de la dispersión en transformadores de difusión mediante atención lineal dispersa ajustable

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

Resumen

Support