Transformer 是怎么训练的：手算一层注意力的梯度下降

在《Transformer 和 FPGA 都是电路》里，我们用一句话概括过模型是怎么训练的：先算出“错误”对每个参数的坡度，再朝下坡的方向挪一小步。这句话没错，但它跳过了所有细节。真正训练一个模型时，屏幕上看到的是这样一串数：

step  0  loss 0.615
step  1  loss 0.594
step  2  loss 0.572
step  3  loss 0.548
step 10  loss 0.336
step 20  loss 0.076
step 40  loss 0.030

step 是第几步，loss 叫作损失，可以先理解成“模型现在错得有多厉害”。它一步步变小，说明模型越来越对。可是 0.615 是怎么算出来的？电脑在每一步里到底做了什么，才让它变成 0.594？为什么前十步降得慢，第十步到第二十步又降得快？

这串数来自一个真的 Transformer，只不过它小到只有 10 个参数。小到这个程度，每一个中间结果都写得下。不过在打开它之前，我们先用一个更小的东西，把“训练”这件事本身弄清楚。

一、热身：一个旋钮

一台只有一个旋钮的机器

设想一台机器，输入一个数 $x$，输出 $y = w \cdot x$。$w$ 是机器上的一个旋钮，可以拧到任何读数。我们希望输入 2 的时候，输出是 6。

你当然一眼就看出 $w$ 应该是 3。但假装我们不知道，只能从某个读数开始，比如 $w = 1$，一点点去拧。神经网络的训练就是这件事，只不过旋钮有成千上亿个，没有人能一眼看出答案。

先要有一个数，说明“现在差多远”。$w = 1$ 时输出是 2，目标是 6，差了 −4。差值可正可负，我们把它平方，得到一个永远不小于 0 的数，叫作损失：

$w = 1$ 时 $L = 16$；拧对了，$w = 3$，$L = 0$。训练的目标就是把 $L$ 拧小。

导数告诉我们往哪边拧

$L$ 是 $w$ 的二次函数，画出来是一条开口向上的抛物线，谷底在 $w = 3$（图 1）。站在 $w = 1$ 这一点，该往左拧还是往右拧？看脚下的坡就知道了。

坡度就是导数。$L = (2w-6)^2$ 的导数是 $L' = 2(2w - 6) \times 2 = 8w - 24$，$w = 1$ 时等于 −16。负数表示往右是下坡：$w$ 调大一点，$L$ 会变小，而且 $w$ 每增加一点点，$L$ 大约减少这一点点的 16 倍。

于是规则很简单：导数是负的就把 $w$ 调大，是正的就调小，调多少和导数成正比。写成式子：

$\eta$ 读作“伊塔”，叫作学习率，管的是一步迈多大。取 $\eta = 0.05$：$w$ 从 1 变成 $1 - 0.05 \times (-16) = 1.8$，损失从 16 降到 5.76。再算一次导数，$8 \times 1.8 - 24 = -9.6$，于是 $w$ 变成 2.28，损失 2.07。这样一步步走下去，$w$ 依次是 2.568、2.741、2.845……越来越接近 3。

这就是梯度下降：算坡度，往下坡挪一步，再算坡度，再挪一步。注意步子是越来越小的，因为越靠近谷底，坡越缓，导数越小。

步子迈多大

学习率不能随便取（表 1）。

这个例子简单到可以算清原因。新的 $w$ 减去 3，等于旧的 $w$ 减去 3 再乘上 $1 - 8\eta$。$\eta = 0.05$ 时乘的是 0.6，离谷底的距离每步缩到六成；$\eta = 0.125$ 时乘的是 0，一步就到；$\eta = 0.25$ 时乘的是 −1，从谷底一边跳到另一边，距离不变，永远在 1 和 5 之间来回；$\eta = 0.3$ 时乘的是 −1.4，每跳一次离得更远。导数只说明脚下这一点的坡，步子太大，就会跨过谷底。

两个旋钮：偏导数

机器再加一个旋钮 $b$，输出变成 $y = w \cdot x + b$。输入还是 2，目标还是 6，损失是 $L = (2w + b - 6)^2$，从 $w = 1$、$b = 0$ 出发。

现在 $L$ 同时取决于两个数。办法是一次只看一个：先假装 $b$ 不动，只对 $w$ 求导；再假装 $w$ 不动，只对 $b$ 求导。这样得到的叫偏导数，写作 $\partial L / \partial w$ 和 $\partial L / \partial b$，读作“L 对 w 的偏导”。

所有参数的偏导数排在一起，叫作梯度。这里的梯度是 $(-16, -8)$，两个旋钮各自往大里拧，$w$ 拧得多一些。

