1.2 注意力机制：从加权检索到万物的基础操作

本章目标：把”注意力”从玄学变成一个你闭眼能写的函数。它是整个大模型时代唯一的计算原语——CV 里它取代了卷积，LLM 里它就是全部，多模态里它负责图文对齐。学完标准：能手写带 causal mask 的多头注意力并与 PyTorch 内置实现对拍，能解释 1/√d 的来历。

1. 为什么需要注意力：RNN 的两大死穴

Transformer 之前序列建模靠 RNN/LSTM：逐 token 读取，用隐状态”记忆”历史。两个致命问题：

无法并行：第 t 步依赖第 t−1 步，GPU 上万个核心干等着看一个序列爬行；
长程依赖衰减：信息要经过 O(N) 次状态更新才能传到远处，距离一长就被冲淡（梯度消失的序列版）。

注意力的解法釜底抽薪：不递归传递信息，让任意位置一步直达任意位置——距离消失，同时整个序列的计算可以一次矩阵乘完成。2017 年《Attention Is All You Need》标题的狂妄就在于此。

2. 注意力的通用形式：一次可微分的”查字典”

想象查字典：你有一个查询 Q（想找什么），字典里每个词条有键 K（索引）和值 V（内容）。硬字典是 argmax 检索（不可导），注意力做的是软检索：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

三步拆解（每步都有几何直觉）：

打分： $QK^\top$ 算出每个 query 和每个 key 的内积——内积 = 对齐程度 = 相关性（0.1 章 §2.1）；
归一化：softmax 把分数变成总和为 1 的权重分布——“把 100% 注意力分配给各位置”；
加权求和：输出 = 各位置的 value 按注意力权重混合。

一句话本质：注意力是按相关性加权的加权平均，而相关性用内积度量，权重用 softmax 归一，全部可导，因此可以端到端学习。 图像块之间做它 = ViT；图像与文本之间做它 = CLIP 式对齐；序列自己内部做它 = GPT。万物皆注意力。

3. 为什么除以 √d_k（面试与理解双高频）

假设 q 和 k 的各分量独立、零均值、方差 1，则内积 $q \cdot k = \sum_{i=1}^{d} q_i k_i$ 的方差是 $d$ （方差相加），标准差 $\sqrt{d}$ 。d=128 时，打分普遍在 ±11 上下波动——softmax 在这么大的输入下会饱和：最大项拿到几乎全部概率，其余趋近 0，梯度消失，且分布对输入噪声极其敏感。

除以 $\sqrt{d_k}$ 把打分的方差拉回 1，softmax 保持”清醒”。这是”方差控制”思想——和初始化（0.3 章 §5）、RMSNorm 是同一个思想在不同位置的应用：让每层输出的数值尺度保持稳定。

4. Self-Attention 与 Multi-Head

Self-attention：Q、K、V 全部来自同一个序列（各自过三个线性投影 $W_Q, W_K, W_V$ ）。每个位置同时是查询者、被查者和内容提供者——序列内的信息交换。
Cross-attention：Q 来自 A 序列，K/V 来自 B 序列。多模态的桥（LLaVA 里文本 query 查图像特征）、翻译里译文查原文，都是它。
Multi-Head：把 d 维投影拆成 h 个 64 维子空间，各自独立做注意力再拼接。动机：一组 QKV 只能表达一种”关注模式”，多头 = 多个视角并行（语法一=head1，指代一=head2……）。与 CNN 的多通道异曲同工：并行多组特征提取器，各学各的。

5. Causal Mask：让 GPT 不会”作弊”

语言模型预测第 t 个 token 时只能看前文。实现：把 $QK^\top$ 矩阵的上三角（未来位置）在 softmax 前填 $-\infty$ ，softmax 后权重为 0：

mask = torch.tril(torch.ones(N, N, dtype=torch.bool))   # 下三角为 True
scores = scores.masked_fill(~mask, float("-inf"))

记住这个机制的名字：解码器里的自回归 = causal mask + 逐位预测。卷3 LLaVA 复现、卷4 自回归生成模型（Janus 路线），全靠它。

6. 手写实现 + 对拍（本章作业的骨架）

import torch, torch.nn as nn, torch.nn.functional as F, math

def attention(q, k, v, mask=None):
    """q,k,v: (B, H, N, D)  mask: (N, N) bool, True=可见"""
    d = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d)          # (B,H,N,N)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    return F.softmax(scores, dim=-1) @ v                     # (B,H,N,D)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        assert d_model % n_head == 0
        self.h, self.dh = n_head, d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)   # 一次投影三份，省两次 GEMM
        self.proj = nn.Linear(d_model, d_model, bias=False)      # 输出融合投影
    def forward(self, x, mask=None):
        B, N, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=2)                    # 各 (B,N,D)
        q, k, v = (t.view(B, N, self.h, self.dh).transpose(1, 2) for t in (q, k, v))  # (B,H,N,dh)
        out = attention(q, k, v, mask)                           # (B,H,N,dh)
        out = out.transpose(1, 2).contiguous().view(B, N, D)
        return self.proj(out)

