我们介绍 Zyphra 分层注意力（ZHA），这是一种稀疏注意力方法，它构建并搜索一棵分层树，以确定每个查询的相关 token 集合。ZHA 通过为路由选择合适的 top-p 阈值，实现了亚二次复杂度，优于现有那些仍维持二次索引器的稀疏注意力方法。我们证明，只需使用几十亿 token 的蒸馏阶段，就能将现有的预训练模型适配到 ZHA，并且 ZHA 在高稀疏度下仍能在具有挑战性的长上下文推理、智能体编码和检索基准上保持性能。我们的性能模型显示，在超长上下文中，ZHA 在速度和计算成本上大幅优于稀疏注意力方法。
为什么高效注意力如此困难？
Transformer 最显著的特性之一，是能够从序列中的任意位置检索信息。
当语言模型生成下一个 token 时，注意力使其能够回顾此前出现的所有内容，并判断哪些信息片段是相关的。例如，编码模型可能需要从数千行之前检索某个函数定义，而推理模型可能需要重新审视长思维链中的某个中间结论。
关键在于，模型不必事先决定哪些信息会重要。每个新查询都可以独立检索它所需要的那部分上下文。
这种灵活性极其强大，但代价也很高昂。
在标准注意力中，每个 token 都要与之前的所有 token 进行比较。这导致二次复杂度。因此，上下文长度翻倍会使注意力总计算量增加约四倍。在生成过程中，每增加一个 token，还需要读取一个不断增大的键值（KV）缓存，其中包含模型先前上下文的表示。
当我们向数千万 token 的上下文推进时，计算成本很快就会变得难以承受。
因此，大量研究都聚焦于让注意力更高效。然而，现有方法往往分为两大类，每一类都有一个重要的局限。
第一种方法是将过去压缩为固定大小的状态。
状态空间模型（SSM）、线性注意力变体和滑动窗口注意力机制，可以通过维护对先前序列的有界表示来实现高度高效的扩展。问题在于，序列中潜在有用信息的数量并不会停止增长。
想象你在阅读一个不断增大的藏书集合，却只被允许保留固定数量的笔记。起初，你的笔记或许能记录大部分重要内容。但最终，新信息必然要覆盖或干扰旧信息，无论这些旧信息日后可能多么有用。
第二种方法是稀疏注意力。
稀疏注意力方法通过为每个查询选择 KV 缓存的一小部分，保留了检索任意信息的能力。模型不必关注之前的所有 token，而可能只关注最相关的那 1%。
然而，这里有个问题： 模型如何确定哪 1% 的 token 是相关的？
许多稀疏注意力方法使用一个轻量级索引器，对之前的所有 token 打分，以找出值得检索的子集。尽管最终的注意力操作是稀疏的，但索引器本身仍必须搜索整个上下文。这实际上把二次计算从注意力转移到了索引器中。让索引器更廉价可以带来大幅的常数因子改进，但底层的扩展性问题依然存在。
一些方法通过索引压缩后的块而非单个 token 来解决这个问题。然而，在单一分辨率下搜索每个块，仍然需要随上下文线性增长的工作量，同样只能在渐近扩展上带来常数因子改进。
这些局限表明，一种有效的注意力替代方案需要同时具备两个特性。
第一， 检索必须保持查询相关性. 不同的问题需要不同的信息，因此我们不能简单地丢弃一组固定的、被认为不重要的 token。
第二， 主动检索的信息量必须允许随上下文长度增长，但增长速度要慢于上下文本身. 固定大小的记忆最终会变得不足，而线性增长的主动记忆又会让我们回到二次注意力。
我们的工作始于这样一个观察：自然语言似乎恰好具有这一特性。当我们考察预训练 LLM 中的注意力分布时，我们发现，占据固定比例注意力概率的 token 数量，往往随序列长度亚线性增长。在我们对长篇小说的实验中，这种关系通常可以用幂律很好地描述。
换句话说，随着文档变长，模型需要检索的信息量会增加，但它未必像文档本身增长得那么快。这提示了一个重要的可能性：或许我们可以在不付出搜索每个 token 代价的情况下，保留注意力的灵活性。
分层注意力：用树搜索上下文
ZHA 背后的关键洞见是，寻找相关信息不应要求逐一检查每一条信息。
设想在一本非常厚的书中查找某个特定段落。你通常不会为了确定该看哪里而读遍每一句话。相反，你可能会先从目录开始，找到一个有希望的章节，浏览其各节，最后定位到相关段落。
这就是分层注意力背后的基本直觉。ZHA 并不在所有先前 token 上维护一个扁平索引，而是构建一棵由 KV cache 的逐级压缩表示组成的树。
构建树
在最底层，我们从模型 KV cache 中的原始键开始。
我们将这些键划分为连续的分组，即 块,并为每个分组构建压缩表示。然后我们将这些摘要分组并再次压缩，递归重复这一过程以构建树的更高层级。
例如，当块大小为四时，一层可概括四个原始键，两层可概括 16 个键，六层则可概括 4,096 个 token 的区域。
其结果是一个层级结构：最高层描述上下文的大片区域，而较低层则提供逐级更精细的细节。
一个重要的挑战是让这些摘要对检索有用。
简单地对一个块中的键取平均，会丢失判断该块是否与特定查询相关所需的大量信息。不同的查询可能关注同一个块中完全不同的方面。
相反，ZHA 使用一小组学习得到的参考查询来构建摘要。 参考查询。每个参考查询都提供了一种考察块中键的不同方式。这些参考查询使 ZHA 能够近似估计未来查询会为该块分配多少总注意力。直观地说，我们不是要求摘要表示一个区域内的所有内容，而是训练它回答一个具体得多的问题：这个区域与我当前试图回答的查询可能有多相关？随后，这些摘要被递归地组合起来，以构建树的更高层级。每个参考查询都提供了一种考察块中键的不同方式。这些参考查询使 ZHA 能够近似估计未来查询会为该块分配多少总注意力。
直观地说，我们不是要求摘要表示一个区域内的所有内容，而是训练它回答一个具体得多的问题：
这个区域与我当前试图回答的查询可能有多相关？
随后，这些摘要被递归地组合起来，以构建树的更高层级，同时原始 KV cache 仍被保留。ZHA 压缩键是为了搜索，而不是永久替换之后可能需要检索的信息。
在树中路由
树构建完成后，检索信息就变成了一种由粗到细的搜索。
当新查询到达时，ZHA 首先将其与树最高层的摘要进行比较。它并不检查这些摘要所代表的全部 token，而是判断哪些区域看起来最相关，并仅展开这些区域。在下一层，它重复这一过程，将查询与所选区域内更细粒度的摘要进行比较。这一过程持续进行，直到搜索到达原始 KV 条目，此时模型对选定的子集执行普通注意力计算。
重要的是，ZHA 并不只是沿着树的一个分支向下走。它可以展开位于上下文中任意位置的多个分支，从而能够从相隔很远的区域检索信息。关于展开哪些分支的决策使用一种名为 top-p 路由top-p 路由不是选择固定数量的节点，而是选择最小的节点组，使其占估计注意力概率的指定比例。
这使得检索量能够自适应于查询。如果查询是在寻找某个单一的具体事实，注意力分布可能会高度集中，ZHA 就只会探索少数几个分支。如果查询需要综合上下文中许多不同部分的信息，注意力分布可能会更加分散，ZHA 就会展开多得多的分支。
自适应检索至关重要。在固定的 top-k 预算下，我们面临一个严格的权衡：k 既要足够大以回答分散的查询，又要在只需一小部分高度集中的子集时避免浪费计算。而 top-p 则允许计算量自然地适应任务。
这如何实现亚二次方扩展？
分层索引最重要的优势在于 我们不再需要以相同的分辨率搜索整个索引在每一层，都可以在不检查其具体内容的情况下，将上下文中的大片区域排除在考虑之外。只要所探索的区域数量随上下文长度呈亚线性增长，路由计算和最终的注意力计算就能在每次查询时保持亚线性。这与那些仍需扫描整个上下文的一维扁平索引的稀疏注意力方法有着本质区别。树结构使得 寻找相关 token 的过程本身是稀疏的.
它也与固定大小的状态有着根本区别。ZHA 不对可存储或检索的信息量施加固定的上限。活跃状态可以随上下文持续增长，同时仍显著小于完整序列。此外，完整的 KV-cache 保留在树的叶级，因此原始的键和值仍可供检索。
将 ZHA 蒸馏到预训练模型中
ZHA 一个特别吸引人的特性是，它不需要从头训练一个新的语言模型。相反，我们可以拿一个已有的预训练 transformer，教一个轻量级的分层索引器去近似它原本的稠密注意力行为。
我们通过一个简短的 KL 蒸馏阶段来实现这一点。在整个训练过程中，原始模型保持冻结。我们观察其稠密注意力如何在上下文的不同区域分配概率，并训练 ZHA 的参考查询，以在树的每一层复现这些概率。
实际上，预训练模型教会了索引器如何搜索自己的记忆。由于唯一可学习的组件是参考查询，因此额外参数的数量极少。在我们的 Qwen3.6-35B-A3B 实验中，ZHA 仅新增 81,920 个可训练参数，约为基座模型参数量的 0.00024%。
我们发现，仅在 32K 上下文长度下用几十亿个 token 进行训练，就足以获得强劲的性能，包括在长得多的上下文上的评估。这一点尤为重要，因为它意味着分层注意力有可能部署在现有的高能力模型之上，而无需重复其昂贵的预训练过程，也无需改变负责其推理和语言能力的权重。
使用 ZHA 时模型质量能否保持？
高效的注意力方法往往主要在相对简单的检索任务上进行评估，例如在长文档中查找单个隐藏事实。然而，这些基准测试并不能完全反映现代推理模型所面临的需求。
因此，我们在更广泛、更具挑战性的任务集合上评估 ZHA，包括 AIME 2026 数学推理、MMLU-Pro、GPQA Diamond、SWE-bench Verified、Terminal-Bench 2.1、RULER 和 LongBench v2。我们将原始稠密模型与两种 ZHA 配置进行比较：宽的单层树和深的六层树。结果令人鼓舞。图 4. 在一系列长上下文推理和检索评估中，包括 SWE-bench 和 Te 等困难的智能体任务，
因此，我们在更广泛、更具挑战性的任务集合上评估 ZHA，包括 AIME 2026 数学推理、MMLU-Pro、GPQA Diamond、SWE-bench Verified、Terminal-Bench 2.1、RULER 和 LongBench v2。
我们将原始稠密模型与两种 ZHA 配置进行比较：宽的单层树和深的六层树。
结果令人鼓舞。
在 RULER-13 在 256K token 下在 256K token 的 RULER-13 上，我们的宽树 ZHA 模型达到 96.2% 的准确率，而原始稠密模型为 96.0%。更深的六层树达到 93.3%，尽管它只关注原始 KV cache 的一小部分。
在 LongBench v2深树模型得分为 63.4%，而稠密注意力为 58.0%，同时在解码期间仅关注约 1.1% 的 KV cache。
而在 SWE-bench Verified这是一项要求严苛的智能体编码基准，涉及真实世界的软件工程任务；ZHA 的性能与稠密基线相当：深树为 67.4%，宽树为 68.4%，而稠密注意力为 67.4%。
我们还与 MiniMax 稀疏注意力（MSA）进行了比较，后者使用由学习到的索引器选出的固定数量块。MSA 索引器仅经历了 KL 预热阶段，使用了 40B token，遵循 MSA 论文中描述的训练配方。
尽管 MSA 在若干推理基准上仍具竞争力，但在长上下文检索上性能下降幅度明显更大。例如在 RULER-13 上，MSA 达到 77.1%，而我们的深树与宽树 ZHA 配置分别为 93.3% 和 96.2%。
在 LongBench v2 上，在关注键比例相近的情况下，我们的深度路由器比 MSA 高出约 12.7 个百分点，尽管索引器所需的训练 token 仅为后者的十分之一。
这些结果尤为重要，因为它们表明分层路由能够在那些远不止检索单个孤立事实的任务上保持性能。
我们还在 Nemotron-3.5-Lightning-30B-A3B 上测试了 ZHA，该模型在其注意力层中不使用位置编码（NoPE）。该方法在推理和智能体编码任务上依然表现良好，不过在 RULER 上的检索性能下降得更为明显，这与其它工作中关于 NoPE 注意力稀疏适配的发现一致。
总体而言，我们的实验表明，分层注意力可以保留稠密注意力的大部分表达能力，同时无需检查大部分上下文。
对树几何结构的鲁棒性以及实现大幅加速的潜力
一个自然的担忧是，分层注意力本身可能会引入新的瓶颈。如果信息必须经过多层压缩传递，那么树越深，准确性不就会越低吗？
这是一个关键问题，因为ZHA的亚二次复杂度扩展最终取决于我们能否在上下文长度增长时增加树的深度。为研究这一点，我们训练并评估了28种不同的树配置，变化参考查询的数量、每个块所表示的token数量以及树的深度。
值得注意的是，我们发现 ZHA 在众多此类选择下都具有高度鲁棒性。在中等稀疏度水平下，许多不同树几何结构下的检索性能相近。更令人惊讶的是，当我们固定整体压缩因子时， 在相同稀疏度水平下，较深的树往往比较浅的树表现更好。.
这表明，通过多个抽象层级逐步细化搜索，可能比做出一次激进的压缩决策更有效。
ZHA 能变得多快？
减少参与计算的 token 数量只是性能表现的一部分。
GPU 的实际性能还取决于构建树、在树的不同层级间路由、选择节点，以及最终对检索到的 token 执行注意力的开销。为研究这些开销，我们构建了一个详细的 GPU roofline 性能模型，将计算、内存带宽、内核启动开销和硬件并行性纳入考量。
我们使用该模型来比较 ZHA 与稠密注意力以及一种带有二次索引器的代表性稀疏注意力方法的预测性能。
结果表明了改变渐近扩展方式为何重要。对于轻量级 ZHA 配置，我们的模型预测，分层注意力在约 48K token 时比稠密注意力更快，在约 480K token 时比 MSA 更快。
超过这些交叉点后，随着上下文增长，优势会继续扩大。在上下文长度为 一千万 token 该模型预测大约可实现以下加速： 32–35× 在稠密和 4.7–5.1× 在 MSA 之上。在 十亿 token，这些预测的加速比达到约 373–417× 和 52–58×分别。
在足够长的上下文中，即便是非常大的常数因子改进，也会变得不如彻底避免二次方扩展那么重要。
通往真正超长上下文的窄路
超越直接的架构与实验结果，我们的工作提出了一个更根本的问题：
要让语言模型有效利用任意长的上下文，究竟需要什么？
人们很容易把长上下文扩展纯粹看作一个计算问题。如果我们能让稠密注意力足够快、足够便宜，也许我们就可以简单地无限扩大上下文窗口。
然而，我们的分析表明，实际情况可能更为微妙。要理解其中的原因，不妨考虑两种相反的失效模式。
注意力变得过于分散
随着上下文增长，稠密注意力必须将查询与越来越多的 token 进行比较。但并非所有这些 token 都包含有用的预测信息。随着更多无关或弱相关的 token 被加入，它们对注意力分布的集体贡献可能变得日益显著。
先前的工作表明，softmax 注意力在大序列长度下会变得分散：概率会扩散到数量不断增多的干扰项上，可能掩盖真正相关的信息。 分散 在大序列长度下：概率会分散到数量不断增多的干扰项上，可能掩盖真正相关的信息。
换句话说，即使我们能够承担在巨大上下文上计算稠密注意力的成本，也未必能保证模型可以从中检索到正确的信息。
这表明稀疏性可能不仅仅在计算效率方面有用。通过过滤掉过去中不相关的部分，稀疏注意力可以起到一种去噪的作用，使模型能够更精确地聚焦于真正重要的信息。
注意力变得过于稀疏
然而，让注意力变得越来越稀疏会带来相反的问题。
如果模型检索到的信息太少，最终它会丢弃有用的上下文。这一点对分层注意力尤为相关，因为随着我们沿树向下，稀疏性会不断累积。
假设我们在每一层都保留固定比例的注意力概率。随着我们增加更多层，经过完整遍历后所保留的原始注意力质量比例可能会逐渐缩小至零。
在极限情况下，这可能导致活跃检索状态实际上坍缩到一个固定大小。
这样一来，我们就会重新制造出与固定状态架构相同的基本局限：随着序列增长，模型将无法再保留越来越多的相关信息。
人们可能会尝试通过在每一层变得越来越保守、随着树越来越深而保留几乎全部注意力概率来解决这个问题。但这会导致相反的问题。如果树的更高层保留了几乎每一个分支，层级结构就不再提供有意义的剪枝，计算成本也会接近稠密注意力的水平。
因此，我们面临一条狭窄的道路，夹在两种不可取的极端之间：
稀疏性太低，注意力在计算上会变得难以承受，并且可能被无关信息淹没。稀疏性太高，模型推理所需的信息又会丢失。
将压缩程度与数据中的预测信息相匹配
我们认为，穿越这一权衡的办法是认识到：适当的压缩量应当由序列本身的统计结构来决定。
因此，一段很长的自然语言序列包含大量冗余。第一百万个 token 仍然可以提供有用的新信息，但它所贡献的额外预测信息量可能远小于较早的 token。
因此，相关信息的总量可以持续增长，而无需随总序列长度线性增长。这正是我们最终想要达到的状态：
- 活跃状态持续增长，避免了固定大小记忆带来的信息瓶颈。
- 它呈亚线性增长，避免了稠密注意力高昂的成本。
- 检索仍然依赖于查询，因此不同的问题可以访问完整历史中不同的部分。
要实现这一状态（渐近地），需要对稀疏性进行比在每个树层级简单应用固定 top-p 阈值更精细的控制。特别是，路由阈值必须随树深度和上下文长度而演变，使累积压缩量能够跟踪预测实际所需的信息。关键在于，这种缩放并非由注意力架构本身所保证。它取决于底层数据的分布。
更广泛的含义是， 真正超长的上下文可能需要比单纯让稠密注意力更快更根本的解决方案.
在足够大的规模下，挑战不仅仅在于存储或处理更多的 token。关键在于确定过去的信息中有多少仍然有用，并精确地检索出这些信息，而不被其他所有内容所淹没。
我们的结果表明，分层注意力为实现这一点提供了有前景的架构基础。通过将查询相关的检索与动态增长、亚线性的活跃状态相结合，ZHA 证明注意力不必局限于固定大小的记忆或二次计算。
我们认为，最令人兴奋的可能性在于：恰当形式的稀疏性或许能够同时应对极长上下文在计算和信息两方面的挑战。
如果我们能够学会让活跃注意力的增长与真实世界序列中预测信息的增长相匹配，那么数亿甚至数十亿 token 的上下文最终可能变得可行。因此，长上下文语言建模的未来或许取决于能否赋予模型一种能力：在每一个尺度上，都能精准找到真正重要的内容。
继续阅读
研究在线向量量化注意力
2026年2月 · 阅读时长8分钟
研究树注意力：面向 GPU 集群长上下文注意力的拓扑感知解码
2024年8月 · 阅读时长3分钟
研究