05. 注意力：每个 token 怎样看上下文

上一篇结束时，“年假怎么申请？”变成了 5 个 4096 维的向量，送进第 1 层。这时“申请”的向量还是词表里那一行，在每句话里都一样，它并不知道自己申请的是年假还是报销。让它知道的，是每层里的自注意力。这一篇先说清注意力要解决什么问题、Q、K、V 从哪里来，再用 4 个 token、每头 2 维的小数字把缩放点积、softmax 和因果掩码完整手算一遍，最后算一算它的计算量和显存为什么随长度平方增长。

读完应能看懂 softmax(QKᵀ/√d + 掩码)·V 这一行公式里每个符号对应什么，能从一个注意力权重矩阵读出“谁看了谁”，也能说清为什么上下文从 8K 拉到 128K，注意力会从配角变成主角。

“申请”需要从前面拿什么

先不管公式，只看需求。模型读到“申请”时，要预测接下来的内容，就需要把前面“年”“假”的信息并到“申请”这个位置上。最简单的办法是把前面所有位置的向量取个加权平均，加到“申请”上：

新的“申请” = w₁ × 年 + w₂ × 假 + w₃ × 怎么 + w₄ × 申请      （w₁ + w₂ + w₃ + w₄ = 1）

问题在权重怎么定。固定权重不行：“年假怎么申请”里“申请”该看“假”，“报销单怎么申请”里该看“报销单”，“申请”自己的位置和前一个词的位置都一样，内容不同。权重必须由内容决定，而且每个位置要能自己决定看谁。注意力就是一套按内容计算这组权重、再做加权平均的方法。

Q、K、V：同一个向量的三种用途

注意力让每个位置的向量 x 经过三个不同的矩阵，得到三个新向量：

• Q（query，查询）：这个位置想找什么信息。
• K（key，键）：这个位置能被怎样找到，相当于贴在自己身上的标签。
• V（value，值）：被选中以后，这个位置交出去的内容。

位置 i 对位置 j 的关注程度，就是 i 的 q 和 j 的 k 的点积：方向越一致，分数越高。分数换算成权重后，按权重把各位置的 v 加起来，就是位置 i 取到的信息。K 和 V 分开，是因为“凭什么被找到”和“找到后给什么”不必相同：“假”可以用“我是一个事项”这个特征被“申请”找到，交出去的却是“具体是什么假”的内容。

三个矩阵都是训练学出来的参数，Q、K、V 里也没有人工写的关键词，上面的“想找什么”“标签”只是帮助记住分工的说法。在 Llama-3.1-8B 里，它们就是第 03 篇参数表里的 q_proj、k_proj、v_proj：

x [n, 4096] ── q_proj (4096 × 4096) → Q [n, 4096] = [n, 32 个头, 128]
            ── k_proj (4096 × 1024) → K [n, 1024] = [n,  8 个头, 128]
            ── v_proj (4096 × 1024) → V [n, 1024] = [n,  8 个头, 128]

4096 维被拆成 32 个头，每头 128 维，每个头独立做一遍注意力；K、V 只有 8 个头，每 4 个查询头共用一组，这是 GQA，下一篇专门讲。这一篇只看一个头。

手算一遍：4 个 token，每个头 2 维

为了能手算，把维度缩到很小：取“年假怎么申请？”的前 4 个 token，每个 token 的输入向量只有 4 维，每个头的 Q、K、V 只有 2 维。为了好理解，给 4 个输入维度起了名字（事项、时间、疑问、动作），真实模型的维度没有这样的含义。所有数都是教学用的，不是 Llama 的真实数值，计算用脚本核对过。

第一步，算 Q、K、V。每个 token 的输入行乘以权重矩阵，例如“申请”的输入是 [0, 0, 0, 1]，只有“动作”一维，所以它的 q 就是 W_Q 的第 4 行 [2, 0]；“假”的输入是 [2, 0, 0, 0]，它的 k 是 W_K 第 1 行乘 2，即 [2, 0]。全部算完：

token	q	k	v
年	[1, 1]	[1, 0]	[1, 0]
假	[0, 2]	[2, 0]	[2, 0]
怎么	[1, 0]	[0, 1]	[0, 1]
申请	[2, 0]	[0, 0]	[0, 1]

第二步，以“申请”为例，用它的 q 和每个位置的 k 做点积，再除以 √d（d 是每个头的维度，这里 d = 2，√2 ≈ 1.414），然后做 softmax：

分数 s_j = q_申请 · k_j / √d
  年：  [2,0]·[1,0] = 2  → 2 / 1.414 = 1.414
  假：  [2,0]·[2,0] = 4  → 4 / 1.414 = 2.828
  怎么：[2,0]·[0,1] = 0  → 0
  申请：[2,0]·[0,0] = 0  → 0