一条会反复用到的规律

把上面两个偏导数拆开看。两个式子前面都有 $2(2w + b - 6)$，它是损失对输出 $y$ 的导数：输出每变一点，损失变多少。$w = 1$、$b = 0$ 时它是 −8，意思是输出应该变大。后面再乘一个数：$w$ 乘的是 2，$b$ 乘的是 1。这两个数正是它们在 $y = w \cdot 2 + b \cdot 1$ 里各自乘着的那个输入。

所以一个参数的偏导数 = 损失对输出的导数 × 这个参数乘着的输入。输入大，这个旋钮对输出的影响就大，分到的“责任”也大；输入是 0，拧它什么也改变不了，偏导数就是 0。

这其实是复合函数求导的一个特例，叫链式法则：损失通过 $y$ 依赖 $w$，那么“$w$ 变一点，损失变多少”就等于“$w$ 变一点，$y$ 变多少”乘上“$y$ 变一点，损失变多少”。中间隔了好几层也是一样，把每一层的变化倍数依次乘起来。像一串齿轮，第一个转一圈，最后一个转几圈，就等于沿途各级传动比相乘。

有了这几样工具，就可以打开那个 10 个参数的 Transformer 了。

二、一个能手算的 Transformer

任务：说出第一个字

先说清楚 Transformer 在做什么。它是今天大语言模型的基本结构，干的事情是“看前面的字，猜下一个字”。我们给它一个小到不能再小的任务：字只有 A 和 B 两个，模型每次读两个字，读到第二个字时，说出第一个字是什么。一共只有四种输入：

• AA → A，AB → A，BA → B，BB → B。

为什么挑这个任务？因为模型站在第二个位置上做预测，手里直接拿着的是第二个字，要答对就必须回头看第一个字。“该看哪里”正是 Transformer 里注意力负责的事；训练好的大模型里，也确实有专门负责“看前一个字”的注意力¹。

一个完整的 Transformer 还有很多别的部件²，我们只留下注意力和最后把结果变成答案的输出层，其余全部去掉。

字怎样变成数

电脑只能处理数，所以每个字先变成三个数。A 记作 $(1, 0, 0)$，B 记作 $(0, 1, 0)$：第一个数表示“是不是 A”，第二个数表示“是不是 B”。第三个数是一个位置标记，只有正在做预测的第二个位置上是 1，相当于告诉模型“这是我自己”。

所以输入 BA 变成两组数：第一个位置 $x_1 = (0, 1, 0)$，第二个位置 $x_2 = (1, 0, 1)$。这一步是固定的，不参与训练。

注意力：先打分，再分配比例

注意力要决定两个位置各看几成。它分三步。

第一步，提问和亮牌。 第二个位置提出一个“问题”，叫查询，记作 $q$；每个位置各亮出一张“牌”，叫键，记作 $k_1$、$k_2$。它们都是用参数和输入算出来的一个数：查询用三个参数 $w_{q1}, w_{q2}, w_{q3}$，键用另外三个参数 $w_{k1}, w_{k2}, w_{k3}$，算法都是“对应位置相乘，再加起来”。比如第二个位置的键：

第二步，打分。 查询和每张牌相乘，得到两个分数 $s_1 = q\,k_1$，$s_2 = q\,k_2$。分数越高，越该看那个位置。

第三步，把分数变成比例。 两个分数要变成两个加起来等于 1 的比例 $a_1$、$a_2$，表示第一个位置看几成、第二个位置看几成。用的办法叫 softmax：先对每个分数取 $e$ 的幂，变成正数，再各自除以总和。

为什么用 $e$ 的幂？因为它永远是正的，分数高的分到的比例就大，而且它的导数就是它自己，后面求导时很方便。

输出：看到了什么，就说什么

按比例把两个字混在一起：看到的东西里，A 占 $h_A$，B 占 $h_B$。BA 这个例子里，第一个字是 B、第二个字是 A，所以 $h_B = a_1$，$h_A = a_2$。

输出层用四个参数 $W_{AA}, W_{AB}, W_{BA}, W_{BB}$ 给 A、B 两个答案打分：

两个分数再用一次 softmax，变成“答 A”和“答 B”的概率 $p_A$、$p_B$。

要训练的参数就是 $w_q$ 的 3 个、$w_k$ 的 3 个、$W$ 的 4 个，一共 10 个。初值我们挑了好算的：$w_{q1} = w_{q2} = w_{q3} = 0.5$；$w_{k1} = w_{k2} = 0$，$w_{k3} = 0.5$；$W_{AA} = W_{BB} = 0.5$，另外两个是 0。真实的训练从随机的小数出发，我们手挑初值，是为了让数字整齐，并且让模型一开始是错的。

