RNN中的注意力机制
和很多人想的不一样,注意力机制不是一诞生就取代了RNN,相反开始RNN和注意力机制是同时存在的。
前面说过RNN的问题是,循环中会把信息反复压缩,这样距离远的token的占比会越来越低,表现为模型输出时忘记前面的内容,为了解决这个问题提出过LSTM,引入长期记忆,但是这个做法效果提升有限,太长了该忘还是忘,而且计算更复杂。
如下图,正常的encoder-decoder模型,生成一个token,比如说y1时,输入只有上一个隐状态s0,以及上一个token,模型对编码阶段的记忆完全由s0提供,但是s0是反复压缩后的结果,记不住前面的x0,x1信息。
于是一个自然的想法是,既然我们想让模型记住前面token的信息,为什么不直接把前面token对应的隐状态传递给当前层?前面每个token的因转台都需要被传递过来,肯定不能拼接这样太长了,考虑做一个加权平均,这个加权平均就是注意力信息,也就是
attni=∑αihiattn_i=\sum \alpha_ih_iattni=∑αihi
这样就能考虑到前面所有token,并且哪个token比较重要,我们可以通过权重αi\alpha_iαi来调整。
最后把这个attniattn_iattni取代sis_isi输入,和yiy_iyi拼接起来,作为生成yi+1y_{i+1}yi+1时的输入
Attention is all you need
更进一步,谷歌论文的指出,我们不需要RNN了,只需要注意力机制就能解决NLP问题,效果比RNN更好,计算还更快。
注意力机制的优点是:
- 从根本上解决了长距离依赖问题
- 并行度远高于RNN,推理,训练大幅加速
最终的benchmark结果表明,纯注意力机制,比RNN得分更高,计算还更快,完爆RNN。
这里的注意力机制,整体仍然和之前RNN的注意力类似,一个token的注意力,是其他所有token的状态的加权平均,需要注意另一个token的话就把需要注意的这个token的权重调大。
不同点在于,这里没有RNN,所以隐状态不是RNN循环层算出来的,而是词向量直接乘上一个矩阵WvW_vWv算出来的。另外注意力权重的具体计算方式是,每个token乘上Wq,WkW_q,W_kWq,Wk,生成q,k向量,一个token计算和其他token的注意力权重,就让这个token的q和其他token的k做点乘,得到的值做一个softmax归一化,就是权重,然后根据这个权重做一个v的加权平均,就是这个token的隐状态
attni=∑αjvj=∑(qikj)vjattn_i=\sum \alpha_jv_j=\sum (q_ik_j)v_jattni=∑αjvj=∑(qikj)vj
这样有两个好处,一是可以直接捕捉长距离依赖,这里的Wq,Wk,WvW_q,W_k,W_vWq,Wk,Wv矩阵都是可训练的;二是注意到我们token的注意力计算不存在依赖关系,可以同时计算,而RNN需要逐个计算出隐状态才能计算注意力
具体来说,就是开始我们可以把全部token拼在一起形成一个矩阵,然后去乘Wq,Wk,WvW_q,W_k,W_vWq,Wk,Wv就能得到Q,K,VQ,K,VQ,K,V矩阵。然后让QKTQK^TQKT相乘就能得到注意力矩阵
然后让QKTQK^TQKT相乘,再softmax就能得到注意力矩阵,最后让注意力矩阵和V矩阵相乘,就能得到每个token的隐状态,然后就能去推理了。整个过程除了要遵循先算QKTQK^TQKT,再算隐状态的这个依赖之外,token之间都是可以贵遇到矩阵乘法计算,然后并行的,而矩阵乘法是一个并行优化非常充分的算子,因此整个注意力计算流程可以实现很高的并行度,这让纯注意力机制比传统RNN有很大的性能优势
最后,来看看形式化公式。转置是因为Q,K,VQ,K,VQ,K,V的形状都是[seqlen,hiddendim][seqlen,hiddendim][seqlen,hiddendim],想要每一行分别做点乘,需要转置。最后QKTQK^TQKT的形状就是[seqlen,seqlen][seqlen,seqlen][seqlen,seqlen]了,可以直接乘上VVV,得到张量形状为[seqlen,hiddendim][seqlen,hiddendim][seqlen,hiddendim],每一行表示一个token的隐状态,这个隐状态就相当于RNN里的hih_ihi了,可以接上普通网络实现下游任务。
Attention(Q,K,V)=Softmax(QKTdk)VAttention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt {d_k}})VAttention(Q,K,V)=Softmax(dkQKT)V
这里softmax之前除以dk\sqrt{d_k}dk是因为如果词向量长度dkd_kdk太大,点乘出来的logits也会太大,经过softmax进一步放到,会导致注意力集中于某几个token,其余token几乎没有权重。所以除以一个正比于词向量长度的量,让logits分布更均匀
多头注意力MHA
有点类似于MoE的专家思路,只能训练一组Wq,Wk,WvW_q,W_k,W_vWq,Wk,Wv矩阵的话,能学习到的特征不够灵活,相同参数量的前提下,不如把大的WqW_qWq拆成多个小的矩阵WqiW_{qi}Wqi,这样得到多组注意力结果,
再拼接起来,隐状态长度还和以前一样,但每组注意力头是分开反向传播的,能学到不同特征,组合起来学习到的特征更全面
LayerNorm
之前的层归一化是用BatchNorm实现的,简单来说就是每一个batch内部平均值和方差,然后进行归一化
但这样在NLP里的问题是,输入shape是[batch,seqlen,hiddendim][batch,seqlen,hiddendim][batch,seqlen,hiddendim],在一个batch内归一化的话,由于seqlen是变长的,我们如果想统一处理,一般会在后面补0到统一长度,但这样再计算均值和方差,可能会把很多占位符都算进去,导致结果失真。
所以对于序列长度可变的问题,改为在每个序列内部归一化,也就是计算每个序列内所有token的词向量的均值,标准差,然后在序列内部归一化,这样计算更简单,结果还不容易失真。另一个好处是,BatchNorm需要区分训练和推理,推理时用训练学会的均值和标准差归一化,训练时计算输入的均值和标准差,但LayerNorm没这个问题,训练和推理都临时计算当前序列的token均值,标准差。
如下图,图中箭头就是归一化的维度
位置编码
仔细分析前面的QKV计算过程,会发现一个问题,我们是批量计算每个token对其他token的注意力的,这里面并不包含token在序列中的位置信息,也就是我们把QKV矩阵里随机交换两行,也不改变注意力的计算结果
但是文本是序列信息,即使组成的token完全相同,顺序不同也会导致语义不同,比如你欠我钱和我欠你钱,显然不一样。
所以我们还要想办法,把位置信息加入进去,这就是位置编码。具体来说,就是每个token生成一个和词向量长度相同的向量,表示这个token在序列中的位置,然后把这个位置向量和词向量相加。
绝对位置编码
比较朴素的方法,把位置编码看成一个二进制数,第i个token的编码就是i的二进制表示。这样的问题是,能处理的最大位置就是2dk2^{d_k}2dk,再长的序列无法编码吗,但是语言模型输入序列长度可能是变化的,如果推理超出训练时的最长序列,就不知道如何编码
三角函数编码
每个token用一组不同相位的三角函数编码,不同位置的token,用波长来区分,第i个token波长取i。
相位的变化规律如下,这里的i不是第i个token,而是词向量里的第i个位置。
这个表示看起来有点奇怪,但是机器学习的一个设计思路是,不要尝试人脑理解模型参数,只要给这个模型足够的表达空间,训练时能自动学会规律就行。事实证明这个编码方式比绝对位置编码要好,首先就解决序列长度有上限这个问题。三角函数是个周期函数,没有值域超过上限的问题,然后区分位置是通过波长进行的,这可以类比为一个时钟,时针分针秒针都是周期函数,值域还一样,我们区分它们的方式是转速,或者说波长。
实际transformer里还更复杂一点,奇数位用sin,偶数位用cos
相位规律如下,也就是词向量里相邻两个位置的sin和cos,相位一样
cos的好处是,对于两个不同位置的token,计算注意力时点乘,根据和差化积能得到,结果只和两个token的位置之差,也就是相对位置有关,和两个token的绝对位置无关。这是符合语言规律的。
整体架构
基于注意力机制的模型整体架构如下,可以看到仍然是一个encoder-decoder架构,处理任务是序列到序列,也就是一般所说的Seq2Seq
编码阶段:
- 接受输入是输入文本
- 先把token加上位置编码,
- 然后经过多头注意力层
- 然后add&norm意思是,一条通路是残差连接,直接接到下一层,同时另一条通路是LayerNorm归一化再连接到下一层
- 最后经过一个FFN,再经过一层add&norm。
- 上面这被称为一个EncoderBlock,这样的Block循环N次,完成编码阶段,结果传给解码器
解码阶段:
- 自回归输出,把生成的输出序列作为输入,先位置编码,经过一个多头注意力和add&norm,这里对自己做注意力,也叫自注意力
- 同时接受编码器的结果,也做一次多头注意力和add&norm,这里注意力计算的不是自己的序列,而是另一个序列,称为交叉注意力。
- 自注意力和交叉注意力结果叠加,传给FFN层
- 上面整体被称为DecoderBlock,也是重复N次,注意是每生成一个Token重复N次,不是一共N次。重复完之后再经过一个线性分类层和softmax,映射到每一个token的概率,决定下一个token是什么。这里有了概率,决定下一个token选什么,是我们第一节里讲过的推理采样,有贪心采样,Beam Search,topk/topp随机采样多种方式。
- 解码阶段什么时候结束,取决于什么时候输出表示结束的特殊token
Attention的并行训练能力
训练时,采用自回归训练模式,也就是给一个无标注文本,截出他的多个前缀,第iii个token的预测训练,就是输入[1,i−1][1,i-1][1,i−1]的前缀,前向传播,然后和实际答案[1,i][1,i][1,i]的前缀对比计算Loss,进行反向传播。
注意到由于整个文本我们都知道了,这个训练方式可以并行,我们可以同时进行i∈[i,1]i∈[i,1]i∈[i,1]的所有token的前向传播,然后和对应答案对比,计算Loss。当然推理阶段肯定还是要逐个token生成
这相比RNN是个巨大进步,RNN训练时是必须串行的,Attention可以并行,大大降低了训练开销。
Masked Attention掩码注意力
注意到前面的架构图里,自注意力不是一般的MHA而是Masked MHA。这是因为如上节所述,训练时我们已经知道每一步,应该预测的下一个token是什么,可以对每一步并行训练。但为了防止模型作弊,直接去看下一个token是什么(下一个是什么在编码器阶段已经见过了),我们需要给注意力做一个掩码,已经预测了前i个token,下一步预测i+1,则只能看到前i个token的注意力结果。
具体来说我们构造一个掩码矩阵,第i行表示第i步推理。计算注意力结果时,1的位置正常计算,0的位置设为负无穷。这样和下一步的交叉注意力相加后,0的位置仍然是一个接近于负无穷的量,然后经过注意力的softmax时,负无穷会被计算为接近于0的值,相当于注意力得分为0了,忽略后面的token。
推理时由于不知道后面的token是什么,没这个作弊问题,不需要掩码。