先看证据
把机制串起来
Transformer模型的核心瓶颈是自注意力机制,它需要查询(query)与每个键(key)进行交互,导致计算量随上下文长度(N)呈二次方增长(O(N^2))。现有稀疏注意力方法试图通过索引器减少计算,但索引器本身可能仍需扫描整个上下文,成本并未真正降低。
ZHA通过递归压缩KV cache中的连续key块来构建一棵摘要树,摘要块越往上层代表越大的上下文区域,底层叶节点保留原始KV信息。利用少量可学习的参考查询,ZHA近似计算每个摘要块对未来查询的注意力潜力。解码时,查询自顶向下遍历树,结合top-p路由策略,选择性地只展开预估注意力占比高的分支,最终仅对选中的原始KV子集计算注意力。
['适用于需要处理极长上下文的任务,尤其是检索和推理。', '通过KL蒸馏适配冻结模型,成本相对较低,易于部署。', '对模型架构的改动主要集中在注意力模块,兼容性较好。', '不适用于对计算速度有极致要求的场景,因其高效性依赖于模型规模和上下文长度。']
注意力瓶颈的新解法,还是计算负担的转移?
大型语言模型(LLM)在处理长上下文时的核心挑战之一,是其标准注意力机制的计算成本会随着序列长度呈二次方增长。这意味着当上下文窗口扩展时,计算资源需求会急剧膨胀,严重限制了模型处理超长文档、复杂代码库或长期对话的能力。业界虽已探索多种稀疏注意力方法,试图减少每个查询与过去所有 token 的交互,但许多方案只是将计算负担从注意力矩阵转移到了一个负责“查找相关 token”的索引器上,而这个索引器本身往往仍需要扫描整个上下文,导致整体计算复杂度并未真正实现次二次突破。
针对这一痛点,Zyphra AI 发布了 Zyphra 分层注意力(ZHA)技术,旨在从根本上解决“如何高效查找相关信息”的问题。ZHA 不仅让最终的注意力计算变得稀疏,更通过引入一个分层树状索引器,使查找过程本身也变得稀疏。这项技术的核心目标是在保持模型性能的同时,实现更具扩展性的次二次计算复杂度,为现有冻结的预训练 LLM 带来处理极长上下文的能力,而无需从头训练庞大的新模型。
机制拆解:树形索引如何高效筛选长上下文
ZHA 的核心在于其独特的分层树状索引器设计,它改变了模型在 KV 缓存中寻找相关信息的范式。首先,ZHA 将连续的 key-value(KV)对划分为若干个块,然后对每个块进行压缩,生成一个摘要 key。这些摘要 key 再被递归地分组和压缩,层层向上构建成一个树结构。在这棵树中,顶层的节点代表着更大的上下文区域,而底层的叶节点则直接对应原始的 KV 缓存条目。
为了评估这些摘要 key 对未来查询的潜在相关性,ZHA 引入了一组数量不多的、可训练的“参考查询”。这些参考查询通过学习,能够为每个块的注意力质量提供切面估计,从而近似出任何未来查询可能给予该块的注意力程度。当解码一个新的 token 时,当前的查询会从树的顶层开始,对顶级节点进行评分。它会选择一小组累计估计注意力概率达到指定阈值(称为 top-p)的节点,然后只展开这些选定节点的子节点,重复这一过程,直到到达原始的 KV 条目。这种自适应的 top-p 路由机制允许 ZHA 根据查询的性质,动态地检索不同数量的上下文,例如,一个简单的事实查找可能只涉及少数分支,而需要综合多个分散信息的查询则会展开更多分支。值得强调的是,这一整套索引器机制是在基础 LLM 保持冻结的情况下,通过与原始模型的稠密注意力分布进行 KL 蒸馏训练而得,只需新增极少量参数(例如在 Qwen3.6-35B-A3B 上仅新增约 8.2 万个参数,占总参数量的 0.00024%),极大地降低了适配成本。
性能验证:实测与模型预测的“次二次”能力
Zyphra AI 在 Qwen3.6-35B-A3B 模型上对 ZHA 进行了广泛评估,展示了其在长上下文任务上的竞争力。通过仅用约 40 亿 token 的 KL 蒸馏,ZHA 能够适配冻结的预训练模型。以下是其在 256K 上下文长度下的部分关键性能指标:
**ZHA 与稠密注意力及 MSA 性能对比(256K 上下文,Qwen3.6-35B-A3B):** | 基准测试 (256K) | 稠密注意力 | ZHA-深树 | ZHA-宽树 | MiniMax Sparse Attention (MSA) | | :------------ | :-------- | :------- | :------- | :--------------------------- | | RULER-13 | 96.0% | 93.3% | **96.2%**| 77.1% | | LongBench v2 | 58.0% | **63.4%**| 62.4% | 49.3% | | SWE-bench Verified | 67.4% | 67.4% | **68.4%**| 63.6% |
*数据来源:Zyphra 自我报告,arXiv 2610.42112,图 5。深树(chunk 4, 6 层),宽树(chunk 256, 1 层),MSA 为 top-16 block of 128,经 40B token 训练。*
从数据中可以看出,ZHA 在多个任务上表现与稠密注意力相当甚至更优。例如,在 256K 上下文的 RULER-13 检索任务中,宽树 ZHA 略优于稠密注意力基线;在 LongBench v2 的长推理任务中,深树 ZHA 显著超出稠密注意力,同时仅访问了约 1.1% 的 KV 键。与 MiniMax Sparse Attention (MSA) 相比,ZHA 的优势更为明显,尤其是在 RULER-13 上取得了远超 MSA 的性能,且 ZHA 索引器的训练数据量远小于 MSA。这些结果表明,ZHA 在保持高准确率的同时,能够实现更高的稀疏度,从而为实现计算效率提升奠定基础。此外,Zyphra 的 GPU roofline 性能模型预测,在轻量配置下,ZHA 的速度优势将在 48K token 时显现(相比稠密注意力),并在 480K token 时超过 MSA;在极端长上下文(如 1,000 万 token)时,理论上甚至能比稠密注意力快 32-35 倍。
实际影响:何时引入 ZHA 值得考量?
对于寻求在实际部署中扩展 LLM 上下文窗口的技术负责人而言,ZHA 提供了一条值得关注的路径。其最直接的影响在于,它允许在不修改或重新训练大型基础模型的前提下,大幅提升模型处理超长输入的能力。这意味着可以利用现有高性能的冻结模型,以较低的开发成本和资源投入,将其应用范围扩展到需要深度理解和处理极长文本的场景,例如法律文档分析、超长篇幅财报解读、复杂代码库的自动修复或企业知识库检索等。
此外,ZHA 的自适应 top-p 路由机制带来了显著的灵活性和效率。当用户提出一个精确的检索查询时,ZHA 可以迅速定位到少数相关信息,而无需处理整个上下文;当任务需要更广泛的背景理解或信息综合时,它又能智能地扩展搜索范围,以满足需求。这种动态的稀疏度管理对于构建响应式、资源高效的 AI 代理和应用至关重要。与那些需要从头训练或对模型架构进行大量修改的稀疏注意力方法相比,ZHA 仅需通过蒸馏训练一个轻量级的索引器,使其成为现有 LLM 基础设施升级的经济有效选项。
局限与权衡:未被压缩的 KV 缓存和模型速度
尽管 ZHA 在理论和基准测试中展现出巨大潜力,但在实际部署前,技术负责人仍需审慎考量其局限性与权衡。首先,ZHA 的树形索引器虽然极大地稀疏了对 KV 缓存的“访问”过程,从而降低了计算复杂度,但它并未从根本上压缩或减少 KV 缓存本身的“存储”需求。原始的 KV 键值对仍需完整地保存在内存中,这意味着随着上下文长度的增加,内存占用依然会线性增长,这在处理极大规模上下文时(例如数百万甚至数十亿 token)可能依然构成挑战。
其次,当前 ZHA 提出的速度优势主要基于理论分析和 GPU roofline 性能模型预测。这些模型预测的数据点,例如在 1,000 万 token 时能比稠密注意力快 32-35 倍,并非实际硬件上的端到端延迟测量结果。ZHA 的性能提升在很大程度上依赖于优化的内核实现,而 Zyphra AI 目前尚未发布相应的优化内核。在缺乏实测数据和生产级优化代码的情况下,将这些预测数字直接转化为实际的生产系统吞吐量和延迟提升,还存在一定的不确定性。此外,“匹配稠密注意力”这一说法也需要限定范围:ZHA 在某些任务上确实与稠密注意力持平或更优,但在其他任务或特定配置下(如深树 ZHA 在 RULER-13 上),性能可能略有下降,这要求在具体应用中进行细致的评估和调优。
技术负责人判断:短期适配与长期潜力
综合来看,Zyphra 分层注意力(ZHA)代表了稀疏注意力领域一个重要的架构进步。它通过解决“如何高效查找相关 token”这一深层问题,为实现真正意义上的次二次注意力提供了一个有前景的框架。对于那些已经部署了大型冻结 LLM 并希望经济高效地扩展其上下文窗口的技术团队来说,ZHA 提供了一种引人注目的短期适配策略,其极低的额外参数量和通过蒸馏实现模型扩展的能力,使其成为一个低风险、高潜力的增强方案。
然而,在将其投入生产环境之前,关键在于密切关注 Zyphra AI 未来是否会发布针对特定硬件的优化内核,以及提供更详尽的实测速度对比数据。只有当这些性能优势能在真实世界的复杂负载下得到验证时,ZHA 的次二次计算复杂度和模型预测的巨大速度提升才能真正转化为实际的运营效益。在短期内,ZHA 值得在非关键性长上下文任务上进行小规模概念验证和性能评估;而从长远看,它有望成为构建下一代超长上下文、高效能 LLM 基础设施的关键一环,推动 LLM 应用边界的进一步拓展。