SparseD：扩散语言模型的稀疏注意力

虽然扩散语言模型（DLM）为自回归模型（AR）提供了一个有前途的替代方案，但现有的开源 DLM 存在推理延迟高的问题。这种瓶颈主要是由于注意力计算所有查询-键对时，其二次复杂度与上下文长度有关。直观地说，为了降低这种复杂性，一个自然的策略是将注意力限制在保留最相关连接的稀疏模式上。这类方法在 AR 中已经非常成熟，因为注意力遵循固定且明确定义的稀疏模式。然而，在 DLM 中，我们观察到不同的稀疏行为：（1）注意力模式因头而异，（2）每个头中的注意力模式在去噪步骤之间保持高度相似，以及（3）早期去噪步骤对生成至关重要。这些发现使得为 AR 设计的稀疏注意力方法与 DLM 基本不兼容，因为它们未能捕捉到特定头的结构，并且在早期去噪步骤中使用时有损生成质量的风险。为了解决这些挑战，我们提出了 SparseD，一种新颖的 DLM 稀疏注意力方法。利用这些观察，SparseD 只需要预先计算特定头的稀疏模式一次，并在所有步骤中重复使用它们。这可以防止在每个去噪步骤中重新计算稀疏模式。同时，SparseD 在早期步骤中使用全注意力，然后在后期切换到稀疏注意力以保持生成质量。总而言之，这些使得 SparseD 成为部署 DLM 进行长上下文应用的实用且高效的解决方案。实验结果表明，SparseD 在 64k 上下文长度和 1,024 个去噪步骤的情况下，实现了无损加速，比 FlashAttention 快 1.50 倍。