三、前向：从输入算到损失

一个样本走一遍

从输入一路算到答案，叫前向计算。拿 BA 来算，正确答案是 B。

查询：$q = 0.5 \times 1 + 0.5 \times 0 + 0.5 \times 1 = 1$。键：第一个位置是 B，$k_1 = w_{k2} = 0$；第二个位置是 A 加位置标记，$k_2 = w_{k1} + w_{k3} = 0.5$。分数 $s_1 = 0$，$s_2 = 0.5$。

softmax：$e^0 = 1$，$e^{0.5} = 1.649$，加起来 2.649，所以 $a_1 = 1 / 2.649 = 0.378$，$a_2 = 0.622$。模型把六成多的注意力放在了第二个字，也就是它自己身上，看错了地方。原因在 $w_{k3}$：它是 0.5，给带着“我自己”标记的那个位置加了分。

看到的东西：$h_A = a_2 = 0.622$，$h_B = a_1 = 0.378$，六成是 A。输出分数 $z_A = 0.5 \times 0.622 = 0.311$，$z_B = 0.5 \times 0.378 = 0.189$。再用 softmax：$p_A = 0.531$，$p_B = 0.469$。模型更倾向于说 A，答错了（图 2）。

损失：为什么用 −ln p

热身里的损失是“差值的平方”。这里的答案是一个字，没有“差多少”可言，用什么当损失？

最直觉的是数一数答对了几道。可这个数没法用来训练：把某个参数从 0.5 改成 0.51，$p_B$ 从 0.469 变成 0.470，答案还是 A，答对的题数一道没变，导数是 0。分数对小的改动毫无反应，就给不出“该往哪边拧”的信息。

所以实际用的是交叉熵损失：

$\ln$ 是以 $e$ 为底的对数。BA 的损失是 $-\ln 0.469 = 0.756$。它有几个好记的点：正确答案的概率是 1 时，$L = 0$；是 0.5，也就是瞎猜时，$L = \ln 2 = 0.693$；概率越接近 0，$L$ 越大，没有上限。$p_B$ 每多一点，损失就少一点，哪怕答案还没翻过来。

另外三个样本也这样算。AA 和 BB 两个字相同，不管注意力放在哪里，看到的都是同一个字，正确答案的概率是 0.622，损失 0.474；AB 和 BA 情况对称，损失也是 0.756。四个数取平均：$(0.474 + 0.756 + 0.756 + 0.474) / 4 = 0.615$，这就是开头打出来的第一个数。

现在的问题和热身一样：10 个旋钮，各该往哪边拧、拧多少，才能让 0.615 变小？

四、梯度：每个参数往哪边调

第一次尝试：挨个动一下

不会求导也有办法：一个一个试。把 $W_{BB}$ 从 0.5 调到 0.51，其余不动，重新算一遍 BA 的损失，0.7563 变成了 0.7542，少了 0.0020。改动 0.01，损失变了 −0.0020，两者之比是 −0.200。这个比值就是 $\partial L / \partial W_{BB}$ 的近似值：负的，说明 $W_{BB}$ 该调大。

再试 $w_{k3}$：从 0.5 调到 0.51，损失多了约 0.00125，比值 +0.125，说明它该调小，和前面“它让模型看错了地方”的判断一致。把 $w_{k3}$ 在一段范围里来回扫，画出损失的样子（图 3），跟热身的抛物线是一回事：损失是一条光滑的曲线，偏导数就是当前这一点的斜率。

图 3 还说明了一件事：就算把 $w_{k3}$ 调到很负、注意力全放在第一个字上，BA 的损失也只降到 0.474，因为输出层的 0.5 太小，答得不够果断。要把损失降下去，“看哪里”和“说什么”两头都得改，10 个参数要一起动。

不足：参数一多就试不动

挨个试，每个参数都要多算一两次前向。10 个参数还算得过来；GPT-3 有 1750 亿个参数，挨个试一遍，就要算 3500 亿次前向，而这只够走一步。

有没有办法算一趟，就得到所有参数的偏导数？

五、反向传播：从损失往回算

思路：沿途的倍数相乘

热身里的链式法则说，隔了好几层，就把每层的变化倍数乘起来。$w_{k3}$ 影响损失的路线是：$w_{k3}$ 改变键 $k_2$，$k_2$ 改变分数 $s_2$，分数改变比例 $a$，比例改变看到的 $h$，$h$ 改变输出分数 $z$，最后 $z$ 改变损失。

