ZHA以分层键摘要树实现次二次稀疏注意力

注意力的二次复杂度增长是限制大型语言模型（LLMs）的长上下文能力的基本瓶颈。各种方法试图规避这一点，但它们要么无法匹配注意力的表达能力（例如，滑动窗口或具有固定大小状态的线性注意力变体），要么只能实现恒定因子的加速（例如，索引器本身按二次方缩放的稀疏注意力方法）。在本文中，我们推导了一种基于分层树结构索引器的渐近次二次稀疏注意力方法，我们将其称为 Zyphra 分层注意力（ZHA）。 ZHA 在 KV cache 上构造一棵递归压缩的键摘要树，未来的查询从上到下遍历以检索键的稀疏子集。至关重要的是，通过迭代地增加树的深度，包括路由和最终注意力的总成本在序列长度上呈次线性增长，从而导致注意力的整体次二次（但超线性）缩放。我们证明，ZHA 可以廉价地应用于现有的预训练模型，只需经过数十亿 token 的简短 KL 蒸馏训练阶段，并保持长上下文检索和具有挑战性的推理、数学和 Agentic 编码评估任务的性能。此外，我们表明我们的方法对于树几何形状具有高度鲁棒性，包括高达 6 层的非常深的树。最后，我们开发了一个详细的性能模型，并表明我们的方法呈次二次复杂度增长，在上下文超过 100 万 token 时显著优于竞争方法，并最终有可能将有效的上下文窗口扩展到数亿或数十亿 token。