Transformer系列：Multi-Head Attenti

关键词：Transfomer，self attention

Transformer Self Attention的作用

Transformer引入Self Attention解决NLP任务，相比于传统的TextCNN, LSTM等模型拥有以下优势

• 1. 解决了传统的RNN无法并行的问题，RNN是自回归模型，下一个RNN单元的计算依赖上一个RNN单元的计算结果，Transformer采用Self Attention每个句子的字/词可以同时输入独立计算
• 2. Transformer能够观察整个句子的每个元素进行语义理解，而TextCNN采用一定尺寸的卷积核只能观察局部上下文，只能通过增加卷积层数来处理这种长距离的元素依赖
• 3. Transformer每一个字/词不仅包含了自身的embedding信息，还自适应地融合和整个句子的上下文的信息，可以实现相同的字/词在不同上下文语境下不同表达，尤其擅长对有强语义关系的数据进行建模

---

Self Attention简介

Self Attention就是自身和自身进行Attention，具体为句子内部的每个字/词之间进行通信，计算出句子中每个字/词和其中一个目标字/词的注意力权重，从而得到目标字/词的embedding表征

self attention示意图

在Transformer中Self Attention采用Scaled Dot-Product Attention（缩放点积注意力），采用向量内积计算两两字/词的相似度，相似度越大注意力权重越大，融合这个词的信息越多

---

Multi-Head Attention网络结构解析

Transformer采用多头注意力机制，模型网络结构如下

Multi-Head Attention

其中h表示头的个数，每个头都包含单独的一个缩放点积注意力以及注意力前的线性映射层，多个头的结果concat，输入到最后的全连接映射层，缩放点积注意力网络结构如下

Scaled Dot-Product Attention

---

(1) Multi-Head Attention流程

Transformer的Multi-Head Attention包含5个步骤：

• 1.点乘: 计算Query矩阵Q、Key矩阵K的乘积，得到得分矩阵scores
• 2.缩放: 对得分矩阵scores进行缩放，即将其除以向量维度的平方根（np.sqrt(d_k)）
• 3.mask: 若存在Attention Mask，则将Attention Mask的值为True的位置对应的得分矩阵元素置为负无穷（-inf）
• 4.softmax: 对得分矩阵scores进行softmax计算，得到Attention权重矩阵attn
• 5.加权求和: 计算Value矩阵V和Attention权重矩阵attn的乘积，得到加权后的Context矩阵

---

(2) 为啥Q,K,V线性变换

Q，K，V是三个矩阵，对原始的输入句子的embedding做线性映射（wx+b，没有激活函数），其中Q和K映射后的新矩阵负责计算相似度，V映射的矩阵负责和相似度进行加权求和。在Transformer的decoder层，Q，K，V对同一个句子进行三次不同的映射，目的是提升原始embedding表达的丰富度，如果有多个头，就有多少套Q，K，V矩阵，他们之间不共享。
如果不引入Q，K而选择直接对原始的embedding做self attention，则计算的相似度是个上三角和下三角对称的

字和字的点积结果

另外如果不引入Q，K，则对角向上的值一定是最大的，因为同一个字相同的embedding是完全重合的，每个字/词必定最关心自己，这是模型不想看到的，因此要引入Q，K。而引入V矩阵主要是提升原始embedding的表达能力

---

(3) 为啥要带有缩放的Scaled Dot-Product Attention

Scaled是缩放的意思，表现在在点乘之后除以一个分母根号下K向量的维度

Attention公式

引入这个分母的作用的防止在Softmax计算中值和值存在过大的差异，导致计算结果为OneHot导致梯度消失。
容易理解除以分母之后整个点乘的结果会变小，可以缓解值和值之间的差异大小，而为什么是除以根号下K向量维度（K,V,Q三个向量维度一样），原因是除以根号下K维度后数据的分布期望和原来一致。举例假设key和query服从均值为0，方差为1的均匀分布, 即D(query)=D(key)=1, 维度大小为64，那么点积后的,我们可以计算他的方差变化

方差变化

因此所有计算出的点积值都除以根号下64似的最终的结果还是符合均值0方差1的分布。
在计算Attention的时候多种策略比如第一种以全连接计算相似性比如GAT中所使用，和第二种类似Transformer的向量内积

两种注意力权重计算方式

其中由于第一种有全连接参数进行学习，还有tanh激活函数压缩，到Softmax的输入是可控的，而第二种随着向量维度的增大，点乘结果的上限越来越高，点乘结果的差异越来越大，因此采用第二种计算Attention权重需要加入scaled

