• 问答
  • 技术
  • 实践
  • 资源
transformer 中: self-attention 部分是否需要进行 mask?精选
来源| 空字符

1 Transformer中的掩码

由于在实现多头注意力时需要考虑到各种情况下的掩码,因此在这里需要先对这部分内容进行介绍。在Transformer中,主要有两个地方会用到掩码这一机制。第1个地方就是在上一篇文章用介绍到的Attention Mask,用于在训练过程中解码的时候掩盖掉当前时刻之后的信息;第2个地方便是对一个batch中不同长度的序列在Padding到相同长度后,对Padding部分的信息进行掩盖。下面分别就这两种情况进行介绍。

1.1 Attention Mask

如图3所示,在训练过程中对于每一个样本来说都需要这样一个对称矩阵来掩盖掉当前时刻之后所有位置的信息。

图 3. 注意力掩码计算过程图

从图3可以看出,这个注意力掩码矩阵的形状为[tgt_len,tgt_len]。在后续实现过程中,我们将通过generate_square_subsequent_mask方法来生成这样一个矩阵。同时,在后续多头注意力机制实现中,将通过attn_mask这一变量名来指代这个矩阵。

1.2 Padding Mask

在Transformer中,使用到掩码的第2个地方便是Padding Mask。由于在网络的训练过程中同一个batch会包含有多个文本序列,而不同的序列长度并不一致。因此在数据集的生成过程中,就需要将同一个batch中的序列Padding到相同的长度。但是,这样就会导致在注意力的计算过程中会考虑到Padding位置上的信息。

图 4. Padding时注意力计算过程图

如图4所示,P表示Padding的位置,右边的矩阵表示计算得到的注意力权重矩阵。可以看到,此时的注意力权重对于Padding位置山的信息也会加以考虑。因此在Transformer中,作者通过在生成训练集的过程中记录下每个样本Padding的实际位置;然后再将注意力权重矩阵中对应位置的权重替换成负无穷便达到了忽略Padding位置信息的目的。这种做法也是Encoder-Decoder网络结构中通用的一种办法。

图 5. Padding掩码计算过程图

如图5所示,对于”我 是 谁 P P“这个序列来说,前3个字符是正常的,后2个字符是Padding后的结果。因此,其Mask向量便为[True, True, True, False, False]。通过这个Mask向量可知,需要将权重矩阵的最后两列替换成负无穷,在后续我们会通过torch.masked_fill这个方法来完成这一步,并且在实现时将使用key_padding_mask来指代这一向量。

到此,对于Transformer中所要用到Mask的地方就介绍完了,下面正式来看如何实现多头注意力机制。

2 实现多头注意力机制

根据前面的介绍可以知道,多头注意力机制中最为重要的就是自注意力机制,也就是需要前计算得到Q、K和V,如图6所示。

图 6. Q、K和V计算过程

然后再根据Q、K、V来计算得到最终的注意力编码,如图7所示:

图 7. 注意力编码计算图

同时,为了避免单个自注意力机制计算得到的注意力权重过度集中于当前编码位置自己所在的位置(同时更应该关注于其它位置),所以作者在论文中提到通过采用多头注意力机制来解决这一问题,如图8所示。

图 8. 多头注意力计算图(2个头)

2.1 定义类MyMultiHeadAttention

综上所述,我们可以给出类MyMultiHeadAttentiond的定义为

class MyMultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0., bias=True):
        super(MyMultiheadAttention, self).__init__()
        """
        :param embed_dim:   词嵌入的维度,也就是前面的d_model参数,论文中的默认值为512
        :param num_heads:   多头注意力机制中多头的数量,也就是前面的nhead参数, 论文默认值为 8
        :param bias:        最后对多头的注意力(组合)输出进行线性变换时,是否使用偏置
        """
        self.embed_dim = embed_dim  # 前面的d_model参数
        self.head_dim = embed_dim // num_heads  # head_dim 指的就是d_k,d_v
        self.kdim = self.head_dim
        self.vdim = self.head_dim
        self.num_heads = num_heads  # 多头个数
        self.dropout = dropout
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim 除以 num_heads必须为整数"
        # 上面的限制条件就是论文中的  d_k = d_v = d_model/n_head 条件
        self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  # embed_dim = kdim * num_heads
        # 这里第二个维度之所以是embed_dim,实际上这里是同时初始化了num_heads个W_q堆叠起来的, 也就是num_heads个头
        self.k_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  # W_k,  embed_dim = kdim * num_heads
        self.v_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))  # W_v,  embed_dim = vdim * num_heads
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        # 最后将所有的Z组合起来的时候,也是一次性完成, embed_dim = vdim * num_heads

在上述代码中,embed_dim表示模型的维度(图8中的d_m);num_heads表示多头的个数;bias表示是否在多头线性组合时使用偏置。同时,为了使得实现代码更加高效,所以Pytorch在实现的时候是多个头注意力机制一起进行的计算,也就上面代码的第17-20行,分别用来初始化了多个头的权重值(这一过程从图8也可以看出)。当多头注意力机制计算完成后,将会得到一个形状为[src_len,embed_dim]的矩阵,也就是图8中多个Z水平堆叠后的结果。因此,第21行代码将会初始化一个线性层来对这一结果进行一个线性变换。

2.2 定义前向传播过程

在定义完初始化函数后,便可以定义如下所示的多头注意力前向传播的过程:

    def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
        """
        在论文中,编码时query, key, value 都是同一个输入, 
        解码时 输入的部分也都是同一个输入, 
        解码和编码交互时 key,value指的是 memory, query指的是tgt
        :param query: # [tgt_len, batch_size, embed_dim], tgt_len 表示目标序列的长度
        :param key:  #  [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param value: # [src_len, batch_size, embed_dim], src_len 表示源序列的长度
        :param attn_mask: # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
                一般只在解码时使用,为了并行一次喂入所有解码部分的输入,所以要用mask来进行掩盖当前时刻之后的位置信息
        :param key_padding_mask: [batch_size, src_len], src_len 表示源序列的长度
        :return:
        attn_output: [tgt_len, batch_size, embed_dim]
        attn_output_weights: # [batch_size, tgt_len, src_len]
        """
        return multi_head_attention_forward(query, key, value, self.num_heads,
                                            self.dropout, self.out_proj.weight, self.out_proj.bias,
                                            training=self.training,
                                            key_padding_mask=key_padding_mask,
                                            q_proj_weight=self.q_proj_weight,
                                            k_proj_weight=self.k_proj_weight,
                                            v_proj_weight=self.v_proj_weight,
                                            attn_mask=attn_mask)

在上述代码中,querykeyvalue指的并不是图6中的Q、K和V,而是没有经过线性变换前的输入。例如在编码时三者指的均是原始输入序列src;在解码时的Mask Multi-Head Attention中三者指的均是目标输入序列tgt;在解码时的Encoder-Decoder Attention中三者分别指的是Mask Multi-Head Attention的输出、Memory和Memory。key_padding_mask指的是编码或解码部分,输入序列的Padding情况,形状为[batch_size,src_len]或者[batch_size,tgt_len];attn_mask指的就是注意力掩码矩阵,形状为[tgt_len,src_len],它只会在解码时使用。

注意,在上面的这些维度中,tgt_len本质上指的其实是query_lensrc_len本质上指的是key_len。只是在不同情况下两者可能会是一样,也可能会是不一样。

2.3 多头注意力计算过程

在定义完类MyMultiHeadAttentiond后,就需要定义出多头注意力的实际计算过程。由于这部分代码较长,所以就分层次进行介绍。

