Improving Video Sparse Attention with Fine-grained Router and Sparse Rebasing

2026-09-26

AI Summary

This study addresses the prohibitive computational overhead of full attention in video diffusion models and the accuracy degradation inherent in conventional sparse alternatives. To this end, it proposes VSA2, a trainable sparse attention mechanism integrated throughout the Diffusion Transformer (DiT) pipeline. VSA2 achieves adaptive computation by dynamically selecting key-value pairs via fine-grained routing, optimizes motion quality through an easy-to-hard curriculum learning strategy, and ensures training stability by introducing checkpoint rebasing. Experimental results demonstrate that VSA2 attains a 4.62× speedup on 720p video generation tasks while halving attention computation costs. Notably, the generated video quality remains comparable to that of full-attention models, effectively reconciling inference efficiency with visual fidelity.

Abstract

We present VSA2, a frontier trainable sparse attention for video DiTs. VSA2 includes a variety of new architectural features and training procedures that we apply across all stages of the DiT development cycle, including pretraining, RL, and inference, to produce a DiT with comparable or better quality than a full attention counterpart. Architecturally, VSA2 introduces a fine-grained router that improves the precision of identifying critical tokens and supports dynamic computation by allowing each query to attend to a variable number of key-value pairs. In training, we identify a Hard-to-Easy Curriculum, where models trained under high sparsity and later evaluated with lower sparsity during inference not only generalize effectively, but also outperform models trained with full attention in motion quality. VSA2 is also flexible: it can replace full attention during the middle of progressive low-to-high resolution pretraining, rebasing early-stage full-attention checkpoints. Experiments show that VSA2 reduces attention computation by half over VSA with lower loss. On 720p videos, it accelerates attention by 8.9x and end-to-end generation by 4.62x compared to the FlashAttention-3 baseline, while achieving comparable or better video quality.

Problem

• Video Diffusion Transformer
• Sparse Attention
• Computational Efficiency
• Video Generation

Innovation

• Sparse Attention
• Video Diffusion Transformer
• Fine-grained Router
• Curriculum Learning
• Sparse Rebasing