PaperlayerView exact revision on arXiv

Source-linked AI summary

VSA: Faster Video Diffusion with Trainable Sparse Attention

Peiyuan Zhang, Yongqi Chen, Haofeng Huang, Will Lin, Zhengzhong Liu, Ion Stoica, Eric Xing, Hao Zhang

arXiv:2505.13389v5cs.CV

TL;DR

Video DiT scaling is constrained by quadratic attention despite concentrated attention mass. VSA replaces full attention with trainable hierarchical sparse attention at training and inference, achieving better scaling and practical speedups while exposing sparsity-level and hardware trade-offs.

  • Problem

    Quadratic 3D attention dominates the cost of scaling video DiTs, whose long sequences contain sparse attention patterns but lack a trainable sparse primitive.

  • Method

    VSA uses pooled coarse attention to predict critical cubes, fine token attention within selected cubes, differentiable combination, and block-sparse GPU execution.

  • Results

    VSA reaches a Pareto point that cuts training FLOPS by 2.53x with no drop in diffusion loss.

  • Takeaways & Limitations

    Trainable sparse attention provides a practical alternative to full attention for scaling video DiTs at both training and inference.

  • Takeaways & Limitations

    Optimal sparsity depends on sequence length and training budget, while the fixed (4, 4, 4) cube size restricts compatible resolutions.

Abstract

from arXiv · show

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.

1 Introduction

Video DiTs are bottlenecked by quadratic attention over very long sequences, motivating VSA, a trainable hierarchical sparse mechanism that predicts critical tokens and computes fine attention within selected tiles. VSA balances localization, model quality, and GPU efficiency while supporting sparse attention during both training and inference.

  • Motivation: Over 100K tokens can arise in a 5-second 720p clip, making attention the primary scaling bottleneck for video DiTs.State-of-the-art video DiTs spend most of their compute on attention for full-resolution, long-sequence training, while inference remains slow.
  • Limitations of Prior Work: Most prior sparse methods apply fixed or profile-derived masks only after full-attention training, leaving training cost unchanged and creating a train-test mismatch.These methods can erode quality when sparsity exceeds a gentle budget, so DiTs still commonly use quadratic 3D attention.
  • VSA: VSA pools (4, 4, 4) token cubes for lightweight coarse attention, selects critical cubes, and performs token-level fine attention only within the top-K selected cubes.A differentiable gate combines the coarse and fine outputs, allowing critical-token selection to be learned end-to-end rather than imposed by heuristics.
  • VSA: VSA maps spatial-temporal cubes to kernel-level tiles, allowing GPU blocks to process each tile densely or skip it under a block-sparse layout.This hardware alignment is intended to preserve efficient execution while applying learned sparsity.
  • Design Trade-offs: Tile size trades off localization and expressiveness against GPU throughput: smaller tiles sharpen sparsity, whereas larger tiles improve arithmetic intensity but blur attention patterns.The reported configuration uses tile size 64 and 87.5% attention sparsity, achieving performance comparable to full attention while maintaining efficient kernel execution.
  • Efficiency Results: The implementation retains 85% of FlashAttention3 MFU, while retrofitting Wan2.1 reduces attention time by 6x and end-to-end inference latency from 31s to 18s on H100.These measurements demonstrate hardware efficiency and inference speedups in an open-source video DiT.

2 Methods

VSA uses hierarchical, cube-based sparse attention to identify critical regions cheaply and perform fine-grained attention in hardware-aligned blocks. Its design balances sparsity, expressiveness, throughput, and adaptation stability across training and inference.

  • Hardware trade-offs: Block-sparse layouts let GPU threadblocks process selected tiles densely or skip them entirely, trading fine-grained expressiveness against hardware efficiency.Smaller tiles improve flexibility and modeling expressiveness but reduce throughput; larger tiles improve throughput while coarsening attention patterns.
  • VSA design: VSA partitions video latents into contiguous cubes and maps each cube to a GPU tile for hardware-aligned sparse computation.The cube dimensions are (Ct, Ch, Cw), with tile size B = Ct × Ch × Cw.
  • Coarse stage: The coarse stage mean-pools cubes, computes cube-level attention, and selects Top-K critical-token cubes without forming the full token-level attention matrix.Selected coarse entries are expanded into B × B blocks for the fine stage.
  • Fine stage: The fine stage applies token-level attention only within selected blocks, while gating vectors combine coarse and fine outputs.The coarse stage contributes less than 1% of total FLOPS, according to the supplied passage.
  • Design exploration: Smaller tile sizes improve performance through finer attention granularity but make attention kernels slower, motivating B = 64 with cube size (4, 4, 4).The paper reports this setting as a favorable trade-off between expressiveness and efficiency.
  • Pretrained adaptation: VSA adapts pretrained full-attention DiTs with annealing that begins equivalently to full attention and gradually reduces K toward the target sparsity.This addresses instability from randomly initialized gating projections and the architectural difference between VSA and full attention.
  • Distillation: VSA is also compatible with distillation, converting Wan-2.1 through brief flow-matching finetuning while preserving the original distillation loss and hyperparameters.The supplied passage describes this as the first reported sparse-attention method compatible with distillation.

