Attention 就是三次矩阵乘

问题:词的意思,要看邻居的脸色

先立问题再上公式。经过 embedding(第 06 章的查表),句子「猫追红色激光点」变成一摞向量 (L, D)。但查表查出的向量是孤立的:「追」的向量里没有「谁在追、追什么」的信息;多义词更惨,「苹果」的向量分不清水果和公司——同一个 token,查同一张表,拿到同一支向量,语境全靠后面挣。要让每个词的表示吸收上下文,机制上需要三件事:每个词得能找到该看的邻居(匹配),决定看多少(权重),然后把邻居的信息搬过来(混合)。Attention 就是这三件事的线性代数化,而且三件事恰好各用一章你学过的数学——这不是巧合,是「匹配、加权、混合」这组语义,在向量世界里的标准译法本来就只有点积、softmax 和线性组合。

第一步:QKᵀ,批发 L² 个点积

在 attention 之前,历史上已经有两种「吸收上下文」的方案,对照着看能看清它赢在哪。RNN 方案:从左到右逐词吞,把上下文压缩进一个固定大小的隐状态——串行(没法并行训练),远距离信息要接力传递几百步(第 16 章矩阵幂的暴政)。卷积方案:每个词只看固定窗口内的邻居(第 04 章滑动模板)——并行是并行了,但「猫」和三十个词之外的「它」要叠很多层才能碰面。Attention 的答案粗暴而优雅:让每个词直接和所有词见面,见面成本一次矩阵乘,距离一律为一步。代价是 L² 的账单——这场「表达力换算力」的交易,就是 2017 年那次架构革命的核心条款。

每个词的向量先兵分三路,分别乘三个待训练的投影矩阵,得到三种「人格」:Q(query,我在找什么)、K(key,我能被什么找到)、V(value,找到我后能带走什么)。用图书馆比喻:Q 是你手里的检索词,K 是每本书的索引标签,V 是书的内容——检索词匹配标签,匹配上了借走内容;第 04 章的向量检索、第 20 章 FFN 的键值记忆,和这里是同一套「以匹配定取舍」的家族语法,只是这次连「检索词怎么写」都交给了训练。三种人格必须分开:「追」作为查询者想找名词,作为被查者希望被「动作相关」的问题命中,两种角色的向量不该是同一支——这就是为什么要三个不同的投影矩阵,而不是拿原向量直接互相点积。

匹配的数学是现成的:第 04 章立过「点积 = 相似度度量衡」。把所有查询摞成 Q(L×d),所有键摞成 K(L×d),那么第 07 章的批发定理说:

一次矩阵乘,L² 个点积批发到手。那个 √d 是唯一的新面孔,来历是第 04 章的高维冷知识:d 维随机向量的点积,典型大小随 √d 增长;不除掉它,d=128 时分数动辄几十,softmax 的指数会把预算全部砸给最高分(温度过低的贫富分化),梯度随之死亡。除以 √d = 把温度校准回「有分辨力但不极端」的区间——一行公式里藏着一整章的高维几何,论文里那句轻描淡写的 "scaled dot-product",你现在知道 scaled 的每一分含义。

第二、三步:softmax 分预算,·V 搬货

分数表逐行过 softmax(axis=-1,第 19 章的语义警告在此站岗):每行变成非负、和为 1 的注意力权重——第 04 章的预算分配,第 03 章的「凸组合配方」。为什么非要 softmax 而不直接用裸分数加权?三个理由:非负(往回搬信息,不该有「负份额」的诡异语义)、归一(输出的尺度不随句长漂移)、可导且赢家放大(高分拿走大头,但人人有份,梯度能流向每一条边——argmax 做不到)。然后第三次矩阵乘:

按行读(第 07 章四种读法之三):O 的第 i 行 = 权重行 · V 矩阵 = 各词 value 的凸组合。凸组合的两条老性质在这里都是设计特性:输出不会跑出各 V 围成的凸包(稳定,第 03 章),权重是当场算出来的(动态配方——对比之下,卷积和 MLP 的权重训练完就焊死,attention 的「权重」每个输入现算,这是它表达力的来源,也是它贵的原因)。至此三步齐了,亲手跑:

