弹性阈值注意力：长上下文解码的学习上下文稀疏性

arXiv:2609.20888v1 公告类型：新 摘要：海量 KV 缓存可能会在长上下文解码过程中导致严重的内存带宽瓶颈。稀疏注意力方法通过选择性加载来缓解这种情况，但这是有代价的：严格的启发式方法会放弃必要的上下文，导致质量下降。我们引入了弹性阈值注意力（ETA），这是一种端到端的可训练架构，可以在不牺牲密集模型质量的情况下实现硬件加速解码速度。 ETA 直接从查询表示中预测动态的上下文阈值，允许模型将密集的上下文分配给困难的检索或推理步骤，同时修剪常规标记。为了从头开始学习此策略而不导致表示崩溃，ETA \emph{在训练期间将}亚阈值 logits 乘法抑制为零，而不是删除它们。针对这种平滑均匀注意力层的训练提供了一个分布式概率库，\textbf{导致初始标记上的局部注意力集中消失}。它还使模型能够在推理时硬修剪无信息的 KV 块，并吸收粗略 GPU 块选择共同承认的附带标记。因此，1.45B 预训练 ETA 模型在训练稀疏度和活动解码密度分别为 $\约 85\%$ 和 $\约 38\%$ 的情况下，可与语言建模、常识推理和长上下文针检索方面的密集注意力相媲美。在推理时，我们在 Triton 中实现了一个自定义解码内核，该内核使用缓存的几何概率边界在 $O(1)$ 时间内筛选 KV 块，在高达 512K 令牌的序列上比 FlashAttention-2 提供高达 $2.5\times$ 的挂钟解码速度。最后，我们引入了一种用于特定领域部署的离线校准算法，该算法冻结每个头的常量阈值以消除预测器开销，从而将注意力计算额外减少 27%\%$。