PaperlayerView exact revision on arXiv

Source-linked AI summary

PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation

Xiaolong Li, Youping Gu, Xi Lin, Weijie Wang, Bohan Zhuang

arXiv:2512.04025v1cs.CVcs.AIcs.LG

TL;DR

Long video sequences make quadratic attention a major efficiency bottleneck, and binary block masks can discard informative context at high sparsity. PSA uses adaptive multi-level pooled KV representations instead of only keeping or dropping blocks. Across video understanding and generation, it reports strong results under low compute budgets, including matching full-attention Video-MME accuracy at approximately 35% of full-attention FLOPs.

  • Problem

    Long video sequences make quadratic attention costly, while binary block-sparse masks can restrict each query’s visible KV region and lose information at high sparsity.

    Source: §1 · pp. 1–2

  • Method

    PSA assigns query blocks adaptive multi-level pooled KV representations, using finer levels for important blocks and coarser levels for less important ones.

    Source: §1 · p. 2

  • Results

    Across video understanding and generation, PSA consistently outperforms BSA-based methods under low compute; on Video-MME it matches full-attention accuracy at approximately 35% full-attention FLOPs.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    PSA generalizes across diverse model architectures and application scenarios under low compute budgets.

    Source: §5 · p. 9

Abstract

from arXiv · show

Attention mechanisms are the core of foundation models, but their quadratic complexity remains a critical bottleneck for scaling. This challenge has driven the development of efficient attention mechanisms, with sparsity emerging as the dominant paradigm. Current methods typically retain or discard entire key-value blocks with binary masks, resulting in substantial information loss under high sparsity. To mitigate this gap, we present Pyramid Sparse Attention (PSA), a versatile module applicable to both video understanding and generation tasks. Instead of binary masking, PSA introduces multi-level pooled KV representations, enabling finer mask granularity. Specifically, each query block dynamically allocates lower pooling levels to critical KV blocks and higher levels to less important ones, creating an informative interpolation between full retention and complete pruning. This design, analogous to fixed-point quantization and classical feature pyramid networks in computer vision, effectively mitigates information loss while preserving computational efficiency under a low compute budget. It works with a native, hardware-friendly kernel that leverages decoupled block-tile design to ensure efficient execution. Across video understanding and generation benchmarks, PSA preserves contextual information and visual fidelity, consistently outperforming or achieving comparable performance over existing sparse attention baselines with superior efficiency-quality trade-offs. Our code and model weights are publicly available at: http://ziplab.co/PSA

1. Introduction

Long video sequences make quadratic attention costly, while binary block masks can discard informative context at high sparsity. PSA replaces hard keep/drop decisions with adaptive multi-level pooled KV representations and reports favorable efficiency-quality trade-offs across video understanding and generation.

1. Introduction

  • Long video attention is a scaling bottleneck: a 720p, 81-frame Wan2.1–14B clip takes nearly two hours to infer on one H20 GPU, with attention exceeding 80% of runtime.Long-context video understanding prefill and video-generation sampling are identified as especially affected.Source: §1 · pp. 1–2
  • PSA assigns finer pooling to important KV blocks, coarser pooling to less important ones, and skips the least important blocks; query-key importance determines each query block’s multi-level mask.This replaces binary retention with finer-grained allocation; the design is likened to feature pyramids and fixed-point quantization.Source: §1 · p. 2
  • At the same compute budget, PSA queries attend to about 70% of KV blocks on average versus 20% for BSA, with outputs closer to full attention.The expanded receptive field yields a superior accuracy-efficiency balance.Source: §1 · p. 2
  • Across tasks, PSA matches full-attention Video-MME accuracy at approximately 35% full-attention FLOPs, outperforms prior block-sparse video-generation methods, and improves TDM-combined VBench scores at 15% compute.The generation result is reported for CogVideoX–5B.Source: §1 · pp. 2–3
  • PSA’s merge–split kernel decouples logical block size from hardware tile size, supports backpropagation, and is compatible with FlashAttention; avoiding token reordering supports use in both understanding and generation.The contribution summary describes efficient GPU utilization for dynamically varying KV block sizes.Source: §1 · p. 3

2. Related Work

Sparse attention reduces computation, but existing approaches trade flexibility against hardware efficiency. Block Sparse Attention uses hardware-friendly coarse blocks, yet its binary keep/drop decisions can lose information at high sparsity.

