Transformer 是怎么知道“他”是小明的

“小明累了，他想睡觉。”

“他”是谁？你不用想就知道是小明。模型得自己找出来。

2017 年以前，翻译模型的主流是 RNN：从左往右一个词一个词地读，每读一个，就把到目前为止的意思压进一个固定长度的向量，交给下一步。读到“他”的时候，“小明”已经在这个向量里被揉过两轮了。句子一长，开头的东西越揉越淡。更麻烦的是第 10 个词得等第 9 个词算完，GPU 再多也只能排队。

《Attention Is All You Need》把这条链拆了：每个词直接去看整句话，自己决定看谁、看多少。“他”和“小明”隔多远都只差一步，所有词还能同时算。这个结构叫 Transformer，后来的 GPT、BERT 都是从它长出来的。

下面一直用这句话。文中说的“词”，在模型里其实是 token，可能是一个字，也可能是半个词，不影响理解。

注意力：带着问题去查

每个词进模型时是一个向量 $x$。注意力先用三个不同的矩阵，把它变成三个新向量：

$$q=xW^Q,\qquad k=xW^K,\qquad v=xW^V.$$

可以当成查资料。$q$（query）是我要找什么；$k$（key）是贴在每个词身上的标签，写着我能提供什么；$v$（value）是标签对上之后，真正被读走的内容。三个矩阵都是训练出来的。

“他”的 $q$ 大概在问“前面哪个是人”，“小明”的 $k$ 写着“我是个人”。两个方向对得越准，“他”就越多地读走“小明”的 $v$。

把这几个向量压到二维画出来。拖动 $q$ 的箭头，盯住右边的 $o$：

整个过程分三步。

打分。 $q$ 和每个 $k$ 做点积。几何上，点积就是 $q$ 落在 $k$ 方向上的影子长度，再乘上 $k$ 的长度。方向越一致，影子越长，分数越高；背过身去，影子落到反方向，分数就是负的。分数还要再除以 $\sqrt{d_k}$，$d_k$ 是向量的维数，这里是 2。为什么要除，下一节说。

分配。 softmax 把五个分数变成五个权重：

$$w_j=\frac{e^{s_j}}{\sum_i e^{s_i}}.$$

权重都是正的，加起来正好是 1。分数高的拿得多，分数低的也不是零。

混合。 用这组权重把五个 $v$ 加权平均，得到 $o$，也就是“他”在这一层的输出。把权重想成质量，$o$ 就是五个 $v$ 的重心。

“他”的 $q$ 对着“小明”，“小明”拿到 57% 的权重。“他”自己的 $v$ 几乎在原点，本来没什么内容；读完这一下，$o$ 被一把拽向“小明”的 $v$，“他”就带上了小明的信息。

再试两件事。把 $q$ 拉长，分数之间的差距跟着放大，权重越来越集中到一个词上，$o$ 贴到那个点上；把 $q$ 缩短，权重趋于平均，$o$ 退回中间。另外，不管怎么拖，$o$ 都跑不出那块灰色区域。它是个加权平均，只能落在几个 $v$ 中间。这件事讲多头时还会用到。

分数为什么要缩小

刚才拉长 $q$，分布就变尖。维数一高，这件事会自己发生。

点积是逐维相乘再相加。论文脚注里算过：假设 $q$、$k$ 的每个分量都是均值 0、方差 1 的独立随机数，点积的方差就是 $d_k$，相加的项越多，结果越容易跑得很大。论文用的 $d_k=64$，分数的标准差就是 8，几个分数之间差个十几二十很正常。

softmax 对差距是指数敏感的。差 20，$e^{20}$ 大约是 5 亿倍，权重几乎全给了最大的那个，其他位置只分到零头。训练就麻烦在这儿：softmax 一旦这样饱和，输入再怎么动，输出都几乎不变，梯度传回来接近零，参数学不动。论文的说法是，softmax 被推进了“梯度极小的区域”。

除以 $\sqrt{d_k}$，方差就回到 1。拖动维数看看两种算法的差别：