关键在于，靠近损失的那几段路是很多参数共用的。10 个参数都要经过“$z$ 怎样影响 $L$”这一段；6 个注意力参数还要共同经过 $h$、$a$、$s$。所以从损失出发往回走，每走一层，就把“损失对这一层的导数”记下来，下一层直接拿来乘。一趟走完，所有参数的偏导数就都有了。这就是反向传播³，它的计算量和一次前向差不多，大约是两倍，跟参数有多少个无关。

下面沿着图 2 的箭头倒着走一遍。每一步只用到热身里那两件事：求导，以及倍数相乘。

逐层往回

第一步：输出分数。 只有两个答案时，$p_B = \dfrac{e^{z_B}}{e^{z_A} + e^{z_B}}$。把损失 $L = -\ln p_B$ 写开，再分别对 $z_A$、$z_B$ 求导，能得到一个很干净的结果：

（推导：$L = -z_B + \ln(e^{z_A} + e^{z_B})$，对 $z_A$ 求导得 $\dfrac{e^{z_A}}{e^{z_A}+e^{z_B}} = p_A$，对 $z_B$ 求导得 $-1 + p_B$。）读法很直接：A 的分数应该降，B 的分数应该升，力度就是模型离正确答案差多远。

第二步：输出层的四个参数。 $z_A = W_{AA}\,h_A + W_{AB}\,h_B$，正是热身里那种“参数乘输入”的形状。用那条规律，偏导数等于“传回来的导数 × 这个参数乘着的输入”：

• $W_{AA}$：$0.531 \times 0.622 = 0.330$；$W_{AB}$：$0.531 \times 0.378 = 0.200$；
• $W_{BA}$：$-0.531 \times 0.622 = -0.330$；$W_{BB}$：$-0.531 \times 0.378 = -0.200$。

$W_{BB}$ 的 −0.200，正好就是前面挨个试出来的那个数。

第三步：看到的东西和注意力比例。 $h_A$ 同时出现在 $z_A$ 和 $z_B$ 里，两条路的影响要加起来：$\partial L / \partial h_A = 0.531 \times W_{AA} + (-0.531) \times W_{BA} = 0.531 \times 0.5 + 0 = 0.265$；同理 $\partial L / \partial h_B = -0.265$。看到的 A 多了不好，B 多了好。在 BA 里 $h_B = a_1$、$h_A = a_2$，所以 $\partial L / \partial a_1 = -0.265$，$\partial L / \partial a_2 = 0.265$：多看第一个字会让损失变小。

第四步：穿过 softmax。 两个比例加起来恒等于 1，一个涨另一个就跌，不能分开算。只有两个位置时，$a_1 = \dfrac{1}{1 + e^{s_2 - s_1}}$，求导可以得到：$s_1$ 每增加一点，$a_1$ 增加 $a_1 a_2$ 倍，$a_2$ 减少同样多。所以

$\partial L / \partial s_2$ 正好相反，是 +0.125。$a_1 a_2$ 这个倍数值得记住：比例是 0.5 和 0.5 时它最大，是 0.25；如果注意力几乎全压在一边，比如 0.999 和 0.001，它就接近 0，梯度传到这里就过不去了。第九节会遇到这种情况。

第五步：查询和键。 $s_1 = q\,k_1$，对 $q$ 求导是 $k_1$，对 $k_1$ 求导是 $q$。所以 $\partial L / \partial q = -0.125 \times k_1 + 0.125 \times k_2 = -0.125 \times 0 + 0.125 \times 0.5 = 0.062$；$\partial L / \partial k_1 = -0.125 \times q = -0.125$，$\partial L / \partial k_2 = 0.125$。

第六步：落到 6 个注意力参数上。 又是“传回来的导数 × 输入”。$q = w_{q1} \times 1 + w_{q2} \times 0 + w_{q3} \times 1$，所以 $w_{q1}$ 和 $w_{q3}$ 的偏导数都是 0.062，$w_{q2}$ 乘着的输入是 0，偏导数是 0。$k_1 = w_{k2}$，$k_2 = w_{k1} + w_{k3}$，所以 $w_{k2}$ 的偏导数是 −0.125，$w_{k1}$ 和 $w_{k3}$ 都是 0.125。$w_{k3}$ 的 0.125，正是前面试出来的那个坡度。

一趟走下来，10 个偏导数全部到手（图 4），每一步都只是求导和乘法。

对账

手算难免出错，我们用两种办法核对⁴。一是前面那种“挨个动一下”，10 个数和手算的结果在四位小数上全部一致。二是交给 PyTorch，一个常用的深度学习工具，它能自动做反向传播。下面这段代码可以跳过，只看输出：

