SWA：滑动窗口注意力与跨层信息传递#

TL;DR#

1. 每个位置都有自己的窗口#

Transformer用 Q 与 K 计算匹配权重，再对 V 加权求和。全因果注意力允许每个位置读取当前位置及之前的全部位置；序列越长，末尾位置读取的历史 K、V 越多。

Sliding Window Attention，简称 SWA，滑动窗口注意力，把每个位置的直接读取范围限制在最近的窗口内。本文讨论用于自回归生成的因果窗口：只读取当前位置及之前的位置。Mistral 7B 使用了这类局部注意力，并配合滚动缓存减少推理开销。[1]

以“小明／来到／公司／开始／工作”为例，假设每个词对应一个 token，位置从 1 到 5。设窗口大小为 $W$，本文的 $W$ 包含当前位置。取 $W=3$，每个位置最多读取“前两个位置＋当前位置”。

第一层分别计算所有位置的新向量：

窗口属于每个查询位置，不是整句话只保留一个窗口。 第一层计算“工作”时，不读取“小明”和“来到”的 K、V；但“小明”“来到”仍会计算自己的输出。序列开头没有足够的前置位置，就使用现有位置，窗口不会绕到句末。

窗口随查询位置移动，彼此重叠。它不会将这句话切成互不联系的固定块。

不同资料和接口对窗口大小的计数可能不同。若参数表示“向左读取的距离”，本例对应左侧距离 2。FlashAttention 的左右窗口参数按距离定义；Q、K 等长时，本文的窗口对应 window_size=(W-1, 0)。[2]

2. 第二层读取第一层的向量#

2.1 第一层仍输出五个位置的向量#

当这句话一起做前向计算时，第一层接收五个输入向量，分别计算后输出：

小明的新向量、来到的新向量、公司的新向量、开始的新向量、工作的新向量。

若每个向量有 $d$ 个数，那么第一层输入和输出的形状都是 $5\times d$。第二层接收第一层输出的全部五行，再为每个位置计算一遍。

每个位置在一层中读取的都是上一层向量。 第一层“工作”读取输入中的“公司”，第二层“工作”才读取第一层输出的“公司”；同层各位置可以并行计算。

2.2 小明通过公司影响第二层的工作#

第一层计算“公司”时，窗口内有“小明／来到／公司”。因此，第一层“公司”的新向量可能已经包含“小明”的信息。

第二层计算“工作”时，窗口仍然是“公司／开始／工作”这三个位置，但读取的是它们的第一层输出向量。其中“公司”的向量已经处理过“小明”的信息，于是形成一条间接依赖路径：

小明的输入向量 → 第一层公司的向量 → 第二层工作的向量。

“公司”的 token ID 没有变成“小明”，变化的是该位置的隐藏向量。图中每个框都是一个向量；同一列中的框属于同一次层输出，箭头表示跨层读取关系。

这是一条结构上允许的依赖路径。模型是否实际保留人物信息，取决于学习到的权重和向量变换；有路径不保证信息被完整保留。Mistral 论文借助堆叠层解释超出单层窗口的信息传播，但理论跨度不等于实际长距离检索能力。[1]

2.3 层数增加时，窗口保持不变#

“工作”在每层都只直接读取上一层的“公司／开始／工作”。随着这些向量经过更多次处理，“工作”可能受到的原始输入影响范围会扩大：

多层扩大的是间接影响范围，每层的窗口仍为 3。 到达序列开头后，范围就不再向前扩展。第三层继续处理第二层的向量，不会重新从原始 token 开始，也不会自动切换成全局 Attention。

一般地，假设序列长度为 $T$，共有 $L$ 层，各层都使用同样的因果窗口，没有全局注意力或其他跨位置操作。FFN 与归一化只在各位置内部计算，残差保留同一位置，所以不会单独扩大位置范围。

一层最多向左扩展 $W-1$ 个位置。对第 $L$ 层的位置 $i$，理论输入范围为：

$$\text{第 }L\text{ 层位置 }i\text{ 的理论输入范围} =[\,\max(1,i-L(W-1)),\;i\,].$$

其覆盖位置数为 $\text{理论覆盖位置数}=\min(i,1+L(W-1))$。若 $W>1$，输入位置 $j$ 到位置 $i$、且 $j\le i$，所需的最小层数为：

$$L_{\min}=\left\lceil\frac{i-j}{W-1}\right\rceil.$$

“小明”在位置 1，“工作”在位置 5，本例得到 $L_{\min}=\lceil(5-1)/(3-1)\rceil=2$。若 $W=1$，每层只能读取自身，无论堆叠多少层都不能从其他位置汇集信息。若 $W\ge T$，局部遮罩与这段序列的全因果遮罩相同。