def multi_head_attention_forward(
  query,  # [tgt_len,batch_size, embed_dim]
  key,  # [src_len, batch_size, embed_dim]
  value,  # [src_len, batch_size, embed_dim]
  num_heads,
  dropout_p,
  out_proj_weight, # [embed_dim = vdim * num_heads, embed_dim]
  out_proj_bias,
  training=True,
  key_padding_mask=None,  # [batch_size,src_len/tgt_len]
  q_proj_weight=None,  # [embed_dim,kdim * num_heads]
  k_proj_weight=None,  # [embed_dim, kdim * num_heads]
  v_proj_weight=None,  # [embed_dim, vdim * num_heads]
  attn_mask=None,  # [tgt_len,src_len] 
  ):
   # 第一阶段: 计算得到Q、K、V
    q = F.linear(query, q_proj_weight)
    #  [tgt_len,batch_size,embed_dim] x [embed_dim,kdim * num_heads] = [tgt_len,batch_size,kdim * num_heads]
    k = F.linear(key, k_proj_weight)
    # [src_len, batch_size,embed_dim] x [embed_dim,kdim * num_heads] = [src_len,batch_size,kdim * num_heads]
    v = F.linear(value, v_proj_weight)
    # [src_len, batch_size,embed_dim] x [embed_dim,vdim * num_heads] = [src_len,batch_size,vdim * num_heads]

在上述代码中,第16-20行所做的就是根据输入进行线性变换得到图6中的Q、K和V。

# 第二阶段: 缩放,以及attn_mask维度判断
   tgt_len, bsz, embed_dim = query.size()  # [tgt_len,batch_size, embed_dim]
    src_len = key.size(0)
    head_dim = embed_dim // num_heads  # num_heads * head_dim = embed_dim
    scaling = float(head_dim) ** -0.5
    q = q * scaling  # [query_len,batch_size,kdim * num_heads]

    if attn_mask is not None:  # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
        if attn_mask.dim() == 2:
            attn_mask = attn_mask.unsqueeze(0)  # [1, tgt_len,src_len]
            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
                raise RuntimeError('The size of the 2D attn_mask is not correct.')
        elif attn_mask.dim() == 3:
            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
                raise RuntimeError('The size of the 3D attn_mask is not correct.')
        # 现在 atten_mask 的维度就变成了3D

接着,在上述代码中第5-6行所完成的就是图7中的缩放过程;第8-15行用来判断或修改attn_mask的维度,当然这几行代码只会在解码器中的Masked Multi-Head Attention中用到。

# 第三阶段: 计算得到注意力权重矩阵
   q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    # [batch_size * num_heads,tgt_len,kdim]
    # 因为前面是num_heads个头一起参与的计算,所以这里要进行一下变形,以便于后面计算。 且同时交换了0,1两个维度
    k = k.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)#[batch_size * num_heads,src_len,kdim]
    v = v.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)#[batch_size * num_heads,src_len,vdim]
    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
    # [batch_size * num_heads,tgt_len,kdim] x [batch_size * num_heads, kdim, src_len]
    # =  [batch_size * num_heads, tgt_len, src_len]  这就num_heads个QK相乘后的注意力矩阵

继续,在上述代码中第1-5行所做的就是交换Q、K、V中的维度,以便于多个样本同时进行计算;第6行代码便是用来计算注意力权重矩阵;其中上contiguous()方法是将变量放到一块连续的物理内存中;bmm的作用是用来计算两个三维矩阵的乘法操作[1]。

需要提示的是,大家在看代码的时候,最好是仔细观察一下各个变量维度的变化过程,笔者也在每次运算后进行了批注。

# 第四阶段: 进行相关掩码操作
  if attn_mask is not None:
        attn_output_weights += attn_mask  # [batch_size * num_heads, tgt_len, src_len]
    if key_padding_mask is not None:
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        # 变成 [batch_size, num_heads, tgt_len, src_len]的形状
        attn_output_weights = attn_output_weights.masked_fill(
            key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf')) 
         # 扩展维度,从[batch_size,src_len]变成[batch_size,1,1,src_len]
        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len,src_len)  
        # [batch_size * num_heads, tgt_len, src_len]