softmax：权重 w_j = e^(s_j) / Σ e^(s_k)
  e^1.414 = 4.113，e^2.828 = 16.919，e^0 = 1，e^0 = 1，合计 23.032
  w = [4.113, 16.919, 1, 1] / 23.032 = [0.179, 0.735, 0.043, 0.043]

softmax 做两件事：取指数让所有权重变成正数、拉开大小差距（分数差 1.414，权重差 4 倍多）；再除以总和，让权重加起来等于 1。

第三步，按权重把 v 加起来：

输出 = 0.179 × [1,0] + 0.735 × [2,0] + 0.043 × [0,1] + 0.043 × [0,1]
     = [1.648, 0.087]

“申请”取到的信息有 73.5% 来自“假”。在这组教学数字里，这是因为 W_Q 把“动作”映射成了“找事项”的方向，W_K 把“事项”特征放在同一个方向上。真实模型里，这种“动词去找它的宾语”的关系是否存在、出现在哪一层哪个头，要看具体的权重；后文会说到，从注意力权重推断模型“在想什么”要很谨慎。

为什么要除以 √d

去掉 √d 公式照样能算，但 d 一大就出问题。假设 q、k 的每个分量都是均值 0、方差 1 的随机数，d 维点积的方差就是 d，标准差是 √d。用随机数模拟一下：d = 2 时点积标准差约 1.4，d = 128（Llama 每头的维度）时约 11.3。

分数的差距动辄十几，softmax 的指数会把它放大成几乎全部权重压在一个位置上。同样用随机数模拟 5 个位置、d = 128：不缩放时，最大的那个权重平均是 0.93；除以 √128 以后，平均是 0.46。前一种情况下 softmax 接近“只选一个”，其他位置的梯度几乎为零，训练时很难学会调整关注对象。除以 √d 把分数的标准差拉回 1 左右，原始 Transformer 论文给出的理由就是这一点。Transformers 里对应 LlamaAttention 的 self.scaling = head_dim ** -0.5。

手算例子里 d = 2，缩放影响不大：不除 √2，“申请”对“假”的权重从 0.735 变成 0.853。

因果掩码：每个位置只能看左边

上面只算了“申请”这一行。实际计算时，4 个位置同时当提问者，分数排成一个 4 × 4 的矩阵：行是提问的 token（Q），列是被看的 token（K）。

decoder-only 模型要求每个位置只能看自己和左边。做法是在分数矩阵的右上角（列号大于行号的位置）加上负无穷，softmax 时 e^(−∞) = 0，这些位置的权重就是 0，其余权重仍然加起来等于 1。这叫因果掩码（causal mask）。

逐行看右边的权重：“年”是第一个 token，只能看自己，权重必然是 1；“假”能看“年”和自己，两个分数都是 0，各得 0.5；“怎么”能看前三个，0.284、0.576、0.140；“申请”能看全部 4 个，就是刚才手算的那一行。

如果去掉掩码，“假”这一行对“怎么”的分数最高（1.414），会把 0.578 的权重放到右边的“怎么”上。训练时这等于偷看答案：模型在“假”这个位置的任务是预测下一个 token 是“怎么”，它却能直接看到“怎么”，学到的只是抄写，到了生成时右边根本没有内容可抄。第 03 篇提到的“训练时一次前向算出所有位置的预测”，正是靠这个掩码才成立。

掩码还带来一个推理上的好处：位置 j 的 K、V 只由它左边的内容决定，后面来了新 token 也不会变，所以可以缓存。这就是 KV Cache 成立的前提。

回到真实形状：一层注意力的完整代码

把一个头的计算推广到 32 个头，再接上输出投影 o_proj，就是 Llama 一层注意力的全部。下面是简化的 PyTorch 风格写法，对应 Transformers 里的 LlamaAttention 和 eager_attention_forward（实际推理多走 FlashAttention 等算子，数学相同）：

def attention(x, positions):                    # x: [n, 4096]
    q = q_proj(x).view(n, 32, 128)              # 32 个查询头
    k = k_proj(x).view(n, 8, 128)               # 8 组 K、V（GQA，第 06 篇）
    v = v_proj(x).view(n, 8, 128)
    q, k = rope(q, k, positions)                # 旋转写入位置信息（第 09 篇）
    k, v = repeat(k, 4), repeat(v, 4)           # 每组 K、V 给 4 个查询头用 → [n, 32, 128]
    s = q @ k.transpose(-1, -2) / sqrt(128)     # 每个头一张 [n, n] 分数表
    s = s + causal_mask                         # 右上角 −∞
    p = softmax(s, dim=-1)                      # 每行加起来等于 1
    out = p @ v                                 # [n, 32, 128]
    return o_proj(out.reshape(n, 4096))         # 32 个头拼回 4096 维，再混合一次

