SLA：微調整可能なスパース線形注意を介した拡散トランスフォーマーにおけるスパース性を超えて

要旨

Diffusion Transformer (DiT) モデル、特にビデオ生成において、長いシーケンス長と二次的な計算複雑性により、アテンションのレイテンシが主要なボトルネックとなっている。本研究では、アテンション重みが、高ランクの大きな重みの小さな部分と、非常に低ランクの残りの重みの2つに分離可能であることを発見した。この発見に基づき、前者にはスパース加速を、後者には低ランク加速を適用することが自然に導かれる。この知見を基に、我々はSLA（Sparse-Linear Attention）を提案する。SLAは、スパースアテンションと線形アテンションを融合させた学習可能なアテンション手法であり、拡散モデルの高速化を実現する。SLAはアテンション重みを、クリティカル、マージナル、ネグリジブルの3つのカテゴリに分類し、クリティカルな重みにはO(N^2)のアテンションを、マージナルな重みにはO(N)のアテンションを適用し、ネグリジブルな重みはスキップする。SLAはこれらの計算を単一のGPUカーネルに統合し、順方向および逆方向の計算をサポートする。SLAを用いたわずかなファインチューニングにより、DiTモデルはアテンション計算を20倍削減し、生成品質を損なうことなく大幅な高速化を実現する。実験結果から、SLAはエンドツーエンドの生成品質を低下させることなくアテンション計算を95%削減し、ベースライン手法を上回る性能を示す。さらに、SLAのための効率的なGPUカーネルを実装し、Wan2.1-1.3Bにおいてアテンション計算で13.7倍、ビデオ生成のエンドツーエンドで2.2倍の高速化を達成した。

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