为什么嵌入尺寸必须能被 MultiheadAttention 中的头数整除？

Question

我正在学习变形金刚。这是MultiheadAttention的 pytorch 文档。在他们的实现中，我发现有一个限制：

\n

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

\n

为什么需要约束：embed_dim must be divisible by num_heads? 如果我们回到方程

\n

\n

假设：\n Q, K,V是n x emded_dim矩阵；所有权重矩阵W是emded_dim x head_dim

\n

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

\n

W^O有尺寸(num_heads*head_dim) x embed_dim

\n

[head_i, ..., head_h] * W^O将成为n x embed_dim输出

\n

我不知道为什么我们需要embed_dim must be divisible by num_heads.

\n

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

\n

Answer 1

据我了解，这是他们为了让事情变得简单而添加的简化。理论上，我们可以像您提出的那样实现模型（类似于原始论文）。在pytorch 文档中，他们简要提到了这一点。

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图层被忽略。请参阅下面的内容，其中显示了第一个编码器（Btach 大小为 32、10 个字、512 个特征）。

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

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")