Transformer(5) - multi-head attention

Transformer(4) - self-attention 结尾留了一个问题：一套 、、 只能学一种“提问方式”，而一个 token 往往需要同时关注好几种关系。这一篇就来看解决办法：multi-head attention（多头注意力）。

解决什么问题

看这句话：“小明把书递给了小红，因为她想看”。为了理解“她”，模型至少要同时搞清楚几件事：“她”指的是谁（小红），“她”后面的动作是什么（想看），想看的是什么（书）。这是几种不同的关系，要去句子里不同的地方找。

问题在于，一套 attention 里每个 token 只有一个 ，也就是只能问一个问题。而且 softmax 算出来的权重加起来等于 1，注意力就这么多，分给了这个就少了那个。如果同时想查好几种关系，只能把好几个问题揉成一个问，得到的答案也是揉在一起的。

还是用通讯录来看

用上一篇的通讯录来看就很清楚。通讯录还是那三条记录， 的 4 个维度表示 [关于张三, 关于李四, 是住址, 是电话]， 的 3 个维度表示 [西安, 138xxxx, 北京]

现在我同时想知道张三的住址和电话。

只有一个问题的时候

只能问一个问题，那就把两个要求合进一个 里：“关于张三”“是住址”“是电话”都要，。按上一篇的步骤算一遍：

除以 再做 softmax，权重是 0.450、0.450、0.100，得到的结果是

住址和电话各拿到了一半。但这两个答案是混在同一个向量里的，而且各自都只拿到了一半的权重。问题越多，每个问题分到的注意力就越少，答案也越模糊。

分成两个问题分别问

如果允许分别问两次，第一次问“张三住在哪”，第二次问“张三电话多少”，上一篇已经算过：

每个问题都有自己独立的一份注意力，各自的答案都很明确。再把两个结果拼在一起：

两个答案各占一段，互不干扰。这就是 multi-head 的思路：与其让一个问题包揽一切，不如分成好几个问题各问各的，最后再把答案汇总起来。

具体怎么做

先看整体的流程，下面再一步一步说

(一) 准备 h 套矩阵

准备 套独立的 、、，每一套叫一个 head（头）。第 个 head 的三个矩阵记作 、、，尺寸和上一篇一样：

每个 head 都有自己的一套矩阵，就意味着每个 head 都能学出自己的“问法”、“标签”和“要传递的内容”。

(二) 每个 head 各自做一遍 attention

每个 head 都拿同一个输入 ，用自己的矩阵算出自己的 、、，然后按上一篇的公式独立的做一遍 attention：

关键在于，每个 head 都是单独做 softmax 的，每个 head 都有完整的一份注意力。一个 head 把注意力全部放在“小红”上，不影响另一个 head 把注意力全部放在“书”上。

(三) 拼起来

把 个 head 的结果按列拼在一起，就像上面通讯录例子里把两个答案拼成一个向量一样：

(四) 乘一个输出矩阵

最后再乘一个输出矩阵 ：

有两个作用

合在一起

输入是 ，输出还是 ，尺寸完全不变，只是每个 token 的向量里都融入了好几种不同角度的上下文信息。

维度怎么分

每个 head 变窄了

self-attention 那篇说过，如果只用一套 attention，最简单的做法是 。用了 个 head 以后，一般把它们平分：

这样拼起来以后 ， 就是一个 的方阵。原版论文里 ，，每个 head 的 ，这就是上一篇提到的那组数字。再看几个模型

总开销几乎不变

分成 个 head 并不会让参数和计算量变成 倍，因为每个 head 都变窄了。以 为例，所有 head 加起来的参数量是

和只用一个 的 head 一模一样。、 也是同理，再加上 的 ，整个 multi-head attention 一共是 个参数。计算量也一样，算分数时每个 head 要做 次乘法， 个 head 加起来是 ，和一个宽的 head 相同。

也就是说，multi-head 是在总预算不变的前提下，把一个大问题拆成了好几个小问题。

head 数也是一种取舍

既然总预算不变，head 越多，每个 head 就越窄。上一篇说过， 太小会“问题问不细”， 太小会“搬不动信息”。所以 head 并不是越多越好，而是在“能同时问几个问题”和“每个问题能问多细”之间折中。从上面的表里也能看到，实际模型里每个 head 的 一般在 64 到 128 之间。

实际实现时是一个大矩阵

虽然公式里写的是 套独立的矩阵，但实际实现时一般把 个 按列拼成一个 的大矩阵，一次乘完，再把结果按列切成 份：

和分开乘的结果完全一样，只是一次大的矩阵乘法在 GPU 上更快。所以看代码时会发现 是一个方阵，算完以后再变形（reshape）成 个 head。

前面那张图为了好理解，把每个 head 画成了一条独立的分支。但实际计算时还是整块算的： 先整体乘一个大的 、、，得到完整的 、、，再按列切给各个 head，如下图所示

训练出来的 head 真的各有分工吗

公式只是给了每个 head “可以问不同问题”的能力，并没有规定谁该问什么。那训练完以后，它们真的会分工吗？有不少研究专门分析过这个问题，结论是确实能找到很多分工明确的 head，但也不是每个都这么清晰

这些分工都不是人为指定的，而是训练过程中自己形成的。不过也要说明，模型里的大多数 head 并没有这么干净的含义，上面这些是从大量 head 里挑出来的典型例子。

顺带一提：GQA

现在的大模型里，常常不是每个 head 都有自己独立的 和 ，而是几个 head 共用一组。比如 Llama 3 8B 有 32 个 head，但只有 8 组 、，每 4 个 head 共用一组。这种做法叫 GQA（grouped-query attention），主要是为了省推理时缓存 、 的内存，等讲 KV Cache 的时候再细说。

小结

multi-head attention 把一套 attention 拆成了 套，每套有自己的 、、，各自问自己的问题、各自做 softmax，最后把结果拼起来，再用 融合并变回 维。每个 head 的宽度一般是 ，所以总参数和计算量几乎不变，却能同时关注好几种不同的关系。

到这里，一个完整的 attention 层就讲完了。但它还不是 Transformer 的全部，一个 Transformer 层里除了 attention，还有前馈网络（FFN）、残差连接和 LayerNorm。下一篇就把这些零件拼起来，组成一个完整的 Transformer block。