DashAttention：可微分和自适应稀疏层次注意力

当前的分层注意力方法，例如 NSA 和 InfLLMv2，根据粗略注意力分数选择前 k 个相关键值 (KV) 块，然后对所选标记应用细粒度的 softmax 注意力。然而，top-k 操作假设任何查询的相关标记的数量是固定的，并且它排除了稀疏阶段和密集阶段之间的梯度流。在这项工作中，我们提出了 DashAttention (Differentiable and Adaptive Sparse Hierarchical Attention)，它利用自适应稀疏 $α$-entmax 变换在第一阶段根据当前查询选择可变数量的块。这反过来又为第二阶段的 softmax 注意力提供了先验，保持整个层次结构完全可微。与其他分层注意力方法相反，我们证明了 DashAttention 是非分散的，这意味着更好的长上下文建模能力。 大语言模型 (LLM) 的实验表明，与 NSA 和 InfLLMv2 相比，DashAttention 达到了与完全注意力相当的精度，稀疏度为 75%，并且具有更好的帕累托前沿，尤其是在高稀疏度情况下。我们还在 Triton 中提供了一种高效的、GPU 感知的 DashAttention 实现，它在推理时的加速速度高达超过 FlashAttention-3。总体而言，DashAttention 提供了一种经济高效的策略来对长上下文进行建模。