AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

Attention 变体 —— 滑动窗口、稀疏、差分 ​

完整注意力是个圈。每个 token 都看每个 token，显存为此买单。四种变体掰弯这个圈的形状，把一半的成本赚回来。

类型： Build 语言： Python 前置要求： 阶段 7 · 02（Self-Attention）、阶段 7 · 03（Multi-Head）、阶段 7 · 12（KV Cache / Flash Attention） 预计时间： ~60 分钟

问题背景 ​

完整注意力在序列长度上花 O(N²) 显存和 O(N²) 算力。对一个 128K 上下文的 Llama 3 70B，那是每层 160 亿个注意力条目，乘以 80 层。Flash Attention（第 12 课）藏掉了 O(N²) 的激活显存，但没改算术成本——每个 token 仍然关注其他每个 token。

三类变体改变注意力矩阵本身的拓扑：

  1. 滑动窗口注意力（SWA）。 每个 token 关注一个固定的邻居窗口，而非整个前缀。显存和算力降到 O(N · W)，其中 W 是窗口。Gemma 2/3、Mistral 7B 的前几层、Phi-3-Long。
  2. 稀疏 / 块注意力。 只有选中的对 (i, j) 被打分；其余被强制为零权重。Longformer、BigBird、OpenAI 稀疏 transformer。
  3. 差分注意力。 用各自的 Q/K 投影算两张注意力图，一张减另一张。干掉那个把权重渗进前几个 token 的"注意力汇（attention sink）"。微软的 DIFF Transformer（2024）。

它们共存。一个 2026 年的前沿模型常把它们混着用：大多数层是 SWA-1024，每第五层是全局完整注意力，还有少数差分头来清理检索。Gemma 3 的 5:1 SWA-对-全局比例是当前的教科书默认。

核心概念 ​

滑动窗口注意力（SWA） ​

位置 i 的每个 query 只关注 [i - W, i]（因果 SWA）或 [i - W/2, i + W/2]（双向）里的位置。窗口外的 token 在分数矩阵里得 -inf。

完整因果：              滑动窗口 (W=4)：
位置 0-7                位置 0-7, W=4
    0 1 2 3 4 5 6 7        0 1 2 3 4 5 6 7
0 | x                0 |  x
1 | x x              1 |  x x
2 | x x x            2 |  x x x
3 | x x x x          3 |  x x x x
4 | x x x x x        4 |    x x x x
5 | x x x x x x      5 |      x x x x
6 | x x x x x x x    6 |        x x x x
7 | x x x x x x x x  7 |          x x x x

对 N = 8192 和 W = 1024，分数矩阵期望上有 1024 × 8192 个非零行——8 倍的削减。

SWA 让 KV 缓存缩小。 每层只需保留 K 和 V 的最后 W 个 token。对一个类 Gemma-3 配置（1024 窗口、128K 上下文），KV 缓存降 128 倍。

质量代价。 纯 SWA transformer 在长程检索上吃力。修法：把 SWA 层和完整注意力层交错。Gemma 3 用 5:1 的 SWA:全局。Mistral 7B 用了一个因果 SWA 栈，信息通过重叠窗口"向前流"——每层把有效感受野扩展 W，L 层之后模型能往回关注 L × W 个 token。

稀疏 / 块注意力 ​

提前挑一个 N × N 的稀疏模式。三种规范形状：

  • 局部 + 跨步（OpenAI 稀疏 transformer）。 关注最后 W 个 token，加上之前每第 stride 个 token。以 O(N · sqrt(N)) 算力同时捕捉局部和长程。
  • Longformer / BigBird。 局部窗口 + 一小组全局 token（如 [CLS]），它们关注所有人也被所有人关注 + 随机稀疏连接。同等质量下经验上 2 倍上下文。
  • Native Sparse Attention（DeepSeek，2025）。 学哪些 (Q, K) 块重要；在 kernel 层跳过零块。兼容 FlashAttention。

稀疏注意力是个 kernel 工程的故事。数学很简单（掩码分数矩阵）；胜利来自从不把零条目加载进 SRAM。FlashAttention-3 和 2026 年的 FlexAttention API 让自定义稀疏模式在 PyTorch 里成为一等公民。

差分注意力（DIFF Transformer，2024） ​

常规注意力有个"注意力汇"问题：softmax 强迫每一行加和为 1，所以不想特别关注什么的 token 就把权重倒在第一个 token（或前几个）上。这偷走了本该给真实内容的容量。

差分注意力通过算两张注意力图并相减来修这点：

A1 = softmax(Q1 K1^T / √d)
A2 = softmax(Q2 K2^T / √d)
DiffAttn = (A1 - λ · A2) V

其中 λ 是一个学到的标量（通常 0.5–0.8）。A1 捕捉真实内容权重；A2 捕捉那个汇。相减抵消了汇，把权重重新分配给相关 token。

报告的结果（微软 2024）：困惑度低 5–10%、同训练长度下有效上下文长 1.5–2 倍、大海捞针检索更锐利。

变体对比 ​

变体算力KV 缓存相比完整的质量生产使用
完整注意力O(N²)每层 O(N)基线每个模型的默认层
SWA（窗口 1024）O(N·W)每层 O(W)-0.1 ppl，配全局层不错Gemma 2/3、Phi-3-Long
局部 + 跨步稀疏O(N·√N)混合类似 SWAOpenAI 稀疏 transformer、Longformer
BigBird（局部 + 全局 + 随机）约 O(N)混合2 倍上下文下追平完整早期长上下文 BERT
Native Sparse（DeepSeek-V3.2）O(N · 激活比例)O(N)在 0.05 ppl 之内DeepSeek-V3.2，2025
差分O(2·N²)O(2N)-5 到 -10% pplDIFF Transformer、2026 早期模型
figure
gqa-kv-sharing

