Sol-Attn: オンザフライ注意機構スパース化による動画生成推論の高速化

要旨

拡散トランスフォーマーは高忠実度の動画生成に不可欠であるが、長いトークン系列によりアテンションが推論の支配的なボトルネックとなる。学習不要の動的スパースアテンションは、選択されたキー・バリューブロックのみを計算することでこのボトルネックを緩和するが、既存手法では以下の2つの理由からアテンションを効率的かつ正確にスパース化することが困難である。(1) 硬直的で予測不可能かつコストのかかるルーティング：プロキシスコアによる上位ランクのブロックを固定割合で選択すると固定予算が課されるのに対し、累積プロキシ確率質量が目標に達するまでブロックを保持すると動的だが不均衡な予算が生じる。いずれもプロキシスコアの計算と具体化に無視できないオーバーヘッドが発生する。(2) 損失を伴う保持・破棄のスパース化：選択されなかったブロックは完全に破棄され、積極的なスパース性の下で精度が低下する。これらの制約から、より低コストな動的予算ルーティングと精度低下の制限が動機となる。本論文では、学習不要のSol-Attn（オンラインアテンションのスパース化）を導入する。これは動的ルーティング、スパース計算、近似補正を単一のオンラインソフトマックスパスに統合し、スパースアテンションにおける精度と効率のより良いトレードオフを実現する。Sol-Attnの中核は、プロキシスコア再利用を伴うオン・ザ・フライのブロックしきい値処理であり、オンラインソフトマックス中にブロックプロキシスコアをしきい値と比較して重要なブロックを選択する。この設計により、プロキシマップを具体化することなく動的で制御可能なブロック予算を実現し、選択されなかったブロックのプロキシスコアを直接再利用してその寄与を近似する。画像・動画生成タスクにおける実験により、Sol-Attnが学習不要のスパースアテンションの品質・効率フロンティアを前進させ、動画生成と編集においてそれぞれ2.1倍および2.3倍のエンドツーエンド高速化を視覚品質を維持しながら実現することを示す。

English

Diffusion transformers are essential for high-fidelity video generation, but long token sequences make attention a dominant inference bottleneck. Training-free dynamic sparse attention alleviates this bottleneck by computing only selected key-value blocks, yet existing methods struggle to sparsify attention both efficiently and accurately for two reasons: (1) Rigid, unpredictable, and costly routing: selecting a fixed fraction of top-ranked blocks by proxy score imposes fixed budgets, whereas retaining blocks to reach a target cumulative proxy probability mass yields dynamic but potentially imbalanced budgets; both incur non-negligible overhead from computing and materializing proxy scores. (2) Lossy keep-or-drop sparsification: unselected blocks are discarded entirely, degrading accuracy under aggressive sparsity. These limitations motivate cheaper dynamic-budget routing while limiting accuracy degradation. In this paper, we introduce training-free Sol-Attn (Sparsifying online attention), which unifies dynamic routing, sparse computation, and approximation correction in a single online-softmax pass, achieving a better accuracy-efficiency trade-off in sparse attention. The core of Sol-Attn is on-the-fly block thresholding with proxy-score reuse, which selects critical blocks by comparing block proxy scores against a threshold during online softmax. This design enables dynamic yet controllable block budgets without materializing the proxy map, while directly reusing the proxy scores of unselected blocks to approximate their contribution. Experiments across image and video generation tasks show that Sol-Attn advances the quality-efficiency frontier of training-free sparse attention, delivering 2.1 times and 2.3 times end-to-end speedups for video generation and editing, respectively, while preserving visual quality.