为此,与其只使用单独一个注意力汇聚,我们可以用独立学习得到的
h
h
h 组不同的线性投影(linear projections) 来变换查询、键和值。然后,这
h
h
h 组变换后的查询、键和值将并行地送到注意力汇聚中。最后,将这
h
h
h 个注意力汇聚的输出拼接在一起,并且通过另一个可以学习的线性投影进行变换,以产生最终输出。这种设计被称为多头注意力(multihead attention)。对于
h
h
h 个注意力汇聚输出,每一个注意力汇聚都被称作一个头(head)。
在实现多头注意力之前,让我们用数学语言将这个模型形式化地描述出来。给定查询
q
∈
R
d
q
\mathbf{q} \in \mathbb{R}^{d_q}
q∈Rdq、键
k
∈
R
d
k
\mathbf{k} \in \mathbb{R}^{d_k}
k∈Rdk和值
v
∈
R
d
v
\mathbf{v} \in \mathbb{R}^{d_v}
v∈Rdv,每个注意力头
h
i
\mathbf{h}_i
hi(
i
=
1
,
…
,
h
i = 1, \ldots, h
i=1,…,h)的计算方法为:
h
i
=
f
(
W
i
(
q
)
q
,
W
i
(
k
)
k
,
W
i
(
v
)
v
)
∈
R
p
v
,
\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},
hi=f(Wi(q)q,Wi(k)k,Wi(v)v)∈Rpv,
其中,可学习的参数包括
W
i
(
q
)
∈
R
p
q
×
d
q
\mathbf W_i^{(q)}\in\mathbb R^{p_q\times d_q}
Wi(q)∈Rpq×dq、
W
i
(
k
)
∈
R
p
k
×
d
k
\mathbf W_i^{(k)}\in\mathbb R^{p_k\times d_k}
Wi(k)∈Rpk×dk和
W
i
(
v
)
∈
R
p
v
×
d
v
\mathbf W_i^{(v)}\in\mathbb R^{p_v\times d_v}
Wi(v)∈Rpv×dv,以及代表注意力汇聚的函数
f
f
f。
f
f
f 可以是之前学习的加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换,它对应着
h
h
h 个头连结后的结果,因此其可学习参数是
W
o
∈
R
p
o
×
h
p
v
\mathbf W_o\in\mathbb R^{p_o\times h p_v}
Wo∈Rpo×hpv:
W
o
[
h
1
⋮
h
h
]
∈
R
p
o
.
\mathbf W_o \begin{bmatrix}\mathbf h_1\\\vdots\\\mathbf h_h\end{bmatrix} \in \mathbb{R}^{p_o}.
Wo⎣⎡h1⋮hh⎦⎤∈Rpo.
基于这种设计,每个头都可能会关注输入的不同部分,可以表示比简单加权平均值更复杂的函数。
import math
import torch
from torch import nn
from d2l import torch as d2l
实现
在实现过程中,我们选择缩放点积注意力作为每一个注意力头。为了避免计算代价和参数代价的大幅增长,我们设定
p
q
=
p
k
=
p
v
=
p
o
/
h
p_q = p_k = p_v = p_o / h
pq=pk=pv=po/h。值得注意的是,如果我们将查询、键和值的线性变换的输出数量设置为
p
q
h
=
p
k
h
=
p
v
h
=
p
o
p_q h = p_k h = p_v h = p_o
pqh=pkh=pvh=po,则可以并行计算
h
h
h 个头。在下面的实现中,
p
o
p_o
po是通过参数 num_hiddens 指定的。