$d_k=64$ 时，不缩放的那一行已经把 99% 给了同一个词。除过之后还是 39%、22%、16%……分数排序一样，分布却还留着余地，梯度也就能传回去。

一整句一起算

刚才只看了“他”。其实五个词同时在做同一件事，各自拿自己的 $q$ 去和所有 $k$ 打分。把五个 $q$ 摞成矩阵 $Q$，五个 $k$ 摞成 $K$，五个 $v$ 摞成 $V$，所有词的注意力就是一个式子：

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.$$

$QK^\top$ 其实就是一张乘法表：行头是每个词的 $q$，列头是每个词的 $k$，格子里填两者的点积。点一个格子看它怎么来的：

表要按行读。第 3 行就是第一张图里的“他”：它读“小明”读了 57%。整张表不对称：“他”读“小明”，用的是“他”的 $q$ 和“小明”的 $k$；反过来是“小明”的 $q$ 和“他”的 $k$，四个不同的向量，分数自然不一样。

每一行各自做 softmax，加起来都是 100%。再拿这一行的权重去加权平均五个 $v$，就是这个词的输出 $o$。

这里没有任何一步需要等别的词先算完。整张表就是两次矩阵乘法，GPU 一次算完；句子有一千个词，就是一千乘一千的表，照样一次算完。RNN 那种必须排队的链条，就这样没了。代价是表的大小按词数的平方涨，句子越长越贵，这是后来很多改进的出发点。

多头：一个平均数装不下两件事

回到第一张图。“他”其实想知道两件事：指的是谁（小明），这人怎么了（累了）。可一组权重只能产出一个 $o$，而 $o$ 只是几个 $v$ 的加权平均。两个都要，就只能各分一部分权重，$o$ 落在两点之间，哪边都不像。

论文里就一句话带过：“With a single attention head, averaging inhibits this.”一个头只做一次平均，没法同时顾上两头。

多头的办法是多准备几套 $W^Q$、$W^K$、$W^V$。每套叫一个头，各自打分、各自分配、各自混合，最后把几个头的输出首尾拼起来，再乘一个矩阵 $W^O$：

$$\begin{aligned} \operatorname{head}_i&=\operatorname{Attention}(XW_i^Q,XW_i^K,XW_i^V),\\ \operatorname{MultiHead}(X)&=\operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_h)W^O. \end{aligned}$$

拼起来，不是平均起来。区别看这张图：

一个头的时候，滑块怎么拖，$o$ 都只能在“小明”和“累了”之间那条线上滑：多拿一点小明，就少拿一点累了。两个头各有一份加起来是 100% 的权重，还各有各的 $W^V$，同一个词在两个头里的 $v$ 也不一样。头 1 的 $o_1$ 贴着小明，头 2 的 $o_2$ 贴着累了，拼起来两样都在，后面的 $W^O$ 想怎么组合都行。

论文的 base 模型用 8 个头。输入是 512 维，每个头先投影到 64 维，在自己的小空间里算注意力，8 个 64 维的输出拼回 512 维。总计算量和一个 512 维的大头差不多，却能同时看 8 个地方。

论文做过对照：同样的计算量，只用一个头，翻译的 BLEU 分数低了 0.9；头太多（32 个，每头 16 维）也会变差。至于每个头学会看什么，没人指定，全靠训练。论文附录画过几个头，有的头会把 its 连到它指代的 Law 上，看起来就像在做指代消解。

位置：注意力分不清顺序

回头看第一张图，图里有“方向”，有“远近”，就是没有“谁在前、谁在后”。

把句子改成“他累了，小明想睡觉”。“小明”的 $q$、$k$、$v$ 一个数都不变，“他”也一样；每个词分给别人的权重、算出来的 $o$，全都原样不动，只是换了个座位。在注意力眼里，这两句话是同一袋词。

所以要把位置告诉它。办法很直接：给每个位置准备一个和词向量一样长的向量，加到词向量上：

$$\text{输入}=\text{词向量}+\text{位置向量}.$$

