Source-linked AI summary
SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning
Jintao Zhang, Kai Jiang, Chendong Xiang, Weiqi Feng, Yuezhou Hu, Haocheng Xi, Jianfei Chen, Jun Zhu
TL;DR
Sparse attention must handle masking failures at high sparsity and preserve generation behavior during adaptation. SpargeAttention2 combines hybrid Top-k+Top-p masking with efficient implementation and velocity-level distillation, achieving 95% sparsity and 16.2× attention speedup while maintaining quality.
Problem
At very high sparsity, Top-k and Top-p can discard important attention information, while standard diffusion-loss fine-tuning can cause behavior drift from data-distribution mismatch.
Method
SpargeAttention2 uses a unified Top-k+Top-p masker, efficient sparse-attention kernels, and velocity-level distillation against a frozen full-attention model.
Results
SpargeAttention2 reaches 95% attention sparsity, 16.2× attention runtime speedup, and up to 4.7× end-to-end video-generation speedup while preserving quality comparable to full attention.
Takeaways & Limitations
Hybrid masking and distillation-style fine-tuning enable high-sparsity acceleration without degrading video-generation quality in the reported experiments.
Takeaways & Limitations
Standard diffusion-loss fine-tuning is sensitive to fine-tuning data quality and distribution mismatch, which can degrade performance even with full attention.
Abstract
from arXiv · showhide
Many training-free sparse attention methods are effective for accelerating diffusion models. Recently, several works suggest that making sparse attention trainable can further increase sparsity while preserving generation quality. We study three key questions: (1) when do the two common masking rules, i.e., Top-k and Top-p, fail, and how can we avoid these failures? (2) why can trainable sparse attention reach higher sparsity than training-free methods? (3) what are the limitations of fine-tuning sparse attention using the diffusion loss, and how can we address them? Based on this analysis, we propose SpargeAttention2, a trainable sparse attention method that achieves high sparsity without degrading generation quality. SpargeAttention2 includes (i) a hybrid masking rule that combines Top-k and Top-p for more robust masking at high sparsity, (ii) an efficient trainable sparse attention implementation, and (iii) a distillation-inspired fine-tuning objective to better preserve generation quality during fine-tuning using sparse attention. Experiments on video diffusion models show that SpargeAttention2 reaches 95% attention sparsity and a 16.2x attention speedup while maintaining generation quality, consistently outperforming prior sparse attention methods.
1. Introduction
SpargeAttention2 addresses failures of existing sparse-attention masking and fine-tuning at high sparsity with a hybrid masker and distillation-style training. It achieves high sparsity and substantial speedups while preserving video-generation quality.
- Video diffusion attention is expensive because long sequences give the attention operator O(N^2) time complexity.
- At sparsity above 90%, Top-k can miss useful context in uniform attention rows, while Top-p can retain attention sinks in skewed rows.
- SpargeAttention2 combines Top-k and Top-p masking with an efficient trainable implementation and distillation-style fine-tuning.
- The method studies why masking rules fail, why trainable attention reaches higher sparsity, and why diffusion-loss fine-tuning can be suboptimal.
- 95% attention sparsity, 16.2× attention speedup, and 4.7× end-to-end generation speedup are achieved without degrading video-generation quality.
2. Preliminaries
The preliminaries formulate standard and block-sparse attention, explain Top-k and Top-p block masking, and describe flow-matching diffusion training. Sparse attention masks low-importance weights to reduce the quadratic attention cost while retaining GPU-compatible block structure.
- 2.1. Block Sparse Attention: Standard attention forms S = QK^T, applies row-wise softmax to obtain P, and produces output O from Q, K, and V.
- 2.1. Block Sparse Attention: The two attention matrix multiplications cost O(N^2d), making attention expensive for large token counts.
- 2.1. Block Sparse Attention: Sparse attention applies a binary mask to remove low-importance weights and skip corresponding score and value computations.
- 2.1. Block Sparse Attention: Block-sparse attention uses all-one or all-zero tiles, aligning masking with GPU tiling for practical speedups.
- 2.2. Masking for Sparse Attention in Diffusion Models: Top-k keeps a fixed fraction of the largest pooled attention positions, whereas Top-p keeps the smallest set whose cumulative probability reaches p%.
- 2.3. Diffusion Loss: Flow matching trains a diffusion model to predict a ground-truth velocity from an interpolated latent, noise, timestep, and text prompt using MSE.
3. Analysis
The analysis identifies distinct failure modes for Top-k and Top-p masking, explains why fine-tuning concentrates attention and reduces sparse-attention error, and shows that diffusion-loss fine-tuning can degrade quality under dataset mismatch.
- Case 1: Top-p performs approximately as well as Top-k+Top-p and better than Top-k on uniform attention-weight matrices.The comparison uses relative L1 distance between sparse- and full-attention outputs at matched sparsity.
- Case 1: Top-k and Top-p fail differently because uniform attention distributions need broader token coverage, whereas skewed distributions can make cumulative thresholds select too few attention-sink tokens.Top-k may miss important probabilities in uniform rows; Top-p may ignore important non-sink probabilities in skewed rows.
- Case 1: Combining Top-k and Top-p addresses both distribution-dependent failures by avoiding fixed-token selection for uniform rows and fixed cumulative thresholds for skewed rows.Top-p supplies broader coverage for relatively uniform rows, while Top-k prevents ineffective attention-sink-only selection in highly skewed rows.
- Case 1: Top-k performs approximately as well as Top-k+Top-p and better than Top-p on skewed attention-weight matrices.In skewed rows, Top-p can satisfy its cumulative threshold using only a few attention-sink probabilities.
- Case 2: After sparse-attention fine-tuning, attention probabilities become more concentrated and the model has smaller attention L1 error at the same 60% sparsity.The analysis links concentration to reduced dropped and renormalization errors; an example changes the distribution from [0.6, 0.2, 0.2] to [0.8, 0.1, 0.1].
- Case 3: Diffusion-loss fine-tuning degrades performance across key metrics when the fine-tuning data differ in quality from closed, high-quality pre-training data.The degradation occurs even with full attention and is attributed mainly to the fine-tuning dataset rather than to full or sparse attention.
4. Method
SpargeAttention2 combines Top-k and Top-p masking to avoid failures on uniform and skewed attention distributions, then adapts sparse attention using velocity distillation against a frozen full-attention model.
- Hybrid Top-k+Top-p Masking: Top-p prevents Top-k from retaining too few useful tokens for relatively uniform attention distributions.Top-k can capture only a small probability fraction when probability mass is spread across many tokens.
- Hybrid Top-k+Top-p Masking: Top-k prevents Top-p from selecting too few tokens around an attention sink in highly skewed distributions.The hybrid mask takes the union of Top-k and Top-p selections.
- Velocity Distillation: Velocity distillation aligns sparse-attention velocity predictions with those of a frozen full-attention teacher instead of using standard diffusion-loss supervision.The fine-tuning data are used only to construct noisy inputs for distillation, reducing drift toward a mismatched data distribution.
- Efficient Implementation: The implementation replaces full attention with a CUDA block-sparse operator that skips masked-out matrix multiplications and softmax computations.The adaptation procedure replaces all attention layers and minimizes the difference between sparse and teacher velocity predictions.
5. Experiments
Experiments evaluate SpargeAttention2 on two Wan2.1 model scales and resolutions, comparing it with prior methods and ablations of masking, training, and objective choices. It maintains strong generation quality under high sparsity while delivering substantial attention and end-to-end speedups.
- Setup: Experiments use Wan2.1-1.3B at 480p and Wan2.1-14B at 720p, trained with 3,000 approximately five-second videos.The study compares representative trainable sparse-attention methods and controls one design component at a time in ablations.
- Setup: Generation quality is evaluated with IQ, OC, AQ, VR, VA, and VT, alongside attention latency and end-to-end generation time.VA and VT denote VQA-a and VQA-t, respectively.
- Effectiveness: Across both model settings, SpargeAttention2 matches or exceeds full attention on generation quality while remaining stable under high sparsity.Prior sparse-attention baselines show noticeable degradation under the same or lower sparsity levels.
- Efficiency: 95% attention sparsity yields a 16.2× attention speedup, while end-to-end generation reaches up to 4.7× speedup.On Wan2.1-1.3B at 480p, attention latency falls from 97s to 6s and end-to-end time from 159s to 68s.
- Efficiency: At 85%–95% sparsity, SpargeAttention2 is the only method reported to combine strong generation quality with substantial efficiency gains.It is 1.8× faster than SLA and more than 4× faster than VSA and VMoBA in the reported comparisons.
- Ablations: The hybrid Top-k+Top-p masker gives the best overall generation quality and alignment across both model scales.Disabling training substantially degrades quality and alignment, while diffusion-loss fine-tuning consistently underperforms velocity distillation.
6. Related Work
Related work distinguishes training-free sparse attention, which applies test-time masks, from trainable methods that adapt sparse attention through training. SpargeAttention2’s Top-k+Top-p hybrid differs from prior hybrid sampling and targets diffusion models.
- Training-Free Sparse Attention: Training-free approaches reduce inference cost by applying attention masks at test time.The cited training-free methods span diffusion and general attention settings.
- Trainable Sparse Attention: Trainable sparse attention methods use pre-training or fine-tuning to reach higher sparsity than training-free approaches.The paper positions SpargeAttention2 within this trainable category.
- Positioning: vAttention combines Top-k with random sampling, whereas SpargeAttention2 combines Top-k with Top-p and is designed for diffusion models.The paper explicitly distinguishes these hybrid masking strategies and application settings.
7. Conclusion
SpargeAttention2 combines hybrid sparse masking, efficient implementation, and distillation-style fine-tuning to achieve high sparsity while preserving video generation quality.
- 95% attention sparsity, 16.2 × attention runtime speedup, and up to 4.7 × end-to-end video generation speedup are achieved.Generation quality remains comparable to full attention, and the method surpasses baselines.
- SpargeAttention2 combines hybrid Top-k and Top-p masking, efficient implementation, and distillation-style fine-tuning.These components target high sparsity and quality preservation in video diffusion models.
A. Hyper-parameters
The experiments use 500 training steps for main results, with model-specific batch sizes and resolutions; ablations use fewer steps for the 14B model and vary masking parameters.
- Main results use 500 training steps, with batch size 64 for Wan2.1-1.3B at 480p and 16 for Wan2.1-14B at 720p.Ablations on the 14B model use 100 training steps to reduce computational cost.
- Top-k = 0.03 and Top-p = 0.2 for Wan2.1-1.3B, while Top-k = 0.03 and Top-p = 0.16 for Wan2.1-14B target approximately 95% sparsity.The corresponding block sizes are bq = 128 and bkv = 64.
- Top-k ablations use Top-k = 0.05, producing approximately 95% sparsity for both model scales.Top-p ablations use 0.4 at 480p and 0.3 at 720p, corresponding to 94% and 93% sparsity.
B. Prompts for Qualitative Visualizations
The qualitative visualizations use prompts spanning animals, urban scenes, landscapes, underwater environments, and a rainy formal scene.
- Prompts for Figure 1: Figure 1 includes a polar bear playing guitar, a teddy bear playing drums in Times Square, a Pacific Coast shoreline, and a clownfish in a coral reef.These prompts cover animal-centered scenes, an urban setting, a coastal landscape, and an underwater environment.
- Prompt for Figure 4: Figure 4 uses an oil-painting prompt depicting a formally dressed couple walking through heavy rain with umbrellas.The scene includes streetlights, puddles, and wet cobblestones.