SLA: Выход за пределы разреженности в диффузионных трансформаторах с помощью настраиваемого разреженно-линейного внимания

Аннотация

В моделях Diffusion Transformer (DiT), особенно для генерации видео, задержка внимания является основным узким местом из-за большой длины последовательности и квадратичной сложности. Мы обнаружили, что веса внимания можно разделить на две части: небольшую долю крупных весов с высоким рангом и оставшиеся веса с очень низким рангом. Это естественным образом подсказывает применение разреженного ускорения к первой части и низкорангового ускорения ко второй. На основе этого наблюдения мы предлагаем SLA (Sparse-Linear Attention) — обучаемый метод внимания, который объединяет разреженное и линейное внимание для ускорения диффузионных моделей. SLA классифицирует веса внимания на критические, маргинальные и незначительные, применяя внимание O(N^2) к критическим весам, O(N) — к маргинальным и пропуская незначительные. SLA объединяет эти вычисления в одном ядре GPU и поддерживает как прямой, так и обратный проходы. Всего за несколько шагов тонкой настройки с использованием SLA модели DiT достигают 20-кратного сокращения вычислений внимания, что приводит к значительному ускорению без потери качества генерации. Эксперименты показывают, что SLA сокращает вычисления внимания на 95% без ухудшения качества сквозной генерации, превосходя базовые методы. Кроме того, мы реализовали эффективное ядро GPU для SLA, которое обеспечивает ускорение вычислений внимания в 13,7 раз и сквозное ускорение генерации видео в 2,2 раза на 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: Выход за пределы разреженности в диффузионных трансформаторах с помощью настраиваемого разреженно-линейного внимания

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

Аннотация

Support