PulseCol：定期刷新列稀疏注意力以加速扩散语言模型

扩散 大语言模型 (dLLM) 中的推理在计算上非常昂贵，因为在没有 KV 缓存的情况下，去噪过程的每一步都必须重复执行完全自注意力。最近的 dLLM 稀疏注意力方法通过块稀疏计算减轻了这种成本，该计算仅在模型性能对粗粒度稀疏近似不太敏感时应用在后续迭代中，但在计算效率和加速方面的改进有限。这激发了更细粒度的稀疏化策略，该策略可以从早期迭代中应用，并利用可重用的稀疏模式，从而进一步提高效率。在这项工作中，我们介绍了 PulseCol，一种用于加速扩散语言模型的定期刷新的列稀疏注意力方法。 PulseCol 用更细粒度的列稀疏结构取代了粗的块级稀疏性，允许更精确地保留重要的注意力交互，同时暴露更大的稀疏性。建立在这个列级公式的基础上，PulseCol 进一步在早期去噪步骤中识别稀疏模式，并在后续迭代中重用它们，仅在少量中间步骤中刷新它们，以跟踪去噪期间稀疏注意力模式的演变。实验表明，与之前的 dLLM 稀疏注意力方法相比，PulseCol 实现了更高的稀疏性和更大的实际加速，同时保持了模型质量。通过针对列稀疏注意力优化的 GPU 内核，PulseCol 在多个上下文长度上比 FlashAttention 提供高达 1.95$\times$ 的端到端加速。