import torch
x1 = torch.tensor([0., 1., 0.])          # B，第 1 个位置
x2 = torch.tensor([1., 0., 1.])          # A，第 2 个位置（带“我自己”标记）
wq = torch.tensor([0.5, 0.5, 0.5], requires_grad=True)
wk = torch.tensor([0.0, 0.0, 0.5], requires_grad=True)
Wo = torch.tensor([[0.5, 0.0], [0.0, 0.5]], requires_grad=True)

q = wq @ x2
k = torch.stack([wk @ x1, wk @ x2])
a = torch.softmax(q * k, dim=0)
h = a[0] * x1[:2] + a[1] * x2[:2]
z = Wo @ h
loss = torch.nn.functional.cross_entropy(z[None], torch.tensor([1]))
loss.backward()
print(loss.item(), wq.grad, wk.grad, Wo.grad, sep="\n")

0.756250262260437
tensor([0.0623, 0.0000, 0.0623])
tensor([ 0.1247, -0.1247,  0.1247])
tensor([[ 0.3303,  0.2003],
        [-0.3303, -0.2003]])

和手算的一样。loss.backward() 那一行，做的就是第一步到第六步。

六、更新：沿着下坡走一步

一步更新

有了 10 个偏导数，更新和热身里一模一样：每个参数减去学习率乘它的偏导数。先只用 BA 这一个样本、取学习率 1 走一步（表 2）。

用新参数把 BA 重新算一遍：注意力 $a_1$ 从 0.378 升到 0.473，$p_B$ 从 0.469 升到 0.625，损失从 0.756 降到 0.470。模型对 BA 的答案翻了过来，从说 A 变成了说 B。

现象：BA 变好了，别的样本变差了

但再看看另外三个样本（表 3）。

AA 的损失从 0.474 涨到 0.777，答案也从对变成了错。原因要回到表 2：变化最大的不是注意力的参数，而是输出层。对 BA 这一个样本来说，最快的下坡路是“不管看到什么，都多说 B”：B 那一行的两个参数调大，A 那一行调小。这条路比把注意力扭过来陡得多（输出层的偏导数是 0.330、0.200，$w_{k3}$ 只有 0.125），可它对 AA、AB 是坏事。

一个样本算出来的坡度，只是这一个样本的坡度。怎样让一步对所有样本都有好处？

七、一批样本：把梯度平均

四个样本各说各的

训练真正要降的是四个样本的平均损失。平均数的导数就是导数的平均，所以把四个样本的偏导数各自算出来，再取平均就行。把它们并排放在一起（表 4），能看出各个样本在“想”什么。

AA 和 BB 对 6 个注意力参数的偏导数全是 0：两个字一样，看哪里都一样，它们对注意力没有意见，只希望输出层更果断地“看到什么就说什么”。

AB 和 BA 在 $w_{k1}$、$w_{k2}$ 上意见正好相反。BA 希望 A 的牌小一点、B 的牌大一点（这样会多看第一个字 B），AB 希望的正好反过来。平均后两者抵消成 0，这是好事：“A 的牌比 B 小”并不是这个任务的规律，只照顾一个样本才会学到它。

它们在 $w_{k3}$ 上意见一致，都是 +0.125，平均后是 0.062。四个样本唯一的共识是：把“我自己”这张牌压低，让模型少看自己。这正是这个任务真正的规律。

平均之后走一步

用平均梯度走一步（学习率还是 1）：$w_{k3}$ 从 0.5 降到 0.438，$W_{AA}$、$W_{BB}$ 从 0.5 升到 0.562，另外两个从 0 降到 −0.062。四个样本里，AA 和 BB 的损失从 0.474 降到 0.429，AB 和 BA 却从 0.756 微微涨到 0.759，平均损失从 0.615 降到 0.594，这就是开头的第二个数。

为什么难的两个样本反而略差？因为这一步主要加强的是输出层的“照抄”：看到什么就更果断地说什么。可 AB、BA 这时看的还主要是自己（$a_1$ 只有 0.397），照抄得越果断，错得越果断。平均梯度保证的是平均损失下降，不保证每个样本都变好。

真实训练的样本有成千上万，不可能每一步都全部算一遍，于是每一步随机抽一小批来算平均梯度，这叫随机梯度下降。一小批的平均只是全体平均的估计，有点偏差，但方向大体是对的，代价却小得多。

八、多步：从平台到收敛

现象：难的样本先停在原地，然后突然下降

一直这样走下去，前 40 步的情况画在图 5 里。

前 6 步，AA、BB 的损失从 0.474 降到 0.260，AB、BA 却几乎原地不动，0.756、0.759、0.755……到第 6 步才到 0.687。这段时间里并非什么都没发生：$w_{k3}$ 每一步都在往下走，0.5、0.438、0.362、0.275……第 6 步降到 −0.028，变成负数；注意力 $a_1$ 跟着从 0.378 慢慢爬到 0.505，越过了一半。

