要約

拡散変換器（DiT）モデル、特に動画生成において、長大なシーケンス長と二次時間計算量の特性から、アテンションの遅延が主要なボトルネックとなっている。本研究では、アテンション重みが大きく高ランクな少数の重みと、低ランクな残りの重みに分離可能であることを発見した。この特性は、前者に対してスパース加速を、後者に対して低ランク加速を適用する自然なアプローチを示唆する。この知見に基づき、本研究ではスパースと線形アテンションを統合した学習可能なアテンション手法「SLA（Sparse-Linear Attention）」を提案する。SLAはアテンション重みを「重要」「僅差」「無視可能」の三つのカテゴリに分類し、それぞれに対してO(N²)のアテンション計算、O(N)のアテンション計算を適用し、無視可能な重みはスキップする。SLAはこれらの計算を1つのGPUカーネルに統合し、順伝搬および逆伝搬の両方をサポートする。SLAをわずか数ステップのファインチューニングで適用するだけで、DiTモデルにおけるアテンション計算量を20倍削減し、生成品質の低下を伴わずに大幅な高速化を実現する。実験の結果、SLAはエンドツーエンドの生成品質を損なわずにアテンション計算量を95％削減し、ベースライン手法を上回る性能を示した。さらに、SLA用の効率的なGPUカーネルを実装した結果、アテンション計算は13.7倍高速化され、Wan2.1-1.3Bモデルにおける動画生成全体の処理速度も2.2倍向上した。