高速なビデオ拡散のための学習可能なスパースアテンション

Faster Video Diffusion with Trainable Sparse Attention

May 19, 2025

著者: Peiyuan Zhang, Haofeng Huang, Yongqi Chen, Will Lin, Zhengzhong Liu, Ion Stoica, Eric P. Xing, Hao Zhang

cs.AI

要旨

ビデオ拡散Transformer（DiT）のスケーリングは、その二次元3Dアテンションによって制限されていますが、ほとんどのアテンションの質量は少数の位置に集中しています。この観察を基に、我々はVSA（Variable Sparse Attention）を開発しました。これは、訓練と推論の両方で完全なアテンションを置き換える、訓練可能でハードウェア効率の高いスパースアテンションです。VSAでは、軽量な粗い段階でトークンをタイルにプールし、高重みの重要なトークンを特定します。細かい段階では、それらのタイル内でのみトークンレベルのアテンションを計算し、ブロック計算レイアウトに従うことでハード効率を確保します。これにより、エンドツーエンドで訓練可能な単一の微分可能なカーネルが得られ、事後のプロファイリングを必要とせず、FlashAttention3のMFUの85%を維持します。我々は、60Mから1.4BパラメータまでのDiTを事前訓練し、大規模なアブレーション研究とスケーリング則実験を行いました。VSAは、拡散損失の低下なしに訓練FLOPSを2.53倍削減するパレート点に到達します。オープンソースのWan-2.1モデルにVSAを適用すると、アテンション時間が6倍速くなり、エンドツーエンド生成時間が31秒から18秒に短縮され、品質は同等です。これらの結果は、訓練可能なスパースアテンションが完全なアテンションの実用的な代替手段であり、ビデオ拡散モデルのさらなるスケーリングの鍵となることを示しています。

English

Scaling video diffusion transformers (DiTs) is limited by their quadratic 3D attention, even though most of the attention mass concentrates on a small subset of positions. We turn this observation into VSA, a trainable, hardware-efficient sparse attention that replaces full attention at both training and inference. In VSA, a lightweight coarse stage pools tokens into tiles and identifies high-weight critical tokens; a fine stage computes token-level attention only inside those tiles subjecting to block computing layout to ensure hard efficiency. This leads to a single differentiable kernel that trains end-to-end, requires no post-hoc profiling, and sustains 85\% of FlashAttention3 MFU. We perform a large sweep of ablation studies and scaling-law experiments by pretraining DiTs from 60M to 1.4B parameters. VSA reaches a Pareto point that cuts training FLOPS by 2.53times with no drop in diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention time by 6times and lowers end-to-end generation time from 31s to 18s with comparable quality. These results establish trainable sparse attention as a practical alternative to full attention and a key enabler for further scaling of video diffusion models.