图解Transformer系列二:Self-Attention(自注意力机制)

1,242 阅读10分钟

一、Transformer中的三种Attention

关于Transformer笔记,预计出以下系列:

  1. Positional Encoding
  2. Self-attention
  3. Batch Norm & Layer Norm
  4. ResNet,残差连接
  5. subword,子词分词器
  6. 组装:Transformer

二、Attention构造

2.1 Attention的基本运作方式

首先,来看RNN这样一个用于处理序列数据的经典模型。

图1: 传统RNN

在RNN当中,tokens是一个一个被喂给模型的。比如在a3的位置,模型要等a1和a2的信息都处理完成后,才可以生成a3。这样的作用机制,使得RNN存在以下几个问题:

(1)Sequential operations的复杂度随着序列长度的增加而增加。

这是指模型下一步计算的等待时间,在RNN中为O(N)。该复杂度越大,模型并行计算的能力越差,反之则反。

(2)Maximum Path length的复杂度随着序列长度的增加而增加。

这是指信息从一个数据点传送到另一个数据点所需要的距离,在RNN中同样为O(N),距离越大,则在传送的过程中越容易出现信息缺失的情况,即数据点对于远距离处的信息,是很难“看见”的。

那么,在处理序列化数据的时候,是否有办法,在提升模型的并行运算能力的同时,对于序列中的每个token,也能让它不损失信息地看见序列里的其他tokens呢?

Attention就作为一种很好的改进办法出现了。

图2: Self-attention

如图,蓝色方框为一个attention模型。在每个位置,例如在a2处产生b2时,attention将会同时看过a1到a4的每个token。此外,每个token生成其对应的输出的过程是同时进行的,计算不需要等待。下面来看attention内部具体的运算过程。

2.2 Attention的计算过程

2.2.1 Self-attention

(1)计算框架

Self-attention的意思是,我们给Attention的输入都来自同一个序列,其计算方式如下:

图3: self-attention计算框架(图片来自李宏毅老师课堂ppt)

这张图所表示的大致运算过程是:

对于每个token,先产生三个向量query,key,value:

  • query向量类比于询问。某个token问:“其余的token都和我有多大程度的相关呀?”
  • key向量类比于索引。某个token说:“我把每个询问内容的回答都压缩了下装在我的key里”
  • value向量类比于回答。某个token说:“我把我自身涵盖的信息又抽取了一层装在我的value里”

以图中的token a2为例:

  • 它产生一个query,每个query都去和别的token的key做“某种方式”的计算,得到的结果我们称为attention score(即为图中的α\alpha )。则一共得到四个attention score。(attention score又可以被称为attention weight)。
  • 将这四个score分别乘上每个token的value,我们会得到四个抽取信息完毕的向量。
  • 将这四个向量相加,就是最终a2过attention模型后所产生的结果b2。

(2)产生query,key和value

下图描述了产生query(q),key(k)和value(v)的过程:

图4: 产生query, key和value

假设batch_size=1,输入序列X的形状为(seq_len = 4, d_model = 6),则对于这串序列,我们产生三个参数矩阵:WQW^{Q}, WKW^{K},WVW^{V}。通过上述的矩阵乘法,我们可以得到最终的结果Q,K,V。

一般来说,WQW^{Q}WKW^{K}都同样使用k_dim,WVW^{V}使用v_dim。k_dim和v_dim不一定要相等,但在transformer的原始论文中,采用的策略是,设num_heads为self-attention的头数,则:

k_dim=v_dim=d_model//num_headsk\_dim = v\_dim = d\_model//num\_heads

上图所绘是num_heads = 1的情况。关于num_heads的概念,在本文的后面会详细解释。

(3)计算attention score

总结一下,到目前为止,对于某条输入序列X,我们有:

{Q=XWQK=XWKV=XWV\left\{\begin{matrix} Q = XW^Q\\ K = XW^K\\ V = XW^V \end{matrix}\right.

现在,我们做两件事:

  • 利用Q和K,计算出attention score矩阵,这个矩阵由图3中的α\alpha组成。
  • 利用V和attention score矩阵,计算出Attention层最终的输出结果矩阵,这个矩阵由图3中的b组成。

记最终的输出结果为Attention(Q,K,V)Attention(Q,K,V),则有:

Attention(Q,K,V)=softmax(QKTdk)VAttention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k} } )V

这个dkd_k就是k_dim,而softmax(QKTdk)softmax(\frac{QK^T}{\sqrt{d_k} } )就是Attention Score矩阵,我们来详细看下这个矩阵的计算过程。

