用可训练的稀疏注意力加速视频扩散模型,训练和推理都更快。
VSA: Faster Video Diffusion with Trainable Sparse Attention
- 设计轻量级粗粒度阶段筛选关键帧,细粒度阶段只在关键区域计算注意力。
- 训练耗时减少2.53倍,生成速度提升6倍,质量基本不变。
- 适合需要高效训练和部署视频生成模型的研究者与开发者。
视频扩散变换器(DiTs)的扩展受限于其三维注意力机制的二次复杂度,尽管大部分注意力权重集中在少数位置。本文提出可训练的稀疏注意力(VSA),在训练和推理中均替代全注意力。VSA包含一个轻量级粗粒度阶段,将序列标记分块并识别高权重的“关键标记”;再通过细粒度阶段仅在这些区块内计算标记级注意力,并采用块计算布局以保证硬件效率。该方法实现单一可微核,端到端训练,无需后期性能分析,且维持85% FlashAttention3 MFU。通过从60M到1.4B参数的大型消融与缩放实验验证,VSA在不降低扩散损失的前提下,将训练浮点运算量减少2.53倍。对开源模型Wan-2.1进行改造后,注意力时间缩短6倍,端到端生成时间由31秒降至18秒,质量相当。结果表明,可训练稀疏注意力是全注意力的实用替代方案,也是视频扩散模型进一步扩展的关键技术。代码将发布于https://github.com/hao-ai-lab/FastVideo。
原文摘要 · Abstract (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 \emph{both} training and inference. In VSA, a lightweight coarse stage pools tokens into tiles and identifies high-weight \emph{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.53$\times$ with no drop in diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention time by 6$\times$ 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. Code will be available at https://github.com/hao-ai-lab/FastVideo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。