把稀疏注意力做对：HiLS-Attention 如何在数学上同时打破效率与效果的天花板
让大模型"读得更长"一直是 Agent、深度推理和海量资料整合等场景的刚需，但标准全注意力机制的计算量随序列长度呈平方级增长，始终是横亘在长上下文建模面前的难题。腾讯混元团队近期开源的 HiLS-Attention，首次在数学层面同时解决了稀疏注意力中 chunk 重要性估计的表达力不足与端到端不可导两大根本问题。
一、问题：全注意力扛不动，稀疏注意力又"不准"
对大语言模型而言，"读得更长"有三座大山：
- 计算量平方级增长：序列翻倍，算力翻四倍。
- 长度外推能力差：训练 8K，测试 32K 就开始崩溃。
- KV Cache 随长度线性膨胀：显存很快就吃不消。
传统的分块稀疏注意力（Chunk-wise Sparse Attention）是一个直觉上很合理的思路：把上下文切成若干个 chunk，每个 query 只挑选最相关的 Top-K 个 chunk 参与注意力计算，其余 KV cache 卸载到 CPU。但实践中的尴尬是------怎么挑 chunk？
早期的做法是用 chunk 内 token 的注意力分数均值做重要性估计，这里有两个核心缺陷：一是均值本身丢失了分布信息，表达能力不足；二是 Top-K 选择操作不可导，梯度无法端到端回传，导致训练阶段和推理阶段的行为不一致。结果就是稀疏注意力在长文本任务上的效果始终追不上全注意力。
二、解法：分层地标稀疏注意力
HiLS-Attention（Hierarchical Landmark Sparse Attention）的核心贡献在于重新设计了三件事：
2.1 地标 token 机制
不再对 chunk 内所有 token 求均值，而是在每个 chunk 内选取少量"地标 token"（Landmark Tokens），用这些地标 token 与 query token 的交互来代表整个 chunk 的重要性。地标 token 在训练过程中通过可学习的投影矩阵动态生成，天然可导。
2.2 分层选择策略
采用两层递进结构：
Layer 1: 粗筛选 → 用地标 token 快速计算 chunk 级重要性分数
Layer 2: 细选择 → 在初筛通过的 chunk 中执行精确注意力计算这种分层设计使得计算复杂度从 O(N²) 降为 O(N·K)，其中 K 为选中的 chunk 数，远小于总 chunk 数。
2.3 端到端可导的 Top-K 近似
HiLS-Attention 使用 Gumbel-Softmax 重参数化技巧来近似 Top-K 选择操作，使得整个选择过程对梯度透明。这意味着模型在训练阶段就能学会"如何挑选重要 chunk"，而非在推理时临时抱佛脚。
用伪代码概括核心逻辑：
python
def hils_attention(Q, K, V, num_chunks, top_k):
# Step 1: 将序列切分为 chunk，每个 chunk 生成地标 token
chunks = split_into_chunks(K, V, num_chunks)
landmarks = [chunk.landmark_projection() for chunk in chunks]
# Step 2: 用地标 token 计算重要性分数
landmark_scores = softmax(Q @ landmarks.T / sqrt(d_k))
# Step 3: Gumbel-Softmax 近似 Top-K 选择（可导）
selected = gumbel_topk(landmark_scores, k=top_k)
# Step 4: 仅在被选中的 chunk 上执行精确注意力
K_selected = concat([chunks[i].K for i in selected])
V_selected = concat([chunks[i].V for i in selected])
return softmax(Q @ K_selected.T / sqrt(d_k)) @ V_selected三、效果：效率与质量首次同时超越
在 345M 至 7B 参数规模上的系统验证显示：
| 指标 | 全注意力 | HiLS-Attention | 提升幅度 |
|---|---|---|---|
| 短文本 PPL | 基准线 | 几乎持平 | 无明显退化 |
| 外推能力 | 8K→32K 崩溃 | 8K→4M（512x） | 免训外推 |
| Prefill 加速（512K） | 1x | 13.5x | 13.5 倍 |
| Decode 加速（512K） | 1x | 15.7x | 15.7 倍 |
| 长上下文检索 | 基准线 | 反超 | 首次超越 |
最令人兴奋的是最后一条：HiLS-Attention 在部分长上下文检索任务上反超了全注意力本身。这意味着稀疏化不仅是"算力不够时的妥协"，稀疏注意力本身可能具备某种正则化效果，让模型在长程依赖建模中更加聚焦。
四、实践启示
- 可以零成本升级现有模型：HiLS-Attention 在短文本上几乎不损失 PPL，意味着已部署的模型可以平滑迁移，无需重新评估短文本基准。
- Agent 场景的直接受益者：多轮对话、长文档 RAG、代码库级理解等需要大量上下文的 Agent 场景，13.5 倍的 prefill 加速意味着从"等不起"变成"实时可用"。
- 开源生态的积极信号：论文（arXiv: 2607.02980）和代码（GitHub: Tencent-Hunyuan/HiLS-Attention）均已公开，社区可以快速集成到主流框架中。
小马
技术标签：#稀疏注意力 #长上下文建模 #HiLS-Attention #注意力机制 #开源