原则上，伪代码是正确的，但不是它是如何实现的。投影和点积注意力可以有效地使用矩阵乘法同时对所有头部进行。

您可以只使用一层密集层进行查询，一层用于键，一层用于值，而不是在头上循环。例如，对于键，你会做一个密集的维度层$h \cdot d_k$然后重塑它。

假设您有批量大小$b$, 序列长度$l$和模型尺寸$d_m$：

• 密集层的输入是形状$b \times l \times d_m$.

• 密集层的输出是有形状的$b \times l \times hd_k$（或者$d_v$分别）。

• 然后你可以重塑查询和键的形状$b \times l \times h \times d_k$.

现在，如果您排列维度，使查询具有形状$b \times h \times l \times d_k$钥匙有形状$b \times h \times d_k \times l$，您可以在最后两个维度进行批量矩阵乘法，并以形状的注意力能量结束$b \times h \times l \times l$. 然后，如果你在最后一个维度做 softmax，你会得到每个 head 和每个查询的注意力分布。

通过与现在具有转置维度的投影值进行批量矩阵乘法$b \times h \times l \times d_v$，你得到加权平均值。所以最后，所有头的“连接”实际上只是另一个张量重塑。

@Jindřich 是完全正确的。这本身不是一个答案，而是一些指向他提到的每个项目的注释转换器中的实现的指针：

密集层的输入是形状$b \times l \times d_m$：

• torch.from_numpy(np.random.randint(1, V, size=(batch, l)))：$b \times l$
• self.lut = nn.Embedding(vocab, d_model); self.lut(x)：$b \times l \times d_m$

密集层的输出具有相同的形状$b \times l \times d_m = b \times l \times hd_k$：

• self.linears = clones(nn.Linear(d_model, d_model), 4)： 这些都是$W^Q,W^K,W^V,W^O$分别，它们的输出是$b \times l \times d_m$
• l(x).view(nbatches, -1, self.h, self.d_k): 将输出转换为$b \times l \times h \times d_k$

然后你可以重塑查询和键的形状$b \times h \times l \times d_k$

• l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2): 将输出转换为$b \times h \times l \times d_k$，为 K、Q 和 V 完成。

现在，如果你改变尺寸......

• scores = torch.matmul(query, key.transpose(-2, -1))：$[b \times h \times l \times d_k]\times[b \times h \times d_k \times l] = [b \times h \times l \times l]$
• scores = scores.masked_fill(mask == 0, -1e9)：面具是$b \times 1 \times 1 \times l$对于编码器层，以及$b \times 1 \times l \times l$对于解码器层（实际上$l-1$但让我们忽略这一点）。两者都可以广播到分数张量，编码器掩码对于所有头部和所有位置都是相同的，而解码器掩码对于每个位置都是不同的，因为它隐藏了后续位置。

然后，如果你在最后一个维度做 softmax ......

• p_attn = F.softmax(scores, dim = -1)：$b \times h \times l \times l$，但现在每个向量总和为 1

通过进行批量矩阵乘法......你得到加权平均值

• torch.matmul(p_attn, value)：$[b \times h \times l \times l] \times [b \times h \times l \times d_k] = [b\times h\times l \times d_k]$，这是我们想要的加权平均值

所以最后，所有头的“连接”实际上只是另一个张量重塑

• x.transpose(1, 2).contiguous()：$[b\times l\times h \times d_k]$
• x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)：$[b\times l\times d_m]$