Demo 用的句子是「猫 追 红色 激光点」,四个 token、每人三维的 Q/K/V 都是手工设定的(真模型里它们由训练得来,这里为了让因果看得见,把「动词找名词」「名词找修饰语」的倾向直接写进了数字)。实验引导:① 按 ①→②→③ 单步走,每一步先猜再看:分数表里哪格该深?softmax 后每行是不是和为 1?② 盯「追」那一行:它同时深看「猫」与「激光点」——动词把主语宾语的信息抓进自己的新向量,readout 里能看到它的输出确实混进了两者的成分;③ 盯「激光点」那行对「红色」的权重:修饰关系被捕捉;④ 复盘:整个过程只有矩阵乘和一个 softmax,而「谁看谁」的智能,全部来自 Q、K 投影矩阵里被训练出来的数字——机制是死的,权重是学的——这句话适用于本书讲过的每一个模块,在这里最扎眼。

写成代码:十行,不藏私

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)   # 第 21 章急救包②
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(X, Wq, Wk, Wv, causal=False):
    Q, K, V = X @ Wq.T, X @ Wk.T, X @ Wv.T    # (L,D)→(L,d) 三种人格
    S = Q @ K.T / np.sqrt(Q.shape[-1])         # (L,L) 第一次矩阵乘
    if causal:
        S = np.where(np.tril(np.ones_like(S)), S, -1e9)  # 下三角通行证
    return softmax(S) @ V                     # 第二、三步

就这些——GPT 心脏的单头版,含因果掩码,十来行。生产实现多出来的部分全是第 19 章的形状杂技(batch 轴、切头、拼头)和第 21 章的数值照顾,数学一分没加。强烈建议把它抄进 REPL,用 L=4、d=3 的小矩阵单步跑一遍,和 Demo 的数字互相印证——「我徒手跑过 attention」这句话,值得拥有。

位置去哪了:序不能靠猜

细心的话你该发现一个漏洞:三次矩阵乘对「词的顺序」完全无感——把输入的行随便洗牌,输出跟着同样洗牌,分数一分不变(用第 10 章置换矩阵的语言:注意力与置换可交换)。「猫追狗」和「狗追猫」在它眼里是同一袋词!所以位置信息必须手动注入,两大流派:加法派——给每个位置配一支「位置向量」加进 embedding(原版 Transformer 的正弦波方案:不同频率的 sin/cos 组合,像给每个位置发一个二进制刻度的模拟版);旋转派(RoPE)——把 Q、K 向量按位置旋转不同角度(第 10 章的旋转小轮子、第 15 章的复数辐角),妙处是两个被旋转过的向量做点积,结果只依赖角度差,即相对距离——「我离你几个词」被编码进了内积的几何里,外推到训练没见过的长度时表现更好,Llama 系全家采用。两派的共同点:位置终究是以向量运算的身份进场的——本书的语言,一个字都不用换。

多头:同一句话,开几路雷达

一套 QKV 只能表达一种「看人的方式」,但语言里同时存在句法关系、指代关系、语义关联。解法朴素:把 D 维切成 H 份(第 19 章那段「切头四连」的 reshape/transpose),每份 d_h = D/H 维,各自独立跑一遍注意力——H 路雷达并行扫描,各扫各的关系——最后把 H 份输出拼回 D 维再过一个投影(第 07 章分块矩阵的拼接)。可解释性研究确实在真模型里找到了分工明确的头:专盯上一个词的、专追指代链的、专找句尾的,还有著名的「归纳头」(induction head)——在上文里找「上次这个词后面跟了谁」,被认为是大模型上下文学习能力的机制雏形之一。「多头」没有新数学,只有形状语言:(B, L, D) → (B, H, L, d_h),在更小的子空间里批量做同一件事;顺带每头的点积维度从 D 降到 d_h,√d 的温度也按头校准,一举两得。