越过一半之后，情况反过来了。模型看到的东西里第一个字占了多数，输出层那份“照抄”的果断从帮倒忙变成了帮忙，两件事互相放大：第 10 步 AB、BA 的损失是 0.504，第 16 步已经是 0.146。开头那串数里，第 10 步 0.336、第 20 步 0.076，降得最快的就是这一段。

再往后，曲线变得很平。10 个偏导数整体的大小（平方和再开方）在第 11 步最大，是 0.228；到第 50 步只有 0.023，第 200 步只剩 0.005。坡越来越缓，每一步挪得越来越少，就像热身里越靠近谷底步子越小。损失从 0.030（第 40 步）慢慢挪到 0.005（第 200 步）。走到这样一片坡度很小、每步几乎不动的地方，就叫收敛。

收敛到了什么

走了 200 步以后，$w_{k3} = -1.94$，$w_{q1}$、$w_{q2}$ 约 0.99，$w_{q3}$ 约 1.48，$W_{AA}$、$W_{BB}$ 是 2.91，另外两个是 −2.41。

把它们代回模型就能读出它学会的规则。查询 $q = 0.99 + 1.48 = 2.47$，是个正数；第一个字的牌是 0，第二个字因为带着“我自己”标记，牌是 −1.94。分数一个是 0、一个是 $2.47 \times (-1.94) = -4.8$，softmax 之后 $a_1 = 0.992$。也就是说，模型学会了**“别看自己，看前一个字”**，再用输出层把看到的字照抄出来。这条规则没有人写过，是四个样本的梯度一步步推出来的。

损失永远不会恰好是 0：要 $p$ 正好等于 1 才行，而 softmax 给不出正好的 1，参数只能越来越大地逼近它。真实训练也不会等到“完全收敛”，而是在预算用完、或者在没见过的数据上的成绩不再变好时停下。

走到的这个地方也不是唯一的。我们换掉手挑的初值，让 10 个参数从随机数出发，跑了 100 次，每次都在 20 到 37 步之内把平均损失降到 0.05 以下，一次也没有失败；可最后的参数各不相同，光是 $w_{k3}$ 就散布在 −1.82 到 1.77 之间。$w_{k3}$ 是正数的那些解，查询恰好是负数，负负得正，同样做到了“别看自己”。梯度下降找到的不是“那个”最优解，而是一大片同样好的解里离起点比较近的一个⁴。

还有一个数是我们随手定的：学习率为什么取 1？

九、学习率：步子迈多大

现象：大了更快，再大就卡死

同一个起点，换四种学习率各跑 60 步（图 6）。

0.1 太小，60 步后平均损失还有 0.469，要两百步左右才能到 0.08。学习率 10 反而最快，3 步就降到 0.031。可一旦加到 30，事情就坏了：第 1 步损失从 0.615 涨到 0.714，第 2 步 0.914，第 3 步跳到 3.79，然后再也下不来。

原因：一步跨过了谷底

热身里已经见过前一半：步子太大，就会跨过谷底，跳到另一边去。这里还多了一件更糟的事。

学习率 30 的第 2 步，$w_{k3}$ 一下变成 7.56，三个查询参数变成 15.2、15.2、30.0。这时 BA 第二个位置的分数是 341，第一个位置是 0，softmax 给出的比例是 $(0, 1)$：模型只看自己。

还记得反向传播第四步里的那个倍数 $a_1 a_2$ 吗？现在它是 0。损失再怎么变，传到分数这里都乘上了 0，6 个注意力参数再也收不到任何信号，注意力永远转不回来。剩下的输出层只能看到第二个字，而第二个字和答案无关，它最好的做法是两个字各给一半概率；可是学习率太大，它每一步都冲过头，在两边来回翻，损失就停在 3.7 上下。

这也是大模型训练里很多做法的来由。GPT-3 的学习率只有 0.00006，开始时还要从 0 慢慢升上去，后面再逐渐降低；偏导数整体太大时，还会按比例缩小再用⁵。这些手段都是为了不让某一步迈得太远，掉进坡度为 0、再也走不出来的地方。

十、放大到真实的 Transformer

我们的模型和 GPT-3 走的是同一套流程：前向算损失，反向算梯度，减一步，重复。差别在下面这张表里。

几处差别值得多说一句。

每个位置都是一道题。 我们的模型只在第二个位置答题；语言模型在每个位置都猜下一个词元，一段 2048 个词元的文字就是 2048 道题，损失是所有题目的平均。所以“一步 320 万道题”就是一步平均了这么多道题的梯度。用 3000 亿除以 320 万粗算，GPT-3 一共走了九万多步。

