Sliding Window Attention

滑动窗口注意力（SWA）将每个token可关注的邻近词元数量限制为固定值；在因果语言模型中，范围为最近的 w 个词元，而不是像完整的自注意力那样覆盖整个序列。注意力计算量随后会随序列长度线性增长，而不是呈二次增长；生成模型只需在其KV 缓存中保留最近 w 个词元的键和值。Longformer 于 2020 年将滑动窗口作为其注意力机制的核心，而 Mistral AI 的 Mistral 7B 于 2023 年将其带入一款广泛使用的开源大型语言模型（LLM），窗口大小为 4,096 个词元。

窗口如何随深度扩大#

一层 SWA 只能看到 w 个 token，但Transformer会堆叠多层。在第 2 层，一个 token 读取的隐藏状态已分别概括了第 1 层中的 w 个 token，因此信息每经过一层大约就能再传播 w − 1 个位置，就像堆叠卷积会扩大感受野一样。这与卷积神经网络（CNN）的原理相同。在 PyTorch中，窗口由布尔注意力掩码表示；将掩码与自身相乘，就能追踪信息的传播范围：

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

对于 16 个 token，该窗口每层计算 58 个分数，而不是 136 个；最后一个 token 的覆盖范围每层增加 3 个 token，从 4 个扩展到 7 个、10 个和 13 个。同一个掩码可以作为 attn_mask 传给 PyTorch 的 scaled_dot_product_attention。对于 Mistral 7B，32 层配合 4,096 个 token 的窗口，理论注意力范围约为 131K 个 token，不过来自很远位置的信息必须经过许多跳才能传递到达。

Mistral 7B 还会将键和值存放在由 w 个槽位构成的循环缓冲区中，把位置 i 写入槽位 i mod w，并覆盖最旧的条目。在包含 32K 个 token 的序列上，这种做法将缓存内存减少到原来的八分之一，且不影响质量；固定的缓存大小也让无论生成运行多久，内存占用都保持不变。

交错使用局部层与全局层#

纯窗口无法一步复制很久之前的信息，因此近期模型会混合使用 SWA 层和完整的全局层。Google 的 Gemma 2 以 1:1 的比例交替使用这两种层，窗口大小为 4,096 个 token；Gemma 3则每五个局部层配一个全局层，窗口仅为 1,024 个 token，因此即使上下文窗口长达 128K 个 token，KV 缓存仍然很小。只有全局层会为整个上下文存储键和值，因此六层中有五层每层只保存 1,024 个 token。

视觉 Transformer也会对高分辨率图像采用同样的混合方式，此时每个图像块都是一个 token。Meta 的 SAM 3在 Ultralytics 包中用于图像分割：其图像编码器的 32 个块中，有 28 个块在互不重叠的 24 × 24 图像块窗口内进行注意力计算；第 7、15、23 和 31 块则在整张 1,008 像素图像的 72 × 72 图像块网格上进行全局注意力计算。

SWA 与相关术语#

• Swin Transformer： 将图像划分为固定且互不重叠的窗口，窗口内的所有 token 共享该窗口；并在层与层之间移动网格，让相邻窗口交换信息。
• 稀疏注意力： 一类更广泛的固定模式。SWA 是其中最简单的一种，其他方法还会加入扩张窗口、全局 token 或随机连接。
• 注意力汇聚点： 序列开头的 token，会吸收多余的注意力权重。对使用完整注意力训练的模型进行推理时，如果应用窗口并移除这些 token，生成过程就会出错；StreamingLLM 会将这些 token 与窗口一同保留在缓存中。
• 可变形注意力： 每个查询会关注少量采样点，并学习这些点的偏移量，因此邻域取决于内容。
• Flash attention： 一种用于精确注意力计算的 GPU 内核。其实现支持窗口大小，因此 SWA 模型会完全跳过窗口之外的块。