ZHA：树状索引器如何实现亚二次注意力

大语言模型的长上下文能力正被注意力的二次方复杂度死死拖住。现有方案要么牺牲表达能力（如滑动窗口或固定状态的线性注意力），要么只带来常数级加速（如索引器本身仍是二次方的稀疏注意力）。Zyphra 提出的 Hierarchical Attention（ZHA）试图从渐近复杂度上破局：用一棵递归压缩的树状索引器，让总成本随序列长度亚线性增长，从而实现整体亚二次（但仍超线性）的注意力。

为什么现有稀疏注意力不够“稀疏”

稀疏注意力的初衷是只计算部分 token 对，但多数方法的索引器（决定哪些 key 被选中）自身就带着二次方开销。换句话说，省下的注意力计算又被路由成本吃掉了。另一类线性注意力变体把历史压进固定大小的状态，长序列下信息必然有损，表达力无法与全注意力对齐。ZHA 的切入点很直接：让索引结构本身的构建与查询成本都随长度亚线性增长，而不是只做常数级优化。

注意力二次方复杂度指标准自注意力中每个 query 都要与全部 key 计算相似度，得到 L×L 的注意力矩阵，计算与显存开销随序列长度 L 的平方增长。当上下文从 8K 扩展到 128K，这部分成本会放大约 256 倍，成为长上下文推理的首要瓶颈。线性注意力则改用核函数或循环状态，把历史信息压缩进固定大小的记忆，复杂度降为线性，但固定状态容量有限，长序列中早期信息会被覆盖或稀释，导致检索类任务性能下降。稀疏注意力试图折中：只计算部分 query-key 对，但多数方法的“索引器”本身仍需对全序列打分，因而保留二次方开销，这正是文中说“省下的注意力又被路由成本吃掉”的含义。

树状索引器的工作机制

ZHA 在 KV cache 上构建一棵递归压缩的 key 摘要树。新来的 query 从树顶向下遍历，逐层筛选，最终只检索一小部分稀疏 key。关键设计在于树的深度可以迭代增长：随着序列变长，树加深，路由与最终注意力的总成本相对序列长度呈亚线性。论文给出的结论是，整体注意力因此获得亚二次但超线性的缩放曲线。

这种层级筛选的思路并不追求一次定位，而是用多层粗到细的摘要逐级缩小候选集。相比平面的稀疏索引，树结构把“找相关 key”这件事的复杂度摊到了对数级的深度上。

KV cache 是自回归推理时缓存的历史 key/value 张量，避免每步重算，其大小随已生成 token 数线性增长，也是长上下文显存占用的主要来源。ZHA 不直接在原始 key 上做平面检索，而是在这份缓存之上额外维护一棵摘要树：叶节点对应原始 key 块，上层节点是其子节点的压缩表示（如均值或学习到的池化向量）。查询时先用 query 与顶层摘要比对，选出若干相关子树，再逐层向下细化，最终只对少量叶节点做精确注意力。这样路由成本取决于树的节点数而非 token 数，树加深时每层候选集指数收缩，从而把“找相关 key”的总代价压到亚线性。

低成本迁移到已有模型

ZHA 的实用价值在于不必从头预训练。论文显示，只需在几千亿 token 量级上做一个简短的 KL 蒸馏训练阶段，就能把 ZHA cheaply 应用到现有预训练模型上，并且在长上下文检索以及推理、数学、智能体编程等评测任务上保持性能。这一点对落地很关键——大多数团队没有预算重训基座模型，短蒸馏加结构替换是更现实的路径。

KL 蒸馏在这里指用原全注意力模型作为教师，让替换成 ZHA 结构的学生模型在相同输入上模仿教师的输出分布，损失为教师与学生下一 token 分布的 KL 散度。相比从头预训练，蒸馏只需少量 token 就能让新结构对齐原有模型的注意力行为，避免灾难性遗忘。文中“几千亿 token 量级”的训练量看似很大，但相对于基座模型通常数万亿 token 的预训练成本仍小一个量级以上，且不需要重新准备大规模语料，因此被称为“cheaply 应用”。这一路径的隐含前提是：ZHA 的树状路由能力可以通过蒸馏从教师的全注意力中学到，而不是必须靠大规模预训练自行涌现。

对树几何的鲁棒性

一个容易被忽视的工程问题是：树的结构参数（分支度、深度）会不会敏感？论文报告该方法对树几何高度鲁棒，甚至能稳定支持深达 6 层的树。这意味着实际部署时不必为不同序列长度精细调参，深树带来的额外路由层数不会破坏效果。

性能模型与长上下文外推

作者构建了详细的性能模型，结论有三层：第一，方法在实测中呈现亚二次缩放；第二，在超过 100 万上下文时显著优于 competing methods；第三，具备把有效上下文窗口扩展到数亿甚至数十亿 token 的潜力。需要说明的是，后两点来自模型推演与实验外推，属于方向性判断而非已落地的产品指标。

从研究脉络看，ZHA 与近期“用层级结构降低长上下文成本”的思路一脉相承，但明确把渐近复杂度作为设计目标，而非只优化常数因子。对关注长上下文推理成本、稀疏注意力改造和模型蒸馏迁移的读者，这篇论文提供了一个值得跟进的架构方向。