技术博客

滑动窗口注意力与全局 token

第 56 / 147 篇 · 2025-10-14 · ai · 注意力机制长文本稀疏注意力

线性注意力和 FlashAttention 分别从"换核"和"省访存"两条路处理 $n^2$。滑动窗口注意力(Sliding Window Attention, SWA)走第三条:承认大多数注意力权重本来就近似为零,稀疏化是被数据原谅的。Mistral 7B 把它作为每一层的默认形态,Longformer 则给出了"窗口 + 全局"的完整方案。

局部性假设:为什么窗口是合理的

经验观察是语言处理强局部:句法依存绝大多数落在相邻几十个 token 内,指代关系的两端离得远,但数量少。注意力矩阵 $\mathrm{softmax}(QK^\top/\sqrt d)$ 中,行 i 的有效质量高度集中在 $j \in [i-w, i]$ 的带状区域——这在可视化里一眼可见。窗口注意力干脆把带外全部 mask 掉:位置 i 只 attend 到自己及之前的 $w$ 个位置(配合 causal mask 时)。

复杂度立刻下来:每行只有 $w$ 个非掩码位置,attention 的 pair 数从 $n^2/2$ 降到 $nw$,FLOPs 与注意力矩阵的物化显存都是 $O(nw)$。窗口不随序列增长,所以长序列下单层成本线性——这就是 Mistral 敢把上下文扩到 32k 而每层 attention 开销几乎不涨的原因(它的实现是 rolling buffer cache, $w=4096$)。

代价要诚实列出。最远的可建模距离是 $w$ 层传播一步:信息要在位置 i 和 i+10000 之间流动,得逐层接力,每层最多前传 $w$ 个 token。$L$ 层堆叠后理论感受野是 $L\times w$——看起来够大,但这是"每层全用满"的上界,实际每层窗口的有效利用远不到 100%,长程依赖在接力中被稀释。窗口注意力学不到"一步到位"的长程检索,这正是它单独用不行、需要全局 token 补位的原因。

全局 token:稀疏里的稠密岛

Longformer(Beltagy et al. 2020)的结构分两类:

  • 滑动窗口(banded):每个位置对左右各 $w/2$ 的窗口算注意力,矩阵是带状。还可加 dilation(窗口内隔位采样)扩大感受野而不增 FLOPs,代价是带内出现空洞。
  • 全局注意力(global):少数被选中的位置(如 <s>/CLS、问题 token、句点分隔符)对所有位置算注意力,且所有位置都对它们算注意力。这两组交互把 $O(nw)$ 的带状世界里凿出 $O(n)$ 的全连接行和列。

全局行的存在改变了信息通路:位置 i 要引用远处的位置 j,不再需要 $w$ 步接力,只需 i → 某个全局 g → j 两步(i attend 到 g,g 的表示里凝聚了全序列信息)。用图论的话说,带状图加少量枢纽节点后直径骤降——这也是后面 BigBird 用随机边保证连通性的同族思想。

账本上,$g$ 个全局 token 的注意力成本是 $2gn$(行加列),与 $nw$ 相加仍是 $O(n)$。实现上全局行不能完全走 SWA 的块状 kernel(带宽不规则),Longformer 的实现把它们拆成两条计算路径再合并——工程上这预示了一个后来反复出现的模式:稀疏注意力的收益在数学上是干净的,在 kernel 上是有杂费的。

窗口注意力:注意力矩阵呈带状 带外全部 mask 复杂度 O(n·w) 加全局 token:两跳可达任意位置 全局 g 所有位置 ⇄ g:O(n) 感受野:L 层接力最多 L×w;经全局 token 两跳覆盖全序列

工程落地:窗口的边界与缓存形态

窗口注意力对推理系统的影响比对训练更大,因为 KV cache 的形状被改变了。全量 KV cache 随序列无限增长,而 SWA 只需保存最近 $w$ 个位置的 K/V——Mistral 的 rolling buffer 就是固定大小的环形缓冲,新 token 覆盖最旧的槽位,显存占用恒定。对比后面要讲的 KV cache 淘汰策略(H2O 等用启发式决定丢谁),SWA 的丢弃规则是位置决定论:谁老丢谁,零开销、零策略,但放弃了模型自己"决定记住谁"的可能。

边界处理有几个坑:窗口跨越文档边界时是否允许 attend 到上一篇的尾部(取决于 packing 策略);RoPE 下窗口内相对位置天然一致,但和全局 token 交互的位置编码没有窗口外的训练先例,Longformer 对全局行用的是单独的一套位置编码方案;混布架构里把窗口层和全局层交替排布(后来 GPT-NeoX、一些混合架构的用法)时,每层窗口大小的配置空间很大,搜起来不便宜。

小结

滑动窗口注意力的合理性建立在语言局部性这个经验事实上,$O(nw)$ 的账干净漂亮。它的短板(感受野靠接力、无显式全局通路)由全局 token 以 $O(n)$ 的杂费补齐,形成 Longformer 式"带状 + 枢纽"的稀疏结构。从系统视角,SWA 还免费附赠了恒定大小的 KV cache——这是它比同复杂度其他稀疏方案在推理侧更受欢迎的直接原因。下一篇把稀疏模式推广到一般形式:当稀疏结构本身需要表达能力保证时,该设计什么样的图。