【动手深度学习-笔记】注意力机制（三）多头注意力

紧接上回：【动手深度学习-笔记】注意力机制（二）注意力评分函数

在实践中，我们希望模型可以基于相同的注意力机制学习到不同的行为，抽取不同的信息（比如长距离依赖关系和短距离依赖关系），再将这些信息组合起来。
为此，与其只使用单独一个注意力汇聚， 我们可以用独立学习得到的$h$组不同的线性投影（linear projections）来变换查询、键和值，并行地送入到注意力汇聚，再将$h$组汇聚结果拼接到一起。
这种设计被称作多头注意力（multihead attention），融合了来自于多个注意力汇聚的不同知识，这些知识的不同来源于相同的查询、键和值的不同的子空间表示。
这里的“头”的概念类似于“通道”的概念，一个头表示一个注意力汇聚。

模型

给定查询$\mathbf{q}\in {\mathbb{R}}^{d_q}$，键$\mathbf{k}\in {\mathbb{R}}^{d_k}$，值$\mathbf{v}\in {\mathbb{R}}^{d_v}$，每个注意力头${\mathbf{h}}_i(i=1,\dots ,h)$的计算方法：
$${\mathbf{h}}_i=f({\mathbf{W}}_i^{(q)}\mathbf{q},{\mathbf{W}}_i^{(k)}\mathbf{k},{\mathbf{W}}_i^{(v)}\mathbf{v})\in {\mathbb{R}}^{p_v},\text{\hspace{1em}(1)}$$
原本$\mathbf{q}$、$\mathbf{k}$、$\mathbf{v}$的维度分别为$d_q$、$d_k$、$d_v$，经过${\mathbf{W}}_i$线性投影为$p_q$、$p_k$、$p_v$维度；
输入到注意力汇聚函数$f$得到汇聚结果${\mathbf{h}}_i\in {\mathbb{R}}^{p_v}$，$f$可以是加性注意力或者是缩放点积注意力等；
然后将得到的${\mathbf{h}}_i$组合，经过${\mathbf{W}}_o\in {\mathbb{R}}^{p_o\times h{p}_v}$进行另一个线性转换，得到最终输出：
$$\begin{array}{r}{\mathbf{W}}_o\left[\begin{array}{c}{\mathbf{h}}_1\\ ⋮\\ {\mathbf{h}}_h\end{array}\right]\in {\mathbb{R}}^{p_o}.\end{array}$$

其中的可学习参数为${\mathbf{W}}_i^{(q)}\in {\mathbb{R}}^{p_q\times d_q}$、${\mathbf{W}}_i^{(k)}\in {\mathbb{R}}^{p_k\times d_k}$、${\mathbf{W}}_i^{(v)}\in {\mathbb{R}}^{p_v\times d_v}$、${\mathbf{W}}_o\in {\mathbb{R}}^{p_o\times h{p}_v}$；

每个头都可能会关注输入的不同部分， 可以表示比简单加权平均值更复杂的函数。

参考

10.5. 多头注意力 — 动手学深度学习 2.0.0-beta1 documentation

进一步

【动手深度学习-笔记】注意力机制（四）自注意力、交叉注意力和位置编码