SLA: Além da Esparsidade em Transformadores de Difusão via Atenção Esparsa-Linear Ajustável

Resumo

Nos modelos Diffusion Transformer (DiT), particularmente para geração de vídeos, a latência de atenção é um grande gargalo devido ao comprimento longo das sequências e à complexidade quadrática. Descobrimos que os pesos de atenção podem ser separados em duas partes: uma pequena fração de pesos grandes com alta classificação e os pesos restantes com classificação muito baixa. Isso sugere naturalmente a aplicação de aceleração esparsa à primeira parte e aceleração de baixa classificação à segunda. Com base nessa descoberta, propomos o SLA (Sparse-Linear Attention), um método de atenção treinável que combina atenção esparsa e linear para acelerar modelos de difusão. O SLA classifica os pesos de atenção em categorias críticos, marginais e negligenciáveis, aplicando atenção O(N^2) aos pesos críticos, atenção O(N) aos pesos marginais e ignorando os negligenciáveis. O SLA combina esses cálculos em um único kernel de GPU e suporta tanto passos diretos quanto reversos. Com apenas alguns passos de ajuste fino usando o SLA, os modelos DiT alcançam uma redução de 20x no cálculo de atenção, resultando em uma aceleração significativa sem perda de qualidade na geração. Experimentos mostram que o SLA reduz o cálculo de atenção em 95% sem degradar a qualidade de geração de ponta a ponta, superando métodos de linha de base. Além disso, implementamos um kernel de GPU eficiente para o SLA, que proporciona uma aceleração de 13,7x no cálculo de atenção e uma aceleração de 2,2x na geração de vídeos de ponta a ponta no 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: Além da Esparsidade em Transformadores de Difusão via Atenção Esparsa-Linear Ajustável

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

Resumo

Support