Zyphra 提出树结构索引器 ZHA,实现亚二次复杂度长上下文注意力
Subquadratic Attention with a Hierarchical Tree-Structured Indexer
Zyphra 的论文提出 ZHA,几步蒸馏就能把现有模型改成亚二次注意力,1M 上下文以上比现有稀疏注意力方法快很多,做长上下文的值得看。
Zyphra 在 arXiv 发布论文,提出基于层级树结构索引器的稀疏注意力方法 ZHA。该方法在 KV cache 上构建递归压缩的 key 摘要树,查询自顶向下遍历检索稀疏 key 子集,使总成本随序列长度亚线性增长。ZHA 只需在几十亿 token 上做短时 KL 蒸馏即可应用到现有预训练模型,在长上下文检索、数学推理和 agentic coding 评测上保持性能。性能模型显示其可扩展到 1M 以上上下文并优于同类方法,理论上可扩展到数亿至数十亿 token 上下文。
Subquadratic Attention with a Hierarchical Tree-Structured Indexer
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.