注意力变体:滑动窗口、稀疏、差分

注意力变体:滑动窗口、稀疏、差分

本节摘要:完整注意力是一个圆——每个 token 看见每个 token,显存为此付代价。四类变体弯折这个圆的形状,省回一半成本。完整注意力在序列长度上花 O(N²) 显存和 O(N²) 算力,128K 上下文的 Llama 3 70B 是每层 160 亿个注意力项乘 80 层;Flash Attention(第 12 节)藏掉了 O(N²) 激活显存,但不改算术成本——每个 token 仍关注每个其他 token。三类变体改变注意力矩阵本身的拓扑:**滑动窗口注意力(SWA)**每 token 只看固定窗口邻居,显存算力降到 O(N·W)(Gemma 2/3、Mistral 7B 首层、Phi-3-Long);稀疏/块注意力只给选中的 (i,j) 对打分(Longformer、BigBird、OpenAI sparse、DeepSeek Native Sparse);差分注意力用两组独立 Q/K 算两张注意力图相减,杀掉向头几个 token 溢出权重的「注意力槽」(微软 DIFF Transformer 2024)。它们共存——2026 前沿模型常混用:多数层 SWA-1024,每五层一次全局,少数差分头清理检索;Gemma 3 的 5:1 SWA:全局比是当前教科书默认。本节带你纯标准库并排实现四种掩码,看清每种换掉了 O(N²) 的哪一部分。

学习目标

阅读完本节,你应当能够:

1. 说清完整注意力为何 O(N²) 且 Flash Attention 只藏激活显存不改算术成本,从而理解变体的动机。
2. 实现滑动窗口掩码:window 参数控制每 token 看多远,KV 缓存随窗口而非序列长度增长,window≥n 退化回完整因果注意力。
3. 区分三类稀疏模式:局部+跨步(OpenAI sparse,O(N·√N))、Longformer/BigBird(局部+全局+随机,等质量 2 倍上下文)、Native Sparse(DeepSeek 学习块级稀疏)。
4. 解释差分注意力如何用 (A1 - λ·A2)V 相减杀掉注意力槽,以及为什么 softmax 强制行和为 1 会制造这个槽。
5. 给出 2026 年何时选哪种:≤16K 全局,>32K 用 SWA+全局 5:1 混合(默认),定制内核用稀疏块,长上下文 RAG 用差分。

一、问题与直觉

完整注意力 O(N²),Flash Attention 藏掉激活显存但不改算术。变体改注意力矩阵本身的拓扑。

滑动窗口注意力(SWA)

位置 i 的查询只关注 [i-W, i](因果 SWA)或 [i-W/2, i+W/2](双向),窗口外得 -inf:

完整因果:              滑动窗口(W=4):
0 1 2 3 4 5 6 7        0 1 2 3 4 5 6 7
0|×                0| ×
1|× ×              1| × ×
2|× × ×            2| × × ×
3|× × × ×          3| × × × ×
4|× × × × ×        4|   × × × ×
5|× × × × × ×      5|     × × × ×
6|× × × × × × ×    6|       × × × ×
7|× × × × × × × ×  7|         × × × ×
​

N=8192, W=1024 时分数矩阵期望有 1024×8192 个非零行——8 倍削减。

KV 缓存随 SWA 缩小:每层只留最后 W 个 token 的 K 和 V。Gemma-3 形状(1024 窗、128K 上下文)KV 缓存降 128 倍。

质量代价:纯 SWA 在长程检索上挣扎。修法:SWA 层与全注意力层交错。Gemma 3 用 5:1 SWA:全局;Mistral 7B 用因果 SWA 堆叠,信息通过重叠窗口「向前流」——每层扩展有效感受野 W,L 层后能关注 L·W 个 token 回。

稀疏/块注意力

预先选 N×N 稀疏模式,三种经典形状:

• 局部 + 跨步(OpenAI sparse):关注最近 W 个 token 加其后每隔 stride 个,O(N·√N) 算力同时捕局部和长程。
• Longformer/BigBird:局部窗口 + 少量全局 token(如 [CLS],关注所有人也被所有人关注)+ 随机稀疏链接,等质量 2 倍上下文。
• Native Sparse(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ᵀ / √d)
A2 = softmax(Q2·K2ᵀ / √d)
DiffAttn = (A1 - λ·A2)·V
​

λ 是学习标量(典型 0.50.8)。A1 捕真实内容权重,A2 捕槽,相减抵消槽、把权重重新分配给相关 token。微软 2024 报告:困惑度低 510%、等训练长度下有效上下文长 1.5~2 倍、针在草堆检索更锐利。

变体对比