---

(4) 为啥要多头

多个头的结果拼接融合，提升特征表征和泛化能力

---

tensorflow代码实现

代码参考attention-is-all-you-need-keras
作者基于tensorflow2和tf.keras，关于Self Attention的代码在MultiHeadAttention类

class MultiHeadAttention():
    # mode 0 - big martixes, faster; mode 1 - more clear implementation
    def __init__(self, n_head, d_model, dropout, mode=0):
        self.mode = mode
        self.n_head = n_head  # 8
        # k的维度，v的维度, q的维度和k一致,因为k,q要计算内积,256/8
        self.d_k = self.d_v = d_k = d_v = d_model // n_head  # 32, d_model为词向量的emb维度
        self.dropout = dropout
        if mode == 0:
            # q,k,v => [None, seq_len, 256]
            # 这个是大矩阵的方案，这个快，这个256包含所有头的线性变换参数w,没有激活函数,8个头统一在一个大矩阵进行线性变换
            self.qs_layer = Dense(n_head * d_k, use_bias=False)
            self.ks_layer = Dense(n_head * d_k, use_bias=False)
            self.vs_layer = Dense(n_head * d_v, use_bias=False)
        elif mode == 1:
            self.qs_layers = []
            self.ks_layers = []
            self.vs_layers = []
            for _ in range(n_head):  # 8个头
                # 保证每个头dense之后的结果拼接和d_model一致
                self.qs_layers.append(TimeDistributed(Dense(d_k, use_bias=False)))
                self.ks_layers.append(TimeDistributed(Dense(d_k, use_bias=False)))
                self.vs_layers.append(TimeDistributed(Dense(d_v, use_bias=False)))
        # 缩放点积注意力
        self.attention = ScaledDotProductAttention()
        # TimeDistributed这个实际上就是一个全连接
        self.w_o = TimeDistributed(Dense(d_model))
        # self.w_o = Dense(d_model)

    def __call__(self, q, k, v, mask=None):
        # 在encoder,q=enc_input,k=enc_input,v=enc_input
        # 在decoder的第一层,q=dec_input, k=dec_last_state, v=dec_last_state
        # 在decoder的第二层,q=decoder第一层的输出, k=enc_output, v=enc_output
        d_k, d_v = self.d_k, self.d_v
        n_head = self.n_head

        if self.mode == 0:
            # [None, seq_len, 256] => [None, seq_len, 256]
            qs = self.qs_layer(q)  # [batch_size, len_q, n_head*d_k]
            ks = self.ks_layer(k)
            vs = self.vs_layer(v)

            def reshape1(x):
                s = tf.shape(x)  # [batch_size, len_q, n_head * d_k]
                # [None, seq_len, 8, 32]
                x = tf.reshape(x, [s[0], s[1], n_head, s[2] // n_head])
                # [8, None, seq_len, 32]
                x = tf.transpose(x, [2, 0, 1, 3])
                # 连续的8个都是同一个原始语句的
                # [8 * batch_size, seq_len, 32]
                x = tf.reshape(x, [-1, s[1], s[2] // n_head])  # [n_head * batch_size, len_q, d_k]
                return x

            # 相当于将for循环头拼接，转化为将for循环放到batch_size里面再整合最后的结果
            qs = Lambda(reshape1)(qs)  # [batch_size, seq_len, 256] => [8 * batch_size, seq_len, 32]
            ks = Lambda(reshape1)(ks)
            vs = Lambda(reshape1)(vs)

            if mask is not None:
                mask = Lambda(lambda x: K.repeat_elements(x, n_head, 0))(mask)
            # head是注意力的输出，attn是注意力权重
            # 如果是大矩阵 [8 * batch_size, seq_len, 32]
            head, attn = self.attention(qs, ks, vs, mask=mask)

            def reshape2(x):
                # 对结果再做整理
                s = tf.shape(x)  # [n_head * batch_size, len_v, d_v]
                # [8, batch_size, seq_len, 32]
                x = tf.reshape(x, [n_head, -1, s[1], s[2]])
                # [batch_size, seq_len, 8, 32]
                x = tf.transpose(x, [1, 2, 0, 3])
                # [batch_size, seq_len, 8 * 32]
                x = tf.reshape(x, [-1, s[1], n_head * d_v])  # [batch_size, len_v, n_head * d_v]
                return x

            head = Lambda(reshape2)(head)
        elif self.mode == 1:
            # 每个头的结果
            heads = []
            # 每个头的注意力权重
            attns = []
            for i in range(n_head):
                # 拿到对应下标的网络
                qs = self.qs_layers[i](q)  # q线性变换  [None, None, 256] => [None, None, 32]
                ks = self.ks_layers[i](k)  # k线性变换  [None, None, 256] => [None, None, 32]
                vs = self.vs_layers[i](v)  # v线性变换  [None, None, 256] => [None, None, 32]
                head, attn = self.attention(qs, ks, vs, mask)
                heads.append(head)
                attns.append(attn)
            # concat [[None, seq_len, 32], [None, seq_len, 32 ...]], Concatenate默认axis=-1,最里面一维合并
            # [None, seq_len, 32 * 8] = [None, seq_len, 256], 最终子注意力产出每个词维度emb是256,和原始的emb维度是一致的
            head = Concatenate()(heads) if n_head > 1 else heads[0]
            attn = Concatenate()(attns) if n_head > 1 else attns[0]

        # 加权求和的结果在做一层全连接,[None, None, 256] => [None. None, 256]
        outputs = self.w_o(head)
        outputs = Dropout(self.dropout)(outputs)
        return outputs, attn

以词的embedding是256为例，其中调用该类的目的是使得输入[batch_size, seq_len, 256]注意力映射为[batch_size, seq_len, 256]的新向量，其中第2位置上的256是8个头的拼接的结果，每个头的embedding维度是32。
其中有两种模式有mode参数控制，默认mode=0走大矩阵方式，该种方式将8个注意头全部平铺在三维输入矩阵的第0维batch_size上，一起进行点乘操作，结果在通过reshape和转置整理为8个头在第2维上的拼接，这种方式计算快。
第二种mode=1是传统的for循环一个一个计算头，再将结果列表进行concat，代码上更清晰一点。
其中点乘计算相似度的ScaledDotProductAttention如下

class ScaledDotProductAttention():
    def __init__(self, attn_dropout=0.1):
        self.dropout = Dropout(attn_dropout)

    def __call__(self, q, k, v, mask):  # mask_k or mask_qk
        # 根号32,向量维度平方根,np.sqrt(d_k)
        # 如果是大矩阵的话，还是32
        temper = tf.sqrt(tf.cast(tf.shape(k)[-1], dtype='float32'))
        # 计算点乘 [None, None, 32] * [None, None, 32]
        # 这个K.batch_dot就是batch0位置不动，1和2位置点乘,相当于tf.matmul(q, tf.transpose(k, [0, 2, 1]))
        # 每个句子内部，每个字和其他字计算一个内积[None, seq_len, 32] * [None, seq_len, 32] => [None, seq_len, seq_len]
        attn = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[2, 2]) / x[2])([q, k, temper])  # shape=(batch, q, k)
        if mask is not None:
            # K.cast(K.greater(src_seq, 0), 'float32') pad=0,非pad=1
            # 将<pad>的置为一个极负的数，使地softmax位置上为0,不把他的特征向量用于加权求和
            mmask = Lambda(lambda x: (-1e+9) * (1. - K.cast(x, 'float32')))(mask)
            attn = Add()([attn, mmask])
        attn = Activation('softmax')(attn)
        attn = self.dropout(attn)
        # 这个地方加权求和
        # [None, seq_len, seq_len] * [None, seq_len, 32] => [None, seq_len, 32] 每个词/句子的最终表达
        # K.batch_dot可以直接改成tf.matmul
        output = Lambda(lambda x: K.batch_dot(x[0], x[1]))([attn, v])
        return output, attn

代码中有一些用到了Keras算子，记录一下

• TimeDistributed

这个就是把一个网络层应用在一个有步长输入矩阵的每一个步长上面，TimeDistributed(Dense(d_k, use_bias=False))相当于原始三维([batch_size, seq_len, emb_size])的[seq_len, emb_size]去做一个Dense全连接，实际上三维可以直接和二维进行全连接，改行代表代表在构建三个线性映射矩阵

• K.batch_dot

代表一个带有batch_size和另一个带有batch_size的矩阵相乘，batch_size不参与计算，axes代表要进行矩阵运算需要匹配的对应维度，axes=[2, 2]表示前一个矩阵的第2维要和后一个矩阵的第2维匹配相等，然后进行相乘，实际上
attn = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[2, 2]) / x[2])([q, k, temper])完全可以替换为一个普通的矩阵相乘，先把矩阵转置一下再矩阵相乘即可
attn1 = tf.matmul(q, tf.transpose(k, [0, 2, 1])) / temper

其他代码解读详情见注释