数据截至 (上游 commit eb980a5c9eea)
注意力机制 — 把「相关」变成一次加权求和
这一章讲三件事: 注意力到底在做哪一个动作;那条公式是怎么一步一步被逼出来的; 以及在它之上加的三个变化——自注意力、掩码、多头——各解决什么。 这是全书的枢纽章。 后面七章里的每一个模型,内部核心都是这一章讲的这一个动作。 只需要第 01 章的一个结论:语义相近的词,词向量也相近。
1. 先看现象:它怎么知道「它」指的是谁
你问模型:
「桌上有一杯水和一本书,它被打翻了。请问被打翻的是什么?」
它答得出来。可这句话里,「它」离「水」隔了六个字,离「书」隔了两个字—— 按距离算,「它」离「书」更近,但答案是「水」。
模型凭什么知道? 因为「打翻」这个词和「水」相关、和「书」不相关。 换句话说,模型必须能在一句话内部,给任意两个位置之间算出一个「有多相关」的数。
这一章讲的就是那个数怎么算出来。
书里给注意力的一句话定义,值得原样记住1:
通过计算查询与键的相关性为真值加权求和,从而拟合序列中每个词同其他词的相关关系。
这句话里有三个陌生词(查询、键、真值),下一节就用一把字典把它们讲清楚。
2. 顶层全景:一条公式,四步推出来
① 我想查什么 → Query(查询) 库里有什么 → Key(键)+ Value(值)
② 「有多像」怎么算 ────────────────→ 点积:两个向量对应位置相乘再相加
③ 一堆分数怎么变成比例 ────────────→ Softmax:压成一组加起来等于 1 的比例
④ 分数太大会把 Softmax 推到极端 ───→ 除以「每个头的维度」的平方根
────────── ──────────────────────────────────────────────────
结果:注意力 = Softmax( Q·Kᵀ / √d ) · V
图说:这四步不是并列的技巧,是一条被逼出来的链——
每一步都在补上一步留下的洞。这一节剩下的篇幅就是把这条链走一遍。
公式看起来吓人,拆开只有三个动作:算相似、变比例、按比例混。
3. 第一步:三个变量,用一把字典就能讲完
这一节解决那三个陌生词。
注意力有三个核心变量:Query(查询值)、Key(键值)、Value(真值)2。
书里的讲法是拿一把普通的字典当起点3:
{ "apple": 10,
"banana": 5,
"chair": 2 }
- 字典的键,就是注意力里的 Key;
- 字典的值,就是注意力里的 Value;
- 你要查的东西,就是 Query。
普通字典的做法:精确匹配
Query 是 "apple",那就直接匹配到键 apple,取出值 10。要求是完全相等4。
注意力的做法:不要求相等,要求「有多像」
现在把 Query 换成 "fruit"。 字典里没有这个键——但你希望
apple 和 banana 都被匹配到,而 chair 不该被匹配到5。
办法是别选一个,而是给每个键分配一个比例,再把这些值按各自的比例混起来 (把每个数乘上自己那份比例再全部相加,这个动作数学上叫加权求和,平均一下就叫加权平均):
| 键 | 分到的比例 | 为什么 |
|---|---|---|
| apple | 0.6 | 和 fruit 很相关 |
| banana | 0.4 | 和 fruit 也相关 |
| chair | 0 | 完全不相关 |
结果 = 10×0.6 + 5×0.4 + 2×0 = 8。
这组比例就叫注意力分数——书里的定义是:「为了查询到 Query,我们应该赋予给每一个 Key 多少注意力」6。
到这里,注意力这个动作已经讲完了。剩下的全部内容,只是在回答一个问题: 那组比例怎么算出来。
换到真实场景是什么样
书里给了一个更贴近实际的例子7:你有一篇新闻报道,想找出报道的时间。
- Query ≈ 「时间」「日期」这类意思的向量;
- Key 和 Value 就是整篇文本;
- 对 Query 和 Key 做运算得到一组比例,这组比例反映了从这个 Query 出发, 文本里每一个标记该分到多少注意力;
- 再拿这组比例和 Value 运算,得到的就是答案。
注意:书里特意括注了一句——为了好懂,这里用文本来表示 Query, 它实际上是一个稠密的向量。 从这里往下,Query、Key、Value 全都是向量或矩阵(矩阵就是把很多个等长的向量摞在一起形成的长方形数表)。
4. 第二步:「有多像」怎么算——点积
这一节回答上一节留下的洞:0.6 和 0.4 是怎么来的。
这里要用到第 01 章挣来的那个性质:经过合理训练,语义相近的词, 它们的词向量在空间里距离更近。
衡量「近」最常见的是欧氏距离,但注意力用的是另一种:点积8。
点积就是两个向量对应位置相乘,再把所有乘积加起来,得到一个数。
书里给的判据很简洁9:
语义相似的两个词,对应词向量的点积应该大于 0;语义不相似的,点积应该小于 0。
为什么选点积不选欧氏距离
书里没解释,我们补一句:
判断(我们的,不是书里的): 选点积的实际理由是它能写成一次矩阵乘法。 一个 Query 对 N 个 Key 的全部点积,等于把 Query 这个行向量乘上「所有 Key 堆起来的矩阵」—— 一步算完,而且正好是显卡最擅长的那种运算。欧氏距离要先做减法再平方再开根, 拼不进一次矩阵乘法里。 如果错,会错在: 如果某种硬件上欧氏距离也能一次算完一大片、代价和点积相同, 这条理由就只剩「历史惯性」了。
一次算完:堆成矩阵
把所有 Key 的词向量竖着堆起来变成一个矩阵(书里写作 K), 那么 Query 和每一个 Key 的相似度,一次矩阵乘法就全出来了10。
q · Kᵀ = [ q·k ₁ , q·k₂ , q·k₃ , … ]
└──── 一次乘法,拿到跟每个键的相似度 ────┘
图说:这是注意力能跑得快的物理原因——所有位置的相似度是同时算出来的,
不像循环网络那样必须一个一个来。
5. 第三步:分数变比例——Softmax
这一节回答:点积出来的是一堆大小不一、还可能是负数的数,怎么变成「比例」。
答案是过一个 Softmax 层,它把任意一串数压成一组都为正、加起来正好等于 1 的比例11。
为什么必须加起来等于 1: 因为下一步要拿它们做加权平均。 这些比例加起来不为 1 的话,输出的数值大小会随着序列长度乱漂,后面的层就没法稳定地接。
到这一步,基本形状已经出来了:
注意力 = Softmax(Q 和 K 的相似度) × V
书里还补了两小步12:每个键对应的值本身可以是一串数(向量),不一定只是一个数; 而且可以一次查询多个 Query——把多个 Query 的向量也堆成矩阵 Q,一次全算完。
6. 第四步:为什么要除以一个数
这一节讲公式里最不直观的那一项,也是最容易被跳过的一项。
先看会出什么问题
Query 和 Key 的维度(第 01 章说过:一串数里有几个数,那个「几」就是维度)如果比较大, 点积的结果数值也会很大。 数值一大,Softmax 就会被推到极端——几乎全部比例集中到一个位置上,其余全成 0。
梯度就是训练时用来决定每个参数该往哪个方向调、调多少的那个量。 书里的说法是:维度较大时「Softmax 放缩时就非常容易受影响,使不同值之间的差异较大, 从而影响梯度的稳定性」13。
Softmax 一旦被推到极端,输出对输入几乎不再变化,梯度就趋近于 0——这一层就学不动了。
怎么解决
把点积结果除以维度的平方根,把数值拉回一个温和的范围14。
维度 d 越大 → 点积的数值越大 → Softmax 越极端 → 梯度越小 → 学不动
└──── 除以 √d,把它拉回来 ────┘
图说:这一项不是数学上的装饰,是让这一层能被训练的前提。
到这里,注意力的核心公式就完整了15。
7. 主走查:给「它」算一遍,每一步都带数
上面四步各自都讲过了,但你还没见过它 们连起来跑一遍。这一节就跑一遍。 用的输入就是第 1 节那句话——「桌上有一杯水和一本书,它被打翻了」, 问题也还是那个:处理到「它」的时候,模型凭什么把「水」取走大半。
先声明:这一节从头到尾的数值全是为演示编的,不是真实数值。 真实模型里每个向量有几百个数,手算不动;这里把它缩到 4 个数,好让你能自己验算。 编的是数值,不是步骤——步骤和真实模型一模一样。
第 0 步:摆出四个向量
第 1 节那句话切完之后,我们只盯三个位置:水、书、打翻;而正在被处理的位置是它。
| 谁 | 它的查询向量 Q | 它的键向量 K | 它的值向量 V |
|---|---|---|---|
| 它(正在处理的位置) | [1, 0, 1, 0] | — | — |
| 水 | — | [3, 1, 3, 2] | [2, 0, 4, 1] |
| 书 | — | [0.6, 2, 0.6, 1] | [0, 3, 0.6, 0] |
| 打翻 | — | [2, 0, 2.8, 1] | [1, 1, 0, 3] |
这四个向量都是同一句话自己算出来的——同一个输入乘三个不同的参数矩阵, 就得到 Q、K、V 三份。这就是第 9 节要讲的「自注意力」,它在主走查上占的就是这一步。
第 1 步:点积—— 「它」和谁更像
按第 4 节的规则:对应位置相乘,再全部加起来。
它·水 = 1×3 + 0×1 + 1×3 + 0×2 = 3 + 3 = 6.0
它·书 = 1×0.6 + 0×2 + 1×0.6 + 0×1 = 0.6+0.6 = 1.2
它·打翻 = 1×2 + 0×0 + 1×2.8 + 0×1 = 2 + 2.8 = 4.8
图说:三个数一出来,「它更像水」这件事就已经是个能比大小的量了。
6.0 是 1.2 的五倍——这就是本章开头那个问题的全部答案。
第 2 步:除以 √d——把数值拉回温和的范围
这里的 d 是每个向量里有几个数,也就是 4;√4 = 2。按第 6 节的规则,三个分数各除以 2:
6.0 ÷ 2 = 3.0 1.2 ÷ 2 = 0.6 4.8 ÷ 2 = 2.4
为什么这一步不是白做的: 拿下一步的算式对照着看就知道—— 不除的话进 Softmax 的是 6.0 和 1.2,差 4.8;除完是 3.0 和 0.6,差 2.4。 差距减半,Softmax 就不会把比例全堆到第一名头上。
第 3 步:Softmax——把三个数压成加起来等于 1 的比例
按第 5 节的规则,先给每个数算一次「以 e 为底的乘方」,再各自除以三者之和:
3.0 → 20.09 0.6 → 1.82 2.4 → 11.02 合计 32.93
水 = 20.09 ÷ 32.93 = 0.61
书 = 1.82 ÷ 32.93 = 0.06
打翻= 11.02 ÷ 32.93 = 0.33 三者相加 = 1.00
图说:这就是第 3 节那张 apple/banana/chair 表里 0.6 / 0.4 / 0 的真身。
那三个数是书里直接给的,这三个是算出来的。
第 4 步:按比例混——「它」这一步的新表示
拿这三个比例去乘各自的值向量,逐位相加:
第 1 位: 0.61×2 + 0.06×0 + 0.33×1 = 1.22 + 0 + 0.33 = 1.55
第 2 位: 0.61×0 + 0.06×3 + 0.33×1 = 0 + 0.18 + 0.33 = 0.51
第 3 位: 0.61×4 + 0.06×0.6 + 0.33×0 = 2.44 + 0.04 + 0 = 2.48
第 4 位: 0.61×1 + 0.06×0 + 0.33×3 = 0.61 + 0 + 0.99 = 1.60
「它」的新表示 = [1.55, 0.51, 2.48, 1.60]
图说:把它和「水」的值向量 [2, 0, 4, 1] 并排看——形状明显更像「水」而不是「书」。
「它」这个位置里,已经掺进了六成的「水」。
这就是模型答得出「被打翻的是水」的物理原因。
这一节到此为止,主走查已经走过四个承重机制:点积、缩放、Softmax、自注意力。 剩下两个——掩码和多头——在第 10、11 节接着用这同一组数走完。
8. 用代码看它有多小
书里用 PyTorch 实现了一遍,核心只有四行16。我们不抄代码,只讲那四行在干什 么:
| 第几步 | 干什么 |
|---|---|
| 1 | 取出键向量的维度(键和值的维度相同) |
| 2 | Query 乘上 Key 的转置(转置就是把一个矩阵的行和列对调过来),再除以维度的平方根 |
| 3 | 过 Softmax,把分数变成加起来等于 1 的比例 |
| 4 | 比例乘 Value,做加权求和 |
中间还可以插一个 dropout——训练时随机把一部分比例丢掉,防止模型死记硬背17。
这一节想让你记住的只有一件事:
注意力机制的核心计算,只有四行。 后面所有的复杂度,都不在这四行里,而在「怎么用它、用几次、给它喂什么」。
9. 变化一:自注意力——三个变量来自同一句话
这一节回答:上面讲的是「拿 A 去查 B」,那「一句话自己查自己」是怎么回事。
先看两种用法的差别
注意力的本质是对两段序列做相似度计算,找出一段序列的每个元素对另一段序列的每个元素的相关度18。 那 Q、K、V 分别从哪儿来,就有两种拼法:
| 拼法 | Q 从哪来 | K、V 从哪来 | 拟合的是什么 |
|---|---|---|---|
| 经典注意力 | 一个序列 | 另一个序列 | 两段序列之间的关系。比如翻译时,Q 来自解码器(负责往下写的那一半)的输入,K、V 来自负责读懂输入的另外那一半的输出19 |
| 自注意力 | 同一个输入 | 同一个输入 | 一句话内部,每个标记和其他所有标记的关系20 |
自注意力具体怎么做
同一个输入,分别乘三个不同的参数矩阵,得到 Q、K、V。 三个矩阵不同,所以算出来的三个东西也不同——虽然它们来自同一句话。
在代码里,自注意力就是把同一个输入当三个参数传进去21:attention(x, x, x)。
为什么这个变化这么重要: 本章开头那个「它指的是谁」的问题 (桌上一杯水和一本书,「它」被打翻了),正是靠自注意力解决的—— 它让句子内部任意两个位置直接连上,不用一站一站传。
10. 变化二:掩码——把「未来」盖住
这一节回答一个反直觉的设计:为什么要故意让模型看不见后面的字。
先看现象:模型是怎么被训练的
模型学的是接龙:给前面的字,猜下一个字。
拿 【BOS】I like you【EOS】 这句话当例子(BOS 和 EOS 是句首、句尾的特殊标记),
学习过程要走四步22:
Step 1:输入【BOS】 → 输出 I
Step 2: 输入【BOS】I → 输出 like
Step 3:输入【BOS】I like → 输出 you
Step 4:输入【BOS】I like you → 输出【EOS】
问题来了:这四步是串行的,必须先做完第 1 步才能做第 2 步。 而注意力相对循环网络的核心优势就是可以并行——这么训练,等于把优势原样丢掉了23。
掩码怎么解决它
把上面四步同时送进模型,每一行把「还不该看到」的位置盖住24:
<BOS> 【MASK】【MASK】【MASK】【MASK】
<BOS> I 【MASK】【MASK】【MASK】
<BOS> I like 【MASK】【MASK】
<BOS> I like you 【MASK】
<BOS> I like you </EOS>
图说:每一行仍然是「只看前面、猜下一个」,但五行可以一次性并行算完。
串行变并行,靠的就是这块遮布。
书里点破了这块遮布的形状:它就是一个和文本序列等长的上三角矩阵25。
实现上的一个小巧思
怎么「盖住」?书里的做法是26:
- 造一个上三角矩阵,该盖的位置填负无穷,其余位置填 0;
- 把这个矩阵直接加到注意力分数上;
- 再过 Softmax。
为什么这样就成了: 负无穷的位置经过 Softmax 之后会被压成 0—— 比例是 0,那个位置的 Value 就完全不参与加权求和27。而下三角区域加了 0,分数不变。
这是一个值得记住的工程手法:「屏蔽某些位置」这个逻辑动作, 被翻译成了「加一个负无穷」这个纯算术动作。 逻辑判断没法在显卡上高效并行,加法可以。
主走查续:在同一组数上加一次掩码
第 7 节那组数还在,这里直接接着用。 在原句「桌上有一杯水和一本书,它被打翻了」里, 「打翻」排在「它」后面——所以只要这一层加了掩码,处理「它」的时候「打翻」就该被盖住。
第 7 节第 2 步算完的三个分数是 3.0 / 0.6 / 2.4。往「打翻」那一位加上负无穷:
加掩码前: 水 3.0 书 0.6 打翻 2.4
加掩码后: 水 3.0 书 0.6 打翻 -inf
└ Softmax:20.09 / 1.82 / 0 合计 21.91
└ 比例: 0.92 / 0.08 / 0.00 ← 「打翻」那一份精确地变成了 0
新表示 = 0.92×[2,0,4,1] + 0.08×[0,3,0.6,0] + 0×[1,1,0,3]
= [1.84, 0.24, 3.73, 0.92]
图说:和第 7 节不加掩码那次 [1.55, 0.51, 2.48, 1.60] 并排看——
「打翻」那 0.33 份被整个抽走,剩下的两份按比例补上。
这三个数把两条路线的差别摆在了同一张纸上: 不盖住未来的那一条在处理「它」时就看得见「打翻」,所以那 0.33 份证据当场就用上了; 盖住未来的那一条在处理「它」时看不见,它得等轮到「打翻」这个位置的时候, 再回头去看「它」——同一份关联,只是被推迟到了后面一步。
这些数仍然是为演示编的,不是真实数值。
用了掩码的自注意力就叫掩码自注意力。 第 04 章会看到, 这一个零件决定了两条路线的全部差别: GPT(就是 OpenAI 那一系列「只管往下接着写」的模型)走的是盖住未来那一条, BERT(第 01 章提过,那类「只负责读懂一句话」的模型)走的是不盖的那一条。
11. 变化三:多头注意力——一次只能拟合一种关系
多头注意力就是把同一段文字同时算好几遍注意力,每一遍关注一种不同的关系,最后把结果拼起来。
这一节回答:为什么要这样把同一件事算很多遍。
先看问题
一次注意力计算,只能拟合一种相关关系28。
可是一句话里的关系有很多种:谁是谁的主语、哪个形容词修饰哪个名词、 哪两个词在指同一个东西……单一的一次注意力很难全面拟合这些关系。
怎么做
同时对同一段文字做多次注意力计算,每次用一组不同的参数矩阵, 所以每次会拟合出不同的关系;最后把多次结果拼接起来, 再乘一个固定的矩阵把它们重新混合一遍(这一步叫线性变换)输出29。
书里说原论文用实验证实了这一点:不同的注意力头确实拟合了语句中不同层次的相关信息30。
实现上的关键技巧
最直观的写法是:n 个头就准备 n 组参数矩阵,各算各的,最后拼起来。 但书里指出这样写时空开销都偏高,实际做法是31:
用三个大矩阵代替 n 组小矩阵。 依据是:「矩阵内积再拼接」等同于「拼接矩阵再内积」。
也就是说,与其做 n 次小矩阵乘法再拼,不如做 1 次大矩阵乘法再切开—— 数学上完全等价,但对显卡友好得多。
先定一个词:隐藏层有多「宽」
下面这条约束里有个新词,它从这里一直用到全书最后一章,先一次说清。
第 01 章说过,神经网络是一层层的加工站,每一层里排着一排单元; 夹在输入和输出中间、不直接露在外面的那些层,就叫隐藏层——「隐藏」两个字没有任何神秘含义。 一层里有多少个单元,就决定了这一层往下送出多少个数。 而 Transformer 有个特点:一个标记从头穿到尾,在每一层里始终被表示成同样多的一串数。
这串数里有多少个数,就叫这台机器的「隐藏层维度」—— 说人话就是这台机器有多「宽」。 隐藏层维度 768,意思就是:从头到尾,每个标记都用 768 个数来表示。
全书只用「隐藏层维度」这一个叫法。 别处还有一堆别名: 模型维度(就是隐藏层维度换了个说法)、每层宽度、d_model、hidden_size—— 说的都是这同一个数,别被叫法晃到。
回到约束本身
有一个硬性约束要记住:隐藏层维度必须是头数的整数倍32。 因为要把输出平均切成头数那么多份,除不尽就切不开。 每个头分到的维度 = 隐藏层维度 ÷ 头数33。
隐藏层维度 768,12 个头
└ 每个头拿 64 维
└ 12 个头各自算一遍注意力,得到 12 个 64 维的结果
└ 拼回 768 维,过一个线性层输出
图说:多头不是「变大」,是「把同样宽度切开,让不同的切片去关注不同的关系」。
参数量和单头基本相同。
主走查收尾:把那 4 维切成 2 个头
768 除以 12 手算不动,所以还用第 7 节那组 4 维的数,切成 2 个头、每头 2 维。 真实模型是 768 切 12、每头 64 维;这里是 4 切 2、每头 2 维——除法是同一个除法。
切法很直白:每个向量的前两位归第 1 个头,后两位归第 2 个头。 注意除的那个数也跟着变了:除以「每个头的维度」的平方根,也就是 √2 ≈ 1.41,不再是 √4 = 2。
第 1 个头(只拿前两位) 它 [1,0]
点积: 水[3,1]→3.0 书[0.6,2]→0.6 打翻[2,0]→2.0
÷1.41: 2.12 / 0.42 / 1.41 → Softmax: 0.60 / 0.11 / 0.29
第 2 个头(只拿后两位) 它 [1,0]
点积: 水[3,2]→3.0 书[0.6,1]→0.6 打翻[2.8,1]→2.8
÷1.41: 2.12 / 0.42 / 1.98 → Softmax: 0.49 / 0.09 / 0.42
图说:同一句话、同一个「它」,两个头给出的比例不一样——
第 1 个头几乎全押在「水」(0.60),第 2 个头把「打翻」抬到了 0.42。
「不同的头关注不同的关系」这句话,落到数上就是这两行的差别。
两个头各自算完 2 维的结果,拼回 4 维,再过一个线性层输出。
这些数同样是为演示编的,不是真实数值。
到这里,这一章六个承重机制在主走查上各占一步: 点积(第 7 节第 1 步)、缩放(第 2 步)、Softmax(第 3 步)、 自注意力(第 0 步:三个向量都 来自同一句话)、 掩码(第 10 节末尾那一次 -inf)、多头(上面这两行)。
12. 作者的判断与证据
这一节把「书里给了理由的」和「书里只是这么做的」分开。
| 说法 | 书里给了什么 | 该怎么看 |
|---|---|---|
| 点积可以衡量语义相似度 | 给了理由——建立在词向量的定义上,并给出了大于 0 / 小于 0 的判据 | 可以采信,但注意这是充分性论证不是必要性论证:书没说为什么不用别的相似度 |
| 除以维度平方根是为了梯度稳定 | 给了机制层面的解释 | 可以采信 |
| 多头能拟合不同层次的关系 | 引了原论文的实验,并给了可视化图 | 有实证支撑,但书里没有复现,也没给数字 |
| 掩码把串行变并行 | 给了完整推演(四步串行 → 五行并行) | 这是全书解释得最漂亮的一段 |
| 「矩阵内积再拼接等同于拼接矩阵再内积」 | 只给了结论,让读者自行模拟 | 结论正确,但书里没证 |
判断(我们的,不是书里的): 这一章最值得带走的,不是公式, 而是那条**「每一步都在补上一步的洞」的推导方式**—— 先有「要算相关性」这个需求,才有点积; 有了点积才发现要归一化(就是把一堆大小不一的分数换算成加起来正好等于 1 的比例),于是有 Softmax; 有了 Softmax 才发现维度一大就崩,于是有缩放。 一个成熟的公式,几乎总是这样一层层补出来的,不是一次想出来的。 以后看到任何一条带着一堆系数的公式,可以试着问:哪一项是主体,哪一项是补丁,补的是什么洞。 如果错,会错在: 有些公式确实是从一个封闭的数学推导里一次性得到的 (比如第 03 章那个用三角函数一次性算出来的、用来标记字词先后顺序的编码),硬套「补丁」框架会看不出它的整体性。
13. 边界与局限
① 书里完全没提的一件大事:注意力很贵
自注意力要算「每个位置对每个其他位置」的相似度,所以它的计算量随序列长度的平方增长。 序列翻一倍,这笔计算量变四倍;第 7 节那三个点积,换成 512 个位置就是 512×512 = 26 万多次。
但要说清一件事,免得后面记岔:随长度平方增长的是「训练时算注意力分数」这笔账, 不是所有和长度有关的账。 生成文字时另有一笔账——把前面每个位置的键和值存下来 反复用的那块地方——它随长度线性增长,翻一倍就是两倍,不是四倍。 两笔账的形状不一样,省法也不一样,第 06 章第 5 节把它们拆开算给你看。
补充(不在书里):2022 年的 FlashAttention 给了一个不改结果、只改访存方式的解法。 先说清一个词:显存(就是显卡自己带的那块内存)——模型跑起来,全靠它装数据。
FlashAttention 把注意力算法做成对显存读写敏感的形式, 办法叫分块——把一整块大矩阵切成一小块一小块,算完一块再算下一块。
这么切是为了少搬东西。芯片上还有一块又小又快的暂存区叫缓存,数据得先从显存搬进缓存才能算; 分块之后,同一份数据在两者之间来回搬的次数大幅减少。
结果是在完全不做近似的前提下明显加速(论文报告 GPT-2 上约 3 倍)34。
书里第 06 章的代码里出现了一个 flash_attn 开关,但正文一个字都没解释它是什么。
② 注意力眼里没有顺序
自注意力对每个标记来说,其他各个位置都是平等的—— 「我喜欢你」和「你喜欢我」在它看来完全一样35。
这是它换来并行的代价。 解法是额外把位置信息编码进去,第 03 章整节讲这件事。
③ 书里这一节的代码有一处对不上
书里第 2.1.6 节给了多头注意力的完整代码,那一段有一处自相矛盾: 建这个部件的时候,代码判断过「这次要不要掩码」,却没把这个判断结果记下来; 等到真正开始算的时候,它又去读那个从来没被记下来的值36。这一读必然出错。
顺带把三个编程词一次说清——后面几章的脚注里还会遇到:
| 词 | 大白话 |
|---|---|
| 类 | 一种零件的图纸。照同一张图纸可以造出很多个同款零件 |
| 构造函数 | 照图纸造零件时先跑一遍的那段设置代码(定尺寸、备好要用的东西) |
| 前向计算 | 零件造好之后,真正把数据从头算到尾的那段代码 |
判断(我们的,不是书里的): 照书里那段代码原样跑,这里会报错。 这类问题在教程型书里很常见——书里的代码块是从可运行工程里摘出来讲解的片段, 摘的时候丢了行。 读法应该是:代码块当伪代码读,要跑就去它的配套仓库拿完整脚本。 如果错,会错在: 如果配套仓库里的对应文件也缺这一行,那就不是「摘漏了」而是真的有 bug—— 我们没有核对仓库,只核对了书。
④ 一个用词上的提醒
书里管注意力的第三个变量叫**「真值」(Value)。 「真值」这个中文词在别处通常指「正确答案」,在这里不是那个意思**—— 它只是「键对应的那个内容」。别把它和训练里的标签混起来。
14. 可带走的
- 注意力的全部动作就一句话:给每个键算一个「有多像」的分数,再按分数把所有值加权平均;
- Query 是你要查的,Key 是索引(拿来和你要查的东西比对的那一栏),Value 是内容——普通字典要求键完全匹配,注意力不要求;
- 「有多像」用点积算,依据是词向量的性质:语义相近则点积为正;
- Softmax 把一串分数压成加起来等于 1 的比例,这样加权平均的数值才稳定;
- 除以维度的平方根不是装饰——不除,维度一大 Softmax 就极端,梯度消失,这一层学不动;
- 核心实现只有四行;后面所有复杂度都在「怎么用它」,不在这四行里;
- 自注意力 = Q、K、V 都来自同一句话,让句内任意两个位置一步相连;
- 掩码 = 把未来位置盖住,目的不是「不让它作弊」,而是让串行的训练变成并行的训练;
- 盖住的手法很妙:加负无穷,过 Softmax 后自动变 0——逻辑判断被翻译成了算术;
- 多头 = 把同样的宽度切开,让不同切片关注不同类型的关系;参数量并不增加;
- 隐藏层维度 = 每个标记在模型内部用多少个数表示,也就是这台机器的「宽度」;它必须能被头数整除,否则切不开;
- 注意力的计算量随序列长度平方增长——这条书里没讲; 但要和另一笔账分开:生成时那块存键和值的地方是线性增长的(第 06 章第 5 节拆开算);
- 注意力眼里没有顺序,位置信息必须另外补。
15. 原文地图
| 主题 | 原书章 | 原文位置 |
|---|---|---|
| 注意力的一句话定义 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:46(搜「拟合序列中每个词同其他词」) |
| 三个核心变量与新闻报道的例子 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:40(搜「Query（查询值）」) · text/06-ch02-transformer.txt:41(搜「新闻报道」) |
| 字典类比:精确匹配 vs 加权混合 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:49(搜「字典为例」) · text/06-ch02-transformer.txt:58(搜「精确的字符串匹配」) · text/06-ch02-transformer.txt:60(搜「fruit」) · text/06-ch02-transformer.txt:67(搜「banana」) |
| 注意力分数的定义 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:73(搜「注意⼒分数」) · text/06-ch02-transformer.txt:74(搜「相关性越」) |
| 点积衡量相似度 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:78(搜「欧式距离」) · text/06-ch02-transformer.txt:80(搜「点积应该⼤于0」) · text/06-ch02-transformer.txt:85(搜「堆叠形成的矩阵」) |
| Softmax 与多 Query | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:86(搜「和为 1 的权重」) · text/06-ch02-transformer.txt:91(搜「⼀次性查询多个」) |
| 缩放:为什么除以根号 d | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:95(搜「梯度的稳定性」) · text/06-ch02-transformer.txt:97(搜「核⼼计算公式」) |
| 四行代码实现 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:114(搜「math.sqrt(d_k)」) · text/06-ch02-transformer.txt:117(搜「dropout」) |
| 经典注意力 vs 自注意力 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:127(搜「相似度计算」) · text/06-ch02-transformer.txt:132(搜「Decoder 的输」) · text/06-ch02-transformer.txt:135(搜「self-attention」) · text/06-ch02-transformer.txt:137(搜「token 对其他所有 token 的关系」) |
| 掩码自注意力:动机与四步串行 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:146(搜「Mask Self-Attention」) · text/06-ch02-transformer.txt:149(搜「历史信息进」) · text/06-ch02-transformer.txt:152(搜「【BOS】I like you【EOS】」) · text/06-ch02-transformer.txt:163(搜「计算效率很低」) |
| 上三角掩码与负无穷技巧 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:165(搜「遮蔽未来信息」) · text/06-ch02-transformer.txt:186(搜「torch.triu」) · text/06-ch02-transformer.txt:197(搜「Softmax 之后会被置为 0」) |
| 多头注意力:动机、实证、实现 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:201(搜「Multi-Head Attention」) · text/06-ch02-transformer.txt:204(搜「原论」) · text/06-ch02-transformer.txt:207(搜「拟合不同层次的相关信息」) · text/06-ch02-transformer.txt:216(搜「时空复杂度均较」) · text/06-ch02-transformer.txt:217(搜「拼接矩阵再内积」) |
| 维度必须被头数整除 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:229(搜「整数倍」) · text/06-ch02-transformer.txt:235(搜「每个头的维度」) |
| 注意力没有顺序信息 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:564(搜「完全相同的」) |
| 代码里 is_causal 没被赋值 | 第二章 Transformer 架构 | text/06-ch02-transformer.txt:253(搜「if is_causal」) · text/06-ch02-transformer.txt:282(搜「self.is_causal」) |