论文

PARK:视频扩散Transformer中稀疏注意力的准确块检索

PARK: Accurate Block Retrieval for Sparse Attention in Video Diffusion Transformers

模型推理推理加速

摘要

扩散Transformer(DiT)已成为视频生成的主要架构,但其效率受到全面注意力的二次复杂度的限制。稀疏注意力通过检索重要的块并仅在其中计算注意力来降低这种成本,但不准确的检索可能会降低生成质量或产生不必要的计算。我们在使用查询和关键块的平均表示检索块的方法中发现了两个检索不匹配:(i)查询端聚合不匹配,其中Softmax之前的平均查询无法保留其各自的注意力偏好;(ii)键端聚类度量不匹配,其中原始键空间中的标准欧几里得聚类可以将当前查询下具有不同QK分数的键分组,因此它们的平均表示可能无法准确表示当前查询如何对各个键进行评分。这些不匹配可能导致块检索不准确。为了解决这些不匹配问题,我们提出了 PARK,一种无需训练的稀疏注意力方法,用于准确的块检索。 PARK 保留每个原始查询,独立地标准化其对关键块的注意力,然后对每个查询块内的这些分布进行平均。它还使用当前查询中的信息在聚类之前转换键,以便将接收相似 QK 分数的键分组在一起。融合的 GPU 内核进一步减少了块检索的开销。在 HunyuanVideo 和 Wan 上的实验表明,PARK 在加速推理的同时提高了块检索精度并保持了生成质量,在比较稀疏注意力方法中实现了最佳质量效率权衡。