効率的な拡散変換器のための学習可能な対数線形スパースアテンション

要旨

拡散トランスフォーマー（DiT）は視覚生成において最先端の性能を達成しているが、その二次関数的なセルフアテンションコストが、長いトークン系列へのスケーリングを根本的に制限している。近年のTop-Kスパースアテンション手法は、トークンをブロック単位の表現に圧縮し、関連するキーブロックの小さな集合を選択することでDiTの計算量を削減するが、依然として（i）圧縮トークンに対する二次関数的な選択コストと、（ii）系列長の増加に伴いモデル品質を維持するために必要なK値の増大という課題に直面している。我々は、この非効率性の原因が単一レベルの設計にあることを明らかにした。単一の粗いレベルでは大域的な構造を表現するには不十分であるためである。本論文では、極めて長いトークン系列のための学習可能なスパースアテンション機構である対数線形スパースアテンション（LLSA）を提案する。LLSAは階層構造を利用することで、選択コストとアテンションコストの両方を二次関数的複雑度から対数線形複雑度に削減する。LLSAは階層的Top-K選択を実行し、前のレベルで見つかったインデックスを用いて段階的にスパースなTop-K選択を適用する。さらに、アテンション計算中に異なる粒度のより少ないトークンを使用しながら大域的な文脈を保持する、階層的KVエンリッチメント機構を導入する。効率的な学習をサポートするため、フォワードパスとバックワードパスの両方でスパースなインデックスのみを使用し、密なアテンションマスクを不要とする高性能なGPU実装を開発した。パッチ化やVAEエンコーディングを使用せずに、高解像度のピクセル空間画像生成に対してLLSAを評価した。LLSAは、256x256ピクセルのトークン系列において、アテンション推論を28.27倍、DiT学習を6.09倍高速化しつつ、生成品質を維持した。この結果は、LLSAが長系列DiTを効率的に学習するための有望な方向性を提供することを示している。コードはhttps://github.com/SingleZombie/LLSA で公開されている。

English

Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost fundamentally limits scaling to long token sequences. Recent Top-K sparse attention approaches reduce the computation of DiTs by compressing tokens into block-wise representation and selecting a small set of relevant key blocks, but still suffer from (i) quadratic selection cost on compressed tokens and (ii) increasing K required to maintain model quality as sequences grow. We identify that their inefficiency is due to the single-level design, as a single coarse level is insufficient to represent the global structure. In this paper, we introduce Log-linear Sparse Attention (LLSA), a trainable sparse attention mechanism for extremely long token sequences that reduces both selection and attention costs from quadratic to log-linear complexity by utilizing a hierarchical structure. LLSA performs hierarchical Top-K selection, progressively adopting sparse Top-K selection with the indices found at the previous level, and introduces a Hierarchical KV Enrichment mechanism that preserves global context while using fewer tokens of different granularity during attention computation. To support efficient training, we develop a high-performance GPU implementation that uses only sparse indices for both the forward and backward passes, eliminating the need for dense attention masks. We evaluate LLSA on high-resolution pixel-space image generation without using patchification and VAE encoding. LLSA accelerates attention inference by 28.27x and DiT training by 6.09x on 256x256 pixel token sequences, while maintaining generation quality. The results demonstrate that LLSA offers a promising direction for training long-sequence DiTs efficiently. Code is available at: https://github.com/SingleZombie/LLSA

効率的な拡散変換器のための学習可能な対数線形スパースアテンション

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

要旨

Support