部分资料把 $W$ 定义为向左的距离，其跨度写作 $LW$；不能将这个表达式与本文“窗口含当前位置”的定义直接混用。

3. 窗口内怎样计算 Attention#

3.1 用窗口遮罩限制匹配范围#

沿用 Transformer 的记号，$H^{(\ell-1)}\in\mathbb R^{T\times d}$ 是进入第 $\ell$ 层的整段隐藏向量。忽略具体归一化顺序，单个头从该层用于注意力投影的输入 $X$ 生成：

$$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V.$$

这里 $X$ 仍有 $T$ 行。SWA 改变的是每个查询允许匹配哪些 K，而不是只为整段序列生成 $W$ 个向量。

用 $\mathcal N_W(i)$ 表示位置 $i$ 允许读取的位置集合，索引从 1 开始：

$$\mathcal N_W(i)=\{j\mid\max(1,i-W+1)\le j\le i\}.$$

对应遮罩将窗口外及未来位置的匹配分数设为 $-\infty$，窗口内保留原分数。设单头的 Query、Key 维度为 $d_k$，定义为：

$$\begin{aligned} M_{ij}^{\mathrm{SWA}}&=\begin{cases} 0,&j\in\mathcal N_W(i),\\ -\infty,&\text{其他位置}, \end{cases}\\ \operatorname{Attention}_{\mathrm{SWA}}(Q,K,V) &=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}+M^{\mathrm{SWA}}\right)V. \end{aligned}$$

softmax 逐行计算，因此每个查询只在自己的窗口内归一化。这个矩阵公式说明数学结果；高效实现应跳过窗口外的计算，而不是先计算完整矩阵再遮住它。

3.2 算出“工作”位置的注意力输出#

下面是假设的一头标量算例，仅用于比较同一组 Q、K、V 在不同遮罩下的输出，不是训练实验。令 $d_k=d_v=1$，“工作”的查询为 $q_5=1$，各位置已投影的 K、V 如下。$\ln$ 是自然对数。

取 $W=3$，“工作”只读取位置 3、4、5。匹配分数由 $q_5k_j/\sqrt{d_k}$ 得到，指数值分别为 1、2、3。因此，窗口内的注意力权重为：

$$\text{工作对公司、开始、工作的权重} =\frac{[\exp(0),\exp(\ln2),\exp(\ln3)]}{1+2+3} =\left[\frac16,\frac26,\frac36\right].$$

用 $a_5^{\mathrm{SWA}}$ 表示这个头在“工作”位置的输出：

$$a_5^{\mathrm{SWA}}=\frac16\times2+\frac26\times4+\frac36\times6 =\frac{14}{3}\approx4.6667.$$

若改用全因果注意力，前 3 个位置的指数值均为 1，位置 4、5 分别为 2、3，归一化分母为 8。同一查询的输出变为：

$$a_5^{\mathrm{full}} =\frac{1\times10+1\times2+2\times4+3\times6}{8} =\frac{38}{8}=4.75.$$

SWA 不只是删掉窗口外的 V，还会在剩余位置之间重新归一化权重。 “小明”的 V 在这次局部加权中不再直接参与；这不等于“工作”在所有层都无法受到“小明”的影响。

$a_5$ 只是注意力头的输出，不是完整层输出或词表概率。多头拼接、输出投影、残差、归一化与 FFN 仍按所选 Transformer 架构执行。SWA 本身不规定必须采用 Post-LN 或 Pre-LN。

4. 计算量与 KV Cache 怎样减少#

4.1 真正跳过窗口外计算，才减少注意力开销#

长度为 $T$ 的序列，全因果注意力共有 $T(T+1)/2$ 个允许读取的位置对。固定窗口为 $W$ 时，第 $i$ 个查询只读取 $\min(i,W)$ 个位置。对单层单头，位置对数量为：

$$\begin{aligned} N_{\mathrm{full}}&=\frac{T(T+1)}2,\\ N_{\mathrm{SWA}}&=\sum_{i=1}^{T}\min(i,W) =TW-\frac{W(W-1)}2\qquad(T\ge W). \end{aligned}$$

本例 $T=5,W=3$，得到 $N_{\mathrm{full}}=15$、$N_{\mathrm{SWA}}=12$。这比较的是同样长度、同样因果方向下的允许位置对，不是拿双向全注意力与因果窗口比较。

当序列远长于固定窗口时，注意力的 QK 匹配与 V 汇集计算由 $O(T^2d_k)$ 降为 $O(TWd_k)$，这里为简洁取 $d_v=d_k$。先算完整 $QK^\top$ 再加局部遮罩，仍会执行稠密匹配，不能获得这项计算量下降。 需要支持局部窗口的实现，例如 FlashAttention 的局部注意力接口。[2]

Q、K、V 投影和 FFN 仍需处理序列各位置；局部窗口也不会使模型层数或参数量自动减少。上述复杂度不等于完整模型的端到端加速比。

