细粒度稀疏注意力：大模型长上下文架构演进

一、背景：长上下文把注意力机制推到了重构临界点

标准 Transformer 的自注意力在训练阶段的时间复杂度是序列长度的平方级；推理解码阶段虽然借助 KV Cache 避免了重复计算，但每生成一个 token，注意力层仍需与全部历史 Key 做点积并加权求和，单步开销随上下文线性增长，KV Cache 的显存占用也随序列长度线性膨胀。当上下文窗口从数千 token 走向数十万乃至更长，全量注意力的计算与访存成本会迅速成为吞吐与成本的主要瓶颈。

与此同时，业界大量分析反复验证了一个经验性结论：对于任意一个 query 位置，真正对其贡献显著注意力的历史位置只占很小的比例。换言之，全量注意力中存在大量“算了但不重要”的冗余计算。既然冗余客观存在，“按需召回”就成为非常自然的架构演进方向。

二、两条主流路线：压缩历史与筛选历史

2.1 线性注意力：把历史压进固定大小的状态

线性注意力的思路是把已经看过的历史压缩进一个固定大小的递归状态中，推理时每处理一个 token 只需对状态做一次常数级更新，单步开销与序列长度无关，显存占用恒定。这类方法近年来通过门控增量更新等改进，状态管理能力已相当成熟，适合维护“长程语义趋势”与“全局主题状态”。

但它本质上是有损压缩：当任务需要精确召回某个具体细节——例如文档中某句话的原文、代码里某个符号的定义位置——固定大小的状态很难无损保留这些信息。这是线性注意力的能力边界，也是它必须与其他机制配合使用的原因。

2.2 细粒度稀疏注意力：先筛选、再精算

另一条路线保留 softmax 注意力的表达能力，但只在被选中的历史子集上计算，其典型实现就是近年受到广泛关注的「Indexer + Top-K」细粒度稀疏注意力：

1. Indexer 打分：一个轻量检索模块为每个历史位置计算与当前 query 的相关性分数；
2. Top-K 筛选：只保留分数最高的 K 个位置（常见取值为数千量级，远小于上下文总长度），通常还会强制保留一个局部滑动窗口，保证近邻上下文永不丢失；
3. 精确计算：在筛选出的子集上执行完整的 softmax 注意力。

“细粒度”的含义在于选择单位是 token 级，而非整块、整页的粗粒度跳过：被选中的位置获得与全量注意力一致的精确计算，被跳过的位置则是明确判断为“不相关”后的主动忽略。对解码阶段而言，注意力层每步只需读取 K 条 KV 而非全部历史，计算量与访存量同时下降，这是长上下文推理降本最直接的收益来源。

2.3 门控设计：让网络自己决定“用多少”

在稀疏注意力层上叠加门控，即在输出注入残差流之前引入可学习的门控参数，控制该层信息的汇入强度。这带来两点工程价值：一是训练更稳定，稀疏化初期的检索噪声不易破坏已有表征；二是层间分工变成可学习目标，便于从旧架构平滑迁移。

三、混合分层架构：务实的折中之道

单靠任何一种注意力机制都很难同时满足“便宜”与“精确”。近年长上下文模型的普遍做法是分层混合：

• 一部分层使用线性注意力：以常数成本维护全局状态，负责“广而粗”的长程依赖；
• 其余层使用门控细粒度稀疏注意力：负责关键事实、指代关系、代码符号等“准而细”的精确召回。

这类架构演进还有一个值得注意的工程细节：升级通常在总参数量与每 token 激活参数量保持不变的前提下完成——线性注意力层原样保留，仅把原本承担精确注意力的层切换为门控稀疏注意力。这样做的收益是：结构兼容部分的预训练投入可以复用，部署侧的参数规划与显存预算基本不动，升级风险可控。

为什么不在所有层上都使用稀疏注意力？原因有三：其一，全文摘要、长文情感判断这类全局聚合任务需要均匀覆盖整个上下文，过度稀疏反而伤害效果；其二，线性注意力层与稀疏注意力层的归纳偏置互补，混合结构比单一机制更鲁棒；其三，全量依赖 Indexer 意味着检索质量成为单点风险，保留部分低成本稠密建模能力是必要的兜底。可以说，”保留线性注意力 + 其余层换用稀疏注意力”是一种兼顾成本、能力与迁移成本的折中方案。

