Subquadratic Attention with a Hierarchical Tree-Structured Indexer

📄 arXiv: 2610.42112v1 📥 PDF

作者: Leon Lufkin, Cutter Dawes, Alexander Speicher, Kamesh Krishnamurthy, Tomás Figliolia, Beren Millidge

分类: cs.CL, cs.AI

发布日期: 2026-10-30

备注: 43 pages

💡 一句话要点

提出Zyphra层次注意力以解决长上下文处理瓶颈

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文处理 稀疏注意力 层次结构 大型语言模型 高效计算 自然语言处理 推理任务

📋 核心要点

  1. 现有的注意力机制在处理长上下文时面临二次扩展的瓶颈，限制了大型语言模型的能力。
  2. 本文提出Zyphra层次注意力（ZHA），通过构建递归压缩的键摘要树实现亚二次稀疏注意力。
  3. 实验结果表明，ZHA在处理超过100万上下文时显著优于现有方法，并能扩展到数亿或数十亿的令牌。

📝 摘要（中文）

注意力机制的二次扩展是限制大型语言模型（LLMs）处理长上下文能力的根本瓶颈。虽然已有多种方法尝试规避这一问题，但大多数要么无法匹配注意力的表达能力，要么仅实现常数因子的加速。本文提出了一种基于层次树结构索引器的渐近亚二次稀疏注意力方法，称为Zyphra层次注意力（ZHA）。ZHA通过构建递归压缩的键摘要树，允许未来查询从树顶向下遍历以检索稀疏的键子集。通过迭代增加树的深度，整体成本在序列长度上以亚线性增长，从而实现注意力的亚二次扩展。我们展示了ZHA可以在现有预训练模型上低成本应用，并在长上下文检索和复杂推理任务中保持性能。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在处理长上下文时的二次扩展瓶颈，现有方法如滑动窗口和稀疏注意力在效率和表达能力上存在不足。

核心思路：提出Zyphra层次注意力（ZHA），通过构建递归压缩的键摘要树，允许查询从树顶向下检索稀疏键子集，从而实现亚二次扩展。

技术框架：ZHA的整体架构包括树结构的构建、键摘要的递归压缩、查询的遍历过程以及最终的注意力计算。主要模块包括树的深度增长和路由机制。

关键创新：ZHA的核心创新在于通过层次树结构实现了注意力的亚二次扩展，显著提升了长上下文处理能力，区别于传统的稀疏注意力方法。

关键设计：在设计中，树的深度可迭代增长，确保在序列长度增加时，整体计算成本以亚线性方式增长，且在训练阶段仅需短时间的KL蒸馏训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Zyphra层次注意力在处理超过100万上下文时显著优于现有方法，整体性能提升幅度达到数倍，且在长上下文检索和复杂推理任务中表现出色，展示了其在实际应用中的巨大潜力。

🎯 应用场景

Zyphra层次注意力的研究具有广泛的应用潜力，尤其在需要处理长文本或复杂推理的自然语言处理任务中，如机器翻译、文本生成和问答系统等。其高效的上下文处理能力将推动大型语言模型在实际应用中的表现，尤其是在需要处理数亿或数十亿令牌的场景中。

📄 摘要（原文）

The quadratic scaling of attention is a fundamental bottleneck limiting the long-context capabilities of large language models (LLMs). A variety of approaches have attempted to circumvent this, but they either fail to match attention's expressiveness (e.g., sliding-window or linear attention variants with fixed-size state) or only achieve constant-factor speedups (e.g., sparse attention methods whose indexers themselves scale quadratically). In this paper, we derive an asymptotically subquadratic sparse attention method based on a hierarchical, tree-structured indexer, which we call Zyphra Hierarchical Attention (ZHA). ZHA constructs a tree of recursively compressed key summaries over the KV cache, which future queries traverse from the top down to retrieve a sparse subset of keys. Crucially, by iteratively growing the tree's depth, the total cost including routing and final attention grows sublinearly in sequence length, resulting in overall subquadratic (but superlinear) scaling for attention. We demonstrate that ZHA can be cheaply applied to existing pretrained models with only a short KL distillation training stage over a few billion tokens, and maintains performance on long-context retrieval and challenging reasoning, mathematical, and agentic coding evaluation tasks. Moreover, we show that our method is highly robust to tree geometry including very deep trees of up to 6 levels. Finally, we develop a detailed performance model and show that our approach scales subquadratically, dramatically outperforms competing methods beyond 1 million context, and ultimately has the potential to extend effective context windows to the hundreds of millions or billions of tokens.