2. Related Work

  • Static sparse patterns are computationally efficient but input-agnostic, risking missed long-range dependencies and sub-optimal performance or unstable generation quality.Examples include sliding windows, attention sinks, and spatiotemporal energy decay patterns.Source: §2 · p. 3
  • Dynamic sparsity creates input-sensitive masks, but fine-grained unstructured operations can cause poor memory access and low GPU utilization.The cited approaches include thresholding attention scores.Source: §2 · p. 3
  • Block Sparse Attention preserves dense computation within hardware-aware blocks, but binary block masks can lose information at high sparsity; PSA instead uses pooled KV levels to widen query receptive fields under the same budget.The related-work passage also notes PSA avoids token reordering and uses an efficient kernel.Source: §2 · p. 3

3. Methodology

PSA replaces binary block keep-or-drop attention with hierarchical pooled KV representations and a multi-level mask that allocates computation across blocks. Its decoupled block-tile kernel supports efficient execution of the resulting heterogeneous blocks.

3. Methodology

  • PSA combines pyramid KV blocks, a multi-level mask generator, and adaptive pyramid attention to allocate computation across pooled key-value representations.The mask specifies which pyramid level to fetch for each query–key block pair.Source: §3 · p. 3, §3.3.2 · p. 5

3.1. Problem Definition and Notation

  • Binary keep-or-drop masks can discard informative regions and limit flexible allocation across blocks of differing importance.PSA addresses this with a multi-level mask whose entries select a pyramid level or skip a block.Source: §3.1 · p. 4, §3.3.2 · p. 5

3.2. Pyramid KV Blocks

  • Pyramid KV blocks become coarser with each level, summarizing broader context while reducing effective KV length by a factor of 2^(h−1).This hierarchy supports coarse-to-fine context under a limited compute budget.Source: §3.2 · p. 4

3.3.1. BLOCK IMPORTANCE ESTIMATION

  • For generation, PSA estimates block importance from the maximum attention score among sampled tokens; for understanding, it uses antidiagonal scoring and adjacent-token similarity checks.The understanding check compares cosine similarity against level-specific thresholds to set a maximum pooling level; generation also applies a Hilbert curve permutation for locality.Source: §3.3.1 · p. 4

3.3.2. MULTI-LEVEL MASK ASSIGNMENT

  • After row-wise normalization and descending sorting, cumulative importance scores are mapped to pyramid levels using thresholds that adapt sparsity to each query’s score distribution.Unlike fixed per-level quotas, the threshold strategy dynamically adjusts sparsity and is reported to yield more accurate attention estimation.Source: §3.3.2 · p. 5

3.4. Adaptive Pyramid Attention

  • At pooling level h, block attention costs O(b_qb_k/2^(h−1)); aggregating across levels gives expected complexity O(ρ̄N^2), with ρ̄ the effective sparsity ratio.The kernel decouples logical block size from execution tile size by flexibly splitting or merging blocks to match hardware-optimized tiles.Source: §3.4 · p. 6, §3.5 · p. 6
  • Adaptive Pyramid Attention fetches each block at its assigned level, giving informative regions finer attention and redundant ones coarser attention at the same compute budget.Pooled tokens aggregate 2^(h−1) original tokens, so attention logits receive a bias of (h−1) ln 2 to maintain the probability distribution after softmax.Source: §3.4 · pp. 5–6

3.5. Hardware-Optimized Kernel Design

  • On NVIDIA H200, the decoupled kernel achieves up to a 10× speedup over a naive PSA implementation while preserving tensor-core utilization for small pooled blocks.It also avoids divergence across pooling levels.Source: §3.5 · p. 6

4. Experiments

Across video generation and understanding, PSA retains quality at high sparsity, and its combination with distillation reports faster denoising without generation-quality loss. Ablations compare multi-level with binary masks and evaluate reordering and mask-generation strategies.

4.1. Experimental Settings

  • Experiments cover open-source video-generation and video-understanding models, use H200 GPUs, and evaluate generation similarity, VBench dimensions, and Video-MME.Sparse attention is applied only during prefill, and PSA's reported sparsity denotes that of a BSA-based method using the same compute budget.Source: §4.1 · p. 7

4.2. Training-Free Video Generation

  • At 0.91 sparsity, PSA preserves sharp details and temporal coherence; under comparable high sparsity, it surpasses SVG, STA, and Sparge on all similarity metrics and most perceptual axes.Against SVG2, PSA has higher Aesthetic and Imaging scores across model sizes while keeping similarity metrics within small margins.Source: §4.2 · p. 8

