摘要

在扩散变换器（Diffusion Transformer, DiT）模型中，尤其是在视频生成任务中，注意力机制的延迟是一个主要瓶颈，这主要源于序列长度较长以及注意力计算具有二次方复杂度。我们发现，注意力权重可被划分为两部分：一小部分高秩的大权重，以及其余大部分低秩的权重。这一现象自然启发我们对第一部分采用稀疏加速，对第二部分采用低秩加速。基于此发现，我们提出了一种可训练的注意力机制——SLA（Sparse-Linear Attention），该方法融合了稀疏注意力与线性注意力，以加速扩散模型的运行。SLA将注意力权重分为三类：关键权重、边缘权重和可忽略权重。对于关键权重，采用复杂度为 $O(N^2)$ 的标准注意力计算；对于边缘权重，采用 $O(N)$ 的线性注意力；而对于可忽略的权重，则直接跳过计算。SLA将上述不同计算方式整合为单一的GPU核函数，并支持前向传播与反向传播。仅通过少量微调步骤引入SLA，DiT模型即可实现注意力计算量减少20倍，显著提升推理速度，同时保持生成质量不受损失。实验结果表明，SLA在不降低端到端生成质量的前提下，使注意力计算量减少了95%，显著优于基线方法。此外，我们实现了一个高效的SLA GPU核函数，在Wan2.1-1.3B模型上实现了注意力计算13.7倍的加速，以及视频生成任务整体2.2倍的端到端加速。