[Paper] PSA：金字塔稀疏注意力用于高效视频理解与生成

Source: arXiv - 2512.04025v1

概览

本文提出了 Pyramid Sparse Attention (PSA)，一种全新注意力模块，能够在保持大部分有用信息的同时，大幅削减视频模型中自注意力的二次方计算成本。通过用多层级池化的键值（KV）表示取代硬二进制掩码，PSA 在速度与精度之间提供了更细粒度的权衡，使其在视频理解（如动作识别）和视频生成（如文本到视频合成）两方面都具备实用性。

主要贡献

• 金字塔式 KV 池化： PSA 不再直接丢弃整个 KV 块，而是为每个块生成多个不同分辨率的池化版本（从低分辨率到高分辨率），并让每个查询自行决定使用哪一级。
• 查询驱动的动态分配： 查询会自动对重要区域使用高分辨率 KV，对不重要的区域使用低分辨率 KV，实现了在完整注意力和激进剪枝之间的“插值”。
• 硬件友好的内核： 作者设计了去耦的块‑瓦实现，能够干净地映射到 GPU/TPU 上，避免了许多稀疏注意力技巧中常见的不规则内存访问。
• 统一用于理解与生成： PSA 在判别式视频任务（如 Kinetics、Something‑Something）和生成式任务（如文本到视频扩散）上均有展示，体现了其通用性。
• 开源发布： 代码、预训练权重以及可直接运行的内核已开源，降低了使用门槛。

方法论

1. 块级注意力基础 – 将输入视频划分为固定大小的查询、键、值块（常见的“块稀疏”设置）。
2. 多层级池化 – 对每个键/值块，PSA 构建一个小金字塔：
   • Level 0：原始（全分辨率）KV。
   • Level 1、2、…：逐步池化（例如平均池化）的版本，降低空间/时间分辨率。
3. 查询驱动的选择 – 轻量评分网络评估每个 KV 块对给定查询块的相关性。根据得分，查询挑选合适的金字塔层级：对“重要”块使用高分辨率，对“非重要”块使用低分辨率。
4. 插值与聚合 – 选中的池化 KV 如有必要进行上采样，然后通过标准的缩放点积注意力公式与查询结合。由于池化是确定性的，梯度能够在所有层级上传递，实现端到端训练。
5. 高效内核 – 实现将块分组为瓦片，每个瓦片在固定计算预算下处理，并利用避免动态内存分配的 CUDA 内核，使 PSA 在普通硬件上运行快速。

结果与发现

• PSA 始终将与密集注意力的差距压缩到 ≤1 % 的绝对误差，同时实现 2–2.5× 的加速和 30–40 % 的内存节省。
• 在定性方面，生成的视频相较其他稀疏方法保留了更清晰的运动边界，且伪影更少。
• 消融实验表明，动态层级选择是性能提升的主要驱动力；若使用单一静态层级池化，则退化至二进制掩码基线的水平。

实际意义

• 更快的视频流水线： 开发者可以将 PSA 插入现有基于 Transformer 的视频模型（如 ViViT、TimeSformer），在不重新设计网络的前提下降低推理延迟。
• 边缘与移动部署： 降低的内存占用使得在 VRAM 受限的设备上运行视频 Transformer 成为可能，为端侧视频分析或 AR/VR 场景打开了大门。
• 成本效益更高的训练： 大型视频扩散模型的每一次前向/反向传播所需 FLOPs 更少，从而可以使用更大的批量或更长的序列，降低训练成本。
• 混合精度友好： PSA 的块‑瓦设计与混合精度（FP16/FP8）训练兼容，符合现代 GPU 流水线，可无缝集成到 PyTorch、TensorFlow 等库中。
• 加速科研： 开源内核为进一步的稀疏性研究提供了基准，例如结合低秩分解或学习式 token 剪枝。

局限性与未来工作

• 粒度受块大小限制： PSA 的效果依赖于所选块尺寸；若块过大，细粒度的时间细节仍可能丢失。
• 池化层级固定： 金字塔层级是预先设定的（如 2×、4× 池化），自适应的池化比例或许能够进一步提升权衡。
• 基准仅限短片段： 实验聚焦于 ≤2 秒的片段，如何扩展到小时级视频或流式场景仍是未解之题。
• 硬件依赖性： 虽然内核对 GPU 友好，但在 CPU 或专用加速器（如 TPU）上的表现可能不同，需要专门的优化。

未来工作可探索 可学习的池化算子、层次化查询路由 以及 与 token‑级剪枝的结合，在进一步提升效率的同时，仍保留高保真视频任务所需的丰富时空线索。

作者

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

论文信息

• arXiv ID: 2512.04025v1
• 分类: cs.CV, cs.AI, cs.LG
• 发布日期: 2025 年 12 月 3 日
• PDF: Download PDF