PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation
Xiaolong Li, Youping Gu, Xi Lin, Weijie Wang, Bohan Zhuang
2025-12-04
Summary
This paper introduces a new way to handle 'attention' in artificial intelligence models, specifically those dealing with video. Attention is how these models focus on the most important parts of the data, but it can be very computationally expensive.
What's the problem?
Traditional attention mechanisms become incredibly slow and require a lot of processing power when dealing with large amounts of data, like long videos. Existing attempts to speed things up often involve simply ignoring large chunks of information, which hurts the model's accuracy. They make a hard choice: keep all the information (slow) or throw a lot of it away (fast but inaccurate).
What's the solution?
The researchers developed 'Pyramid Sparse Attention' (PSA). Instead of completely keeping or discarding information, PSA uses a more nuanced approach. It creates different levels of detail for the information, kind of like how image pyramids work in computer vision. Important information gets kept at a high level of detail, while less important information is summarized. This allows the model to focus on what matters without losing too much crucial data, and it's designed to work efficiently with existing computer hardware.
Why it matters?
PSA offers a better balance between speed and accuracy compared to other methods for speeding up attention. This is important because it allows AI models to process longer and more complex videos without needing massive amounts of computing power, opening the door for more advanced video understanding and generation applications.
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