具有可训练稀疏注意力的快速视频扩散

视频扩散变压器（DiTs）的扩展受到其二次三维注意力机制的限制，尽管大部分注意力集中在一小部分位置上。我们基于这一观察提出了VSA，一种可训练的、硬件高效的稀疏注意力机制，该机制在训练和推理过程中替代了全注意力机制。在VSA中，一个轻量级的粗略阶段将标记汇聚成块，并识别出高权重的关键标记；精细阶段仅在这些块内计算标记级别的注意力，并遵循块计算布局以确保硬件效率。这导致了一个单一的可微内核，可以端到端地进行训练，无需事后分析，并且能够维持FlashAttention3 MFU的85%性能。我们通过预训练参数从6000万到14亿的DiTs模型进行了大规模的消融研究和扩展定律实验。VSA达到了一个帕累托点，在不降低扩散损失的情况下将训练FLOPS减少了2.53倍。对开源Wan-2.1模型进行改造后，注意力时间加快了6倍，端到端生成时间从31秒降至18秒，同时保持了相当的质量。这些结果确立了可训练稀疏注意力作为全注意力的一种实用替代方案，并为进一步扩展视频扩散模型提供了关键支持。