English

VSA: Faster Video Diffusion with Trainable Sparse Attention

Computer Vision and Pattern Recognition 2025-10-29 v5

Abstract

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.

Keywords

Cite

@article{arxiv.2505.13389,
  title  = {VSA: Faster Video Diffusion with Trainable Sparse Attention},
  author = {Peiyuan Zhang and Yongqi Chen and Haofeng Huang and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},
  journal= {arXiv preprint arXiv:2505.13389},
  year   = {2025}
}

Comments

Accepted by Neurips 2025

R2 v1 2026-07-01T02:22:35.556Z