4.2 每层缓存最近窗口所需的 K、V#

训练或整段 prefill 可以一起计算各位置；逐 token 推理只计算新位置的隐藏向量，并读取 KV Cache 中已有的历史 K、V。缓存按层保存：第 $\ell$ 层的历史 K、V 来自该层当时的输入投影，并非所有层共用一份 embedding。

对纯因果 SWA，位置 $i$ 只需要最近的窗口，因此每层可以使用最多 $W$ 行 K、V 的滚动缓存。以 $W=3$ 为例，缓存槽编号为 0、1、2，将绝对位置 $i$ 写入 $\text{缓存槽}(i)=(i-1)\bmod3$：

处理位置 4 时，位置 1 已不在它及未来查询的直接窗口中，可以被覆盖。物理槽顺序不一定是时间顺序，读取时仍须按真实位置正确排列或计算索引。Mistral 的参考实现分别维护缓存槽与绝对位置。使用 RoPE 时，位置编码也不能简单改成缓存槽编号。[3]

删除旧 K、V 不等于删除所有旧信息。 旧位置的信息可能已经进入较新位置的隐藏向量，进而影响后续层的 K、V。但固定层数下这种依赖仍受理论范围限制，不会因为持续生成就自动获得无限记忆。

4.3 缓存容量与输入截断是不同操作#

设层数为 $L$，每层 KV 头数为 $n_{\mathrm{kv}}$，每头维度为 $d_h$，每个缓存元素占 $b$ 字节。对一条序列、不计额外索引与分配开销，保留窗口的缓存容量为：

$$\text{SWA KV Cache 字节数} =2L\min(T,W)n_{\mathrm{kv}}d_hb.$$

系数 2 分别对应 K 和 V。全历史缓存将其中的 $\min(T,W)$ 换为 $T$；这里固定 KV 头数，单独比较窗口作用，不把 GQA 减少 KV 头数的收益混入其中。

例如 $L=32,n_{\mathrm{kv}}=8,d_h=128,b=2$，历史长度 $T=32{,}768$、窗口 $W=4{,}096$ 时，按上述公式计算：

这里采用二进制单位，$1\ \mathrm{MiB}=2^{20}$ 字节、$1\ \mathrm{GiB}=2^{30}$ 字节，缓存容量之比为 $T/W=8$。这是指定配置的容量计算，不是实测显存占用，也不包括模型权重、激活和临时工作区。

滚动缓存按顺序处理完整前缀，只淘汰以后不再直接读取的 K、V。将原始输入直接裁成最后 $W$ 个 token，再从头计算，并不与滚动缓存等价：重新计算会丢失这些位置先前已汇集的更早信息。

长 prompt 可以分块 prefill，但不能在计算块内所有查询之前，只留下整个块结束时的最后 $W$ 行。块首查询与块尾查询需要不同的窗口；实现必须同时处理历史缓存和当前块，并施加正确的局部因果遮罩。[1] [3]

5. 局限性#

远距离直接读取被局部传播替代。 全因果注意力在一层中即可让“工作”读取“小明”的 K、V；$W=3$ 的纯窗口结构至少需要 2 层才能建立本例的依赖路径。中间变换可能没有保留需要的信息，理论覆盖范围不能替代检索或推理任务的实际评测。

固定深度的纯窗口结构存在有限输入范围。 超过 $L(W-1)$ 的位置距离，没有依靠这些局部层传播到当前输出的路径。若模型混用全局层、其他稀疏连接或额外记忆，其依赖范围与缓存规则需要重新分析。

窗口大小改变模型实际执行的计算。 将一个按全注意力训练的模型在推理时直接改成 SWA，可能改变输出，不能仅凭缓存节省认定质量不变。窗口应作为架构和训练设置的一部分验证；Mistral 整体模型的成绩也不能单独证明 SWA 在其他模型上具有相同收益。

参考资料#

1. Jiang 等，2023，Mistral 7B，§2、图 1–3：窗口注意力、跨层传播、滚动缓存与分块 prefill。本文统一采用“$W$ 包含当前位置”的精确计数，手算句子与数值为教学构造。
2. FlashAttention 官方 README：How to use FlashAttention，flash_attn_qkvpacked_func 与 flash_attn_func 的 window_size 说明。左右窗口按距离计数；Q、K 长度不同时还需按接口处理位置偏移。
3. Mistral 官方参考实现：cache.py，BufferCache、CacheView、unrotate：缓存容量、绝对位置、滚动槽与分块遮罩。此处引用具体机制，不将代码默认配置视为所有 Mistral 模型的共同设置。

延伸阅读#

• 稀疏注意力：选择与压缩机制比较：不同稀疏方案如何选择参与注意力计算的信息？