如图5,计算attention score的主流方式有两种,在transformer的论文中,采用的是dot-product(因为不需要额外再去训练一个W矩阵,运算量更小),因此我们来重点关注一下dot-product。

图5: 计算attention score的两种方式(图片来自李宏毅老师课堂ppt)

更确切地说,论文中所采用的是scaled dot-product,因为乘上了因子dk\sqrt{d_k}。在softmax之后,attention score矩阵的每一行表示一个token,每一列表示该token和对应位置token的α\alpha值,因为进行了softmax,每一行的α\alpha值相加等于1。

图6: scaled-dot-product

之所以进行scaling,是为了使得在softmax的过程中,梯度下降得更加稳定,避免因为梯度过小而造成模型参数更新的停滞。下面我们通过数学证明,来解释这个结论。为了表达方便(也为了和论文的标识保持一致),我们把k_dim写成dkd_{k},同理v_dim写成dvd_{v},S表示softmax函数,假设在做softmax之前,紫色矩阵里的每一个值为αij\alpha_{ij}^{*},则有:

αij=S(αij)=eαijj=1dkeαij\alpha_{ij}=S(\alpha_{ij}^{*}) = \frac{e^{\alpha_{ij}^{*}}}{ {\textstyle \sum_{j=1}^{d_k}e^{\alpha_{ij}^{*}}} }

聚焦到紫色矩阵的某一行,对于其中某个j{j}' ,我们有:

S(αij)αij=S(αij)(1S(αij))S(αij)αij=S(αij)S(αij),jj\begin{aligned} \frac{\partial S(\alpha_{i{j}'}^{*})}{\partial \alpha_{i{j}'}^{*}} &= S(\alpha_{i{j}'}^{*})(1-S(\alpha_{i{j}'}^{*}))\\ \frac{\partial S(\alpha_{i{j}'}^{*})}{\partial \alpha_{ij}^{*}} &= -S(\alpha_{i{j}'}^{*})S(\alpha_{ij}^{*}), \qquad j\ne {j}' \end{aligned}

从上面可以看出:

  • αij\alpha_{i{j}'}^{*}相对于同一行其他的αij\alpha_{ij}^{*}更大的时候,S(αij)S(\alpha_{i{j}'}^{*})趋近于1,S(αij)S(\alpha_{ij}^{*})趋近于0,此时以上的两个结果都趋近于0。
  • αij\alpha_{i{j}'}^{*}相对于同一行其他的αij\alpha_{ij}^{*}更小的时候,S(αij)S(\alpha_{i{j}'}^{*})趋近于0,S(αij)S(\alpha_{ij}^{*})趋近于1,此时以上的两个结果都趋近于0。

