PSA: Atención Piramidal Dispersa para la Comprensión y Generación Eficiente de Videos

PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation

December 3, 2025

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

cs.AI

Resumen

Los mecanismos de atención son el núcleo de los modelos base, pero su complejidad cuadrática sigue siendo un cuello de botella crítico para la escalabilidad. Este desafío ha impulsado el desarrollo de mecanismos de atención eficientes, surgiendo la dispersión (sparsity) como paradigma dominante. Los métodos actuales normalmente retienen o descartan bloques clave-valor completos con máscaras binarias, lo que resulta en una pérdida sustancial de información bajo alta dispersión. Para mitigar esta brecha, presentamos Pyramid Sparse Attention (PSA), un módulo versátil aplicable tanto a tareas de comprensión como de generación de video. En lugar de enmascaramiento binario, PSA introduce representaciones KV agrupadas (pooled) multinivel, permitiendo una granularidad de máscara más fina. Específicamente, cada bloque de consulta asigna dinámicamente niveles de agrupamiento más bajos a bloques KV críticos y niveles más altos a los menos importantes, creando una interpolación informativa entre la retención completa y la poda total. Este diseño, análogo a la cuantización de punto fijo y las redes clásicas de pirámide de características (feature pyramid networks) en visión por computadora, mitiga eficazmente la pérdida de información mientras preserva la eficiencia computacional con un presupuesto bajo de cálculo. Funciona con un kernel nativo y compatible con el hardware que aprovecha un diseño desacoplado de bloque-baldosa (block-tile) para garantizar una ejecución eficiente. En diversos benchmarks de comprensión y generación de video, PSA preserva la información contextual y la fidelidad visual, superando consistentemente o logrando un rendimiento comparable con las líneas base de atención dispersa existentes, con mejores compensaciones entre eficiencia y calidad. Nuestro código y los pesos del modelo están disponibles públicamente en: http://ziplab.co/PSA

English

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