动手构建 ​

见 code/main.py。我们实现一个因果掩码对比器，在一个玩具序列上把完整、SWA、局部+跨步、差分注意力并排展示。

第 1 步：完整因果掩码（基线） ​

python
def causal_mask(n):
    return [[0.0 if j <= i else float("-inf") for j in range(n)] for i in range(n)]

第 07 课的基线。下三角；对角线以上零权重。

第 2 步：滑动窗口因果掩码 ​

python
def swa_mask(n, window):
    M = [[float("-inf")] * n for _ in range(n)]
    for i in range(n):
        lo = max(0, i - window + 1)
        for j in range(lo, i + 1):
            M[i][j] = 0.0
    return M

一个参数——window。window >= n 时，你恢复完整因果注意力。window = 1 时，每个 token 只关注自己。

第 3 步：局部 + 跨步稀疏掩码 ​

python
def strided_mask(n, window, stride):
    M = [[float("-inf")] * n for _ in range(n)]
    for i in range(n):
        lo = max(0, i - window + 1)
        for j in range(lo, i + 1):
            M[i][j] = 0.0
        for j in range(0, i + 1, stride):
            M[i][j] = 0.0
    return M

稠密局部窗口加上回溯到序列起点的每第 stride 个 token。感受野随层数增加按 log 步长增长。

第 4 步：差分注意力 ​

python
def diff_attention(Q1, K1, Q2, K2, V, lam):
    A1 = softmax_causal(Q1 @ K1.T / sqrt_d)
    A2 = softmax_causal(Q2 @ K2.T / sqrt_d)
    return (A1 - lam * A2) @ V

两次注意力通过，用一个学到的混合系数相减。代码里我们对比单注意力和差分注意力的注意力汇热力图，看着汇塌掉。

第 5 步：KV 缓存大小 ​

打印 N = 131072 时每种变体的每层缓存大小。SWA 和稀疏变体降 10–100 倍。差分翻倍。清醒地付你的显存账单。

实际使用 ​

2026 年生产模式：

python
from transformers import AutoModelForCausalLM
# Gemma 3 以 5:1 混合 SWA（window=1024）和全局层。
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-27b-it")
# print(model.config.sliding_window, model.config.layer_types)

PyTorch 2.5+ 里的 FlexAttention 接受一个掩码函数：

python
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def swa_pattern(b, h, q_idx, kv_idx):
    return (q_idx - kv_idx < 1024) & (q_idx >= kv_idx)

mask = create_block_mask(swa_pattern, B=batch, H=heads, Q_LEN=n, KV_LEN=n)
out = flex_attention(q, k, v, block_mask=mask)

这会编译成一个自定义 Triton kernel。对常见模式，速度在 FlashAttention-3 的 10% 之内，而掩码函数是个 Python 可调用对象。

什么时候选哪个：

  • 纯完整注意力 —— 约 16K 上下文以内的每一层，或检索质量至上时。
  • SWA + 全局混合 —— 长上下文（>32K）、训练和推理受显存约束。32K 以上的 2026 默认。
  • 稀疏块注意力 —— 自定义 kernel、自定义模式。留给专门负载（检索、音频）。
  • 差分注意力 —— 任何注意力汇污染会伤害的负载（长上下文 RAG、大海捞针）。

拿去用 ​

见 outputs/skill-attention-variant-picker.md。这个 skill 会根据目标上下文长度、检索需求和训练/推理算力画像，为一个新模型挑选注意力拓扑。

练习 ​

  1. 简单。 跑 code/main.py。验证 window=4 的 SWA 把每行最后 4 个 token 之外的一切置零。验证 window=n 逐 bit 复现完整因果注意力。
  2. 中等。 在第 07 课的收官项目之上实现 window=1024 的因果 SWA。在 tinyshakespeare 上训 1,000 步。验证损失相比完整注意力退化多少？峰值显存降多少？
  3. 困难。 在收官模型里实现 Gemma-3 风格的 5:1 层混合（5 个 SWA、1 个全局）。在同等参数下，把损失、显存和生成质量和纯 SWA、纯全局基线对比。
  4. 困难。 实现每头一个学到的 λ 的差分注意力。在一个合成检索任务（一根针、2,000 个干扰项）上训练。在同等参数下，测检索准确率相比单注意力基线的情况。

关键术语 ​

术语大家嘴上怎么说实际是什么意思
滑动窗口注意力（SWA）"局部注意力"每个 query 关注它最后 W 个 token；KV 缓存缩到 O(W)。
有效感受野"模型能往回看多远"在窗口为 W 的 L 层 SWA 栈里，最多 L × W 个 token。
Longformer / BigBird"局部 + 全局 + 随机"带几个始终关注的全局 token 的稀疏模式；早期长上下文方法。
Native Sparse Attention"DeepSeek 的 kernel 把戏"学块级稀疏；在 kernel 层跳过零块同时保住质量。
差分注意力"两张图，一张减"DIFF Transformer：从第一张注意力图里减去 λ 倍的第二张来抵消注意力汇。
注意力汇（Attention sink）"权重渗到 token 0"softmax 归一化强迫行加和为 1；无信息的 query 把权重倒在位置 0。
FlexAttention"掩码即 Python"PyTorch 2.5+ 的 API，把任意掩码函数编译成 FlashAttention 形状的 kernel。
层类型混合"5:1 SWA-对-全局"在栈里交错稀疏和完整注意力层，以更低显存保住质量。

延伸阅读 ​