💡 注意力槽的本质:softmax 是「软」归一化,永不出精确零。当一个 token 在前缀里找不到强相关对象时,它必须把概率撒在某处——往往撒在位置 0(序列开头)。这是为什么长上下文检索里,无关位置反而被位置 0 抢走权重。差分注意力的减法是直接抵消这个偏置的巧妙解。

二、从零实现

完整代码见原课程 phases/07-transformers-deep-dive/15-attention-variants/code/main.py。实现一个因果掩码比较器,在玩具序列上并排展示完整、SWA、局部+跨步、差分注意力。

Step 1:完整因果掩码(基线)

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

第 07 节的基线,下三角,对角线上方零权重。

Step 2:滑动窗口因果掩码

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 只看自己。

Step 3:局部+跨步稀疏掩码

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,感受野随层数对数增长。

Step 4:差分注意力

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
​

两次注意力 pass,用学习混合系数相减。代码里对比单 vs 差分的注意力槽热力图,看槽如何塌缩。

Step 5:KV 缓存大小

打印 N=131072 下每种变体的每层缓存大小。SWA 和稀疏变体降 10~100 倍,差分翻倍——有意识地付你的显存账。

设计要点:SWA 堆叠的有效感受野是 L·W——这是关键洞察。单层 SWA 看不远,但 L 层重叠窗口后信息能「接力」传 L·W 个 token 回。Gemma 3 的 5:1 混合就是这个原理的工程化:5 层 SWA 把局部信息聚合,第 6 层全局注意力把它们与远距离信号融合。

三、框架对比:Gemma 3 与 FlexAttention

2026 生产模式:

from transformers import AutoModelForCausalLM
# Gemma 3 在 5:1 比例混 SWA(窗=1024)和全局层
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-27b-it")
# print(model.config.sliding_window, model.config.layer_types)
​

PyTorch 2.5+ 的 FlexAttention 接受掩码函数:

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 callable。

何时选哪种:

• 纯全局——≤16K 上下文,或检索质量至上的每层。
• SWA + 全局混合——长上下文(>32K),训练推理 memory-bound。2026 年 32K 以上的默认。
• 稀疏块——定制 kernel、定制模式,留给专用负载(检索、音频)。
• 差分——注意力槽污染伤人的负载(长上下文 RAG、针在草堆)。

四、可复用产物

原课程产出 outputs/skill-attention-variant-picker.md:一个选择器 Skill,给定目标上下文长度、检索需求、训练推理算力画像,为新模型选注意力拓扑。

五、练习

1. (Easy) 跑 code/main.py,验证 window=4 的 SWA 把每行最后 4 个 token 外全清零,window=n bit 一致复现完整因果注意力。
2. (Medium) 在第 14 节毕业项目上加因果 SWA(window=1024),tinyshakespeare 上训 1000 步。验证损失相对全局退化多少?峰值显存降多少?
3. (Hard) 在毕业项目里实现 Gemma-3 式 5:1 层混合(5 SWA + 1 全局),等参下对比损失、显存、生成质量。
4. (Hard) 实现每头学习 λ 的差分注意力,在合成检索任务(1 针 + 2000 干扰)上训,测等参下相对单注意力基线的检索准确率。

本节要点回顾

1. 完整注意力 O(N²):Flash Attention 只藏激活显存不改算术成本——这是变体存在的根本动机。
2. SWA = O(N·W):每 token 只看窗口内,KV 缓存降到 O(W);window≥n 退化回完整因果,window=1 只看自己。
3. SWA 堆叠有效感受野 L·W:单层不远但多层重叠接力;Gemma 3 的 5:1 SWA:全局是工程化默认。
4. 三类稀疏:局部+跨步(OpenAI sparse O(N·√N))、Longformer/BigBird(局部+全局+随机,2× 上下文)、Native Sparse(DeepSeek 学习块稀疏,kernel 跳零块)。
5. 稀疏本质是 kernel 工程:数学简单(掩分数),赢点在不把零载入 SRAM;FlexAttention 让掩码成 Python callable。
6. 差分注意力 (A1-λA2)V:两张图相减杀掉注意力槽(softmax 强制行和为 1 制造的偏置),困惑度降 510%、有效上下文长 1.52 倍。
7. 注意力槽:softmax 永不出精确零,无关 token 把概率撒在位置 0,长上下文检索被抢权重。
8. 2026 选择:≤16K 全局,>32K 用 SWA+全局 5:1,定制用稀疏块,长上下文 RAG 用差分。

下一节也是本章最后一节,我们讲投机解码——小模型起草、大模型验证,把自回归生成的串行延迟砍掉一半以上,是 2026 年每个生产推理栈的标配。