PyTorch nn.MultiHead() 参数理解
之前一直是自己实现MultiHead Self-Attention程序，代码段又臭又长。后来发现Pytorch 早已经有API nn.MultiHead()函数，但是使用时我却遇到了很大的麻烦。
首先放上官网说明：
M u l t i H e a d ( Q , K , V ) = C o n c a t ( h e a d 1 , … , h e a d h ) W O w h e r e h e a d i = A t t e n t i o n ( Q W i Q , K W i K , V W i V ) MultiHead(Q,K,V)=Concat(head_1,…,head_h)W_O\quad where\ head_i=Attention(QW_i^Q,KW_i^K,VW_i^V) MultiHead(Q,K,V)=Concat(head1,…,headh)WOwhere headi=Attention(QWiQ,KWiK,VWiV)
# 模型初始化
torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None)
'''
embed_dim – 嵌入向量总长度.
num_heads – 并行的head数目，即同时做多少次不同语义的attention.
dropout – dropout的概率.
bias – 是否添加偏置.默认: True.
'''
# 模型运算
forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None)
'''
Inputs:
query: (L, N, E) where L is the target sequence length, N is the batch size, E is the embedding dimension.
key: (S, N, E) , where S is the source sequence length, N is the batch size, E is the embedding dimension.
value: (S, N, E) where S is the source sequence length, N is the batch size, E is the embedding dimension.
key_padding_mask: (N, S)(N,S) , ByteTensor, where N is the batch size, S is the source sequence length.
attn_mask: 2D mask (L, S)(L,S) where L is the target sequence length, S is the source sequence length. 3D mask (N*num_heads, L, S)(N∗num_heads,L,S) where N is the batch size, L is the target sequence length, S is the source sequence length.
Outputs:
attn_output: (L, N, E)(L,N,E) where L is the target sequence length, N is the batch size, E is the embedding dimension.
attn_output_weights: (N, L, S)(N,L,S) where N is the batch size, L is the target sequence length, S is the source sequence length.
'''
值得注意的一点是，query，key，value的输入形状一定是 [sequence_size, batch_size, emb_size]
官网例子：
>>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
>>> attn_output, attn_output_weights = multihead_attn(query, key, value)
embed_dim, num_heads参数
但是我执行程序却报错了：
A = torch.arange(1,25).view(4,3,2)
A = A.float()
self_attn = torch.nn.MultiheadAttention(embed_dim=2, num_heads=4, dropout=0.0)
res,weight = self_attn(A,A,A)
报错信息：
self_attn = torch.nn.MultiheadAttention(embed_dim=2, num_heads=4, dropout=0.0)
File "E:\Anaconda3\lib\site-packages\torch\nn\modules\activation.py", line 740, in __init__
assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
AssertionError: embed_dim must be divisible by num_heads
切到源码看不懂，而且我用pycharm 一直ctrl+鼠标进入不了最底层的代码，只有前几层的代码。(求相关领域大佬教教我)
经过自己的尝试，nn.MultiheadAttention(embed_dim, num_heads)中的要满足两点约束：
- embed_dim == input_dim ，即query，key，value的embedding_size必须等于embed_dim
- embed_dim%num_heads==0
上面的约束，也就说明了在使用nn.MultiheadAttention(embed_dim, num_heads)时， num_heads不是我们想设多少就设定多少。
我的看法：
nn.MultiheadAttention(embed_dim, num_heads) 中的embed_dim 是输入的embeddingsize，即query输入形状(L, N, E)的E数值，nn.MultiheadAttention 想要实现的是无论head 数目设置成多少，输出的向量大小都是不变的。写完这句话，发现自己表达能力真是弱，自己都看不懂。可以结合下面公式理解。
M u l t i H e a d ( Q , K , V ) = C o n c a t ( h e a d 1 , … , h e a d h ) W O w h e r e h e a d i = A t t e n t i o n ( Q W i Q , K W i K , V W i V ) Q ∈ R ( L , N , E ) , K ∈ R ( S , N , E ) , V ∈ R ( S , N , E ) , W i ∈ R ( E , E / h ) , W O ∈ R ( E , E ) MultiHead(Q,K,V)=Concat(head_1,…,head_h)W_O\quad where\ head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)\\ Q\in R^{(L, N, E)},\ K\in R^{(S, N, E)},\ V\in R^{(S, N, E)},\ W_i\in R^{(E,E/h)},W_O\in R^{(E,E)} MultiHead(Q,K,V)=Concat(head1,…,headh)WOwhere headi=Attention(QWiQ,KWiK,VWiV)Q∈R(L,N,E), K∈R(S,N,E), V∈R(S,N,E), Wi∈R(E,E/h),WO∈R(E,E)
attn_mask参数
self-attention公式: SA ( Q , K , V ) = s o f t m a x ( Q K T d k ) V \operatorname{SA}(Q, K, V)=softmax\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V SA(Q,K,V)=softmax(dkQKT)V
Q K T Q K^{T} QKT生成权重分布，但是在应用中有些位置的权重是不可见的，比如在时间序列中，第t天时，我们并不知道t+1天之后的信息。这时就需要传入attn_mask参数，屏蔽这些不合理的权重。attn_mask要求是booltensor，某个位置true表示掩盖该位置。
import torch
import torch.nn as nn
A = torch.Tensor(5,2,4)
nn.init.xavier_normal_(A)
print(A)
# tensor([[[ 0.3688, 0.0391, 0.2048, -0.0906],
# [-0.0654, 0.1193, -0.1792, 0.0470]],
#
# [[ 0.0812, -0.4180, -0.1353, -0.2670],
# [ 0.0433, 0.1442, 0.1733, 0.0535]],
#
# [[ 0.2352, -0.3314, -0.0238, 0.4116],
# [ 0.1062, 0.5122, 0.1572, -0.2991]],
#
# [[ 0.3381, 0.4004, -0.1936, -0.1553],
# [-0.0168, 0.5914, 0.7389, -0.1740]],
#
# [[ 0.0446, -0.1739, -0.2020, 0.2580],
# [-0.0109, 0.0854, 0.2634, -0.4735]]])
M = nn.MultiheadAttention(embed_dim=4, num_heads=2)
attention_mask = ~torch.tril(torch.ones([A.shape[0],A.shape[0]])).bool()
print(attention_mask)
# tensor([[False, True, True, True, True],
# [False, False, True, True, True],
# [False, False, False, True, True],
# [False, False, False, False, True],
# [False, False, False, False, False]])
attn_output, attn_output_weights=M(A,A,A, attn_mask=attention_mask)
print(attention_mask)
# tensor([[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
# [0.5067, 0.4933, 0.0000, 0.0000, 0.0000],
# [0.3350, 0.3276, 0.3374, 0.0000, 0.0000],
# [0.2523, 0.2511, 0.2549, 0.2417, 0.0000],
# [0.2004, 0.1962, 0.2039, 0.1981, 0.2013]],
#
# [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
# [0.5025, 0.4975, 0.0000, 0.0000, 0.0000],
# [0.3325, 0.3312, 0.3363, 0.0000, 0.0000],
# [0.2535, 0.2429, 0.2633, 0.2404, 0.0000],
# [0.2002, 0.1986, 0.2008, 0.1976, 0.2028]]], grad_fn=)