PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation

2025-12-03

AI Summary

The quadratic computational complexity of attention mechanisms severely hinders their scalability in video understanding and generation. Existing sparse attention methods rely on binary block masks, leading to substantial information loss under high sparsity. To address this, we propose Pyramid Sparse Attention (PSA), which constructs hierarchical key-value (KV) representations via multi-level pooling and dynamically assigns queries to appropriate pooling levels based on their importance—enabling fine-grained, continuous sparsity that preserves critical contextual information while drastically reducing computation. PSA employs a hardware-efficient decoupled block-tile kernel and an importance-aware KV block hierarchy. Experiments across diverse video tasks demonstrate that PSA significantly outperforms state-of-the-art sparse attention methods under tight computational budgets, achieving superior long-range context modeling and higher visual fidelity—thereby delivering a better efficiency–quality trade-off.

Technology Category

• Computer Vision: Video Understanding & Activity Analysis
• Machine Learning: Hardware-aware ML
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Security and Privacy: Large-scale security measurements
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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

Problem

• Addresses quadratic complexity bottleneck in attention mechanisms
• Mitigates information loss from binary masking in sparse attention
• Enables efficient video understanding and generation with finer granularity

Innovation

• Multi-level pooled KV representations for finer granularity
• Dynamic allocation of pooling levels to KV blocks
• Hardware-friendly kernel with decoupled block-tile design