Transformer self-attention源码及原理理解

自注意力计算公式：

在公式(1)中Q(query)是输入一个序列中的一个token，K(key)代表序列中所有token的特征。
$QK^{T}$ 可以得到当前token与序列中其他token的相关性。
在论文原文中 $d_{model}$ =512，表示每个token用512维特征表示（序列符号的embedding长度）。 $d_{k}$ = $d_{model}\div h$ =64，表示每个头的大小为64。

自注意力机制的pytorch实现：

def attention(query, key, value, mask=None, dropout=None):"Compute 'Scaled Dot Product Attention'"d_k = query.size(-1)scores = torch.matmul(query, key.transpose(-2, -1)) \/ math.sqrt(d_k)if mask is not None:scores = scores.masked_fill(mask == 0, -1e9)p_attn = F.softmax(scores, dim = -1)if dropout is not None:p_attn = dropout(p_attn)return torch.matmul(p_attn, value), p_attn

多头注意力机制的pytorch实现如下：

class MultiHeadedAttention(nn.Module):def __init__(self, h, d_model, dropout=0.1):"Take in model size and number of heads."super(MultiHeadedAttention, self).__init__()assert d_model % h == 0# We assume d_v always equals d_kself.d_k = d_model // hself.h = hself.linears = clones(nn.Linear(d_model, d_model), 4)self.attn = Noneself.dropout = nn.Dropout(p=dropout)def forward(self, query, key, value, mask=None):"Implements Figure 2"if mask is not None:# Same mask applied to all h heads.mask = mask.unsqueeze(1)nbatches = query.size(0)# 1) Do all the linear projections in batch from d_model => h x d_k query, key, value = \[l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)for l, x in zip(self.linears, (query, key, value))]#这段代码首先使用zip函数，将self.linears和(query, key, value)这两个列表打包成一个元组列表，其中每个元组包含一个线性层对象和一个输入张量#对遍历的每一个Linear层，对query key value分别计算，结果放在query key value中输出# 2) Apply attention on all the projected vectors in batch. x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)# 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous() \.view(nbatches, -1, self.h * self.d_k)return self.linears[-1](x)