为什么在MultiheadAttention中，嵌入维度必须能够被头数整除？

Question

我正在学习Transformer。这里是PyTorch文档中MultiheadAttention的介绍。在它们的实现中，我发现有一个限制：

 assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

为什么要求约束条件：embed_dim必须被num_heads整除？如果我们回到这个方程：

假设： Q，K，V 是 n x emded_dim 的矩阵；所有的权重矩阵 W 是 emded_dim x head_dim 的。

那么，拼接 [head_i, ..., head_h] 将会是一个 n x (num_heads*head_dim) 的矩阵；

W^O 的大小为 (num_heads*head_dim) x embed_dim。

[head_i, ..., head_h] * W^O 将变成一个 n x embed_dim 的输出。

我不知道为什么我们需要 embed_dim必须被num_heads整除。

假设我们有 num_heads=10000，结果是相同的，因为矩阵乘法将吸收这个信息。

- jason

Answer 1

据我了解，这是他们为了简化而添加的一个简化操作。理论上，我们可以按照您提出的方式实现模型（类似于原始论文）。 在pytorch documention中，他们简要提到了这一点。

Note that `embed_dim` will be split across `num_heads` (i.e. each head will have dimension `embed_dim` // `num_heads`)

此外，如果您查看Pytorch实现，您会发现与最初提出的模型相比有一些不同（在我看来进行了优化）。例如，他们使用MatMul而不是Linear，并且忽略了Concat层。请参考下面显示的第一个编码器（批量大小为32，10个单词，512个特征）。

附言： 如果您需要查看模型参数（如上图所示），这是我使用的代码。

import torch
transformer_model = torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=1,num_decoder_layers=1,dim_feedforward=11)  # change params as necessary
tgt = torch.rand((20, 32, 512))
src = torch.rand((11, 32, 512))
torch.onnx.export(transformer_model, (src, tgt), "transformer_model.onnx")

- Wenuka

Answer 2

当您有一个seq_len x emb_dim（即20 x 8）的序列，并且您想使用num_heads=2时，该序列将沿着emb_dim维度分割。因此，您会得到两个20 x 4的序列。您希望每个头具有相同的形状，如果emb_dim不能被num_heads整除，则无法实现这一点。例如，取一个序列20 x 9和num_heads=2。那么你会得到20 x 4和20 x 5，它们的尺寸不同。

- Theodor Peifer

---

如果seq_len x emb_dim是20 x 9，且num_heads=2，选择head_dim=77作为例子，那么我们可以得到head_i是一个20 x 144的矩阵。因此，[head_1, head_2]是20 x 288，我们仍然可以选择W^O为288 x 9。最终我们可以得到20 x 9的结果。我的观点是，我们也可以将emb_dim映射到任何长度，并使用W^O将其投影回emb_dim。为什么需要将emb_dim分成偶数长度？谢谢。 - jason