一句话总结

腾讯 HY 团队、上海科技大学及合作者提出了分层地标稀疏（HiLS）注意力机制，一种分块稀疏注意力机制，通过让每个查询独立关注检索到的块，并通过与语言建模损失端到端优化的学习检索分数融合输出，从而在域内匹配全注意力，同时外推至训练长度的 $64\times$64× 倍以上，检索精度达 $90\%$90%，打破效率与性能的典型权衡，实现高效、超长上下文的大语言模型。

核心贡献

• HiLS Attention 是一种分块稀疏注意力机制，可在语言建模损失下端到端学习块选择，从而支持原生稀疏训练。
• 在域内上下文长度上，其性能与全注意力相当或更优，并且可外推至超过训练上下文长度 64 倍，检索精度达 90%。
• 现有的全注意力模型可通过轻量级持续预训练转换为 HiLS Attention，在保持域内性能的同时获得超长上下文外推能力，从而打破效率与性能的权衡。

引言

在大语言模型中扩展上下文窗口对于长程 agent 任务、复杂推理和大规模信息整合至关重要，但全注意力的二次复杂度、较差的长文本外推能力以及不断增长的 KV 缓存开销使其不切实际。分块稀疏注意力方法通过仅关注少数相关块提供了一种有前景的替代方案，然而它们尚未达到全注意力的质量，特别是在长上下文检索中，因为它们依赖弱化的非参数块摘要（如均值池化），并且缺乏对选择过程的端到端优化。作者提出了 HiLS-Attention，一种原生稀疏注意力机制，它利用全注意力块质量的一阶泰勒展开从地标 token 中提取富有表达力的块摘要，然后将这些摘要嵌入分层 softmax 分解中，使得块选择在语言建模损失下端到端训练。该设计实现了准确、可学习的块检索，在困惑度上与全注意力相当，在长上下文基准上超越全注意力，并支持极端长度外推（例如，仅在 8K token 上预训练后，在 4M token 下保持超过 90% 的针尖检索准确率）。

方法

作者提出了 HiLS-Attention，一种完全可微的稀疏注意力机制，通过学习块级摘要来克服传统块稀疏注意力（BSA）的表达力局限性。整体设计应对两个核心问题：如何利用紧凑摘要可处理地近似 log-sum-exp 块质量，以及如何将其与语言建模目标端到端地学习。

在朴素的 BSA 中，序列被划分为大小为 $S$S 的统一块。对于每个查询 token $i$i，一个大小为 $W$W 的局部滑动窗口与块边界对齐，注意力被限制在前 $K$K 个远距离块上，这些块按块质量 $Z_{i,c} = \sum_{j \in \mathcal{T}_c} \exp(s_{i,j})$Zi,c​=∑j∈Tc​​exp(si,j​) 降序选择，其中 $s_{i,j}$si,j​ 是标准点积注意力 logit。这虽然实现了精确选择，但计算真实块质量需要评估块内所有 token 级别的 logit，从而抵消了稀疏性的计算节省。实际中，使用均值池化键的方法隐式地通过 logit 的均值来近似块质量，但 log-sum-exp 函数仅在 logit 近乎均匀时表现为均值，在某个 logit 占主导时表现为最大值。这种不匹配可能导致块排序错误，并遗漏重要 token。

HiLS‑Attention 用可学习的替代量替换了非参数化的块质量。第一步是通过泰勒展开对 log‑sum‑exp 进行线性化，得到如下形式的块级替代分数：