# 对拍验证
x = torch.randn(2, 16, 128)
mha = MultiHeadAttention(128, 8); mha.eval()
mine = mha(x)
q, k, v = mha.qkv(x).split(128, dim=2)
ref = F.scaled_dot_product_attention(
    q.view(2,16,8,16).transpose(1,2), k.view(2,16,8,16).transpose(1,2), v.view(2,16,8,16).transpose(1,2))
ref = ref.transpose(1,2).reshape(2,16,128)
print((mha.proj(ref) - mine).abs().max())   # 应 < 1e-6

工程常识：真实训练用 F.scaled_dot_product_attention——它自动分派到 FlashAttention 内核。FlashAttention 你不需要实现，只需记住它的思想：不做优化数学，优化显存读写（避免把巨大的 N×N 注意力矩阵写进显存再读出来，用分块重算换取 10 倍级速度）——这是”系统工作”在大模型里值钱程度的最好例子。

7. 代价与对比

复杂度 $O(N^2 d)$ ：注意力矩阵 N×N 是显存和时间的大头。长上下文、高分辨率图像（patch 数多）、视频（帧数爆炸）的成本都从这里来——卷3/4 里所有”长上下文""多帧”问题，根源都是这个平方项。
与卷积的对比表（理解架构选择的地图）：

	卷积	Self-Attention
感受野	局部，逐层扩大	全局，一步到位
参数	共享，先验强	投影矩阵，先验弱
数据效率	小数据占优	大数据占优
位置信息	结构自带	必须额外注入（→1.3 位置编码）
多模态亲和	弱（模态间无天然接口）	强（都是 token 序列）

8. 自测题

注意力的三步是什么？各自的作用？（打分/归一化/加权求和）
1/√d 的方差推导为什么假设零均值独立？（中心极限式的和：分量方差相加，d 项 → 方差 d）
Self 和 Cross attention 的区别一句话？（Q 和 K/V 是否同源）
为什么多头而不是一个大的头？（多组子空间并行学不同关注模式；总计算量不变）
causal mask 填 −∞ 而不是 0，为什么？（softmax 后 0 权重需要 −∞ 分数；填 0 会产生非零注意力泄漏未来信息）
注意力的平方复杂度在哪一步产生？它如何影响视频生成的帧数上限？（ $QK^\top$ 的 N×N 矩阵；帧数↑→N↑→平方爆炸）
FlashAttention 优化的是什么资源？（显存读写 IO，不是浮点计算）

9. 动手作业

跑通第 6 节对拍代码；然后给 MultiHeadAttention 加 causal mask 支持，用 is_causal=True 的内置实现再对拍。
可视化注意力：随机初始化 + 训练一个 2 层小模型在玩具序列任务上（如反转数字序列），用 matplotlib 画 softmax(scores) 热力图，观察注意力”看见了什么”。
把第 6 节的 MHA 包进 Transformer block（残差 + LayerNorm + MLP），就是 1.3 章的开头——提前搭好骨架。
挑战：读 FlashAttention 论文第 1-2 节，用三句话讲清”分块 + 重算”为什么省显存。

图解 · 一次注意力的完整数据流

flowchart TB
  X["输入 x (B×N×D)"] --> QKV["单个 Linear 投影出 3D 维"]
  QKV --> S["拆成 Q / K / V (B×H×N×d_h)"]
  S --> SC["scores = QKᵀ/√d_h  →  (B×H×N×N)"]
  SC --> M["causal mask：未来位置填 −∞"]
  M --> SM["softmax（每行归一成权重）"]
  SM --> AV["加权求和 scores @ V"]
  AV --> CO["concat 所有头 (B×N×D)"]
  CO --> PJ["输出投影 W_O"]

平方成本就藏在 scores 那个 N×N 矩阵里——N 翻倍它变四倍：这是 2.3 的 KV cache 账、2.8 的 prefill 瓶颈、3.3 的 token 压缩，全部问题的同一源头。

案例：两个 token 的迷你注意力（全手算）

设 D=2、单头，序列 [“the”, “cat”]，投影后 $q_1=[1,0]$ ， $k_1=[1,1],\ k_2=[0,1]$ ，V 取值同 K。

位置 1（有 causal mask）只看自己 → 权重 $[1, 0]$ ；位置 2 看两个： $s_{2,1}=1{\cdot}1+0{\cdot}1=1$ ， $s_{2,2}=0{\cdot}1+1{\cdot}1=1$ → softmax 各 0.5 → 输出 $o_2 = 0.5\,v_1 + 0.5\,v_2 = [0.5,\ 1]$ 。

你刚刚手动执行了 GPT 的一层。四句话总结注意力：mask 决定能看谁，内积决定多想看，softmax 归一成百分比，加权求和取材料。

10. 延伸资源

李沐《论文精读》Transformer 一期——逐段精读 Attention Is All You Need
Jay Alammar《The Illustrated Transformer》——图解第一名
Karpathy《Let’s build GPT: from scratch》——本笔记 1.3/2.3 章的实拍版

云镜检索

稍后阅读