把稀疏注意力做对：HiLS-Attention 如何在数学上同时打破效率与效果的天花板
把稀疏注意力做对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-AttentionHierarchical Landmark Sparse Attention的核心贡献在于重新设计了三件事2.1 地标 token 机制不再对 chunk 内所有 token 求均值而是在每个 chunk 内选取少量地标 tokenLandmark 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而非在推理时临时抱佛脚。用伪代码概括核心逻辑defhils_attention(Q,K,V,num_chunks,top_k):# Step 1: 将序列切分为 chunk每个 chunk 生成地标 tokenchunkssplit_into_chunks(K,V,num_chunks)landmarks[chunk.landmark_projection()forchunkinchunks]# Step 2: 用地标 token 计算重要性分数landmark_scoressoftmax(Q landmarks.T/sqrt(d_k))# Step 3: Gumbel-Softmax 近似 Top-K 选择可导selectedgumbel_topk(landmark_scores,ktop_k)# Step 4: 仅在被选中的 chunk 上执行精确注意力K_selectedconcat([chunks[i].Kforiinselected])V_selectedconcat([chunks[i].Vforiinselected])returnsoftmax(Q K_selected.T/sqrt(d_k)) V_selected三、效果效率与质量首次同时超越在 345M 至 7B 参数规模上的系统验证显示指标全注意力HiLS-Attention提升幅度短文本 PPL基准线几乎持平无明显退化外推能力8K→32K 崩溃8K→4M512x免训外推Prefill 加速512K1x13.5x13.5 倍Decode 加速512K1x15.7x15.7 倍长上下文检索基准线反超首次超越最令人兴奋的是最后一条HiLS-Attention 在部分长上下文检索任务上反超了全注意力本身。这意味着稀疏化不仅是算力不够时的妥协稀疏注意力本身可能具备某种正则化效果让模型在长程依赖建模中更加聚焦。四、实践启示可以零成本升级现有模型HiLS-Attention 在短文本上几乎不损失 PPL意味着已部署的模型可以平滑迁移无需重新评估短文本基准。Agent 场景的直接受益者多轮对话、长文档 RAG、代码库级理解等需要大量上下文的 Agent 场景13.5 倍的 prefill 加速意味着从等不起变成实时可用。开源生态的积极信号论文arXiv: 2607.02980和代码GitHub: Tencent-Hunyuan/HiLS-Attention均已公开社区可以快速集成到主流框架中。小马技术标签#稀疏注意力 #长上下文建模 #HiLS-Attention #注意力机制 #开源
)