多出来的部件，大多是为了让梯度传得动。 我们的梯度往回只走了六步，96 层的模型要往回走几百步，一路连乘，很容易乘到接近 0，或者大得离谱。所以真实的 Transformer 给梯度留了一条可以直通的近路，还在每层把数值拉回固定范围，FPGA 那篇讲过这一点。至于每一步反向传播本身，做的仍然是第五节那几件事：求导，乘上传回来的导数。

每个参数有自己的步长。 我们的做法里，10 个参数共用一个学习率，偏导数大的步子就大；输出层就是因此先跑在了注意力前面。GPT-3 用的方法叫 Adam，它记着每个参数最近的偏导数有多大，按这个调整各自的步长，让偏导数一直很小的参数也能走得动⁶。所以它的学习率和我们这里的“1”不能直接比大小。

十一、与现实中的优化问题相比

现象：现实里找最优解很难，这里却很顺

现实世界里，找一个问题的最优解往往很难。一个面包配方要烤出来尝了才知道好不好，一种药要做完实验才知道有没有效，一张课表的排法多到排不完。可在 Transformer 里，1750 亿个参数，靠热身里那个简单的规则，一步步就走到了很低的地方。为什么这里就这么顺？

先要纠正一点：Transformer 找到的并不是最优解。第八节的 100 次随机起点各自停在了不同的地方；GPT-3 训练结束时的损失也远不是 0。训练要的只是“足够低”，而足够低的地方有很多。不过现实中很多问题也只要求够好，就算只求够好依然很难。所以差别不在目标，而在找解的条件上。

原因：坡度是不是现成的

第四节那个“挨个动一下”的办法，其实就是现实中找解的常态。不知道配方和口味之间的公式，就只能改一点、烤一炉、尝一口，只看结果，不知道坡度，这叫黑箱优化。Transformer 的损失是我们自己写出来的公式，每一步运算都能求导，反向传播一趟就把所有参数的坡度一起算出来。

这个差别在小问题上看不出来。我们让一个只看损失值的随机搜索去解本文那个 10 个参数的问题：每次随机改几个参数，变好就留下，变差就退回。从同一个初值出发跑 20 次，中间那一次用了 77 次计算，就把平均损失降到 0.05 以下；梯度下降要 27 步，每步一次前向加一次反向，大约相当于 81 次前向。两者差不多。

把同一个任务放大，差距才显出来（表 6）。字从 2 个加到 64 个，序列也一样长，参数从 36 个增加到 5120 个。随机搜索这一边换成更聪明的版本，步子会根据成败自动调整；梯度这一边用 Adam。

只有 36 个参数时，随机搜索反而比梯度省；到 192 个参数，它已经要多花一倍；到 5120 个参数时，梯度两次分别用了 368 步和 243 步，随机搜索两次都用完了 30 万次计算，损失停在 0.667 和 0.113，没有达标。梯度需要的步数随参数增加涨得很慢，因为每一步都同时拿到所有参数的方向；随机搜索每次只能知道一个“好还是不好”，参数越多，碰巧改对的机会越小，需要的次数大致跟着参数个数涨，还会卡在半路下不来。照这个趋势推到 1750 亿个参数，随机搜索毫无希望。现实中的问题还要再乘上一次试验的代价：这里一次前向不到一毫秒，可烤一炉面包要一小时，做一轮药物实验要几个月。

另外两个条件：试错便宜，参数连续

试错便宜、可重复、能撤回。 一次前向在电脑里算，结果完全确定，算错了把参数改回去就行，可以重复亿万次。现实中的试错有误差，要花时间和钱，有的还不可逆。这并不是说训练不花代价：GPT-3 的训练用了约 3×10²³ 次浮点运算，只是这些代价交给了机器，而且能拆成无数个可以同时进行的小步骤。

参数是连续的，好解连成一片。 课表、路线、电路开关这类问题，变量只能一格一格地换，换一格结果就跳一下，坡度无从谈起，FPGA 那篇讲的就是这件事。Transformer 的参数可以取任意小数，每一步运算都能求导，损失随参数平滑变化；参数又远多于“恰好够用”，前面 100 次随机起点都能走到低处，靠的就是这一点。

反过来看

现实中的问题一旦凑齐这些条件，同样可以用梯度来解。工程师设计机翼、桥梁结构或者光学镜片的形状时，会把物理过程写成能求导的计算，一次算出目标对成千上万个设计参数的导数，它和反向传播是同一类计算。

