pytorch-多头注意力(维度分析)重要

阅读目录

• 多头注意力
  • 模型
  • 总结
• Attention、Self-attention、Multi-headed Self-attenion
• 实现
• Multihead Attention中维度变化分析

 

回到顶部

多头注意力

在实践中，当给定相同的查询、键和值的集合时，我们希望模型可以基于相同的注意力机制学习到不同的行为，然后将不同的行为作为知识组合起来，捕获序列内各种范围的依赖关系（例如，短距离依赖和长距离依
赖关系）。因此，允许注意力机制组合使用查询、键和值的不同子空间表示（representation subspaces）可能是有益的。为此，与其只使用单独一个注意力汇聚，我们可以用独立学习得到的h组不同的线性投影（linear projections）来变换查询、键和值。然后，这h组变换后的查询、键和值将并行地送到注意力汇聚中。最后，将这h个注意力汇聚的输出拼接在一起，并且通过另一个可以学习的线性投影进行变换，以产生最终输出。这种设计被称为多头注意力（multihead attention）(Vaswani et al., 2017)。对于h个注意力汇聚输出，每一个注意力汇聚都被称作一个头（head）。下图展示了使用全连接层来实现可学习的线性变换的多头注意力。

模型

在实现多头注意力之前，让我们用数学语言将这个模型形式化地描述出来。给定查询$q \in R^{d_q}$、键$k \in R^{d_k}$和值$v \in R^{d_v}$，每个注意力头$h_i（i = 1, . . . , h）$的计算方法为：

$$h_i = f(W^{(q)}_iq, W^{(k)}_ik,W^{(v)}_iv) \in R^{p_v}$$

其中，可学习的参数包括$W^{(q)}_i \in R^{p_q×d_q}$、$W^{(k)}_i \in R^{p_k×d_k}$和$W^{(v)}_i \in R^{p_v×d_v}$，以及代表注意力汇聚的函数f。f可以的加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换，它对应着h个头连结后的结果，因此其可学习参数是 $W_o \in R^{p_o×hp_v}$：

总结

多头注意力机制现在的使用是非常广泛的。为什么需要比较多的head呢?可以想成相关这件事情在做Self-attention的时候，就是用q去找相关的k，但是相关这件事情有很多种不同的形式，有很多种不同的定义，所以我们不能只有一个q，应该要有多个q，不同的q负责不同种类的相关性。
我们应在怎么做呢？首先对于这个$q_i$我们分别乘两个矩阵变成$q^{i,1}$和$q^{i,2}$。这个可以理解为两种不同的相关性。之后q,k,v都要有两个：

用第一个head:

用第二个head:

将这两个接起来，然后通过一个trannsform，也就是乘上一个矩阵，得到$b_i$传到下一层去。

回到顶部

Attention、Self-attention、Multi-headed Self-attenion

如上图所示，最底层的输入$x_1,s_2,x_3.....,x_T$,表示输入的序列数据，比如，$x_1$可以代表某个句子的第一个词所对应的向量。首先，通过嵌入层（可选）将它们进行初步的embedding，得到$a_1,a_2,a_3.....,a_T$。然后，使用三个矩阵$W^Q,W^K,W^v$分别与之相乘，得到$q_i,k_i,v_i,i \in (1,2,3...T)$。上图显示了与输入$x_1$所对应的输出$b_1$是如何得到的。即：

• 利用$q_1$分别与$k_1,k_2,k_3...k_T$计算向量点积，得到$a_{1,1},a_{1,2},a_{1,3}....a_{1,T}$（从数值上看，$a_{i,i}$还不一定是0-1之间的数，还需经过softmax处理）;
• 将$a_{1,1},a_{1,2},a_{1,3}....a_{1,T}$输入softmax层,从而得到均在0-1之间的注意力权重值：$\widehat {a_{1,1}},\widehat {a_{1,2}},\widehat {a_{1,3}}....\widehat {a_{1,T}}$。分别于对应位置上的$v_1,v_2,v_3....v_T$相乘。然后求和，这样便得到了与输入的$x_1$所对应的输出$b_1$。
  同样地，与输入的$x_2$所对应的输出$b_2$，也根据类似过程获得，只是此时是利用与$b_2$对应的$q_2$分别与$k_1,k_2,k_T$计算向量点积，主要过程如下图所示：
  
  其他输入的计算过程以此类推，如下图所示：
  
  
  对于输入的序列$x_1,x_2,x_3....x_T$来说，与RNN/LSTM的处理过程不同，Self-attention机制能够并行对$x_1,x_2,x_3....x_T$进行计算，这大大提高了对$x_1,x_2,x_3....x_T$特征进行提取（即获得$b_1，b_2，b_3....b_T$）的速度。结合上述Self-attention的计算过程，并行计算的原理如下图所示:
  
  由上图可以看到，通过对输入$I$分别乘以矩阵$W^{Q},W^{K},W^{V}$，我们便得到了三个矩阵$Q,K,W$，然后通过后续计算得到注意力矩阵$\widehat {\alpha}$，进而得到输出$O$。