4.3. Distillation

  • At 85% sparsity, PSA combined with TDM achieves the highest VBench score (0.826), surpassing both the 4-step distilled baseline and the original 50-step model.The combination also achieves a 30.2× denoising time speedup over the original 50-step model, without generation-quality loss.Source: §4.3 · p. 8, Table 2 · p. 8

4.4. Training-Free Video Understanding

  • On Video-MME, PSA achieves its best overall performance at 0.65 sparsity, matching or exceeding full attention across most categories and improving Medium- and Long-Video scores.It remains competitive on Short Video.Source: §4.4 · p. 8, Table 3 · p. 8

4.5. Ablation Study

  • Multi-level masking improves PSNR and SSIM and lowers LPIPS versus binary masking despite higher sparsity; threshold-based reordering reconstructs better than quantile-based reordering.The comparisons use Wan2.1-1.3B at 480p in the training-free setting; the first 25% of sampling steps use no sparse attention.Source: §4.5 · pp. 8–9, Table 4 · p. 8, Table 5 · p. 8

5. Conclusion

PSA uses multi-level KV pooling to preserve information under low compute budgets and consistently outperforms existing block-sparse attention methods across video understanding and generation. The appendix supplies implementation details, ablations, and visual comparisons.

5. Conclusion

  • PSA consistently outperforms existing block-sparse attention methods under low compute budgets across video understanding and generation.The conclusion also reports strong generalization across model architectures and application scenarios.Source: §5 · p. 9

A. Appendix

  • The appendix details PSA modules, including the cosine-similarity pooling constraint, block-size variants, and importance-pooling operators.Source: A. Appendix · p. 12
  • The appendix reports additional ablations on cosine similarity, block size, pooling operators, and multi-level allocation across video understanding and generation.Source: A. Appendix · p. 12
  • The appendix provides extended qualitative examples for training-free video-generation evaluation using the main paper's layout.Source: A. Appendix · p. 12

B. Implementation Details of PSA

PSA can constrain its importance-driven multi-level mask using within-block cosine similarity, limiting how coarsely each key block is pooled. The appendix specifies the threshold-based eligibility rule and experimental settings.

B.1. Cosine Similarity Based Pooling Constraint

  • Internally similar key blocks may use coarser pyramid levels, while blocks with heterogeneous tokens should remain at finer levels.Source: §B.1 · p. 12
  • A block is eligible for level h only when its similarity at the corresponding pooling stride exceeds that level's threshold.The thresholds determine each block's maximum admissible level.Source: §B.1 · p. 12
  • The final mask combines the similarity constraint with the importance-driven mask and prevents blocks from exceeding their similarity-permitted coarseness.Source: §B.1 · p. 12
  • Experiments use four pyramid levels; the similarity thresholds differ between Video-MME and the video-generation PSA (sim) preset.Setting all thresholds to −1 disables the constraint, while setting them to 1 collapses the pyramid to one fine-grained level.Source: §B.1 · p. 12

B.2. Block Size and Pooling Variants

PSA separates the block sizes used for importance and mask computation from the kernel's hardware tile sizes. The appendix also describes a mean-pooling ablation and the video-generation evaluation setup.

B.2. Block Size and Pooling Variants

  • Query and key block sizes set the granularity of importance scores and masks, independently of kernel hardware tile sizes.The decoupled block-tile design permits moderate block sizes while retaining relatively high tensor-core utilization.Source: §B.2 · p. 12
  • The mean-pooling ablation replaces max pooling with an average over sampled attention scores while keeping the other components unchanged.The corresponding results are reported in Table 8.Source: §B.2 · p. 12
  • Unless noted otherwise, appendix video-generation experiments use Wan2.1-1.3B at 480p, enable sparse attention after 25% of sampling steps, and report FLOP-equivalent sparsity relative to full attention.Source: §B.2 · p. 13

C.1. Cosine Similarity (Video Understanding)

On Video-MME, adding the cosine-similarity pooling constraint improves PSA's overall score under a matched sparsity budget. The comparison uses sparse variants at approximately 0.65 FLOP-equivalent sparsity.