同一个词放在不同位置，输入就不一样了，算出来的 $q$、$k$、$v$ 也跟着变。

位置向量长什么样？论文用的是不同频率的正弦和余弦。我觉得最好懂的看法，是把它当成一排钟：每两维是一根指针，位置每往后走一格，每根指针转过一个固定的角度，越靠后的指针转得越慢。拖动位置看看：

最快的指针 6 步一圈，相邻两个位置一眼就能分开；最慢的要两万来步才转一圈，在一句话的范围里几乎不动，管的是离得很远的位置。就像读时间要同时看秒针、分针和时针，单看哪一根都会重复，一组指针合起来，每个位置都对应一组独一无二的读数。

它还有个好性质：从任意位置往后挪 3 步，每根指针都只是多转一个固定的角度，跟从哪儿出发无关。“往前数 3 个”就成了一个固定的旋转。论文猜这能让模型更容易学会按相对位置去找，再加上它也许能推广到比训练时更长的句子，就选了它。论文也试过直接学一组位置向量，效果几乎一样（BLEU 25.7 比 25.8）。

原论文的位置编码公式

$$\begin{aligned} \operatorname{PE}(pos,2i)&=\sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right),\\ \operatorname{PE}(pos,2i+1)&=\cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right). \end{aligned}$$

$pos$ 是从 0 开始的位置，$d_{\text{model}}$ 是向量维数，$i$ 是第几对维度，也就是第几根指针。分母里的 $10000^{2i/d_{\text{model}}}$ 决定这根指针转多慢，波长从 $2\pi$ 到 $10000\cdot2\pi$ 按等比排开。

编码器和解码器

论文做的是翻译。人工翻译大概是这样：先把原文通读一遍，再一边写译文，一边回头看原文。Transformer 也分成这两半：

编码器负责读原文。六层叠起来，每层两步：先做一次自注意力，让每个词去读整句；再过一个前馈网络，每个位置单独做一次两层的非线性变换（512 维升到 2048 再降回来）。两步都会把输入加回到输出上（残差连接），再做一次 LayerNorm。读完，原文的每个位置都有了一个带上下文的向量。

解码器负责写译文，也是六层，每层三步。第一步是自注意力，读自己已经写出来的部分，而且只能往前看。第二步是交叉注意力：$q$ 来自解码器，$k$ 和 $v$ 来自编码器的最终输出，公式和自注意力一模一样，只是查的是原文。第三步还是前馈网络。

最后一层的输出乘上一个矩阵、再过一次 softmax，就是词表里每个词作为下一个词的概率。

训练一次算完，生成只能一个一个来

把“小明累了，他想睡觉。”翻译成 Xiaoming is tired, he wants to sleep.。

训练时，标准译文是现成的。把它整体右移一位，开头补一个 <bos>，当作解码器的输入；每个位置的任务，是预测下一个词。输入 he 的位置，要预测的是 wants。

可是如果每个位置都能看到整句输入，he 那一格往后看一眼，wants 就躺在下一个位置上，等于抄答案。所以解码器的自注意力要加因果掩码：把“往后看”的分数设成 $-\infty$，softmax 之后权重就是 0。

训练时，十个位置的输入都是现成的，掩码保证谁也看不到自己的答案，于是一次前向就能把十个位置的预测全算出来，和标准答案逐个对比。

生成时没有标准答案。一开始只有 <bos>，算出 Xiaoming 之后，要把它接到输入末尾，才能算下一个。后面的输入就是前面的输出，只能一步一步来。能同时算很多个位置，不等于能同时知道还没写出来的词。

今天的大模型基本只留下了解码器这一半，生成的方式和这里一样。一次回复从头到尾怎么一个 token 一个 token 地冒出来，我在《从按下回车到最后一个字》里跟过一遍。

参考资料

• Vaswani et al. Attention Is All You Need, 2017。第 3 节是注意力、多头、位置编码和整体结构；表 3 是头数、维数和位置编码的对照实验；附录有几张注意力的可视化。