SLA：通过可微调的稀疏线性注意力，在扩散Transformer中超越稀疏性

在扩散 Transformer (DiT) 模型中，尤其是在视频生成方面， 由于序列长度长且具有二次复杂度，注意力延迟是主要瓶颈。我们发现注意力权重可以分为两部分：一小部分高秩的大权重和其余的低秩权重。这自然地暗示着对第一部分应用稀疏加速，对第二部分应用低秩加速。基于这一发现，我们提出了 SLA（稀疏线性注意力），一种可训练的注意力方法，它融合了稀疏和线性注意力以加速扩散模型。SLA 将注意力权重分为关键、边缘和可忽略类别，对关键权重应用 O(N^2) 注意力，对边缘权重应用 O(N) 注意力，并跳过可忽略的权重。SLA 将这些计算整合成一个 GPU 内核，并支持前向和后向传播。通过仅进行几次使用 SLA 的微调，DiT 模型在注意力计算上实现了 20 倍的缩减，从而在不损失生成质量的情况下显著加速。实验表明，SLA 在不降低端到端生成质量的情况下将注意力计算减少了 95%，优于基线方法。此外，我们为 SLA 实现了一个高效的 GPU 内核，该内核在 Wan2.1-1.3B 上实现了 13.7 倍的注意力计算速度提升和 2.2 倍的端到端视频生成速度提升。