C.1. Cosine Similarity (Video Understanding)

  • The Video-MME comparison includes a dense FlashAttention2 baseline and two PSA variants, with sparse variants matched at approximately 0.65 FLOP-equivalent sparsity.Source: §C.1 · p. 13
  • 0.654 versus 0.647 overall Video-MME score for PSA with versus without the cosine-similarity constraint, at the same sparsity budget.Source: §C.1 · p. 13
  • The PSA w/o sim variant uses the importance-based multi-level mask, while PSA w/ sim adds the similarity-based cap.Source: §C.1 · p. 13

C.2. Cosine Similarity (Video Generation)

This ablation compares PSA presets with and without a cosine similarity constraint in video generation. Under comparable sparsity budgets, the constraint improves the quality–sparsity trade-off by limiting coarse pooling of KV blocks whose tokens are not sufficiently similar.

C.2. Cosine Similarity (Video Generation)

  • PSA (sim) uses thresholds T = {0.5, 0.7, 0.8, 0.9} and Ts = {0.7, 0.65, 0.6} to keep FLOP-equivalent sparsity near 0.8 despite local reversion to finer levels.Its progressive allocation avoids coarse pooling when neighboring tokens are insufficiently similar.Source: §C.2 · p. 13
  • PSA (no-sim) disables the cosine cap with Ts = {−1, −1, −1}, retains more blocks at the finest level h = 1, and has sparsity 0.79 versus 0.80 for PSA (sim).Its thresholds are T = {0.7, 0.8, 0.9, 0.9}, and allocation depends only on importance scores.Source: §C.2 · pp. 13–14
  • PSA (1-level) sets Ts = {1, 1, 1}, forcing every kept block to Lj = 1 and eliminating pooling, so it has the highest compute cost among the three presets.It serves as an upper bound on reconstruction quality for this importance pattern.Source: §C.2 · p. 14
  • Under comparable sparsity budgets, PSA (sim) achieves a better quality–sparsity trade-off than PSA (no-sim), while PSA (1-level) offers a quality reference at higher compute cost.The cosine cap avoids assigning coarse pooling to KV blocks whose tokens are not sufficiently similar.Source: §C.2 · p. 14

C.3. Block Size and Importance Pooling

At fixed sparsity, PSA's generation quality is robust to max- versus mean-based importance pooling, while smaller logical blocks improve quality by isolating high-importance regions more precisely.

C.3. Block Size and Importance Pooling

  • The ablation varies logical block size (bq, bk) and the importance-estimation pooling operator under a fixed sparsity budget; the default importance score uses max-pooling over sampled token-pair attention weights.The mean-based variant changes the outer max operator.Source: §C.3 · p. 14
  • At fixed sparsity 0.8, smaller blocks such as (32, 32) consistently yield better generation quality by more precisely isolating high-importance regions and reducing approximation errors.This comparison holds the sparsity budget fixed.Source: §C.3 · p. 14
  • At sparsity 0.8, max-pooling yields slightly higher PSNR/SSIM at block sizes (128, 128) and (64, 64), while mean-pooling performs comparably at (32, 32).The ablation replaces the outer max with mean while keeping the other listed components unchanged.Source: §C.3 · p. 14

C.4. Multi-Level Allocation Strategy

At a fixed compute budget, balanced allocation across multiple pyramid levels outperforms distributions concentrated on either very coarse or a few fine blocks. The section also provides additional qualitative video-generation comparisons.

C.4. Multi-Level Allocation Strategy

  • Five presets share a compute budget but differ in allocation across levels h = 1, 2, 3 and a dropped level h = 0; the ablation varies allocation while holding other settings fixed.The reported setting uses Wan2.1-1.3B at 480p, about 0.25× full-attention FLOPs, and sparse attention after the first 25% of sampling steps.Source: §C.4 · p. 15, Table 9 · p. 15
  • PSA-1 performs worse when it spends compute on a few fine blocks and drops the rest, while PSA-2 degrades across metrics when it keeps all blocks at one coarse level.The former restricts each query block's effective receptive field; the latter lets overly coarse KV representations dominate reconstruction degradation.Source: §C.4 · p. 15
  • At the same FLOP-equivalent cost, PSA-3/4/5 achieve higher PSNR/SSIM and lower LPIPS than skewed allocations, with PSA-3 selected as the default strategy.Finding the optimal allocation remains open.Source: §C.4 · p. 15, Table 10 · p. 15

D. Additional Visual Comparisons

Loading 2512.04025v1…