一、Transformer中的三种Attention
关于Transformer笔记,预计出以下系列:
- Positional Encoding
- Self-attention
- Batch Norm & Layer Norm
- ResNet,残差连接
- subword,子词分词器
- 组装: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=XWVleft{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,因为乘上了因子dksqrt{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αij∗∑j=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′∗)(1−S(αij′∗))∂S(αij′∗)∂αij∗=−S(αij′∗)S(αij∗),j≠j′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 jne {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=1dkqijkijalpha_{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[(qi−E[qi])(ki−E[ki])]=E[qiki]−E[qi]E[ki]=0begin{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[((qi−E(qi))2]=E[qi2]−E2[qi]=E[qi2]=1begin{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)=1begin{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个WQ,WK,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
也就是说,WQ,WK,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)
参考
本文正在参加人工智能创作者签约季