PSA: 効率的な映像理解と生成のためのピラミッド型スパースアテンション

要旨

注意機構は基盤モデルの核心をなすが、その二次計算量はスケーリングにおける重大なボトルネックとして残っている。この課題が効率的な注意機構の開発を促し、スパース性が主要なパラダイムとして台頭してきた。現在の手法は一般に、バイナリマスクを用いてキー・バリューブロック全体を保持または破棄するため、高スパース性下では情報損失が大きくなる。この問題を緩和するため、本論文では映像理解と生成の両タスクに適用可能な多目的モジュールであるPyramid Sparse Attention（PSA）を提案する。PSAはバイナリマスクの代わりに、マスクの細かい粒度を実現する多段階プーリングKV表現を導入する。具体的には、各クエリブロックが重要なKVブロックには低いプーリングレベルを、重要度の低いブロックには高いレベルを動的に割り当て、完全保持と完全剪定の間の情報豊かな補間を生成する。この設計は、固定小数点量子化や計算機ビジョンにおける古典的特徴ピラミッドネットワークに類似し、低計算予算下で計算効率を保ちつつ情報損失を効果的に軽減する。PSAは分離されたブロック・タイル設計を活用したハードウェアフレンドリーなネイティブカーネルで動作し、効率的な実行を保証する。映像理解と生成のベンチマークにおいて、PSAは文脈情報と視覚的忠実度を保持し、優れた効率性と品質のトレードオフを達成しつつ、既存のスパース注意ベースラインを一貫して上回る、または同等の性能を達成する。実装コードとモデル重みはhttp://ziplab.co/PSAで公開している。

English

Attention mechanisms are the core of foundation models, but their quadratic complexity remains a critical bottleneck for scaling. This challenge has driven the development of efficient attention mechanisms, with sparsity emerging as the dominant paradigm. Current methods typically retain or discard entire key-value blocks with binary masks, resulting in substantial information loss under high sparsity. To mitigate this gap, we present Pyramid Sparse Attention (PSA), a versatile module applicable to both video understanding and generation tasks. Instead of binary masking, PSA introduces multi-level pooled KV representations, enabling finer mask granularity. Specifically, each query block dynamically allocates lower pooling levels to critical KV blocks and higher levels to less important ones, creating an informative interpolation between full retention and complete pruning. This design, analogous to fixed-point quantization and classical feature pyramid networks in computer vision, effectively mitigates information loss while preserving computational efficiency under a low compute budget. It works with a native, hardware-friendly kernel that leverages decoupled block-tile design to ensure efficient execution. Across video understanding and generation benchmarks, PSA preserves contextual information and visual fidelity, consistently outperforming or achieving comparable performance over existing sparse attention baselines with superior efficiency-quality trade-offs. Our code and model weights are publicly available at: http://ziplab.co/PSA

PSA: 効率的な映像理解と生成のためのピラミッド型スパースアテンション

PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation

要旨

Support