$$\hat{s}_{i,c} = \frac{\mathbf{q}_i^\top \mathbf{k}'_c}{\sqrt{d}} + b'_c,$$s^i,c​=d​qi⊤​kc′​​+bc′​,

其中 $\mathbf{k}'_c$kc′​ 和 $b'_c$bc′​ 由与块相关联的可学习查询 $\mathbf{q}'_c$qc′​ 计算得到。具体而言，$\mathbf{q}'_c$qc′​ 在块 $c$c 的键上诱导注意力分布；$\mathbf{k}'_c$kc′​ 是这些键的加权和，$b'_c$bc′​ 是该分布的熵。熵项在 log‑sum‑exp 的均匀和对数主导区域之间自适应插值。为实例化 $\mathbf{q}'_c$qc′​，作者在每个块末尾附加一个特殊的地标 token，并使用其查询向量。得到的 $(\mathbf{k}'_c, b'_c)$(kc′​,bc′​) 对作为熵校准的压缩键。然后，通过按 $\hat{s}_{i,c}$s^i,c​ 排序并保留前 $K$K 个块进行块选择，配分函数估计为 $\hat{\mathcal{Z}}_i = \sum_{c \in \mathcal{I}_i} \exp(\hat{s}_{i,c}) + Z_{i,\mathrm{swa}}$Z^i​=∑c∈Ii​​exp(s^i,c​)+Zi,swa​。由于每个 token 的路由代价为 $O(N/S)$O(N/S)，总复杂度保持为 $O(N^2/S)$O(N2/S)。

为使地标查询可学习，作者将注意力质量分解为块内归一化项和块间质量项。对于选中的 token $j$j，注意力权重被重写为

$$w_{i,j} = \frac{\exp(s_{i,j})}{Z_{i,\mathrm{c}(j)}} \times \frac{Z_{i,\mathrm{c}(j)}}{\mathcal{Z}_i} \approx \frac{\exp(s_{i,j})}{Z_{i,\mathrm{c}(j)}} \times \frac{\hat{Z}_{i,\mathrm{c}(j)}}{\hat{\mathcal{Z}}_i},$$wi,j​=Zi,c(j)​exp(si,j​)​×Zi​Zi,c(j)​​≈Zi,c(j)​exp(si,j​)​×Z^i​Z^i,c(j)​​,

其中第二个因子是学习到的替代质量。这种分层 softmax 分解使得语言建模损失的梯度可以流入地标参数，从而对预测更有用的块获得更大的质量。经验表明，这种端到端学习的替代量优于朴素 BSA。

为让 HiLS‑Attention 在现代架构中有效，引入了若干实践设计选择。首先，作者发现标准 RoPE 在稀疏注意力下会降低困惑度；他们采用 HoPE，仅保留旋转周期不超过预训练上下文长度的 RoPE 维度，其余用 NoPE 替代。其次，一个低秩查询校准模块（Q‑Cal）通过添加学习修正 $\Delta\mathbf{q}_i = \mathbf{W}^{\mathrm{up}}\mathbf{W}^{\mathrm{down}}\mathbf{h}_i$Δqi​=WupWdownhi​ 来调整原始查询 $\mathbf{q}_i$qi​ 以进行块级评分，其中 $\mathbf{h}_i$hi​ 是隐藏状态，投影矩阵的秩 $r \ll d_{\mathrm{model}}$r≪dmodel​。这显著提升了困惑度和长度外推能力。第三，对于使用分组查询注意力（GQA）的模型，同一组内的查询头必须共享相同的检索块集以获得内核效率。为保持头级别的灵活性，先按头计算块权重，然后对组内各头取最大值进行聚合，再从组级分数中选择前 $K$K 个块。另一种无地标变体每层使用一个共享的可学习查询，其域内性能相当，但外推能力较弱。

在硬件方面，每次处理一个查询 token 时，朴素的稀疏注意力内核难以充分利用 Tensor Core。HiLS‑Attention 改为将 $M$M 个相邻查询 token 分组，并在它们所选块的并集上计算注意力。由于相邻查询常常检索高度重叠的块，这种批处理将计算形状变为 $(M \times G, d) \times (d, S)$(M×G,d)×(d,S)，其中 $G$G 是每个 GQA 组的查询头数。高效的 Tensor Core 利用仅需 $M \times G \ge 16$M×G≥16，降低了对大 $G$G 的依赖，并通过重用加载的键/值块减少了冗余内存访问。

训练方面，提供了两种持续训练策略。在地标 token 调优中，冻结所有基础模型参数，仅训练地标 token 嵌入和低秩查询适配器矩阵（少于 $1\%$1% 的参数），仅需数十亿 token 即可达到与基础模型相当的性能。全参数调优联合更新所有参数，随机初始化地标嵌入和适配器，其余继承自基础模型；该策略与 HoPE 位置编码结合时尤为有效，可最大化长度泛化。

实验

在 3.45 亿参数模型上的小规模实验表明，HiLS-Attention 在困惑度和上下文检索上与全注意力相当，同时优于其他稀疏注意力方法，压缩改善了检索，且地标 token 对长度外推至关重要。扩展到 14 亿参数从头训练证实，HiLS-Attention 保留了短上下文性能，并提供了更强的外推能力，而在 70 亿参数上的持续预训练取得了与全注意力相当或更好的 LongBench 结果。消融实验证明了 HoPE 位置编码、低秩查询校准和地标 token 的重要性。推理效率基准测试显示，在超过 16K 上下文长度后，HiLS-Attention 比全注意力快得多，预填充和解码延迟呈次线性扩展，块重叠分析验证了内核的相邻查询重用策略。

在 3.45 亿参数的小规模研究中，HiLS-Attention 在从 64 到 512K token 的每个评估上下文长度上均取得最优或次优的困惑度，展现出强大的长上下文外推能力。其性能始终接近全注意力，而其他稀疏方法在较长上下文下往往退化。使用 RoPE 的全注意力在短上下文下优于 HoPE 变体，而 HSA-UltraLong 的困惑度变得由滑动窗口主导，限制了其长上下文有效性。HiLS-Attention 在所有上下文长度（64 至 512K）上均取得最优或次优困惑度，优于所有其他稀疏注意力基线。使用 RoPE 的全注意力在短上下文（最多 512 token）下比使用 HoPE 的全注意力具有更低的困惑度，使其成为短序列任务更强的基线。HSA-UltraLong 的困惑度由滑动窗口主导，其性能接近独立的滑动窗口注意力即可说明这一点。

在 8K 训练长度下，使用 RoPE 或 HoPE 的全注意力实现了完美的针尖检索，而除 HiLS-Attention 外的所有稀疏方法均未达标。应用 HoPE 位置嵌入使全注意力能够外推至 16K 和 32K，但大多数稀疏方法无法泛化到 8K 以上。HiLS-Attention 是唯一能匹配全注意力域内精度，并在与 HoPE 结合时展现强大长度外推能力的原生稀疏方法。Full-Attn RoPE 和 Full-Attn HoPE 在 8K 时均达到 100% 的单针检索（S-N），而稀疏基线如 NSA（49%）和 SWA（18%）明显落后。Dash-Attention-RoPE 在 8K 时达到 92% 的 S-N，但在 16K 时骤降至 0%，说明外推的脆弱性。Full-Attn HoPE 在 16K 时保持 99% 的 S-N 和 97% 的多键多查询（MK-MQ），在 32K 时仍部分保持性能，之后在 128K 时失败。SWA-RoPE 和 NSA-RoPE 在所有扩展长度上检索率几乎为零，证实了在没有改善位置处理的情况下，稀疏性严重限制了长上下文能力。

HiLS-Attn-HoPE 在上下文从 1M 增长到 4M token 时保持了强劲的 S-N 和 MK-MQ 分数，S-N 仅从 100 降至 96，MK-MQ 从 97 降至 89。VT 分数降幅较大，同一范围内从 53 降至 43。S-N 在所有上下文长度上保持在 100 附近，仅从 1M 到 4M 下降了 4 个点。MK-MQ 在 2M 时降至 87，随后在 4M 时回升至 89，而 VT 从 1M 时的 53 稳步降至 4M 时的 43。

当扩展到 256K 训练时，结合 HoPE 的 HiLS-Attention 在长评估长度（从 256K 到 1M token）上始终取得比全注意力更低的困惑度。稀疏方法在原始 RoPE 基频下保持稳定，而使用 HoPE 的全注意力在 512K 时崩溃。增大 RoPE 基频有助于全注意力，但 HiLS-Attention 仍具微弱优势，仅增加 0.6% 的额外参数。使用 HoPE 的 HiLS-Attention 在 256K 及以上评估长度下困惑度低于全注意力，差距在 1M token 时达 0.53。使用原始 RoPE 基频的 Full-Attention HoPE 在 256K 以上失效，512K 时困惑度超过 100，而 HiLS 变体保持在 7.55。将 RoPE 基频从 $10^4$104 增加到 $10^7$107 将全注意力在 256K 的困惑度从 7.87 降至 7.49，但 HiLS-Attn-HoPE 仍略胜一筹（7.45）。HiLS-Attention 相比全注意力仅增加 0.6% 的参数量。

在 3.45 亿参数模型上以 256K 上下文长度持续训练表明，使用大 RoPE theta（$1\times10^7$1×107）的 HiLS-Attn-HoPE 在长达 1M token 的所有长度上保持近乎完美的针尖检索，而全注意力 RoPE 基线在 256K 以上急剧退化。HiLS-Attn-HoPE 的变量追踪准确率也远高于基线，在每个上下文长度下均保持在 50 以上。使用较小 RoPE theta（$1\times10^4$1×104）的全注意力模型完全无法泛化，凸显了分层注意力与位置编码缩放相结合的重要性。HiLS-Attn-HoPE 从 8K 到 1M token 保持近乎完美的 S-NIAH（针尖检索）分数，而全注意力 RoPE 在 512K 时降至 41，1M 时降至 2。HiLS-Attn-HoPE 的变量追踪（VT）准确率在所有上下文长度下始终高于 50，而全注意力 RoPE 在长上下文下接近零。当 RoPE theta 降至 $1\times10^4$1×104 时，全注意力模型甚至在短长度下也失败，但使用更大 theta 的 HiLS-Attn-HoPE 成功，这凸显了分层注意力和 RoPE 缩放之间的协同作用。多键多查询（MK-MQ）准确率在两种方法上均较低，但 HiLS-Attn-HoPE 在 512K 和 1M 时有所改善，从 1 和 0 升至 13 和 5。

实验在长达 4M token 的上下文长度上评估了 HiLS-Attention 的语言建模和针尖检索性能。HiLS-Attention 在困惑度和检索精度上始终与全注意力相当或更优，尤其在与 HoPE 位置嵌入和缩放 RoPE 基频结合时，而其他稀疏方法在长上下文下严重退化。它展现出强大的长度外推能力，在多达 1M token 下保持近乎完美的检索，并在 256K 及以上长度下困惑度低于全注意力，仅增加 0.6% 的参数量，凸显了分层注意力与位置编码缩放之间的协同作用。