总结起来,即当αij\alpha_{i{j}'}^{*}相对于其他结果过大或者过小时,都会造成softmax函数的偏导趋近于0(梯度过低)。在这种情况下,整个模型在backprop的过程中,经过softmax之后,就无法继续传播到softmax之前的函数上,造成模型参数无法更新,影响了模型的训练效率。

那么αij\alpha_{i{j}}^{*}是怎么计算来的呢?通过前面的讲解可以知道:

αij=qkT=j=1dkqijkij\alpha_{ij}^{*}=qk^{T} = \sum_{j=1}^{d_{k}}q_{ij}k_{ij}

假设向量q和k中的每一个元素都是相互独立,均值为0,方差为1的随机变量,那么易知αij\alpha_{i{j}}^{*}的均值也为0,方差为dkd_{k}dkd_k较大,意味着不同αij\alpha_{i{j}}^{*}间值的差距也很大,这就导致了上面所说的梯度消失的问题。

证明这个结论:

已知E[qi]=E[ki]=0,Var[qi]=Var[ki]=1E[q_{i}] = E[k_{i}] = 0, Var[q_{i}] = Var[k_{i}] = 1

证明:E[i=1dkqiki]=0E[\sum_{i=1}^{d_{k}}q_{i}k_{i}] = 0,Var[i=1dkqiki]=dkVar[\sum_{i=1}^{d_{k}}q_{i}k_{i}] = d_k

因为qiq_{i}kik_{i}相互独立,则它们的协方差为0,也就是:

Cov(qi,ki)=E[(qiE[qi])(kiE[ki])]=E[qiki]E[qi]E[ki]=0\begin{aligned} Cov(q_{i}, k_{i}) &= E[(q_{i}-E[q_{i}])(k_{i}-E[k_{i}])] \\ & = E[q_ik_i]-E[q_i]E[k_i] \\ & = 0 \end{aligned}

因此,我们有E[qiki]=E[qi]E[ki]=0E[q_ik_i] = E[q_i]E[k_i] = 0

对于方差,有:

Var(qi)=E[((qiE(qi))2]=E[qi2]E2[qi]=E[qi2]=1\begin{aligned} Var(q_i) &= E[((q_i-E(q_i))^2] \\ & = E[q_i^2] -E^2[q_i]\\ & = E[q_i^2]\\ & = 1 \end{aligned}

Var(ki)=E[ki2]=1Var(k_i) = E[k_i^2] = 1

因此有:

Var(qiki)=E[(qiki)2]E2[qiki]=E[qi2]E[ki2](E[qi]E[ki])2=Var(qi)Var(ki)=1\begin{aligned} Var(q_ik_i) &= E[(q_ik_i)^2] - E^2[q_ik_i]\\ & = E[q_i^2]E[k_i^2] - (E[q_i]E[k_i])^2 \\ & = Var(q_i)Var(k_i)\\ & = 1 \end{aligned}

则有:

E[i=1dkqiki]=i=1dkE[qiki]=0E[\sum_{i=1}^{d_{k}}q_{i}k_{i}]=\sum_{i=1}^{d_{k}}E[q_ik_i] =0

Var(i=1dkqiki)i=1dkVar(qiki)=dkVar(\sum_{i=1}^{d_{k}}q_{i}k_{i})\sum_{i=1}^{d_{k}}Var(q_ik_i) = d_k

2.2.2 Masked Attention

有时候,我们并不想在做attention的时候,让一个token看到整个序列,我们只想让它看见它左边的序列,而要把右边的序列遮蔽(Mask)起来。例如在transformer的decoder层中,我们就用到了masked attention,这样的操作可以理解为模型为了防止decoder在解码encoder层输出时“作弊”,提前看到了剩下的答案,因此需要强迫模型根据输入序列左边的结果进行attention。

Masked的实现机制其实很简单,如图:

图7: Masked Attention

首先,我们按照前文所说,正常算attention score,然后我们用一个MASK矩阵去处理它(这里的+号并不是表示相加,只是表示提供了位置覆盖的信息)。在MASK矩阵标1的地方,也就是需要遮蔽的地方,我们把原来的值替换为一个很小的值(比如-1e09),而在MASK矩阵标0的地方,我们保留原始的值。这样,在进softmax的时候,那些被替换的值由于太小,就可以自动忽略不计,从而起到遮蔽的效果。

举例来说明MASK矩阵的含义,每一行表示对应位置的token。例如在第一行第一个位置是0,其余位置是1,这表示第一个token在attention时,只看到它自己,它右边的tokens是看不到的。以此类推。

2.2.3 Multihead Attention

在图像中,我们知道有不同的channel,每一个channel可以用来识别一种模式。如果我们对一张图采用attention,比如把这张图的像素格子拉平成一列,那么我们可以对每个像素格子训练不同的head,每个head就类比于一个channel,用于识别不同的模式。

而在NLP中,这种模式识别同样重要。比如第一个head用来识别词语间的指代关系(某个句子里有一个单词it,这个it具体指什么呢),第二个head用于识别词语间的时态关系(看见yesterday就要用过去式)等等。

图8展示了multihead attention的运作方式。设头的数量为num_heads,那么本质上,就是训练num_heads个WQWK,WVW^Q,W^K, W^V个矩阵,用于生成num_heads个Q,K,VQ, K, V结果。每个结果的计算方式和单头的attention的计算方式一致。最终将生成的b连接起来生成最后的结果。图9详细展示了8个head的矩阵化的运算过程,由于拆分成了多头,则此时有

k_dim=v_dim=d_model//num_headsk\_dim = v\_dim = d\_model // num\_heads

也就是说,WQWK,WVW^Q,W^K, W^V的维度变为(d_model,d_model//num_heads)(d\_model, d\_model // num\_heads)。按照这个规则拆分后,多头的运算量和原来单头的运算量一样。同时在图9中,在输出部分出现了一个WOW^O矩阵,这个矩阵用于将拼接起来的多头输出转换为最终总输出

图8: Multihead Attention(图片来自李宏毅老师课堂ppt)

图9: 8头Attention矩阵化计算过程(图片来自:来自:jalammar.github.io/illustrated…

将每个head上的attention score分数打出,可以具象化地感受每个head的关注点,以入句子"The animal didn't cross the streest because it was too tired"为例,可视化代码可点此(存在Google colab上,需要翻墙)。

图10: 单头attention可视化(图片来自:jalammar.github.io/illustrated…

如图10,颜色越深表示attention score越大,我们构造并连接五层的attention模块,可以发现it和animal,street关系密切。现在我们把8个头全部加上去,参见图11。

图11: 8头attention(图片来自:来自:jalammar.github.io/illustrated…

如图11,一种颜色表示一个头下attention score的分数,可以看出,不同的头所关注的点各不相同。

三、Attention代码实践

这里提供一个Mutihead Attention的python实现方法,它可以快速帮助我们了解一个attention层的计算过程,同时可以很方便地打出中间步骤。Tensorflow和Pytorch的源码里有更为工业化的实现方式,包加速运算、引入bias,自定义维度等等。

import numpy as np
import torch
from torch import Tensor
from typing import Optional, Any, Union, Callable
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time


class MultiHeadedAttention(nn.Module):
    def __init__(self, 
                num_heads: int, 
                d_model: int, 
                dropout: float=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        # Assume v_dim always equals k_dim
        self.k_dim = d_model // num_heads
        self.num_heads = num_heads
        self.proj_weights = clones(nn.Linear(d_model, d_model), 4) # W^Q, W^K, W^V, W^O
        self.attention_score = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, 
                query:Tensor, 
                key: Tensor, 
                value: Tensor, 
                mask:Optional[Tensor]=None):
        """
        Args:
            query: shape (batch_size, seq_len, d_model)
            key: shape (batch_size, seq_len, d_model)
            value: shape (batch_size, seq_len, d_model)
            mask: shape (batch_size, seq_len, seq_len). Since we assume all data use a same mask, so
                  here the shape also equals to (1, seq_len, seq_len)
        
        Return:
            out: shape (batch_size, seq_len, d_model). The output of a multihead attention layer
        """
        if mask is not None:
            mask = mask.unsqueeze(1)
        batch_size = query.size(0)
        
        # 1) Apply W^Q, W^K, W^V to generate new query, key, value
        query, key, value \
            = [proj_weight(x).view(batch_size, -1, self.num_heads, self.k_dim).transpose(1, 2)
                for proj_weight, x in zip(self.proj_weights, [query, key, value])] # -1 equals to seq_len
        
        # 2) Calculate attention score and the out
        out, self.attention_score = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) "Concat" output
        out = out.transpose(1, 2).contiguous() \
             .view(batch_size, -1, self.num_heads * self.k_dim)

        # 4) Apply W^O to get the final output
        out = self.proj_weights[-1](out)
        
        return out



def clones(module, N):
        "Produce N identical layers."
        return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
    
    
def attention(query: Tensor, 
              key: Tensor, 
              value: Tensor, 
              mask: Optional[Tensor] = None, 
              dropout: float = 0.1):
    """
    Define how to calculate attention score
    Args:
        query: shape (batch_size, num_heads, seq_len, k_dim)
        key: shape(batch_size, num_heads, seq_len, k_dim)
        value: shape(batch_size, num_heads, seq_len, v_dim)
        mask: shape (batch_size, num_heads, seq_len, seq_len). Since our assumption, here the shape is
              (1, 1, seq_len, seq_len)
    Return:
        out: shape (batch_size, v_dim). Output of an attention head.
        attention_score: shape (seq_len, seq_len).

    """
    k_dim = query.size(-1)

    # shape (seq_len ,seq_len),row: token,col: that token's attention score
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(k_dim)
        
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e10)

    attention_score = F.softmax(scores, dim = -1)

    if dropout is not None:
        attention_score = dropout(attention_score)
        
    out = torch.matmul(attention_score, value)
    
    return out, attention_score # shape: (seq_len, v_dim), (seq_len, seq_lem)


if __name__ == '__main__':
    d_model = 8
    seq_len = 3
    batch_size = 6
    num_heads = 2
    # mask = None
    mask = torch.tril(torch.ones((seq_len, seq_len)), diagonal = 0).unsqueeze(0)
    
    input = torch.rand(batch_size, seq_len, d_model)
    multi_attn = MultiHeadedAttention(num_heads = num_heads, d_model = d_model, dropout = 0.1)
    out = multi_attn(query = input, key = input, value = input, mask = mask)
    print(out.shape)

参考

  1. jalammar.github.io/illustrated…
  2. speech.ee.ntu.edu.tw/~hylee/ml/m…
  3. colab.research.google.com

本文正在参加人工智能创作者签约季