32 个头各算出 128 维，拼起来正好 4096 维，o_proj（4096 × 4096）再把各头的结果混合，输出 [n, 4096]，加回残差流。整个计算里，顺序只出现在两处：因果掩码决定每个位置能看到哪些 token，RoPE 把位置写进 q 和 k。去掉 RoPE，点积只按内容比较，“申请”能看到的 4 个 token 换个排列，它得到的结果完全不变。第 09 篇从这一点讲起。

长度的代价：分数表随长度平方增长

注意力每个头都要算一张 n × n 的分数表，n 是序列长度。手算例子 n = 4，一共 16 个分数；小林那句话 n = 5，25 个。资料助手如果把整本员工手册放进上下文，n 就是几万到十几万，这张表会迅速变大。

先看计算。第 n 个位置的 token 要和前面 n 个位置算点积（QKᵀ），再对 n 个 v 加权求和（PV），每个头各约 2 × n × 128 次运算，32 个头、两步加起来：

每层每 token 注意力计算 ≈ 4 × n × 4096 FLOPs
32 层，n = 8,192：   4 × 8,192 × 4,096 × 32   ≈ 4.3 GFLOPs
32 层，n = 131,072： 4 × 131,072 × 4,096 × 32 ≈ 68.7 GFLOPs
对照：每 token 和权重相乘的部分 ≈ 15.0 GFLOPs（第 03 篇），与 n 无关

单个 token 的注意力计算随 n 线性增长；一整段 n 个 token 一起处理（Prefill），每个 token 平均看 n/2 个位置，总量约 2 × n² × 4096 × 32，随 n 平方增长。按这个口径估算，8K 的 Prefill 里注意力约占固定部分的 14%，32K 时约 57%，128K 时是固定部分的 2.3 倍。上下文越长，注意力越从配角变成主角，这也是 Llama 3 把长上下文训练放到最后阶段的原因之一（见训练系列的预训练）。

再看显存。如果把分数表整张写出来，每层 32 个头、每个 n × n、BF16 每个 2 字节：8K 时每层 4 GiB，32K 时 64 GiB，128K 时 1024 GiB。长度翻 4 倍，表涨 16 倍，任何显卡都放不下。实际推理用的 FlashAttention 把计算切成小块，在 GPU 片上的高速存储（SRAM）里边算边做 softmax，从不写出整张表：结果和标准注意力完全一样，计算量也没少，省掉的是这张平方大小的中间表和读写它的时间。细节见第 12 篇和推理系列的注意力算子。

真正随长度线性增长、每一步都要读的，是 KV Cache：Llama-3.1-8B 每个 token 在 32 层共存 128 KiB 的 K、V，128K 上下文就是 16 GiB。第 06 篇讲的 GQA、MLA，第 12、13 篇讲的滑动窗口、稀疏注意力和线性注意力，都是在这两笔账上做文章。

注意力权重能说明什么

注意力矩阵很直观，常被画成热力图，用来说明“模型看了哪里”。使用时要知道它的边界。

第一，softmax 强制每行权重加起来等于 1。当前面没有真正有用的内容时，权重也必须分出去。实际模型里常见大量权重堆在第一个 token 上，这个现象叫注意力汇聚（attention sink），2023 年 StreamingLLM 的研究系统地描述了它：第一个 token 往往并不重要，只是充当“没处可放时的落脚点”。如果做长文本推理时为了省显存把开头的 token 丢掉，模型输出可能明显变差，原因就在这里（第 12 篇）。看热力图时，第一列特别亮不代表模型在认真读开头。

第二，权重高不等于对结果影响大。一个位置的权重可能很高，但它的 v 很小，或者被后面的 o_proj 和 FFN 抵消。2019 年 “Attention is not Explanation” 和随后 “Attention is not not Explanation” 两篇论文的争论，结论大致是：注意力权重提供了一些线索，但不能单独当作模型做决定的解释，需要配合干预实验（比如遮掉某个位置看输出怎样变）才能下结论。手算例子里“申请”看“假”是我们特意设计的数字；在真实模型的 32 层 × 32 个头里，某个头确实承担了类似的功能，需要实验证据才能说。

第三，长上下文会稀释权重。n 很大时，即使相关的那一段分数最高，几万个不相关位置的小权重加起来也可能压过它。“能放进 128K”和“在 128K 里准确找到那一句”是两回事，资料助手直接塞整本手册之前，应该用自己的长文档实测找得准不准。

注意力让每个位置用自己的 q 去和所有可见位置的 k 比较，比较结果经缩放和 softmax 变成权重，再按权重取走各位置的 v；因果掩码保证只看左边，这既让训练能并行、也让 K、V 可以缓存；代价是每个头一张随长度平方增长的分数表。Llama 每层有 32 个这样的头，却只有 8 组 K、V，下一篇讲多头注意力与 KV 头：为什么要多个头，KV 头从 32 减到 8 甚至压缩成低维向量，省了什么，又丢了什么。