AB-Attention：让视频扩散模型以更少注意力计算逼近全量结果

视频扩散模型的生成质量不断提升，但其推理成本也随视频帧数、空间分辨率和 Transformer 层数快速增长。尤其是 3D 全注意力需要让大量视频 token 两两交互，二次复杂度成为生成速度的重要瓶颈。稀疏注意力可以减少计算，却常常在速度与画质之间做出明显妥协。

AB-Attention 关注的核心问题，不只是“保留多少注意力块”，而是如何在高稀疏率下更准确地估计哪些内容真正重要。现有动态块稀疏方法通常依据粗粒度摘要为键值块打分，或者通过全局聚类寻找相似 token。这些方案可能忽略块内部的差异：一个 tile 中也许只有少数键与当前查询高度相关，但传统方法往往只能整块保留或整块删除。此外，被删除的块不会参与 softmax，近似结果容易偏离稠密注意力。

论文提出的第一部分是 Intra-Block Clustering Selector（IBCS）。它只在每个查询块内部对查询进行聚类，不改变 token 的排列，从而以较低成本估计不同键值块的贡献。与单一块级摘要相比，这种做法能够保留查询块内部更细的需求差异，同时避免全局聚类带来的额外开销。

第二部分是 Sampling with Inclusion-Probability Weighting（SIP）。AB-Attention 不再把低分块简单截断，而是按照重要性估计值为候选块分配采样概率，并使用包含概率的倒数进行重新加权。这样，一次采样得到的是对稠密注意力的统计估计，未被直接选中的块不会在概念上被永久判定为无关。该设计尤其适合在激进稀疏设置下缓解注意力质量下降。

在硬件实现方面，论文还设计了配套内核：将小于常规 tile 的键值块收集并组合成完整的 Tensor Core tile，在不增加 Tensor Core 计算成本的前提下进行更细粒度的选择。也就是说，算法层面的块筛选与 GPU 矩阵计算单元的执行形态被同时考虑，而不是只优化理论上的 FLOPs。

实验覆盖 Wan2.1 和 HunyuanVideo。结果显示，AB-Attention 可跳过 89% 至 92% 的注意力计算；在相同对比中，基线方法需要计算约为其 1.8 至 3.6 倍的注意力量。论文报告其相对最强基线最高提升 4.4 dB PSNR，并实现最高 2.23 倍端到端加速。

这项工作的意义在于，它把视频扩散注意力加速从“粗略截断”推进到“带校正的近似计算”：既利用稀疏性降低成本，又尽量保留稠密注意力的数值特征。需要注意的是，当前结论来自特定视频模型和实验设置，实际收益仍会受到硬件、序列长度、稀疏率及生成任务的影响。若其实现能够适配更多扩散 Transformer，可能为高分辨率和长视频生成提供更具实用性的推理优化路径。

来源：Hugging Face Daily Papers