对于在Transformer及BERT模型中用到的Multi-headed Self-attention结构与之略有差异，具体体现在：如果将前文中得到的整体看做一个“头”，则“多头”即指对于特定的 来说，需要用多组与之相乘，进而得到多组。如下图所示：

如上图所示，以右侧示意图中输入的$a_1$为例，通过多头（这里取head=3）机制得到了三个输出$b^{1}_{head},b^{2}_{head},b^{3}_{head}$,为了获得与对应的输出，在Multi-headed Self-attention中，我们会将这里得到的$b^{1}_{head},b^{2}_{head},b^{3}_{head}$进行拼接（向量首尾相连），然后通过线性转换（即不含非线性激活层的单层全连接神经网络）得到。对于序列中的其他输入也是同样的处理过程，且它们共享这些网络的参数。
注意其中$W^{q},W^{k},W^{v}$是可以学习的参数。

回到顶部

实现

在实现过程中通常选择缩放点积注意力作为每一个注意力头。为了避免计算代价和参数代价的大幅增长，我
们设定$p_q = p_k = p_v = p_o/h$。值得注意的是，如果将查询、键和值的线性变换的输出数量设置为$p_qh = p_kh =p_vh = p_o$，则可以并行计算h个头。在下面的实现中，$p_o$是通过参数num_hiddens指定的。

import math
import torch
from torch import nn
from d2l import torch as d2l

注意这里的维度变化是为了增加并行性。

#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状:
        # (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状:
        # (batch_size，)或(batch_size，查询的个数)
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数，
        # num_hiddens/num_heads)
    #    print(queries.shape,keys.shape,values.shape)
        
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)
        
      #  print(queries.shape,keys.shape,values.shape)
        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

为了能够使多个头并行计算，上面的MultiHeadAttention类将使用下面定义的两个转置函数。具体来说，transpose_output函数反转了transpose_qkv函数的操作。

#@save
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

下面使用键和值相同的小例子来测试我们编写的MultiHeadAttention类。多头注意力输出的形状是（batch_size，num_queries，num_hiddens）。

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape

queries,keys,values的维度
torch.Size([2, 4, 100]) torch.Size([2, 6, 100]) torch.Size([2, 6, 100])
queries,keys,values的维度
torch.Size([10, 4, 20]) torch.Size([10, 6, 20]) torch.Size([10, 6, 20])
torch.Size([2, 4, 100])

这个最后输出的维度为$batch\_size \times tar\_len \times hidden\_size$

• 多头注意力融合了来自于多个注意力汇聚的不同知识，这些知识的不同来源于相同的查询、键和值的不同的子空间表示。
• 基于适当的张量操作，可以实现多头注意力的并行计算。

回到顶部

Multihead Attention中维度变化分析

1.Input： Encoder Multihead Attention 输入的 query, key, value 是相同的，都是经过了word embedding和pos embedding之后的source sentence，其维度为$batch\_size \times sr\_len \times hidden\_size$。由于有num_heads个头需要并行计算，首先query, key, value分别经过一个线性变换，再将数据split给num_heads个头分别做注意力查询，即：
query:
$batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}$
key:
$batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}$
value:
$batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}$

由于query, key, value 是相同的，因此有 sr_len_q = sr_len_k = sr_len_v
2.DotProductAttention： num_heads 个头的计算是并行的，即：

Encoder Multihead Attention中在计算softmax之前对 key 进行了 mask，目的是消除 padding 的影响。事实上，padding不仅对key有影响，对query也有影响，但在实际代码中mask仅针对key，而没有针对query。其实最原始代码是既有key mask，也有query mask的，但后来作者将query mask删去了，因为在最后计算 loss 的时候对 padding 位置的 loss 进行mask，也可达到相同的效果。

假设 batch_size = num_heads = 1，sr_len_q = sr_len_k = 6，source sentence 的最后两个位置是padding，那么Encoder Multihead Attention 中的 mask 为：

$$\begin{pmatrix} 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ \end{pmatrix}$$

即只对 key 的 padding 位置进行了 mask

3.Output： 需要将上面输出的num_heads个头的结果堆叠之后，再做一个线性变换：

$$batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden_size}{num_heads}$$

$$reshape$$

$$\Downarrow$$

$$batch\_size \times sr\_len\_q \times hidden\_size$$

$$线性变换$$

$$\Downarrow$$

$$batch\_size \times sr\_len\_q \times hidden\_size$$

可以看一下这个图：