Transformer模型详解
Transformer
Transformer真是个细节满满的框架呢,大三读到根本不敢看,考研复试前看了看,以为懂了其实差得还远,两个多月前看了,还是一知半解,如今终于经过细细分析,算是知道了Transformer的基本概貌,如有错误,希望可以得到您宝贵的建议。
1 整体架构
Transformer 是一种深度学习模型架构,它在2017年由Vaswani等人在论文Attention is All You Need中提出。它的设计初衷是为了改进序列到序列(Seq2Seq)任务的处理方式,比如机器翻译等自然语言处理任务。与之前的模型如RNN或LSTM相比,Transformer引入了全新的机制——自注意力机制(Self-attention mechanism),使得模型能够并行化训练,并且更好地捕捉长距离依赖关系,从而更好地建模文本的序列结构。随着研究的深入,Transformer逐渐被应用到其他领域。在计算机视觉中, 研究人员将图像分割成多个patches,并且将每个patch视为一个token,并将众多图像块视为一个序列(考虑到图像的复杂空间结构,采用可学习的位置编码)。在语音识别、语音合成中,Transformer可以被用于建模语音信号的时间结构,在生物信息学中,蛋白质也可以看作氨基酸序列,因此也可以用Transformer来建模…
Transformer模型主要依靠自回归(Auto Regression)的方式完成序列到序列(Seq2Seq)任务。自回归意味着生成序列中的下一个元素时依赖于之前已经生成的元素,具体来说,编码器(Encoder)将输入序列,编码为高维语义表示,不参与生成过程,解码器利用编码器输出和已生成的历史token预测下一个token
机器翻译任务中Transformer自回归过程的伪代码如下:
1 初始化输入序列和最大输出序列的长度
2 将输入字符串转换嵌入向量,添加位置编码后再输入编码器,得到其高维上下文表示
3 初始化起始字符串"<SOS>"
4 循环直到当前已生成的字符串包含结束标记 <EOS> \text{<EOS>} <EOS>或达到输出最大序列长度
5 将初始字符串转换嵌入向量,添加位置编码后再输入解码器
6 解码器结合编码器给出的输入序列高维表示和自身当前输入预测下一个词的概率分布
7 从预测出的概率分布中采样下一个词,添加到已生成的字符串上
8 将新的已生成字符串转换为嵌入向量,添加位置编码后再输入解码器
2 编码器和解码器
编码器由 N = 6 N=6 N=6 个相同的编码器层堆叠而成,每个编码器层包含两个子层:
-
第一个子层是多头自注意力机制层;
-
第二个子层是逐位置的全连接前馈网络层。
在每个子层的外部使用残差连接,随后进行层归一化,最终的输出为
L a y e r N o r m ( x + S u b l a y e r ( x ) ) (1) \mathrm{LayerNorm}(\mathbf{x} + \mathrm{Sublayer}(\mathbf{x}))\tag{1} LayerNorm(x+Sublayer(x))(1)
其中 S u b l a y e r ( x ) \mathrm{Sublayer}(\mathbf{x}) Sublayer(x)是该子层本身实现的函数,如MHA或FFN为了使这些残差连接可行,Transformer模型中的所有子层以及嵌入层的输出维度都被设定为 d m o d e l = 512 d_{model}=512 dmodel=512
解码器由 N = 6 N=6 N=6 个相同的解码器层堆叠而成,每个解码器层包含三个子层:
- 第一个子层是多头自注意力机制层;
- 第一个子层是多头交叉注意力机制层;
- 第三个子层是逐位置的全连接前馈网络层。
与编码器类似,在解码器每个子层外部都添加了残差连接,随后进行层归一化。我们还对解码器堆栈中的自注意力子层应用了掩码机制,以防止某个位置关注未来的位置。
3 嵌入层和输出层
待翻译的自然语句在传入嵌入层之前首先要进行分词(tokenization),即转换为词元(token)。在得到分词结果后,进一步将其映射为token ID,上述两个步骤,可以通过分词器(tokenizer)来完成。
嵌入层(embedding layer)接受 x t o k e n I D ∈ R 1 × n \mathbf{x}_{\mathrm{tokenID}}\in\mathbb{R}^{1\times n} xtokenID∈R1×n作为输入,线性变换矩阵 W E ∈ R n × d m o d e l \mathbf{W}^{E}\in\mathbb{R}^{n\times d_\mathrm{model}} WE∈Rn×dmodel将其映射为张量 x ∈ R 1 × d m o d e l \mathbf{x}\in\mathbb{R}^{1\times d_{\mathrm{model}}} x∈R1×dmodel,也称为词向量(Word Embedding),即
x = x t o k e n I D W E (2) \mathbf{x}=\mathbf{x}_{\mathrm{tokenID}}\mathbf{W}^{E}\tag{2} x=xtokenIDWE(2)
线性变换矩阵 W E \mathbf{W}^{E} WE可以在训练过程中被学习
在softmax之前的线性变换层和两个嵌入层(输入嵌入层和输出嵌入层)共享权重,即 W O = ( W E W E ) T ∈ R d m o d e l × n \mathbf{W}^O=(\mathbf{W}^E\vphantom{\mathbf{W}^E})^T\in\mathbb{R}^{d_\mathrm{model}\times n} WO=(WEWE)T∈Rdmodel×n,这样保证了输入和输出的词向量空间一致,增强模型一致性。设最后一个解码器的输出为 x ∈ R 1 × d m o d e l \mathbf{x}\in\mathbb{R}^{1\times d_\mathrm{model}} x∈R1×dmodel,则线性层输出为
o = x W O = x ( W E W E ) T = [ x w 1 T , ⋯ , x w n T ] (3) \mathbf{o}=\mathbf{x}\mathbf{W}^{O}=\mathbf{x}(\mathbf{W}^E\vphantom{\mathbf{W}^E})^T=[\mathbf{x}\mathbf{w}_{1}^T,\cdots,\mathbf{x}\mathbf{w}^T_{n}]\tag{3} o=xWO=x(WEWE)T=[xw1T,⋯,xwnT](3)
上式表示解码器输出和每个原始token的语义相似度,这些相似度可以被 s o f t m a x \mathrm{softmax} softmax归一化,进而用于选择最大概率的下一个词。最终的训练结果总是趋向让语义相似的词向量在嵌入空间中方向上位置相近。
下图将预训练语言模型BERT(Bidirectional Encoder Representations from Transformers)的嵌入向量投影到二维平面进行可视化
4 位置编码
Transformer采用无顺序结构的注意力模型,它不像RNN那样天然具有序列顺序,所以必须显式地加入位置信息,否则模型无法区分token的位置。
对序列的嵌入表示 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_\mathrm{model}} X∈Rn×dmodel,其位置编码是矩阵 P E ∈ R n × d m o d e l \mathbf{PE}\in\mathbb{R}^{n\times d_\mathrm{model}} PE∈Rn×dmodel,序列中位置 k k k上的token的第 2 i 2i 2i和 2 i + 1 2i+1 2i+1个嵌入维度的位置编码分别为
P E ( k , 2 i ) = sin ( k 10000 2 i d m o d e l ) P E ( k , 2 i + 1 ) = cos ( k 10000 2 i d m o d e l ) (4) \begin{align*} PE(k,2i)&=\sin\left(\frac{k}{10000^{\frac{2i}{d_{\mathrm{model}}} }}\right)\\ PE(k,2i+1)&=\cos\left(\frac{k}{10000^{\frac{2i}{d_{\mathrm{model}}} }}\right) \end{align*}\tag{4} PE(k,2i)PE(k,2i+1)=sin(10000dmodel2ik)=cos(10000dmodel2ik)(4)
d m o d e l d_{\mathrm{model}} dmodel总为偶数,可取 i = 1 , ⋯ , d m o d e l 2 − 1 i=1,\cdots,\frac{d_{\mathrm{model}}}{2}-1 i=1,⋯,2dmodel−1,当 i i i的值较低时,即嵌入维度低时,位置编码随 k k k变化快,距离接近的token的相同嵌入维度的位置编码差异较大,使得模型能够容易地区分出短距离内的不同位置,能够提供精细的局部时序结构,对相距较远的token,由于周期性波动快,在该维度上也容易出现接近的位置编码值,从而有利于捕捉重复模式;当 i i i的值较高时,即嵌入维度较高时,位置编码位置编码随 k k k变化慢,位置相距很远的token相同嵌入维度的位置编码也很接近,这样可以捕捉文本中更广泛的信息流和长距离依赖关系,对相距较近的token,嵌入编码十分接近,容易合并局部语义信息。这种多尺度设计使Transformer能同时兼顾全局与局部的信息,从而更全面地理解语言结构。
位置编码的第一个性质是:
[ P E ( k + k 0 , 2 i ) P E ( k + k 0 , 2 i + 1 ) ] = [ cos ( ω i k 0 ) sin ( ω i k 0 ) − sin ( ω i k 0 ) cos ( ω i k 0 ) ] [ P E ( k , 2 i ) P E ( k , 2 i + 1 ) ] (5) \begin{bmatrix} PE(k+k_0,2i)\\ PE(k+k_0,2i+1) \end{bmatrix}= \begin{bmatrix} \cos\left(\omega_i k_0\right)&\sin\left(\omega_i k_0\right)\\ -\sin\left(\omega_i k_0\right)&\cos\left(\omega_i k_0\right) \end{bmatrix}\begin{bmatrix} PE(k,2i)\\ PE(k,2i+1) \end{bmatrix} \tag{5} [PE(k+k0,2i)PE(k+k0,2i+1)]=[cos(ωik0)−sin(ωik0)sin(ωik0)cos(ωik0)][PE(k,2i)PE(k,2i+1)](5)
其中 ω i = 1 10000 2 i d m o d e l \omega_i=\frac{1}{10000^{\frac{2i}{d_{\mathrm{model}}} }} ωi=10000dmodel2i1,这说明对相同的偏移 k 0 k_0 k0,不管起始点 k k k在哪里,都会导致 2 i 2i 2i和 2 i + 1 2i+1 2i+1维度上的位置编码向量经过逆时针旋转 w i k w_i k wik弧度的线性变换。
位置编码的第二个性质是:记第 j j j个token的位置编码向量为 P E ( j , ⋅ ) = [ P E ( j , 1 ) , ⋯ , P E ( j , d m o d e l ) ] \mathbf{PE}(j,\cdot)=\begin{bmatrix}PE(j,1),\cdots,PE(j,d_{\mathrm{model}})\end{bmatrix} PE(j,⋅)=[PE(j,1),⋯,PE(j,dmodel)],则
P E ( k + k 0 , ⋅ ) T P E ( k , ⋅ ) = ∑ i = 1 d m o d e l 2 − 1 cos ( ω i k 0 ) (6) \mathbf{PE}(k+k_0,\cdot)^{T}\mathbf{PE}(k,\cdot)=\sum_{i=1}^{\frac{d_{\mathrm{model}}}{2}-1}\cos(\omega_i k_0)\tag{6} PE(k+k0,⋅)TPE(k,⋅)=i=1∑2dmodel−1cos(ωik0)(6)
上式说明位置编码之间的内积只依赖偏移量 k 0 k_0 k0,可以绘制出位置编码内积随相对位置的变化
由图可见,随着相对距离 k 0 k_0 k0的增大, P E ( k + k 0 , ⋅ ) \mathbf{PE}(k+k_0,\cdot) PE(k+k0,⋅)和 P E ( k , ⋅ ) \mathbf{PE}(k,\cdot) PE(k,⋅)的内积呈减小趋势
也可以从注意力角度理解位置编码。设查询向量为 q i = ( x i + P E i ) W Q \mathbf{q}_i=(\mathbf{x}_i+\mathbf{PE}_i)\mathbf{W}^Q qi=(xi+PEi)WQ,键向量为 k j = ( x j + P E j ) W K \mathbf{k}_j=(\mathbf{x}_j+\mathbf{PE}_j)\mathbf{W}^K kj=(xj+PEj)WK,其中 x i , x j ∈ R 1 × d m o d e l \mathbf{x}_{i},\mathbf{x}_{j}\in\mathbb{R}^{1\times d_\mathrm{model}} xi,xj∈R1×dmodel, W Q , W K ∈ R d m o d e l × d k \mathbf{W}^{Q},\mathbf{W}^{K}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{k}} WQ,WK∈Rdmodel×dk,具体定义可见下文, P E i = P E ( i , ⋅ ) \mathbf{PE}_i=\mathbf{PE}(i,\cdot) PEi=PE(i,⋅),在注意力机制中,需要计算
q i k j T = ( x i + P E i ) W Q W K W k T ( x j + P E j ) T = x i W Q W K W k T x j + x i W Q W K W k T P E j + P E i W Q W K W k T x j T + P E i W Q W K W k T P E j T (7) \begin{align*} \mathbf{q}_i\mathbf{k}_{j}^{T} &=(\mathbf{x}_i+\mathbf{PE}_i)\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}(\mathbf{x}_j+\mathbf{PE}_j)^{T}\\ &=\mathbf{x}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{x}_{j}+\mathbf{x}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{PE}_j+\mathbf{PE}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{x}_{j}^T+\mathbf{PE}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{PE}_j^{T}\\ \end{align*}\tag{7} qikjT=(xi+PEi)WQWKWkT(xj+PEj)T=xiWQWKWkTxj+xiWQWKWkTPEj+PEiWQWKWkTxjT+PEiWQWKWkTPEjT(7)
其中 x i W Q W K W k T x j \mathbf{x}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{x}_{j} xiWQWKWkTxj衡量两个token的语义相似度,例如猫和狗语义相似度较高, x i W Q W K W k T P E j \mathbf{x}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{PE}_j xiWQWKWkTPEj衡量第 i i i个token是否应该关注第 j j j个位置,例如动词更加关注宾语所在的位置, P E i W Q W K W k T x j T \mathbf{PE}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{x}_{j}^T PEiWQWKWkTxjT衡量第 j j j个token是否应该关注第 i i i个位置, P E i W Q W K W k T P E j T \mathbf{PE}_i\mathbf{W}^Q\mathbf{W}^K\vphantom{\mathbf{W}^k}^{T}\mathbf{PE}^{T}_j PEiWQWKWkTPEjT衡量两个token间的相对距离。从注意力角度同样可以说明低维嵌入和高维嵌入的作用,方法是令 x i + P E i = [ x i l + P E i l , x i h + P E i h ] \mathbf{x}_{i}+\mathbf{PE}_{i}=[\mathbf{x}^{l}_{i}+\mathbf{PE}^{l}_{i},\mathbf{x}^{h}_{i}+\mathbf{PE}^{h}_{i}] xi+PEi=[xil+PEil,xih+PEih],其中 x i l + P E i l \mathbf{x}^{l}_{i}+\mathbf{PE}^{l}_{i} xil+PEil表示低维词向量和嵌入, x i h + P E i h \mathbf{x}^{h}_{i}+\mathbf{PE}^{h}_{i} xih+PEih表示高维词向量和嵌入,随后做法与 ( 5 ) (5) (5)相同。
5 注意力机制
注意力函数是从查询(query)和键-值对(key-value pairs)到输出的映射,其中query可以是单个查询向量 q \mathbf{q} q,也可以是一组查询向量 q i , i = 1 , ⋯ , m \mathbf{q}_{i},i=1,\cdots,m qi,i=1,⋯,m,而key和value必须成对出现,且二者均有多个,即 ( k i , v i ) , i = 1 , ⋯ , n (\mathbf{k}_{i},\mathbf{v}_{i}),i=1,\cdots,n (ki,vi),i=1,⋯,n。query表示token对上下文的关注点,key表示token能提供的信息,注意力函数的输出通过对所有value的加权求和计算得到,其中每个value所对应的权重,是通过query与对应key之间的相似度函数计算出来的。
5.1 缩放点积注意力机制
在Transformer中,注意力机制体现为多头注意力机制,而缩放点积注意力机制(Scaled Dot-Product Attention)是Transformer中多头注意力机制的基础。
在缩放点积注意力机制中,同时对一组查询向量进行计算,并采用矩阵运算的方式来实现注意力机制
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V (8) \mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}\tag{8} Attention(Q,K,V)=softmax(dkQKT)V(8)
其中 Q , K ∈ R n × d k , V ∈ R n × d v \mathbf{Q},\mathbf{K}\in\mathbb{R}^{n\times d_{k}},\mathbf{V}\in\mathbb{R}^{n\times d_{v}} Q,K∈Rn×dk,V∈Rn×dv, d k d_k dk表示查询空间和键空间的维度, d v d_v dv表示值空间的维度,二者都是模型可调整的超参数,为分析 ( 2 ) (2) (2),可以将三者展开
Q = [ q 1 T ⋮ q n T ] K = [ k 1 T ⋮ k n T ] V = [ v 1 T ⋮ v n T ] (9) \mathbf{Q}=\begin{bmatrix} \mathbf{q}_{1}^T\\ \vdots\\ \mathbf{q}_{n}^T\\ \end{bmatrix} \qquad \mathbf{K}=\begin{bmatrix} \mathbf{k}_{1}^T\\ \vdots\\ \mathbf{k}_{n}^T\\ \end{bmatrix} \qquad \mathbf{V}=\begin{bmatrix} \mathbf{v}_{1}^T\\ \vdots\\ \mathbf{v}_{n}^T\\ \end{bmatrix}\tag{9} Q= q1T⋮qnT K= k1T⋮knT V= v1T⋮vnT (9)
首先计算
Q K T = [ q 1 T k 1 q 1 T k 2 ⋯ q 1 T k n q 2 T k 1 q 2 T k 2 ⋯ q 2 T k n ⋮ ⋮ ⋱ ⋮ q n T k 1 q n T k 2 ⋯ q n T k n ] (10) \mathbf{Q}\mathbf{K}^T=\begin{bmatrix} \mathbf{q}_{1}^T\mathbf{k}_{1}&\mathbf{q}_{1}^T\mathbf{k}_{2}&\cdots&\mathbf{q}_{1}^T\mathbf{k}_{n}\\ \mathbf{q}_{2}^T\mathbf{k}_{1}&\mathbf{q}_{2}^T\mathbf{k}_{2}&\cdots&\mathbf{q}_{2}^T\mathbf{k}_{n}\\ \vdots&\vdots&\ddots&\vdots\\ \mathbf{q}_{n}^T\mathbf{k}_{1}&\mathbf{q}_{n}^T\mathbf{k}_{2}&\cdots&\mathbf{q}_{n}^T\mathbf{k}_{n}\\ \end{bmatrix}\tag{10} QKT= q1Tk1q2Tk1⋮qnTk1q1Tk2q2Tk2⋮qnTk2⋯⋯⋱⋯q1Tknq2Tkn⋮qnTkn (10)
上述矩阵的维度是 Q K T ∈ R n × n \mathbf{Q}\mathbf{K}^T\in\mathbb{R}^{n\times n} QKT∈Rn×n,根据该矩阵,可以计算注意力权重矩阵 s o f t m a x ( Q K T d k ) \mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right) softmax(dkQKT),其中, s o f t m a x \mathrm{softmax} softmax是在矩阵的列维度上进行,即在使每个query对应一组归一化的注意力权重,用于加权value向量,具体来说,注意力权重矩阵的第 ( i , j ) (i,j) (i,j)行表示第 i i i个query对第 j j j个key-value对分配的注意力占比,第 i i i个query对所有key-value对分配的注意力占比和为 1 1 1
下面分析除以 d k \sqrt{d_k} dk的原因
对 Q K T \mathbf{Q}\mathbf{K}^{T} QKT矩阵中的某个的元素 q k T = ∑ i = 1 d k q i k i \mathbf{q}\mathbf{k}^{T}=\sum_{i=1}^{d_{k}}q_i k_i qkT=∑i=1dkqiki,假设 q , k \mathbf{q},\mathbf{k} q,k自身的各维度之间服从独立同分布,均值都为 0 0 0,方差为 D ( q ) , D ( k ) D(q),D(k) D(q),D(k),二者彼此的各维度之间也相互独立,此时有 D ( q k T ) = d k D ( q ) D ( k ) D(\mathbf{q}\mathbf{k}^{T})=d_{k}D(q)D(k) D(qkT)=dkD(q)D(k)。可以看到,随着键/值空间维度 d k d_k dk的增加, D ( q k T ) D(\mathbf{q}\mathbf{k}^{T}) D(qkT)的方差也在增加。而对于 Q K T \mathbf{Q}\mathbf{K}^{T} QKT矩阵中的同一行 [ q i k 1 T , ⋯ , q i k n T ] [\mathbf{q}_i\mathbf{k}_{1}^{T},\cdots,\mathbf{q}_i\mathbf{k}_{n}^{T}] [qik1T,⋯,qiknT],方差增大意味着行内取值的稀疏化程度更高,因此会有较高概率取得离群的正数大值,下面说明这样的一行在经过 s o f t m a x \mathrm{softmax} softmax后会引发梯度消失问题。
向量 x = [ x 1 , ⋯ , x n ] T \mathbf{x}=[x_1,\cdots,x_n]^T x=[x1,⋯,xn]T,记 y ^ = s o f t m a x ( x ) = [ y 1 ^ , ⋯ , y d ^ ] T \hat{\mathbf{y}}=\mathrm{softmax}(\mathbf{x})=[\hat{y_1},\cdots,\hat{y_d}]^T y^=softmax(x)=[y1^,⋯,yd^]T,如果其中某个元素的值 x i x_i xi很大,那么 s o f t m a x ( x ) \mathrm{softmax}(\mathbf{x}) softmax(x)会将绝大部分的概率值分配给其所对应的 y ^ i \hat{y}_i y^i,也即 y i → 1 , y j → 0 , j ≠ i y_i\rightarrow1,y_j\rightarrow0,j\neq i yi→1,yj→0,j=i,又因为
∂ s o f t m a x ( x ) ∂ x = d i a g ( y ^ ) − y ^ y ^ T = [ y ^ 1 0 ⋯ 0 0 y ^ 2 ⋯ 0 ⋮ ⋮ ⋱ ⋮ 0 0 ⋯ y ^ d ] − [ y ^ 1 2 y ^ 1 y ^ 2 ⋯ y ^ 1 y ^ d 0 y ^ 2 2 ⋯ y ^ 2 y ^ d ⋮ ⋮ ⋱ ⋮ y ^ d y ^ 1 y ^ d y ^ 2 ⋯ y ^ d 2 ] (11) \frac{\partial\mathrm{softmax}(\mathbf{x})}{\partial\mathbf{x}}=\mathrm{diag}(\mathbf{\hat{y}})-\mathbf{\hat{y}}\mathbf{\hat{y}}^T=\begin{bmatrix} \hat{y}_1&0&\cdots&0\\ 0&\hat{y}_2&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ 0&0&\cdots&\hat{y}_d\\ \end{bmatrix}-\begin{bmatrix} \hat{y}_1^2&\hat{y}_1\hat{y}_2&\cdots&\hat{y}_1\hat{y}_d\\ 0&\hat{y}_2^2&\cdots&\hat{y}_2\hat{y}_d\\ \vdots&\vdots&\ddots&\vdots\\ \hat{y}_d\hat{y}_1&\hat{y}_d\hat{y}_2&\cdots&\hat{y}_d^2\\ \end{bmatrix}\tag{11} ∂x∂softmax(x)=diag(y^)−y^y^T= y^10⋮00y^2⋮0⋯⋯⋱⋯00⋮y^d − y^120⋮y^dy^1y^1y^2y^22⋮y^dy^2⋯⋯⋱⋯y^1y^dy^2y^d⋮y^d2 (11)
故 ∂ s o f t m a x ( x ) ∂ x → 0 \frac{\partial\mathrm{softmax}(\mathbf{x})}{\partial\mathbf{x}}\rightarrow0 ∂x∂softmax(x)→0,即遇到了梯度消失的问题。矩阵 Q K T \mathbf{Q}\mathbf{K}^{T} QKT某一行的行向量即可看做存在利离群大正值元素的 x \mathbf{x} x,其数值由初始输入经过权重可学习的线性变换得到,故直接对 Q K T \mathbf{Q}\mathbf{K}^{T} QKT矩阵的每行取 s o f t m a x \mathrm{softmax} softmax会导致权重更新困难。
通过除以 d k \sqrt{d_k} dk,可使得 D ( q k T d k ) = D ( q ) D ( k ) D\left(\frac{\mathbf{q}\mathbf{k}^{T}}{\sqrt{d_k}}\right)=D(q)D(k) D(dkqkT)=D(q)D(k),即不再依赖键/值空间维度,从而使得行内稀疏化程度降低,避免经过 s o f t m a x \mathrm{softmax} softmax导致的梯度消失问题。
5.2 多头注意力机制
Transformer中,多头注意力机制(Multi-Head Attention,MHA)可分为多头自注意力机制(Multi-Head Self Attention)和交叉注意力机制(Multi-Head Cross Attenrion)。
自注意力机制是在同一序列内部所有位置之间进行注意力计算。每个位置会同时关注该序列中的其他所有位置,从而获取整个序列的上下文信息,实现对全局依赖关系的建模。交叉注意力机制则是在两个不同序列之间进行注意力计算。通常用于解码器中,其中目标序列(如已生成的历史词)中的每个位置会关注源序列(如编码器输出的输入表示)中的所有位置,从而融合两个序列的信息,实现跨序列的上下文交互。
多头注意力允许模型从不同的表示子空间中学习信息,即每个注意力头可以关注输入序列中不同类型的依赖关系,从而提升模型的整体表达能力和灵活性。同时,虽然使用多个头会增加参数数量,但由于各个头之间是相互独立的,因此可以很好地利用GPU进行并行计算。
多头自注意力机制层的输入是一个张量 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_{\mathrm{model}}} X∈Rn×dmodel,多头交叉注意力机制层的输入包括来自编码器的张量 X e n c ∈ R n × d m o d e l \mathbf{X}_{\mathrm{enc}}\in\mathbb{R}^{n\times d_{\mathrm{model}}} Xenc∈Rn×dmodel和来自解码器的张量 X d e c ∈ R n × d m o d e l \mathbf{X}_{\mathrm{dec}}\in\mathbb{R}^{n\times d_{\mathrm{model}}} Xdec∈Rn×dmodel。多头自注意力机制层和多头交叉注意力机制层的输出都可以总体地写为
M u l t i H e a d ( X ) = C o n c a t ( h e a d 1 , ⋯ , h e a d h ) W O (12) \mathrm{MultiHead}(\mathbf{X})=\mathrm{Concat}(\mathrm{head}_1,\cdots,\mathrm{head}_h)\mathbf{W}^O\tag{12} MultiHead(X)=Concat(head1,⋯,headh)WO(12)
多头自注意力机制的每个注意力头可以计算为
h e a d i = A t t e n t i o n ( X W i Q , X W i K , X W i V ) (13) \mathrm{head}_i=\mathrm{Attention}(\mathbf{X}\mathbf{W}^{Q}_{i},\mathbf{X}\mathbf{W}^{K}_{i},\mathbf{X}\mathbf{W}^{V}_{i})\tag{13} headi=Attention(XWiQ,XWiK,XWiV)(13)
多头交叉注意力机制的每个注意力头可以计算为
h e a d i = A t t e n t i o n ( X e n c W i Q , X d e c W i K , X d e c W i V ) (14) \mathrm{head}_i=\mathrm{Attention}(\mathbf{X}_{\mathrm{enc}}\mathbf{W}^{Q}_{i},\mathbf{X}_{\mathrm{dec}}\mathbf{W}^{K}_{i},\mathbf{X}_{\mathrm{dec}}\mathbf{W}^{V}_{i})\tag{14} headi=Attention(XencWiQ,XdecWiK,XdecWiV)(14)
其中 W i Q , W i K ∈ R d m o d e l × d k , W i V ∈ R d m o d e l × d v \mathbf{W}^{Q}_{i},\mathbf{W}^{K}_{i}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{k}},\mathbf{W}^{V}_{i}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{v}} WiQ,WiK∈Rdmodel×dk,WiV∈Rdmodel×dv为对输入张量的可学习的线性变换矩阵, X W i Q , X W i K , X W i V \mathbf{X}\mathbf{W}^{Q}_{i},\mathbf{X}\mathbf{W}^{K}_{i},\mathbf{X}\mathbf{W}^{V}_{i} XWiQ,XWiK,XWiV或 X e n c W i Q , X d e c W i K , X d e c W i V \mathbf{X}_{\mathrm{enc}}\mathbf{W}^{Q}_{i},\mathbf{X}_{\mathrm{dec}}\mathbf{W}^{K}_{i},\mathbf{X}_{\mathrm{dec}}\mathbf{W}^{V}_{i} XencWiQ,XdecWiK,XdecWiV即为多缩放点积注意力机制中的query,key和value, h e a d i ∈ R n × d v \mathrm{head}_i\in\mathbb{R}^{n\times d_{v}} headi∈Rn×dv, C o n c a t ( h e a d 1 , ⋯ , h e a d h ) ∈ R n × h d v \mathrm{Concat}(\mathrm{head}_1,\cdots,\mathrm{head}_h)\in\mathbb{R}^{n\times hd_{v}} Concat(head1,⋯,headh)∈Rn×hdv表示在第一个维度上拼接 h h h个不同的注意力头得到的结果。拼接结果再乘以变换矩阵 W O ∈ R h d v × d m o d e l \mathbf{W}^O\in\mathbb{R}^{hd_{v}\times d_{\mathrm{model}}} WO∈Rhdv×dmodel即可得到MHA层最终输出 M u l t i H e a d ( X ) ∈ R n × d m o d e l \mathrm{MultiHead}(\mathbf{X})\in\mathbb{R}^{n\times d_{\mathrm{model}}} MultiHead(X)∈Rn×dmodel
5.3 掩码自注意力机制
5.3.1 填充掩码注意力机制
Transformer的输入批量中包含可变长度的句子,为了能够更好的一次性处理这些长度不同的句子,需要对Batch内的句子进行填充,使这些句子的长度相同。但是在计算注意力的时候,这些填充的部分不应该参与计算,所以要在进行softmax之前要把 Q K T d k \frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d_k}} dkQKT的填充部分进行Mask(值设置为 − ∞ -\infty −∞),使之在经过softmax后无限接近 0 0 0,上述过程称为填充掩码注意力机制(Padding Mask Head Attention)
所有注意力层均须要进行填充Mask,设注意力层的输入是 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_{\mathrm{model}}} X∈Rn×dmodel, n n n是序列统一长度,首先计算 Q = X W Q , K = X W K , V = X W V \mathbf{Q}=\mathbf{X}\mathbf{W}^{Q},\mathbf{K}=\mathbf{X}\mathbf{W}^{K},\mathbf{V}=\mathbf{X}\mathbf{W}^{V} Q=XWQ,K=XWK,V=XWV,其中 W Q , W K ∈ R d m o d e l × d k , W V ∈ R d m o d e l × d v \mathbf{W}^{Q},\mathbf{W}^{K}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{k}},\mathbf{W}^{V}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{v}} WQ,WK∈Rdmodel×dk,WV∈Rdmodel×dv是该注意力头的线性变换权重,加入填充掩码后,注意力头的输出可以计算为
P a d d i n g M a s k A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k + M p a d d i n g ) V (15) \mathrm{PaddingMaskAttention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}+\mathbf{M}_{\mathrm{padding}}\right)\mathbf{V}\tag{15} PaddingMaskAttention(Q,K,V)=softmax(dkQKT+Mpadding)V(15)
其中为该注意力头的线性变换矩阵, M p a d d i n g ∈ R n × n \mathbf{M}_{\mathrm{padding}}\in\mathbb{R}^{n\times n} Mpadding∈Rn×n称填充掩码矩阵,其中
M p a d d i n g ( i , j ) = { − ∞ 第 j 个 t o k e n 被填充 0 第 j 个 t o k e n 没有被填充 (16) \mathbf{M}_{\mathrm{padding}}(i,j)=\left\{ \begin{array}{ll} -\infty&第j个\mathrm{token}被填充\\ 0&第j个\mathrm{token}没有被填充 \end{array} \right.\tag{16} Mpadding(i,j)={−∞0第j个token被填充第j个token没有被填充(16)
填充掩码矩阵中,当某个token被填充时,直接将一整列元素的值都置为负无穷,表示所有查询都不关注该token,在实际应用中,负无穷可以用绝对值很大的负数代替。
5.3.2 因果掩码注意力机制
在传统的循环神经网络(Recurrent Neural Network, RNN)中,序列是按时间步逐步自回归生成的,这种方式称为串行处理,效率很低。为了充分利用GPU的能力加速训练,Transformer并不是真的每步都重新运行一次前向传播,而是一次性把整个目标序列输入解码器,然后并行地计算每个位置的输出。
在上述并行机制下,由于一次性输入了整个目标序列到解码器,故对解码器的自注意力层采用因果掩码自注意力机制(Causal Mask Self Attention),来防止模型在计算自注意力时看到未来的信息,以维持因果关系。因果掩码同样通过对 Q K T d k \frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d_k}} dkQKT的不含对角线的上三角部分的元素进行Mask来实现。
设解码器自注意力层的输入是 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_{\mathrm{model}}} X∈Rn×dmodel,首先计算 Q = X W Q , K = X W K , V = X W V \mathbf{Q}=\mathbf{X}\mathbf{W}^{Q},\mathbf{K}=\mathbf{X}\mathbf{W}^{K},\mathbf{V}=\mathbf{X}\mathbf{W}^{V} Q=XWQ,K=XWK,V=XWV,其中 W Q , W K ∈ R d m o d e l × d k , W V ∈ R d m o d e l × d v \mathbf{W}^{Q},\mathbf{W}^{K}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{k}},\mathbf{W}^{V}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{v}} WQ,WK∈Rdmodel×dk,WV∈Rdmodel×dv是该解码器自注意力头的线性变换权重,加入填充掩码后,解码器自注意力头的输出可以计算为
C a u s a l M a s k A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k + M c a u s a l ) V (17) \mathrm{CausalMaskAttention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}+\mathbf{M}_{\mathrm{causal}}\right)\mathbf{V}\tag{17} CausalMaskAttention(Q,K,V)=softmax(dkQKT+Mcausal)V(17)
其中 M c a u s a l ∈ R n × n \mathbf{M}_{\mathrm{causal}}\in\mathbb{R}^{n\times n} Mcausal∈Rn×n称为因果掩码矩阵掩码,其定义为
M c a u s a l ( i , j ) = { − ∞ j > i 0 j ⩽ i (18) \mathbf{M}_{\mathrm{causal}}(i,j)= \left\{ \begin{array}{ll} -\infty&j>i\\ 0&j\leqslant i \end{array} \right.\tag{18} Mcausal(i,j)={−∞0j>ij⩽i(18)
多头掩码自注意力机制中,每个注意力头都按照 ( 2 ) (2) (2)计算。在推理阶段,由于生成过程是串行进行的,每次只基于当前已生成的部分来预测下一个token,因此不需要使用掩码。实际上,解码器自注意力层需要同时考虑填充掩码和因果掩码,其掩码矩阵为 M p a d d i n g + M c a u s a l \mathbf{M}_{\mathrm{padding}}+\mathbf{M}_{\mathrm{causal}} Mpadding+Mcausal,同样将其加到 Q K T d k \frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d_k}} dkQKT上即可。
6 逐位置前馈神经网络
逐位置前馈神经网络(Positional-wise Feed-Forward Network, Positional-wise FFN)由两个线性变换,及中间夹着的一个ReLU激活函数组成,通过引入组合线性与非线性变换,Positional-wise FFN层加强了每个token表示的建模能力。设Positional-wise FFN层的输入张量是 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_{\mathrm{model}}} X∈Rn×dmodel,每一行 x ∈ R 1 × d m o d e l \mathbf{x}\in\mathbb{R}^{1\times d_{\mathrm{model}}} x∈R1×dmodel都是一个token表示。Positional-wise FFN层对序列中每一个位置的token,都独立地应用相同的前馈神经网络
F F N ( x ) = m a x ( 0 , x W 1 + b 1 ) W 2 + b 2 (19) \mathrm{FFN}(\mathbf{x})=\mathrm{max}(\mathbf{0},\mathbf{x}\mathbf{W}_{1}+\mathbf{b}_{1})\mathbf{W}_{2}+\mathbf{b}_{2}\tag{19} FFN(x)=max(0,xW1+b1)W2+b2(19)
其中 W 1 ∈ R d m o d e l × d f f , W 2 ∈ R d f f × d m o d e l \mathbf{W}_{1}\in\mathbb{R}^{d_{\mathrm{model}}\times d_{\mathrm{ff}}},\mathbf{W}_{2}\in\mathbb{R}^{d_{\mathrm{ff}}\times d_{\mathrm{model}}} W1∈Rdmodel×dff,W2∈Rdff×dmodel, b 1 ∈ R d m o d e l , b 2 ∈ R d f f \mathbf{b}_{1}\in\mathbb{R}^{d_{\mathrm{model}}},\mathbf{b}_{2}\in\mathbb{R}^{d_{\mathrm{ff}}} b1∈Rdmodel,b2∈Rdff, d f f d_{\mathrm{ff}} dff是Positional-wise FFN的中间输出维度。因为每个token共享权重,故对完整输入 X \mathbf{X} X的处理可以直接通过张量乘法进行,而无需进行张量拼接。
Positional-wise FFN层和普通全连接层的区别在于普通全连接层的单个样本输入是 X ∈ R C \mathbf{X}\in\mathbb{R}^{C} X∈RC, C C C是样本特征维度,其变换对象是样本,Positional-wise FFN层的单个样本输入是 X ∈ R n × d m o d e l \mathbf{X}\in\mathbb{R}^{n\times d_{\mathrm{model}}} X∈Rn×dmodel,其变换对象是样本的每个token
7 层归一化
在介绍层归一化(Layer Normalization)之前,首先回顾批归一化(Batch Normalization)。在神经网络训练过程中,每一层的输入分布会随着前一层参数的变化而变化,这种现象称为内部协变量偏移(Internal Covariate Shift),该现象会导致这种现象会导致后续层需要不断适应新的输入分布,使得模型训练不稳定,收敛变慢。BatchNorm通过标准化每一层的输入,使其分布趋于稳定,从而减少内部协变量偏移,加速训练并提高模型稳定性。
BatchNorm层的输入张量是 X ∈ R B × C \mathbf{X}\in\mathbb{R}^{B\times C} X∈RB×C,BatchNorm层独立地计算每个特征维度 c = 1 , ⋯ , C c=1,\cdots,C c=1,⋯,C在整个Batch上的均值和方差
μ c = 1 B ∑ i = 1 B X i , c σ c 2 = 1 B ∑ i = 1 B ( X i , c − μ c ) 2 (20) \begin{align*} \mu_{c}=\frac{1}{B}\sum_{i=1}^{B}X_{i,c}\qquad\sigma^2_{c}=\frac{1}{B}\sum_{i=1}^{B}(\mathbf{X}_{i,c}-\mu_{c})^2 \end{align*}\tag{20} μc=B1i=1∑BXi,cσc2=B1i=1∑B(Xi,c−μc)2(20)
其中 X i , c X_{i,c} Xi,c表示批量上第 i i i个样本的第 c c c个维度的取值。得到均值和方差后,BatchNorm对输入张量进行标准化,使得每个特征维度在当前Batch下均值为 0 0 0,方差为 1 1 1,并通过可学习的参数 γ \gamma γ和 β \beta β进行缩放和平移
Y i , c = X i , c − μ c σ c 2 + ϵ ∗ γ + β (21) \mathbf{Y}_{i,c}=\frac{X_{i,c}-\mu_c}{\sqrt{\sigma_c^2+\epsilon}}*\gamma+\beta\tag{21} Yi,c=σc2+ϵXi,c−μc∗γ+β(21)
在BatchNorm中,认为不同输入样本在某特征维度上的取值是从固定的取值分布中的多次采样,初始情况下该分布会受到内部协变量偏移影响,通过标准化规范当前Batch下该特征维度上的取值的均值为和方差可以起到加速收敛的效果。在自然语言处理任务中,输入通常是是批量数据 { X i } i = 1 , ⋯ , B \{\mathbf{X}_{i}\}_{i=1,\cdots,B} {Xi}i=1,⋯,B,批量中的每个样本 X i ∈ R n i × d m o d e l \mathbf{X}_{i}\in\mathbb{R}^{n_i\times d_{\mathrm{model}}} Xi∈Rni×dmodel表示一个长度为 n i n_i ni的可变长序列(在Transformer中填充到统一长度),其每个token向量维度是 d m o d e l d_{\mathrm{model}} dmodel,一个token向量称为一层。由于批量中的不同样本的同一位置上的取值几乎没有关联,不可看做从固定的取值分布中的采样,此时标准化失去意义,对不同类别的数据做标准化,可能破坏原始特征结构,故不能应用BatchNorm
与BatchNorm不同,LayerNorm对每个样本的每一层特征进行标准化处理,让每个token的特征向量尺度一,从而保证了注意力机制的稳定性。采用LayerNorm是为了保证自注意力机制的稳定性,我们需要独立地计算批量中每个样本的每个token在其 d m o d e l d_{\mathrm{model}} dmodel的维特征空间上的均值和方差
μ i , j = 1 d m o d e l ∑ d = 1 d m o d e l X i , j , d σ i , j 2 = 1 B ∑ i = 1 B ( X i , j , d − μ i , j ) 2 (22) \begin{align*} \mu_{i,j}=\frac{1}{d_{\mathrm{model}}}\sum_{d=1}^{d_{\mathrm{model}}}X_{i,j,d}\qquad\sigma^2_{i,j}=\frac{1}{B}\sum_{i=1}^{B}(X_{i,j,d}-\mu_{i,j})^2 \end{align*}\tag{22} μi,j=dmodel1d=1∑dmodelXi,j,dσi,j2=B1i=1∑B(Xi,j,d−μi,j)2(22)
其中 X i , j , d X_{i,j,d} Xi,j,d表示批量上第 i i i个样本的位置 j j j上的词向量的第 d d d个维度的取值。得到均值和方差后,LayerNorm同样对输入张量进行标准化,使得每个样本的每个token的词向量具有零均值和单位方差,并通过可学习的参数 γ \gamma γ和 β \beta β进行缩放和平移
Y i , j , d = γ X i , j , d − μ i , j σ i , j 2 + ϵ + β (23) \mathbf{Y}_{i,j,d}=\gamma\frac{X_{i,j,d}-\mu_{i,j}}{\sqrt{\sigma_{i,j}^2+\epsilon}}+\beta\tag{23} Yi,j,d=γσi,j2+ϵXi,j,d−μi,j+β(23)
由于LayerNorm是基于单个样本的特征维度进行标准化,因此非常适合应用于自然语言处理等需要处理变长序列的任务中。与BatchNorm相似,也可以认为LayerNorm解决了token特征向量的内部协变量偏移问题。
需要说明,token特征向量模长也具有意义,LayerNorm确实损失了模长的直接信息,但这是设计者有意为之的选择,目的在于提升训练稳定性和模型泛化能力。而模长的重要性,会通过其他路径,如线性层等,重新学习回来,也就是说,模长是可以通过上下文、参数学习间接编码的,而不是直接保留最初模长。