Transformer 里凑不齐这些条件的部分，也就和现实一样难。模型要多少层、学习率取多少、训练数据怎么搭配，这些都没有导数可求，只能一次一次地试，而试一次就是完整训练一个模型。本文的学习率 1、10、30 就是试出来的；GPT-3 的论文同时训练了 8 个大小不同的模型⁵，大模型团队普遍先在小模型上把这些选择试好，再放大到大模型。

所以，不是 Transformer 里的问题天生简单，而是人们把它造成了梯度能解的样子：损失写得出公式，每一步都能求导，计算便宜到可以重复亿万次。难的部分并没有消失，它转到了“怎样把问题造成这个样子”上。

十二、回到开头

现在可以把开头那串数逐行读一遍了。它就是下面这段代码跑出来的，同样可以跳过代码，只看后面的解释⁴：

import torch
emb = {"A": [1., 0, 0], "B": [0, 1., 0]}
X = torch.tensor([[emb[s[0]], [a + b for a, b in zip(emb[s[1]], [0, 0, 1.])]]
                  for s in ["AA", "AB", "BA", "BB"]])
Y = torch.tensor([0, 0, 1, 1])                       # 目标：第一个字

wq = torch.tensor([0.5, 0.5, 0.5], requires_grad=True)
wk = torch.tensor([0.0, 0.0, 0.5], requires_grad=True)
Wo = torch.tensor([[0.5, 0.0], [0.0, 0.5]], requires_grad=True)
opt = torch.optim.SGD([wq, wk, Wo], lr=1.0)

def model(X):
    q = X[:, 1] @ wq                                 # 每个样本一个查询（第 2 个位置）
    k = X @ wk                                       # 每个位置一个键
    a = torch.softmax(q[:, None] * k, dim=1)         # 注意力比例
    h = (a[:, :, None] * X[:, :, :2]).sum(1)         # 按比例混合看到的字
    return h @ Wo.T                                  # 两个答案的分数

for step in range(41):
    loss = torch.nn.functional.cross_entropy(model(X), Y)
    if step in (0, 1, 2, 3, 10, 20, 40):
        print(f"step {step:2d}  loss {loss.item():.3f}")
    opt.zero_grad()
    loss.backward()
    opt.step()

第 0 步的 0.615，是四个样本损失 0.474、0.756、0.756、0.474 的平均。loss.backward() 算出的是表 4 最后一列那 10 个平均偏导数，opt.step() 把它们乘上学习率 1 从参数里减掉，于是第 1 步变成 0.594，其中 AA、BB 变好了，AB、BA 还略差一点。第 2、3 步的 0.572、0.548 降得慢，是因为注意力还在慢慢转，出力的主要是输出层。第 10 步的 0.336 和第 20 步的 0.076 之间降得最快，那是注意力越过一半之后，“看对地方”和“照抄看到的字”互相放大的阶段。第 40 步的 0.030 之后，坡已经很缓，模型基本学会了“看前一个字”。

小结

1. 训练在做什么？ 把一个叫损失的数拧小。损失要对参数的每一点变化都有反应，才算得出坡度：“答对几道”不行，交叉熵 $-\ln p_{\text{正确答案}}$ 可以。
2. 梯度下降是什么？ 算出损失对每个参数的偏导数，也就是坡度，每个参数减去学习率乘它的偏导数，重复。一个旋钮时，它就是沿着抛物线一步步走到谷底。
3. 反向传播做了什么？ 从损失往回，把沿途每一层的变化倍数依次相乘，一趟算出所有参数的偏导数，不用挨个去试。最常用的一条：参数的偏导数等于传回来的导数乘上它乘着的输入，输入是 0，偏导数就是 0。
4. 为什么要一批样本一起算？ 只按一个样本走，会学到只对它有用的捷径；按一批样本的平均走，互相矛盾的意见抵消，剩下的是共同的规律。
5. 为什么要走很多步，曲线为什么先平后陡？ 坡度只在脚下可靠，所以每次只走一小步。注意力没转过来之前，难的样本几乎不动；越过一半以后，“看哪里”和“说什么”互相放大，损失迅速下降；之后坡越来越缓，这就是收敛。
6. 学习率为什么要小心？ 太小走得慢，太大会跨过谷底，甚至跳进坡度为 0 的地方再也出不来。
7. 为什么这里找解比现实中容易？ 找到的不是唯一的最优解，而是一大片好解里的一个；更关键的是，损失是写出来的公式，一趟反向传播就拿到所有参数的坡度，试错又便宜。只看结果的随机搜索，代价随参数个数上涨，到大模型的规模就完全走不动。现实问题凑齐这些条件也能用梯度解；Transformer 里凑不齐的部分，比如层数和学习率，也只能靠试。

参考资料