SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning

1. Problem Statement

Video diffusion models suffer from high computational costs due to the $O(N^2)$ complexity of the attention mechanism, especially given the long sequence lengths inherent in video generation. While sparse attention methods exist to mitigate this, current approaches face three critical challenges:

1. Failure of Standard Masking Rules at High Sparsity: Existing methods rely on either Top-k (keeping the top $k$ tokens) or Top-p (keeping tokens until cumulative probability reaches $p$).
   • Top-k Failure: In rows with uniform probability distributions, Top-k retains too few tokens, discarding valuable context.
   • Top-p Failure: In rows with highly skewed distributions (e.g., dominated by "attention sinks"), Top-p may satisfy the threshold with only a few uninformative tokens, dropping crucial context.
2. Suboptimal Fine-Tuning Objectives: Most trainable sparse attention methods fine-tune models using the standard diffusion loss on real-world data. However, open-source video models often lack access to their original, high-quality pre-training datasets. Fine-tuning on lower-quality or distribution-mismatched data using standard diffusion loss causes behavior drift, degrading generation quality even when using full attention.
3. Efficiency vs. Quality Trade-off: Achieving very high sparsity (e.g., >90%) typically results in significant degradation in visual quality and temporal coherence.

2. Methodology: SpargeAttention2

The authors propose SpargeAttention2, a trainable sparse attention framework designed to achieve high sparsity without compromising generation quality. It consists of three core components:

A. Hybrid Top-k + Top-p Masking

To address the failure modes of individual masking rules, the authors propose a unified masking strategy that combines both Top-k and Top-p.

• Mechanism: For each row of the attention weight matrix, the mask is the union of the tokens selected by Top-k and those selected by Top-p.
  • If the distribution is uniform, Top-p ensures enough tokens are kept to capture the probability mass.
  • If the distribution is skewed, Top-k ensures that the "attention sink" does not dominate the selection, forcing the inclusion of other informative tokens.
• Result: This hybrid approach provides robust masking performance across diverse attention weight distributions, maintaining accuracy even at extreme sparsity levels.

B. Velocity Distillation Fine-Tuning

To solve the data distribution mismatch and behavior drift problem, the authors replace the standard data-driven diffusion loss with a Velocity Distillation Loss.

• Teacher-Student Setup:
  • Teacher: The original pre-trained full-attention model (frozen).
  • Student: The model being adapted to use SpargeAttention2.
• Objective: The student model is trained to match the velocity predictions ($u_{sparse}$) of the teacher model ($u_{full}$) given the same noisy latent inputs ($x_t$), timestep ($t$), and text prompts ($c_{text}$).
• Loss Function:
  $$L_{VD} = \mathbb{E}[\| u_{sparse}(x_t, c_{text}, t) - u_{full}(x_t, c_{text}, t) \|^2]$$
• Benefit: This approach decouples the adaptation process from the quality of the fine-tuning dataset. The model learns to preserve the original generation dynamics of the full-attention teacher, preventing degradation even when fine-tuning data is limited or lower quality.

C. Efficient Kernel Implementation

The method utilizes a block-sparse attention implementation built on FlashAttention.

• It partitions tensors into tiles and applies the hybrid mask at the block level.
• The CUDA kernel efficiently skips masked matrix multiplications and softmax computations, ensuring practical speedups on modern GPUs.

3. Key Contributions

1. Theoretical Analysis: The paper identifies and analyzes why Top-k and Top-p fail under specific distribution conditions (uniform vs. skewed) and demonstrates that trainable sparse attention can achieve higher sparsity by concentrating the attention distribution during training.
2. Hybrid Masking Strategy: Introduction of a simple yet effective Top-k+Top-p union masking rule that stabilizes sparse attention at very high sparsity levels.
3. Distillation-Based Adaptation: A novel fine-tuning objective using velocity distillation that aligns sparse attention models with full-attention teachers, solving the data distribution mismatch problem common in open-source video model adaptation.
4. State-of-the-Art Performance: SpargeAttention2 achieves unprecedented sparsity while maintaining or exceeding the quality of full-attention baselines.

4. Experimental Results

The method was evaluated on Wan2.1 video diffusion models (1.3B and 14B parameters) across 480p and 720p resolutions.

• Sparsity & Speed:
  • Achieved 95% attention sparsity.
  • 16.2× speedup in attention runtime compared to full attention.
  • 4.7× end-to-end video generation speedup (e.g., reducing 14B 720p generation time from ~3043s to ~650s).
• Generation Quality:
  • SpargeAttention2 matched or exceeded the full-attention baseline on all metrics (Imaging Quality, Aesthetic Quality, Vision Reward, VQA accuracy).
  • In contrast, prior methods (VSA, VMoBA, SLA) showed noticeable degradation in quality at similar sparsity levels.
• Ablation Studies:
  • Removing the hybrid masking (using only Top-k or Top-p) led to quality drops.
  • Removing training (using a frozen sparse attention) resulted in severe quality degradation.
  • Replacing velocity distillation with standard diffusion loss caused performance to drop, confirming the necessity of the distillation objective.

5. Significance

SpargeAttention2 represents a significant leap in efficient video generation. By addressing both the algorithmic limitations of masking strategies and the optimization challenges of fine-tuning, it enables the deployment of high-resolution video diffusion models on more accessible hardware. The ability to achieve 95% sparsity without quality loss makes real-time or near-real-time high-fidelity video generation feasible, bridging the gap between theoretical efficiency and practical application.