进一步,在上述代码中第2-3行便是用来执行图3中的步骤;第4-8行便是用来执行图5中的步骤,同时还进行了维度扩充。

  attn_output_weights = F.softmax(attn_output_weights, dim=-1)# [batch_size * num_heads, tgt_len, src_len]
    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
    attn_output = torch.bmm(attn_output_weights, v)
    # Z = [batch_size * num_heads, tgt_len, src_len]  x  [batch_size * num_heads,src_len,vdim]
    # = # [batch_size * num_heads,tgt_len,vdim]
    # 这就num_heads个Attention(Q,K,V)结果

    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
    # 先transpose成 [tgt_len, batch_size* num_heads ,kdim]
    # 再view成 [tgt_len,batch_size,num_heads*kdim]
    attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)

    Z = F.linear(attn_output, out_proj_weight, out_proj_bias)
    # 这里就是多个z  线性组合成Z  [tgt_len,batch_size,embed_dim]
    return Z, attn_output_weights.sum(dim=1) / num_heads  # 将num_heads个注意力权重矩阵按对应维度取平均

最后,在上述代码中第1-3行便是用来对权重矩阵进行归一化操作,以及计算得到多头注意力机制的输出;第13行代码便是用来对多个注意力的输出结果进行线性组合;第15行代码用来返回线性组合后的结果,以及多个注意力权重矩阵的平均值。

2.4 示例代码

在实现完类MyMultiHeadAttention的全部代码后,便可以通过类似如下的方式进行使用。

if __name__ == '__main__':
    src_len = 5
    batch_size = 2
    dmodel = 32
    num_head = 1
    src = torch.rand((src_len, batch_size, dmodel))  # shape: [src_len, batch_size, embed_dim]
    src_key_padding_mask = torch.tensor([[True, True, True, False, False],
                                         [True, True, True, True, False]])  # shape: [src_len, src_len]

    my_mh = MyMultiheadAttention(embed_dim=dmodel, num_heads=num_head)
    r = my_mh(src, src, src,key_padding_mask = src_key_padding_mask)

在上述代码中,第6-11行其实也就是Encoder中多头注意力机制的实现过程。同时,在计算过程中还可以打印出各个变量的维度变化信息:

进入多头注意力计算:
  多头num_heads = 1, d_model=32, d_k = d_v = d_model/num_heads=32
  query的shape([tgt_len, batch_size, embed_dim]):torch.Size([5, 2, 32])
   W_q 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
    Q  的shape([tgt_len, batch_size,kdim * num_heads]):torch.Size([5, 2, 32])
 ----------------------------------------------------------------------
   key 的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
   W_k 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
    K  的shape([src_len,batch_size,kdim * num_heads]):torch.Size([5, 2, 32])
 ----------------------------------------------------------------------
  value的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
   W_v 的shape([embed_dim,vdim * num_heads]):torch.Size([32, 32])
    V  的shape([src_len,batch_size,vdim * num_heads]):torch.Size([5, 2, 32])
 ----------------------------------------------------------------------
  ***** 注意,这里的W_q, W_k, W_v是多个head同时进行计算的. 因此,Q,K,V分别也是包含了多个head的q,k,v堆叠起来的结果 *****
  多头注意力中,多头计算结束后的形状(堆叠)为([tgt_len,batch_size,num_heads*kdim])torch.Size([5, 2, 32])
  多头计算结束后,再进行线性变换时的权重W_o的形状为([num_heads*vdim, num_heads*vdim  ])torch.Size([32, 32])
  多头线性变化后的形状为([tgt_len,batch_size,embed_dim]) torch.Size([5, 2, 32])


相关推荐

pytorch nn.Transformer 的 mask 理解

  • 0
  • 0
  • 1161
收藏
暂无评论