第 11 章 稀疏注意力¶
11.1 有效的注意力¶
如果把训练收敛后的注意力权重画成图，常见形状是：绝大多数位置接近零，真正亮起来的只有少数几类
- 最近的一段（局部语法、指代、代码缩进）；
- 序列开头附近（系统提示、角色设定、Attention Sink，见 11.4）；
- 少量语义相关的远处位置（攻略里那句「周一闭馆」、代码里的函数定义）。
若每个 query 只访问约 \(k\) 个 key（或等价地，只访问总计约 \(k\) 个 token 的块），注意力计算从
当 \(k\ll n\)，可以看到，稀疏化会大大减少计算量
11.2 滑动窗口¶
位置 \(i\) 只和最近 \(w\) 个 token 做注意力（区间 \([i-w+1,\ i]\)）。Prefill 从 \(O(n^2)\) 降到 \(O(nw)\)；推理若丢掉窗外 KV（rolling），单层缓存从 \(O(n)\) 降到 \(O(w)\)。
当上下文几乎都落在最近 \(w\) 里——短多轮对话、总长不超过窗口的文档，滑动窗口和全注意力一致，节省算力和缓存
当文本总长度超过滑动窗口大小，较早的 token 虽然参与了特征计算，但当前位置的 query 无法直接读取窗口外的 K、V。远距离信息不能通过注意力直接引用，仅能依靠多层残差微弱传递，精确事实类细节容易丢失。 工程上很少单独依靠纯滑动窗口处理超长上下文。常用的弥补手段：保留最开头少量 token 的 KV 作为 Attention Sink，固定系统提示不随窗口滑动；采用混合架构，部分层使用滑动窗口，少数层启用全局或检索式注意力捕获远处信息；训练阶段就内置滑动窗口与 sink 机制。
12.3 全局 token¶
全局 token 是 Longformer 提出的思路：在序列中插入少量特殊 token，这类 token 可以访问全部位置，同时所有 token 也都能读取全局 token，充当整个序列的信息枢纽。该设计天然适合自带标题、[CLS]这类锚点的文档任务。但对于自回归对话模型，需要额外规划全局 token 的放置策略，实现成本更高，因此不如滑动窗口普及。
12.4 Attention Sink¶
Attention Sink（StreamingLLM）面向流式长文本推理。如果采用滑动窗口缓存，直接丢弃窗口之外最开头 token 的 KV 缓存，生成质量会明显下滑。工程上的解决办法是永久保留开头少量 token 对应的 KV，再搭配滑动窗口使用。直观解释是 softmax 归一化要求权重总和为 1，模型倾向将无明确匹配目标的注意力概率分配给起始 token。完整机制同时和训练分布、位置编码相关。实践中，保留少量 sink token 的流式生成，性能明显优于单纯保留最近窗口内的 KV。
注意：这属于推理阶段的 KV 缓存策略，不需要重新预训练模型。
12.5 按块稀疏¶
理想稀疏注意力允许每个 query 独立挑选语义最相关的 key。该方案理论简洁，但在 GPU 上很难充分发挥硬件性能：随机离散的内存访问无法适配矩阵计算单元对连续内存的偏好，理论上减少了浮点运算量，但实际运行延迟改善有限。
因此工业界主流的稀疏方案，普遍采用按块筛选，把长度为 \(n\) 的序列按固定步长 \(B\) 切成连续段：
常见 \(B=64\) 或 \(128\)（有的实现更小，如第 16 章 QSA 的微块约 4 token）。筛选与访存都以整块为单位：选中一块，就读这一段连续的 K、V；没选中，整块跳过。
玩具例子：\(n=16\)、\(B=4\)，行程助手上下文切成 4 块——
| 块 | token 下标 | 内容（示意） |
|---|---|---|
| 0 | 0–3 | 系统提示开头 |
| 1 | 4–7 | 「周一闭馆」那句攻略 |
| 2 | 8–11 | 用户改行程 |
| 3 | 12–15 | 当前要续写的位置 |
选中块 1，就读下标 4–7 整段连续 K、V，再在块内做普通注意力。
import torch
n, d_h, B = 16, 8, 4
n_block = n // B
q, k, v = (torch.randn(n, d_h) for _ in range(3))
k_blocks = k.view(n_block, B, d_h) # [4, 4, d_h]
v_blocks = v.view(n_block, B, d_h)
selected = [0, 1] # 手写：sink + 块 1
k_sel = k_blocks[selected].reshape(-1, d_h) # [2B, d_h]
v_sel = v_blocks[selected].reshape(-1, d_h)
score = (q[-1:] @ k_sel.T) * (d_h ** -0.5) # 当前 query × 选中块
out = score.softmax(dim=-1) @ v_sel
全注意力是 q @ k.T 得到 [n, n]；这里每条 query 只对 \(|S|\) 个块、共 \(|S|\,B\) 个 key 算，量级 \(O(n\cdot |S|\cdot B\cdot d_h)\)。
12.6 可学习检索式稀疏¶
滑动窗口的底层假设，是有效信息基本都集中在近期上下文。一旦关键事实藏在很远的片段，窗口外的 token 就无法直接读取。可学习检索式稀疏注意力换了思路：让模型根据内容自行挑选需要读取的上下文块；挑选动作保持块粒度，方便 GPU 并行。
我们看下NSA 这类三分支架构是典型代表
压缩
把序列切块后，k.view(n_block, B, d_h)，再对块内做聚合，压缩成 1 个或少数几个代表 key。最简单实现是取均值；实际代码多用可学习池化或者小型 MLP，把一块里 B 个 key 压缩成\(\tilde k_b\in\mathbb R^{d'}\)，\(d'\)一般小于原始维度。整条长序列就被压缩成\(\tilde K\in\mathbb R^{n_{\mathrm{block}}\times d'}\)，长度是块数\(n/B\)，不再是原始 token 长度n。
打分与选择
粗筛阶段只对块代表打分，不会遍历每一个 token。拿当前 query 和所有\(\tilde K\)做点积，得到长度等于块总数的分数，选出 top-k 块下标，对应前面代码里的selected：
# 接上一节：k_blocks [n_block, B, d_h]
k_rep = k_blocks.mean(dim=1) # [n_block, d_h]，均值压缩仅作示意
block_score = (q[-1] @ k_rep.T) * (d_h ** -0.5) # 逐个块打分
k_sel_idx = block_score.topk(2).indices # k=2，选出块索引
# 取出选中块对应的原始token k_blocks[k_sel_idx]，再做精细token级注意力
粗筛代价按块计算：\(O(n_{\mathrm{block}}\,d')=O((n/B)\,d')\)，远低于遍历全部 token 的\(O(nd_h)\)。高维、开销大的精细注意力，只在被选中的块内部执行，token 总量大约\(k\cdot B\)。 NSA 是从压缩支路的注意力结果导出块分数；第五部分要介绍的QSA 单独训练一个轻量 indexer。两者底层逻辑一致：对全部块粗打分，选出 top-k。
滑动窗口分支 最近一段上下文永远走精细注意力，不走检索流程。局部语法、就近指代这类基础依赖，不用靠检索能否命中远端块来兜底。三路分支最后通过门控融合输出。
块选择模块必须和主干模型联合训练；如果直接外挂固定掩码，会和预训练的数据分布错位。漏检问题不可逆：只要存放答案的块没能进入 top-k，后续块内注意力再精细，也读不到目标信息。大海捞针这类长距离事实任务对此尤其敏感，平均评测指标往往掩盖这种单点失败。
对比纯滑动窗口：它可以读到远端信息，但很多实现依旧保留完整历史 KV，显存开销不一定下降。对比全局全注意力：精细计算只限制在选中块 + 局部窗口，避开 Prefill 阶段完整\(O(n^2)\)复杂度。最终效果好坏，核心取决于选块精度。选得准，效果接近全局注意力；一旦频繁漏检，表现甚至不如滑动窗口 —— 滑窗至少保证最近一段上下文一定可用。
核心问答¶
Q1：稀疏注意力和第 7 章的 GQA 是一回事吗？ 不是。GQA 减的是 KV 的份数（缓存）；稀疏减的是参与计算的 token / 块数（计算）。正交，可同时用。
Q2：Attention Sink 为什么重要？ 经验上，训练后的模型常把开头位置当成稳定的注意力落点；流式推理若突然拿掉，等于改变模型见惯的分布。softmax 泄压是有用直觉，不是唯一解释。做滑窗流式时，保留少量开头 KV，通常比只保留最近窗口稳。
Q3：为什么稀疏要按块，不能每个 token 自由选？ 自由选访存太随机，矩阵乘单元吃不饱，纸面 FLOPs 换不成实际延迟下降。按块牺牲一点选择粒度（整块进、整块出），换连续访存和真实加速。
Q4：训好的稠密模型能不能直接改成稀疏推理？ 可以尝试，但会掉点，幅度看任务。滑窗加 sink 是少数常被用作免训练近似的方案；检索式稀疏基本需要训练配合。
Q5：稀疏注意力省显存吗？ 看哪一种。滑窗配 rolling cache 会省；普通检索式稀疏通常不省全量历史 KV；跨层复用等还要另算。读论文时把「少读 KV」和「少存 KV」分开核对。
Q6：稀疏能解决「中间找不到」吗？ 只部分相关。它可能减少干扰，也可能因漏块制造新的盲区。质量仍要靠数据、位置编码，以及第 14 章里是否保留足够的精确回看层。