3 Experiments

Experiments show that VSA’s effectiveness depends on adaptive sparsity, tile granularity, and hardware-aware kernel design. Across scaling and sparse-adaptation settings, VSA preserves quality while reducing computation and inference time.

  • Ablation Studies: VSA’s data-dependent patterns outperform fixed local patterns and full attention across the examined training settings.The advantage reverses for prior sparse methods under extended training, while VSA remains stronger than both fixed-pattern methods and full attention.
  • Ablation Studies: Smaller tiles reduce model loss through finer attention granularity but can substantially lower GPU efficiency.Decreasing tile size from 256 × 256 to 64 × 16 significantly reduces MFU; 64 × 16 tiles are 2.26× slower despite slightly better performance.
  • Ablation Studies: Average pooling outperforms max pooling and convolutional pooling for VSA’s coarse stage.The convolutional approach also causes training instability.
  • Sparse Adaptation & Distillation: VSA preserves generation quality under sparse adaptation, including 91.2% attention sparsity for Wan-1.3B and 90% sparsity for Wan-14B.The Wan-1.3B VBench results are close to full-attention finetuning, while human evaluation reports preserved quality for Wan-14B.
  • Kernel Performance: Nearly 7× kernel speedup over FlashAttention-3 and over 6× speedup after coarse-stage computation translate to 2–3× inference speedups on Wan-1.3B and Hunyuan.VSA’s implementation retains 85% of FA3’s MFU.

4 Related Work

Sparse attention in language models has largely used fixed patterns, while video DiTs face more urgent quadratic-attention costs because their sequences are much longer and training uses full-resolution data. Recent video methods mostly apply sparsity after dense pretraining, motivating trainable sparse attention as a core video-DiT design.

  • LLM sparse attention has seen limited uptake because most training uses short sequences under the “train-short, adapt-long” paradigm.Sliding-window variants remain a prominent practical example.
  • Video DiTs face longer sequences and devote most compute to full-resolution, long-sequence training, making quadratic attention costly during both training and inference.A 100K-token context corresponds to only about five seconds of video.
  • Prior video-DiT sparse-attention methods generally retrofit sparsity at inference after models are pretrained with full attention.This differs from using sparsity as a training primitive.
  • VSA is positioned as trainable sparse attention for video DiTs rather than a post-hoc fix, while DSV uses a multi-stage, profiler-based training design.The paper presents these approaches as distinct design choices for incorporating sparsity during training.

5 Limitation and Conclusion

VSA is a trainable, hardware-efficient sparse attention mechanism for scaling video DiTs. Across model sizes and compute budgets, it matches full-attention performance at lower training cost, while retaining stated resolution and sparsity-selection boundaries.

  • VSA matches full-attention performance at 2.53× lower training cost across 60M–1.4B models and budgets up to 4 × 10^21 FLOPS.It also achieves 85% MFU of FlashAttention3 and reduces Wan2.1-1.3B end-to-end latency by 1.7×.
  • VSA currently uses fixed (4, 4, 4) cubes, requiring video latent dimensions divisible by 4.The paper notes that generating a slightly larger latent and cropping can address compatible-resolution restrictions.
  • The optimal sparsity level remains an open question requiring scaling laws that explicitly account for sparsity, model size, and training compute.

A Qualitative Examples

The qualitative examples show how Wan-1.3B outputs evolve during finetuning after switching from full attention to VSA. Frames are sampled either repeatedly at the video midpoint or uniformly across the video.

  • Figure 6(a) samples the same middle frame at each video step, while Figure 6(b) uniformly samples four frames across the video.
  • Wan-1.3B initially exhibits artifacts after switching from full attention to VSA, then gradually recovers coherent video generation during training.The examples use validation videos at selected training steps with K = 32.

B Pseudocode of VSA

The pseudocode implements VSA by tiling inputs, computing coarse attention over pooled blocks, retaining top-k blocks, applying masked fine attention, and combining coarse and fine outputs before untiling.

  • The algorithm keeps only the top-k blocks from coarse attention to construct the sparse selection pattern.The pseudocode sets block = 64 and topk = 32.
  • The coarse stage averages q, k, and v within blocks, computes dense block-level attention, and repeats the coarse output across each block.
  • The fine stage computes token-level QK attention, masks unselected positions, applies softmax, and multiplies by v.
  • The implementation assumes cube size (4,4,4) and video size (16,32,32), with tiling and untiling rearranging tokens between cube and sequence layouts.These operations can be moved outside individual attention layers.

