Transformer 原理 2
上一篇把 Q、K、V 和多头注意力拆开了,但那是局部。这一篇把整条链路走完——文字进去,下一个 Token 出来,中间每一步的形状变化都标清楚;最后再按 GPT-3 175B 的真实配置,把「1750 亿参数」这个数自己算一遍。
整套图用的是 GPT-3 175B 的示例尺寸:
1 | 96 个 Transformer 层 |
1. 全流程总览
1 | 文字 |
一句话概括:
每一层先让 Token 之间交流,再让每个 Token 自己加工;连续做 96 轮之后,预测下一个 Token。
2. 输入是怎么变成 [3, 12288] 的
第一步:Token 化
拿三个 Token 举例:
1 | 人、类、简 |
Token 化以后先变成编号:
1 | 人 → Token ID 1 |
此时只是整数编号,还没有任何语义。
第二步:Embedding 查表
拿编号去 Embedding 参数表里查对应的向量,每个 Token 查到一个 12288 维向量:
1 | Token 0 → 12288 维向量 |
组成矩阵:
1 | X:[3, 12288] |
3 行是 3 个 Token,每行 12288 个数字是一个 Token 的表示。
第三步:加位置信息
只用 Embedding 的话,模型只知道有哪些 Token,不知道先后顺序。而这两句:
1 | 我喜欢你 |
Token 差不多,顺序一换意思就变了。所以要给每个 Token 加上位置向量:
1 | 最终输入向量 = Token Embedding + Position Embedding |
注意是对应位置相加,不是拼接:
1 | 12288 维 + 12288 维 = 12288 维 |
所以加完位置信息,矩阵仍然是 [3, 12288],形状没变。
3. 一层里发生了什么
一个完整的 Transformer 层,主要做两件事:
1 | 1. 多头自注意力:Token 之间交换信息 |
先看多头注意力。
4. 一个注意力头内部
在当前示例里:
1 | 输入:[3, 12288] |
一个注意力头拥有自己的 Wq、Wk、Wv,它们把每个 12288 维向量映射成 128 维:
1 | Q = X·Wq |
所以一个头得到:
1 | Q:[3, 128] |
用 Q 和 K 算相关性
Q 中的每个 Token 都要与 K 中的所有 Token 比较,3 个 Token 就得到一个 [3, 3] 的相关性矩阵:
1 | K0 K1 K2 |
每一行表示:当前这个 Token 应该分别关注三个 Token 多少。
正式计算写成矩阵乘法就是:
1 | Q × Kᵀ |
还会除以 √128,防止数值太大。
因果 Mask
对 GPT 这种自回归模型,还要再加一层因果 Mask:
1 | 第 1 个 Token:不能看后面 |
这样训练时模型就不能偷看后面的答案。
Softmax 与加权求和
相关性数字还要过一遍 Softmax,让每一行变成一组加起来等于 1 的权重:
1 | 0.1 + 0.7 + 0.2 = 1 |
然后拿这些权重对 V 加权求和。比如第 0 个 Token 的新向量:
1 | 新向量 0 |
因此一个头最终输出:
1 | Head1 输出:[3, 128] |
这就是一次信息聚合——每个 Token 根据注意力权重,吸收其他 Token 的 V 信息。
5. 96 个头拼接与 Wo 融合
上面只是 Head 1。一层里有 96 个头:
1 | Head 1 → [3, 128] |
每个头都有自己学到的 Wq、Wk、Wv,因此能从不同的投影空间观察同一段输入。它们可能分别更擅长捕捉:
1 | 指代关系 |
(同样地,这只是帮助理解,模型并没有被规定「某个头必须负责哪一种关系」。)
因为每个头输出 128 维:
1 | 96 × 128 = 12288 |
所以把 96 个头横向拼接:
1 | [3, 128] × 96 |
对每个 Token 来说:
1 | 头 1 的 128 维 + 头 2 的 128 维 + …… + 头 96 的 128 维 = 12288 维 |
拼起来以后,96 个头的信息还只是并排放着,所以再经过输出矩阵 Wo:
1 | Wo:[12288, 12288] |
计算:
1 | 多头拼接结果 × Wo |
输出仍然是 [3, 12288]。Wo 的作用就是把 96 个头提供的信息重新混合起来。
完整的注意力过程:
1 | 输入 X |
6. 残差连接与 LayerNorm
上面这些手绘图为了简化,基本都没画完整。
实际的层里还会把原来的输入保留下来:
1 | 注意力输出 + 原输入 |
这叫残差连接,可以防止深层网络把原始信息完全冲掉,也有利于训练。
还会用 LayerNorm 稳定数值。GPT-3 论文说明其沿用了 GPT-2 风格的预归一化结构。
先简单记成一句话:
1 | Attention 不是直接覆盖原向量, |
7. FFN:12288 → 49152 → 12288
注意力让 Token 之间交流完之后,每个 Token 分别进入 FFN:
1 | 12288 维 |
即:
1 | 12288 → 12288 × 4 → 12288 |
矩阵形状:
1 | W1:[12288, 49152] |
图上多出来的 +1 或 12289,是课程把偏置 b 合并进矩阵一起画出来的结果。
FFN 对 3 个 Token 分别处理:
1 | Token 0 → FFN → 新 Token 0 |
但它们是共用同一层的 W1、W2,不是各有一套。输出仍然是 [3, 12288]。
为什么要既有 Attention 又有 FFN
1 | Attention: |
所以一个 Transformer 层就是:
1 | Token 之间交流 |
FFN 之后同样有残差连接和归一化。
8. 从最后一个 Token 到 50257 个候选
96 层做完,如果输入是 3 个 Token,最后仍然得到:
1 | [3, 12288] |
但这三个向量已经不是最开始的普通词向量,而是包含了整段上下文信息的向量。
为什么只取第三条
因为当前的任务是:根据前面已经输入的所有 Token,预测紧接着的下一个 Token。
最后那个位置的向量已经通过注意力吸收了前面所有允许看到的信息,所以取:
1 | 最后一个 Token 的最终向量:[12288] |
去预测下一个 Token。
训练时其实每个位置都可以同时预测它后面的 Token;但实际生成下一个 Token 时,我们只关心最后一行。
Linear 映射到整个词表
最后一个 Token 是 12288 维,但模型要回答的是:词表里每一个 Token 成为下一个 Token 的可能性有多大?所以经过一个输出 Linear:
1 | 12288 维 |
按课程用的词表大小:
1 | 12288 → 50257 |
得到 50257 个分数:
1 | Token 0: 5.2 |
这些原始分数叫 logits,还不是概率。
Softmax 变成概率
Softmax 把 50257 个 logits 转成加起来等于 1 的概率:
1 | 「是」: 30% |
然后按解码策略挑一个:
1 | 贪心: 选概率最高的 |
50257 是什么
50257 是词表大小,也就是模型一共认识的 Token 种类数。
不是说模型只认识 50257 个「单词」,而是有 50257 个可选的 Token 编号:
1 | Token ID:0 |
一个 Token 可能是:一个完整单词(apple)、单词的一部分(ing)、一个汉字或汉字的一部分、标点符号、空格加单词,或者特殊标记。比如:
1 | un + believe + able |
这三个部分可能分别是词表里的一个 Token。
为什么刚好是 50257:
1 | 50256 个普通 Token |
其中 <|endoftext|> 可以表示文本结束等特殊含义。
9. 自回归生成与 KV Cache
选出下一个 Token 之后,把它追加到原文后面,再跑一遍模型:
1 | 人工智能是 → 预测「一种」 |
这就是自回归生成:
1 | 一次只生成一个 Token |
实际推理时会用 KV Cache:把前面 Token 已经算出的 K、V 缓存起来,避免每生成一个新 Token 就把全部历史从头算一遍。
10. 1750 亿参数是怎么算出来的
GPT-3 175B 的关键配置:
1 | 层数 L = 96 |
一层 Transformer 有多少参数
一层包含多头注意力(Wq、Wk、Wv、Wo)、FFN(W1、W2)、偏置,以及两个 LayerNorm。
虽然有 96 个头,但把所有头合并起来看,矩阵形状就是:
1 | Wq:[12288, 12288] |
为什么是 12288 × 12288?因为每个头输出 128 维、共 96 个头:
1 | 96 × 128 = 12288 |
一个矩阵的参数量:
1 | 12288 × 12288 = 150,994,944 |
四个矩阵:
1 | 4 × 150,994,944 = 603,979,776 ≈ 6.04 亿 |
再加注意力部分的偏置:
1 | Wq、Wk、Wv 偏置:3 × 12288 |
注意力总参数:
1 | 603,979,776 + 49,152 = 604,028,928 |
FFN 部分:
1 | W1:12288 × 49152 = 603,979,776 |
两个 LayerNorm,各有一个缩放参数 γ 和一个偏移参数 β:
1 | 一个 LayerNorm:2 × 12288 = 24,576 |
一层合计:
1 | 注意力: 604,028,928 |
可以简写成公式:
1 | 每层参数 = 12d² + 13d |
代入 d = 12288:
1 | 12 × 12288² + 13 × 12288 = 1,812,099,072 |
其中:
1 | 4d²:Attention 的 Wq、Wk、Wv、Wo |
96 层一共
1 | 1,812,099,072 × 96 = 173,961,510,912 ≈ 1739.62 亿 |
光 96 个 Transformer 层,就已经接近 1740 亿了。
加上 Embedding
词嵌入——50257 个 Token,每个一个 12288 维向量:
1 | 50257 × 12288 = 617,558,016 ≈ 6.176 亿 |
位置嵌入——上下文长度 2048,每个位置也是 12288 维:
1 | 2048 × 12288 = 25,165,824 ≈ 0.252 亿 |
96 层之后还有一个最终 LayerNorm:
1 | 2 × 12288 = 24,576 |
输出层要不要再算一次
最后确实要做 12288 → 50257 的映射,看起来应该再有:
1 | 12288 × 50257 = 617,558,016 |
但按课程的说法,GPT-3 沿用 GPT-2 架构,输入 Token Embedding 和最终输出矩阵是权重共享的——同一个矩阵,输入时用来查 Embedding,输出时用小隐藏向量乘它的转置。所以这 6.176 亿已经在 Token Embedding 里算过,不再重复计算。
(这一点 §11 会再说明,它不是没有争议的。)
全部相加
1 | 96 个 Transformer 层:173,961,510,912 |
换算:
1 | 174,604,259,328 = 174.604259328B = 1746.04259328 亿 |
四舍五入:
1 | ≈ 175B |
按 GPT-3 的公开架构算,得到约 1746.04 亿参数。官方把它四舍五入称为 1750 亿。
可以背的简化公式
忽略很小的偏置和 LayerNorm:
1 | 总参数 |
其中 12d² 拆开是:
1 | 4d²:注意力 |
代入 GPT-3:
1 | 96 × 12 × 12288² = 173,946,175,488 |
1 | ≈ 1745.89 亿 |
11. 几处需要再精确一点的地方
一、位置编码:课件画的是正弦余弦,GPT-3 用的是学出来的
图上那张 PE 公式(sin、cos、10000^(2i/d_model))来自 2017 年原始 Transformer 的固定正弦位置编码——它是一个确定性函数,没有参数。
但 GPT-2、GPT-3 用的是可学习的位置嵌入表,GPT 原论文写得很直接:
We used learned position embeddings instead of the sinusoidal version proposed in the original work.
这两件事不能同时成立。而它正好解释了 §10 里为什么会有 2048 × 12288 = 25,165,824 个位置参数——正弦编码是固定的,没有参数可算;只有学出来的嵌入表才有参数量。
所以看到那张公式图时,把它当成「原始 Transformer 的设计」来看,不要和 GPT-3 的尺寸混成一个具体实现。课程里其实也提醒了这一点。
二、权重共享这件事是有争议的
§10 里说「GPT-3 沿用 GPT-2 架构、输入 Embedding 和输出矩阵共享」,所以最后那 6.176 亿不重复计算。这个说法网上分两派:
- GPT-2 确实是共享的,OpenAI 的 GPT-2 实现里输出 logits 直接用输入词嵌入矩阵
wte的转置,这是可以查证的。 - 但 GPT-3 的论文并没有写明这一点。做可解释性研究的一派(Neel Nanda 等)认为 GPT-3 恰恰没有共享权重,理由包括 GPT-3 和 GPT-J 在这类现象上表现一致。
如果没共享,总量就要多 6.176 亿,约 1752 亿。
两边都四舍五入到 175B,所以官方那个数不受影响。但要清楚:「1746.04 亿」这个精确值只在权重共享成立的前提下才对。
三、4× 是 GPT-3 的配置,不是通用规则
FFN 中间层 49152 = 4 × 12288,这个 4 倍关系在 GPT-2/GPT-3 这一代里是标准配置,但别把它当成所有大模型的规则——后来很多模型用的是其他比例(8/3 倍、或者用门控结构)。算参数量时以模型自己的配置为准。
12. 小结
一、整条链路可以压成一句话: Embedding 把文字变成数字;位置编码告诉模型顺序;Attention 负责 Token 之间交流;FFN 负责每个 Token 自己加工;96 层负责反复深化;Linear + Softmax 负责从词表里选出下一个 Token。
二、形状走一圈会回到原点。 输入 [3, 12288],注意力中间拆成 96 个 [3, 128] 再拼回 [3, 12288],FFN 中间扩到 [3, 49152] 再压回 [3, 12288]。每一层的输入输出形状相同,所以才能摞 96 层。
三、位置信息是相加,不是拼接。 加完形状不变,仍是 [3, 12288]。
四、因果 Mask 是自回归的关键。 它让训练时每个位置只能看自己和前面,从而可以一次并行地监督所有位置,同时又不泄露答案。
五、输出只取最后一个 Token 的向量。 训练时每个位置都在预测,但生成时只关心最后一行——它已经吸收了前面全部允许看到的信息。
六、50257 是词表大小,不是单词数。 一个 Token 可能是半个单词、一个汉字,或者标点。
七、1750 亿参数的算法:L × 12d² + Vd + n_ctx d。 96 层吃掉 99% 以上(约 1739.6 亿),其余是词嵌入 6.176 亿、位置嵌入 0.252 亿。参数几乎全在层里,而层里的参数又几乎全是 d² 那两个矩阵堆出来的。
八、三个要留意的点: 课件画的正弦位置编码属于原始 Transformer,GPT 系列用的是学出来的嵌入表;权重共享在 GPT-3 上是存疑的;FFN 的 4 倍不是通用规则。
参考
- Attention Is All You Need(2017,正弦位置编码出自这里):https://arxiv.org/abs/1706.03762
- Language Models are Few-Shot Learners(GPT-3,配置见 Table 2.1):https://arxiv.org/abs/2005.14165
- GPT-2 论文(learned position embeddings、预归一化):https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf
- The Illustrated GPT-2:https://jalammar.github.io/illustrated-gpt2/
- 上一篇:Transformer 原理 —— https://wywzixin.github.io/llm/transformer/
Transformer 原理 2