四、机制细节与示意实现

工程上落地「Indexer + Top-K」有几个关键设计点：

• Indexer 保持轻量：通常采用低维投影加打分头的结构，可与主注意力共享部分表示，避免检索模块本身成为新的开销大头；
• 候选集合的构成：局部滑窗必选 + 全局 Top-K，兼顾近邻语境与远程关键信息；
• 训练策略：业界常见做法是先以稠密注意力预热，再引入 Indexer 联合训练，并让检索分数与主注意力分布保持一致的监督信号，避免出现“检索分布漂移”导致的能力回退；
• 推理系统优化：Top-K 之后每层的计算规模固定、访存模式规整，非常利于算子融合与 KV Cache 分层卸载（例如把低频访问的历史 KV 放到更大容量的存储层级）。

下面的伪代码描述了核心流程：

# 门控细粒度稀疏注意力（Indexer + Top-K）示意
def gated_sparse_attention(x, kv_cache, top_k=2048, window=512):
    q = W_q(x)                                  # 当前 query
    # 1) Indexer 为全部历史位置打分
    scores = indexer.score(q, kv_cache.keys)    # [seq_len]
    # 2) 候选集合 = 局部窗口 + 全局 Top-K
    selected = union(last_window(window),
                     topk(scores, top_k))
    # 3) 仅在选中子集上执行标准注意力
    attn  = softmax(q @ kv_cache.keys[selected].T / sqrt(d))
    ctx   = attn @ kv_cache.values[selected]
    # 4) 门控控制该层输出注入残差流的强度
    return x + gate(ctx)

三类机制的复杂度对比如下（K 为每层选取的历史位置数，n 为上下文长度）：

维度	全量注意力	线性注意力	细粒度稀疏注意力
训练复杂度	O(n²)	O(n)	O(n·K)
解码单步开销	O(n)	O(1)	O(K)
历史信息利用	全量精确	有损压缩	选中位置精确
精确召回能力	最强	较弱	接近全量
KV Cache	全量常驻	常数状态	全量存放、按需读取

五、落地评估清单与工程建议

如果团队正在评估或自研类似架构，建议从三个维度验收：

1. 能力回归：覆盖长文档问答、多跳检索、长文本关键信息定位、代码仓库级理解等典型长上下文任务，重点观察稀疏化后精确召回是否回退；
2. 效率指标：对比首 token 延迟、增量解码吞吐、长序列显存峰值与 P99 延迟，确认理论复杂度优势转化为实际收益；
3. 稳定性观察：持续跟踪 Indexer 命中率、门控权重分布，并在全局聚合类任务上设置退化保护。

工程化层面，通用的云上能力即可支撑完整闭环：以弹性 GPU 算力完成稀疏化继续训练与评测，用对象存储统一管理长文本数据集与 checkpoint，通过容器平台编排评测流水线，再借助日志服务与监控告警持续观察长序列推理延迟与检索命中率。这类基础设施能力属于通用类目，不构成额外的选型障碍。

六、选型建议与小结

• 适合采用混合架构的场景：超长上下文 Agent、仓库级代码理解、海量文档分析等对长程精确召回与推理成本同时敏感的业务；
• 可以暂缓的场景：上下文长度有限、算力预算充足时，全量注意力方案仍是稳妥选择；
• 关键认知：稀疏化不是“算不起之后的妥协”，而是把算力集中花在关键位置上的架构表达。Indexer 的质量决定召回上限，门控与混合分层决定系统鲁棒性，三者共同构成细粒度稀疏注意力落地的完整拼图。

从“保留线性注意力层、替换其余注意力层”这一演进路径可以看出：长上下文架构设计的主线正在从“单一机制做到极致”转向“多种机制按层分工”。对开发者而言，理解 Indexer + Top-K 的筛选逻辑与混合架构的折中思路，比记住任何具体模型的名字都更有长期价值。