Technical Summary: Veda - Scalable Video Diffusion via Distilled Sparse Attention

1. Problem Statement

Scaling Diffusion Transformers (DiTs) for high-resolution, long-duration video generation is currently bottlenecked by the quadratic computational and memory complexity ($O(N^2)$) of self-attention mechanisms relative to the spatiotemporal token sequence length. While sparse attention offers a theoretical path to alleviate this, existing methods face a critical trade-off:

• Static Patterns: Methods using pre-defined sparse patterns (e.g., sliding windows) lack adaptability to the dynamic, head-specific attention structures learned by DiTs.
• Dynamic Selection: Methods that learn dynamic masks (e.g., VSA, VMOBA) often suffer from inaccurate estimation. They rely on implicit supervision via the diffusion objective and coarse statistics (e.g., mean pooling), leading to misalignment with the true geometry of full attention.
• The Core Issue: Empirical evidence shows that at high sparsity levels (≥90%), these existing approaches degrade generation quality, producing structured artifacts such as spatial warping, "water-ripple" patterns, and temporal flickering. The paper posits that the degradation is not caused by sparsity itself, but by the failure of sparse masks to align with the tile-wise geometry of full attention.

2. Methodology: Veda Framework

Veda is a distilled sparse attention framework designed to preserve the tile-wise structure and ranking of full attention even under extreme sparsity. It treats tile selection as an explicit reconstruction problem rather than an implicit byproduct of the diffusion objective.

2.1. Distilled Tile Scoring

Instead of relying on the diffusion loss to shape sparsity, Veda employs a lightweight estimator to explicitly learn tile-level attention scores from a full-attention backbone.

• Target Construction: A reference "Oracle" mask is constructed by applying max-pooling to the full-attention matrix over query-key tile regions. Max-pooling is chosen over averaging to preserve salient, high-frequency signal peaks that are often diluted by background noise.
• Statistic-Aware Estimator: To reconstruct these scores, Veda uses a TripPool descriptor for each tile, concatenating Average, Maximum, and Minimum statistics. This enriches the tile representation beyond simple mean pooling.
• Head-Specific Projections: The estimator uses distinct MLP projections for each attention head to capture heterogeneous dependency patterns.
• Optimization: The model is trained using a row-wise KL-divergence distillation objective ($L_{distill}$) to align the predicted tile-score distribution with the full-attention reference. Crucially, a stop-gradient operation is applied to the backbone features feeding the estimator. This decouples mask learning from feature learning, preventing the disruption of the pre-trained generative manifold.

2.2. Head-Aware Tiling

Recognizing that attention heads exhibit significant heterogeneity in spatiotemporal dependencies (some capturing local spatial interactions, others long-range temporal dependencies), Veda abandons uniform tiling strategies.

• Configuration Search: For each layer and head, Veda searches for an optimal tiling configuration $\pi_{l,h} = (p_t, p_h, p_w)$ that factorizes the hardware tile size.
• Offline Selection: The optimal configuration is selected offline on a calibration set by minimizing the approximation error between the full-attention output and the sparse-attention output, ensuring the tiling strategy matches the specific structural needs of each head.

2.3. Hardware-Efficient Implementation

To translate theoretical FLOP reductions into practical wall-clock speedups, Veda implements a custom tile-skipping sparse attention kernel using the ThunderKittens DSL.

• Asynchronous Execution: The kernel leverages NVIDIA Hopper's Tensor Memory Access (TMA) and Warp Specialization. It decouples data movement from computation using a producer-consumer paradigm, where producer warps fetch non-contiguous key/value tiles into shared memory while consumer warps execute tensor core operations.
• Efficiency: This design hides memory latency behind dense matrix math, achieving approximately 80% of the Memory FLOPs Utilization (MFU) of FlashAttention-3.
• Training Efficiency: A specialized TileLang kernel generates ground-truth heatmaps in two passes, allowing for sparse supervision (random subsets of query tiles) to reduce training overhead without compromising performance.

3. Key Contributions

1. Empirical Insight: The paper demonstrates that generation quality in sparse video diffusion is determined by the alignment of the sparse mask with the tile-wise geometry of full attention, rather than the sparsity ratio itself.
2. Distilled Sparse Attention: Veda introduces a framework that formulates tile selection as an explicit reconstruction problem, using a statistics-aware estimator and head-specific projections to minimize estimation error.
3. Head-Aware Tiling: A novel strategy that assigns distinct spatiotemporal tile factorizations to different attention heads, addressing the heterogeneity of dependency patterns in Video DiTs.
4. Hardware-Optimized Kernel: A custom tile-skipping kernel that achieves near-linear scaling with sequence length and high hardware utilization, converting algorithmic sparsity into tangible latency gains.

4. Experimental Results

Experiments were conducted on large-scale video diffusion models, including Waver-T2V-12B and Wan2.1-T2V-14B, across resolutions up to 720P and sequence lengths up to 241 frames.

• Speedup: On Waver-T2V-12B (720P, 241 frames), Veda achieves a 5.1× end-to-end speedup (reducing sample time from 19.4 to 3.8 minutes) and a 10.5× self-attention speedup. The attention overhead is reduced from 92% to 50%.
• Scalability: Speedups increase with sequence length. At 245K tokens, Veda scales near-linearly (309.1ms) compared to the quadratic growth of Full Attention (1576.5ms), yielding a 5.1× speedup.
• Quality Preservation:
  • Human Evaluation: On Waver-bench 1.0, Veda at 90% sparsity achieves perceptual parity with Full Attention. At 95% sparsity, it significantly outperforms the state-of-the-art VSA method (even when VSA operates at lower sparsity).
  • Artifact Reduction: Unlike baselines that suffer from water-ripple patterns and temporal flickering at high sparsity, Veda maintains high visual fidelity and temporal coherence.
  • Quantitative Metrics: VBench evaluations show Veda maintains subject consistency and motion smoothness comparable to Full Attention baselines, even at 95% sparsity.

5. Significance and Claims

The paper claims that Veda successfully addresses the scalability bottleneck of high-resolution video generation without compromising structural integrity. By reframing sparse attention as a distillation problem rather than a heuristic search or implicit learning task, Veda enables aggressive sparsity (up to 95%) that was previously unattainable without severe quality degradation.

The authors emphasize that the gains are not merely theoretical; the custom kernel ensures that sparsity translates to practical wall-clock acceleration. The framework is presented as a modular, hardware-aware solution that scales favorably with spatiotemporal resolution, making high-fidelity, long-video generation feasible on current hardware constraints. The work suggests that future improvements could focus on tighter kernel fusion, adaptive sparsity strategies across timesteps, and temporal caching of tile scores.