SLA: Oltre la sparsità nei Transformer di diffusione tramite Attenzione Sparse-Lineare fine-tunabile

Abstract

Nei modelli Diffusion Transformer (DiT), in particolare per la generazione di video, la latenza dell'attenzione rappresenta un collo di bottiglia significativo a causa della lunghezza delle sequenze e della complessità quadratica. Abbiamo osservato che i pesi dell'attenzione possono essere separati in due parti: una piccola frazione di pesi grandi con rango elevato e i restanti pesi con rango molto basso. Ciò suggerisce naturalmente di applicare un'accelerazione sparsa alla prima parte e un'accelerazione a basso rango alla seconda. Sulla base di questa scoperta, proponiamo SLA (Sparse-Linear Attention), un metodo di attenzione addestrabile che combina attenzione sparsa e lineare per accelerare i modelli di diffusione. SLA classifica i pesi dell'attenzione in categorie critiche, marginali e trascurabili, applicando un'attenzione O(N^2) ai pesi critici, O(N) ai pesi marginali e saltando quelli trascurabili. SLA unisce questi calcoli in un singolo kernel GPU e supporta sia il passaggio in avanti che quello all'indietro. Con pochi passi di fine-tuning utilizzando SLA, i modelli DiT ottengono una riduzione di 20x nel calcolo dell'attenzione, ottenendo un'accelerazione significativa senza perdita di qualità nella generazione. Gli esperimenti dimostrano che SLA riduce il calcolo dell'attenzione del 95% senza degradare la qualità della generazione end-to-end, superando i metodi di base. Inoltre, implementiamo un kernel GPU efficiente per SLA, che produce un'accelerazione di 13.7x nel calcolo dell'attenzione e un'accelerazione end-to-end di 2.2x nella generazione di video su 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: Oltre la sparsità nei Transformer di diffusione tramite Attenzione Sparse-Lineare fine-tunabile

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

Abstract

Support