装配完成:一个 Transformer 块的全家福

把 20、21、22 三章的零件拧到一起,大模型的「一层」全貌如下,每个词你都认识:

翻译:归一化(第 05/20 章,摁住范数)→ 注意力(本章,token 之间沿 L 轴交流)→ 加回残差流(第 03 章的流水账);再归一化 → FFN(第 20 章,每个 token 独立沿 D 轴精加工)→ 再加回。交流一拍、加工一拍,轮流往残差流里存款,如此三十二层,首尾各接一次 embedding 查表(第 06 章)和输出投影(第 01 章的 32000 路打分)。这就是 GPT、Llama、Claude 的骨架全图——没有一个部件在本书 24 章之外。另有一个变体叫交叉注意力(cross-attention):Q 来自一个序列、K/V 来自另一个(翻译模型里「译文查原文」,多模态里「文字查图像块」),机制一字不差,只是换了交流对象。

因果掩码与 KV 缓存:生成时代的两件套

因果掩码:GPT 生成第 i 个词时不许偷看未来,实现是在 softmax 前把分数表的上三角(j > i 的格子)填上 −∞(softmax 后精确为 0)——第 10 章那只下三角通行证兽的完全体。KV 缓存:逐词生成时,历史 token 的 K、V 不随新词改变,存下来复用,每步只算新词的 Q 去点积历史 K——这就是「为什么长对话越聊显存越涨」的形状账(缓存是 (层数, L, D) 级的张量,L 一直在长),也是第 07 章「少搬数据」哲学的推理版。两件套都不是新数学:一个是掩码矩阵,一个是缓存策略,但它们决定了你每次调 LLM API 的延迟和账单——「首 token 慢」(要把整段 prompt 的 KV 算一遍,俗称 prefill)与「后续 token 快」(每步只增量算一列),这条你肉眼可见的体验曲线,就是 KV 缓存的形状在说话。

L² 大账单的省钱流派

那张 (L, L) 分数表是整个架构最贵的家当,围绕它已经形成了几个省钱流派,全部能用本书语言一句话概括:滑动窗口注意力(Mistral 等)——每个词只看最近 w 个邻居,分数表变成带状稀疏矩阵(第 10 章稀疏兽),L² 降到 L·w,远距离靠层数叠出来(和卷积殊途同归);稀疏/块状注意力——按预设模式(跳跃、全局锚点)只算部分格子;低秩流派(Linformer 等)——赌注意力矩阵本身低秩(第 12 章!),先把 K、V 沿 L 轴投影到 k 维再算,L² 变 L·k。三派加上正文说的线性注意力,构成「长上下文军备竞赛」的主要战线——每条战线的武器,都是你行囊里的旧零件:稀疏、低秩、结合律。

小结

Attention = 三次矩阵乘夹一个 softmax:QKᵀ/√d 批发相似度(点积 + 高维温度校准),softmax 逐行分预算(凸组合配方),·V 按配方搬运(动态加权平均);它相对 RNN 与卷积的胜利,是「人人直连一步到位」对「串行接力/固定窗口」的胜利,代价是 L² 账单,滑窗、稀疏、低秩、结合律四路省钱流派因此而生;Q、K、V 是同一向量的三种受训人格,位置靠正弦向量或 RoPE 旋转注入,多头是切子空间开多路雷达,因果掩码是下三角通行证,KV 缓存是形状账上的 memoization;一个 Transformer 块 = 归一化、交流(Attn)、存款、归一化、加工(FFN)、再存款,三十二层循环——全图之内,无一零件出自本书之外。大模型的心脏,你已经全程手摸过一遍——公式、代码、Demo 三个入口都走通了。下一章收官应用篇:把「低秩」这条全书最长的伏笔——从外积薄片到 LoRA——一次收线,顺便看看量化和 embedding 算术,收完那条线,就该画那张随书带走的地图了。