C Experimental Details

This section documents the detailed experimental setups for the results presented in Section 3.

  • The section provides detailed experimental setups for Section 3 results.
  • Its purpose is to specify how the reported experiments were conducted.
  • The section serves as an experimental-reference component for the paper’s Section 3 findings.

C.1 Model Architecture

The experiments follow Wan2.1’s architecture, using separate video self-attention and text cross-attention, with sparsity applied only to self-attention.

  • C.1 Model Architecture: Experiments use the Wan2.1 architecture with video self-attention and textual cross-attention.Sparse attention is applied only to self-attention layers.
  • C.1 Model Architecture: Sparse attention replaces full attention only in the video-token self-attention layers.Cross-attention remains the mechanism for injecting textual information.
  • C.1 Model Architecture: The ablation model has 120M parameters and uses GPT-NeoX initialization, UMT5-XXL text encoding, and Wan2.1 VAE tokenization.

C.2 Ablation Experiments Setup

Ablations evaluate sparse attention under long video contexts, compare multiple attention designs, and estimate compute using parameter, token, and attention dimensions.

  • C.2 Ablation Experiments Setup: Long sequences of shape 61 × 512 × 512 are used because attention dominates computation at this scale.Short sequences are considered insufficiently challenging for demonstrating sparse-attention effectiveness.
  • C.2 Ablation Experiments Setup: The baseline grid searches batch sizes {512, 1024, 2048} and learning rates {5 × 10^-5, 1 × 10^-4, 2 × 10^-4, 6 × 10^-4}.
  • C.2 Ablation Experiments Setup: Training uses a fixed compute budget of 4.5 × 10^20 FLOPs for the 120M-parameter ablation model.Each ablation job takes around 10 hours on 64 Nvidia H200 GPUs.
  • C.2 Ablation Experiments Setup: Spatial-temporal attention alternates spatial and temporal attention across layers.
  • C.2 Ablation Experiments Setup: The Spatial-Full variant uses four spatial layers and one full-attention layer every five layers because temporal attention can become overly sparse.For latent shape (16, 32, 32), temporal attention accounts for less than 1% of full 3D-attention FLOPs.
  • C.2 Ablation Experiments Setup: Compress KV applies 2 × 2 × 2 average pooling to keys and values, reducing attention FLOPs by 8× while keeping queries at full resolution.
  • C.2 Ablation Experiments Setup: Strided Window attention uses Wt = 2 for spatial windows and Ws = 8 for temporal windows to expand token interactions.
  • C.2 Ablation Experiments Setup: Conv Pooling uses a 3D convolution with kernel and stride (4, 4, 4), matching the block size.

C.5 Sparse Adaptation Setup

Sparse adaptation bridges full and sparse attention through a gradual sparsity schedule and is evaluated by finetuning Wan models on synthetic videos.

  • C.5 Sparse Adaptation Setup: Training begins with full attention for 50 steps before progressively reducing attended cubes.The schedule reduces Top-K by 4 every 50 steps until reaching Top-K = 32.
  • C.5 Sparse Adaptation Setup: Wan-1.3B finetuning uses 80,000 synthetic videos at 448 × 832 resolution and 61 frames.Training runs for 4,000 steps on 32 H200 GPUs with per-GPU batch size 1 and gradient accumulation 2.
  • C.5 Sparse Adaptation Setup: Wan-14B finetuning uses final sparsity 0.9 and 200,000 synthetic videos at 768 × 1280 resolution and 77 frames.Training runs for 4,000 steps on 64 H200 GPUs with global batch size 64.

C.6 Sparse Distillation Setup

The sparse distillation setup replaces the attention module with VSA while holding the score models and DMD-related hyperparameters fixed. Runtime profiling identifies coarse-stage overhead as a limitation for shorter sequences, while longer sequences remain the primary target.

  • Sparse Distillation Setup: The generator is initialized from the pretrained base model, its attention module is replaced with VSA at sparsity=0.8, and the real-score and fake-score models remain unchanged.The setup uses 64 H200 GPUs, per-GPU batch size 1, and runs for 12 hours.
  • Sparse Distillation Setup: All DMD-related hyperparameters remain fixed, including denoising steps, generator update ratio, and real-score guidance scale.
  • Coarse Stage Runtime: For shorter sequences, coarse-stage overhead is more pronounced, with Top-K selection dominating the profiled runtime.Fusing attention scaling, softmax, and Top-K operations yields only modest improvements.
  • Coarse Stage Runtime: The coarse-stage overhead becomes negligible at longer sequence lengths, which are the primary target of the method.Further coarse-stage acceleration is left for future research.
Loading 2505.13389v5…