Transformer详解（附代码）

引言

 $\mathrm{Transformer}$模型是$\mathrm{Google}$团队在$2017$年$6$月由$\mathrm{Ashish}\mathrm{Vaswani}$等人在论文《$\mathrm{Attention}\mathrm{Is}\mathrm{All}\mathrm{You}\mathrm{Need}$》所提出，当前它已经成为$\mathrm{NLP}$领域中的首选模型。$\mathrm{Transformer}$抛弃了$\mathrm{RNN}$的顺序结构，采用了$\mathrm{Self}$-$\mathrm{Attention}$机制，使得模型可以并行化训练，而且能够充分利用训练资料的全局信息，加入$\mathrm{Transformer}$的$Seq2seq$模型在$\mathrm{NLP}$的各个任务上都有了显著的提升。本文做了大量的图示目的是能够更加清晰地讲解$\mathrm{Transformer}$的运行原理，以及相关组件的操作细节，文末还有完整可运行的代码示例。

注意力机制

 $\mathrm{Transformer}$中的核心机制就是$\mathrm{Self}$-$\mathrm{Attention}$。$\mathrm{Self}$-$\mathrm{Attention}$机制的本质来自于人类视觉注意力机制。当人视觉在感知东西时候往往会更加关注某个场景中显著性的物体，为了合理利用有限的视觉信息处理资源，人需要选择视觉区域中的特定部分，然后集中关注它。注意力机制主要目的就是对输入进行注意力权重的分配，即决定需要关注输入的哪部分，并对其分配有限的信息处理资源给重要的部分。

Self-Attention

 $\mathrm{Self}$-$\mathrm{Attention}$工作原理如上图所示，给定输入$\mathrm{word}\mathrm{embedding}$向量$a^1,a^2,a^3\in {\mathbb{R}}^{d_l\times 1}$，然后对于输入向量 a i , i ∈ { 1 , 2 , 3 } a^i,i\in \{1,2,3\} ai,i∈{1,2,3}通过矩阵$W^q\in {\mathbb{R}}^{d_k\times d_l},W^k\in {\mathbb{R}}^{d_k\times d_l},W^v\in {\mathbb{R}}^{d_l\times d_l}$进行线性变换得到$\mathrm{Query}$向量$q^i\in {\mathbb{R}}^{d_k\times 1}$，$\mathrm{Key}$向量$k^i\in {\mathbb{R}}^{d_k\times 1}$，以及$\mathrm{Value}$向量$v^i\in {\mathbb{R}}^{d_l\times 1}$，即 { q i = W q ⋅ a i k i = W k ⋅ a i , i ∈ { 1 , 2 , 3 } v i = W v ⋅ a i \left\{\begin{aligned}q^i&=W^q \cdot a^i\\k^i&=W^k \cdot a^i,\quad i\in\{1,2,3\}\\v^i&=W^v \cdot a^i\end{aligned}\right. ⎩ ⎨ ⎧​qikivi​=Wq⋅ai=Wk⋅ai,i∈{1,2,3}=Wv⋅ai​如果令矩阵$A=(a^1,a^2,a^3)\in {\mathbb{R}}^{d_l\times 3}$，$Q=(q^1,q^2,q^3)\in {\mathbb{R}}^{d_k\times 3}$，$K=(k^1,k^2,k^3)\in {\mathbb{R}}^{d_k\times 3}$，$V=(v^1,v^2,v^3)\in {\mathbb{R}}^{d_l\times 3}$，则此时则有$$\{\begin{array}{rl}Q& =W^q\cdot A\\ K& =W^k\cdot A\\ V& =W^v\cdot A\end{array}$$接着再利用得到的$\mathrm{Query}$向量和$\mathrm{Key}$向量计算注意力得分，论文中采用的注意力计算公式为点积缩放公式 α l i = ( q i ) ⊤ ⋅ k l d k = d k d k ∑ n = 1 d k k n l ⋅ q n i , i , l ∈ { 1 , 2 , 3 } \alpha^{i}_l=\frac{(q^i)^{\top}\cdot k^l}{\sqrt{d^k}}=\frac{\sqrt{d^k}}{d^k}\sum\limits_{n=1}^{d^k}k^l_n\cdot q^i_n,\quad i,l \in \{1,2,3\} αli​=dk ​(qi)⊤⋅kl​=dkdk ​​n=1∑dk​knl​⋅qni​,i,l∈{1,2,3}论文中假定$\mathrm{Key}$向量$k^l=(k_1^l,k_2^l,k_3^l)$的元素和$\mathrm{Query}$向量$q^i=(q_1^i,q_2^i,q_3^i)$的元素独立同分布，且令均值为$0$，方差为$1$，则此时注意力向量$a^i\in {\mathbb{R}}^{3\times 1}$的第$l$个分量${\alpha }_l^i$的均值为$0$，方差$1$具体的计算公式如下 E [ α l i ] = d k d k ∑ n = 1 d k E [ k n l ] ⋅ E [ q n i ] = 0 , i , l ∈ { 1 , 2 , 3 } V a r [ α l i ] = 1 d k ∑ n = 1 d k V a r [ k n l ] ⋅ V a r [ q n i ] = 1 , i , l ∈ { 1 , 2 , 3 } \begin{aligned}\mathbb{E}\left[\alpha^i_l\right]&=\frac{\sqrt{d^k}}{d^k}\sum\limits_{n=1}^{d^k}\mathbb{E}\left[k^l_n\right]\cdot \mathbb{E}\left[q^i_n\right]=0,\quad i,l \in \{1,2,3\}\\ \mathrm{Var}\left[\alpha^i_l\right]&=\frac{1}{d^k}\sum\limits_{n=1}^{d^k}\mathrm{Var}\left[k^l_n\right]\cdot \mathrm{Var}\left[q^i_n\right]=1,\quad i,l \in \{1,2,3\}\end{aligned} E[αli​]Var[αli​]​=dkdk ​​n=1∑dk​E[knl​]⋅E[qni​]=0,i,l∈{1,2,3}=dk1​n=1∑dk​Var[knl​]⋅Var[qni​]=1,i,l∈{1,2,3}​令注意力分数矩阵$\Lambda =({\alpha }^1,{\alpha }^2,{\alpha }^3)\in {\mathbb{R}}^{3\times 3}$，则有$$\Lambda =\frac{K^{\top }\cdot Q}{\sqrt{d^k}}$$注意分数向量${\alpha }^i$经过$\mathrm{softmax}$层得到归一化后的注意力分布${\beta }^i$，即为 β j i = e α j i ∑ n = 1 3 e α n i , i , j = { 1 , 2 , 3 } \beta^i_j = \frac{e^{\alpha^{i}_j}}{\sum\limits_{n=1}^3e^{\alpha^{i}_n}},\quad i,j=\{1,2,3\} βji​=n=1∑3​eαni​eαji​​,i,j={1,2,3}最后利用得到的注意力分布向量${\beta }^i$和$\mathrm{Value}$矩阵$V$获得最后的输出$b^i\in {\mathbb{R}}^{d_l\times 1}$，则有 b i = ∑ l = 1 3 β l i ⋅ v l , i ∈ { 1 , 2 , 3 } b^i=\sum\limits^{3}_{l=1}\beta^{i}_l \cdot v^{l},\quad i \in \{1,2,3\} bi=l=1∑3​βli​⋅vl,i∈{1,2,3}令输出矩阵$B=(b^1,b^2,b^3)\in {\mathbb{R}}^{d_l\times 3}$，则有$$B=\mathrm{Attention}(Q,K,V)=V\cdot \mathrm{softmax}\left(\frac{K^{\top }\cdot Q}{\sqrt{d^k}}\right)$$

Multi-Head Attention

 $\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$的工作原理与$\mathrm{Self}$-$\mathrm{Attention}$的工作原理非常类似。为了方便图解可视化将$\mathrm{Multi}$-$\mathrm{Head}$设置为$2$-$\mathrm{Head}$，如果$\mathrm{Multi}$-$\mathrm{Head}$设置为$8$-$\mathrm{Head}$，则上图的 q i , k i , v i , i ∈ { 1 , 2 , 3 } q^i,k^i,v^i,i\in\{1,2,3\} qi,ki,vi,i∈{1,2,3}的下一步的分支数为$8$。给定输入$\mathrm{word}\mathrm{embedding}$向量$a^1,a^2,a^3\in {\mathbb{R}}^{d_l\times 1}$，然后对于输入向量 a i , i ∈ { 1 , 2 , 3 } a^i,i\in \{1,2,3\} ai,i∈{1,2,3}通过矩阵$W^q\in {\mathbb{R}}^{d_k\times d_l},W^k\in {\mathbb{R}}^{d_k\times d_l},W^v\in {\mathbb{R}}^{d_l\times d_l}$进行第一次线性变换得到$\mathrm{Query}$向量$q^i\in {\mathbb{R}}^{d_k\times 1}$，$\mathrm{Key}$向量$k^i\in {\mathbb{R}}^{d_k\times 1}$，以及$\mathrm{Value}$向量$v^i\in {\mathbb{R}}^{d_l\times 1}$。然后再对$\mathrm{Query}$向量$q^i$通过矩阵$W^{q1}\in {\mathbb{R}}^{d_m\times d_k}$和$W^{q2}\in {\mathbb{R}}^{d_m\times d_k}$进行第二次线性变换得到$q^{i1}\in {\mathbb{R}}^{d_m\times 1}$和$q^{i2}\in {\mathbb{R}}^{d_m\times 1}$，同理对$\mathrm{Key}$向量$k^i$通过矩阵$W^{k1}\in {\mathbb{R}}^{d_m\times d_k}$和$W^{k2}\in {\mathbb{R}}^{d_m\times d_k}$进行第二次线性变换得到$k^{i1}\in {\mathbb{R}}^{d_m\times 1}$和$k^{i2}\in {\mathbb{R}}^{d_m\times 1}$，对$\mathrm{Value}$向量$v^i$通过矩阵$W^{v1}\in {\mathbb{R}}^{\frac{d_l}{2}\times d_l}$和$W^{v2}\in {\mathbb{R}}^{\frac{d_l}{2}\times d_l}$进行第二次线性变换得到$v^{i1}\in {\mathbb{R}}^{\frac{d_l}{2}\times 1}$和$v^{i2}\in {\mathbb{R}}^{\frac{d_l}{2}\times 1}$,具体的计算公式如下所示： { q i h = W q h ⋅ W q ⋅ a i k i h = W k h ⋅ W k ⋅ a i , i = { 1 , 2 , 3 } , h = { 1 , 2 } v i h = W v h ⋅ W v ⋅ a i \left\{\begin{aligned}q^{ih}&=W^{qh}\cdot W^{q} \cdot a^i\\ k^{ih}&=W^{kh}\cdot W^{k} \cdot a^i,\quad i=\{1,2,3\},\quad h=\{1,2\}\\v^{ih}&=W^{vh}\cdot W^{v} \cdot a^i\end{aligned}\right. ⎩ ⎨ ⎧​qihkihvih​=Wqh⋅Wq⋅ai=Wkh⋅Wk⋅ai,i={1,2,3},h={1,2}=Wvh⋅Wv⋅ai​令矩阵$$\begin{array}{ll}{Q}^1=(q^{11},q^{21},q^{31})\in {\mathbb{R}}^{d_m\times 3}& Q^2=(q^{12},q^{22},q^{32})\in {\mathbb{R}}^{d_m\times 3}\\ K^1=(k^{11},k^{21},k^{31})\in {\mathbb{R}}^{d_m\times 3}& K^2=(k^{12},k^{22},k^{32})\in {\mathbb{R}}^{d_m\times 3}\\ V^1=(v^{11},v^{21},v^{31})\in {\mathbb{R}}^{\frac{d_l}{2}\times 3}& V^2=(v^{12},v^{22},v^{32})\in {\mathbb{R}}^{\frac{d_l}{2}\times 3}\end{array}$$此时则有$$\begin{array}{ll}{Q}^1=W^{q1}\cdot W^q\cdot A& Q^2=W^{q2}\cdot W^q\cdot A\\ K^1=W^{k1}\cdot W^k\cdot A& K^2=W^{k2}\cdot W^k\cdot A\\ V^1=W^{v1}\cdot W^v\cdot A& V^2=W^{v2}\cdot W^v\cdot A\end{array}$$对于每个$\mathrm{Head}$利用得到对于$\mathrm{Query}$向量和$\mathrm{Key}$向量计算对应的注意力得分，其中注意力向量${\alpha }^{ih}$的第$l$个分量的计算公式为 α l i h = ( q i h ) ⊤ ⋅ k l h , i ∈ { 1 , 2 , 3 } , h ∈ { 1 , 2 } , l ∈ { 1 , 2 , 3 } \alpha^{ih}_l=(q^{ih})^{\top}\cdot k^{lh},\quad i\in\{1,2,3\},h\in\{1,2\},l\in\{1,2,3\} αlih​=(qih)⊤⋅klh,i∈{1,2,3},h∈{1,2},l∈{1,2,3}令注意力分数矩阵${\Lambda }^1=({\alpha }^{11},{\alpha }^{21},{\alpha }^{31})$，${\Lambda }^2=({\alpha }^{12},{\alpha }^{22},{\alpha }^{32})$，则有$${\Lambda }^1=\frac{(K^1{)}^{\top }\cdot Q^1}{\sqrt{d_m}},{\Lambda }^2=\frac{(K^2{)}^{\top }\cdot Q^2}{\sqrt{d_m}}$$注意分数向量${\alpha }^{ih}$经过$\mathrm{softmax}$层得到归一化后的注意力分布${\beta }^{ih}$，即为 β j i h = e α j i h ∑ n = 1 3 e α n i h , i , j = { 1 , 2 , 3 } , h = { 1 , 2 } \beta^{ih}_j = \frac{e^{\alpha^{ih}_j}}{\sum\limits_{n=1}^3e^{\alpha^{ih}_n}},\quad i,j=\{1,2,3\}, h=\{1,2\} βjih​=n=1∑3​eαnih​eαjih​​,i,j={1,2,3},h={1,2}对于每一个$\mathrm{Head}$利用得到的注意力分布向量${\beta }^{ih}$和$\mathrm{Value}$矩阵$V^h$获得最后的输出$b^{ih}\in {\mathbb{R}}^{\frac{d_l}{2}\times 1}$，则有 b i h = ∑ l = 1 3 β l i h ⋅ v l h , i ∈ { 1 , 2 , 3 } , h ∈ { 1 , 2 } b^{ih}=\sum\limits^{3}_{l=1}\beta^{ih}_l \cdot v^{lh},\quad i \in \{1,2,3\}, h\in\{1,2\} bih=l=1∑3​βlih​⋅vlh,i∈{1,2,3},h∈{1,2}两个$\mathrm{Head}$的$b^{ih}$的向量按照如下方式拼接在一起，则有$$B=\left(\begin{array}{lll}{b}^{11}& b^{21}& b^{31}\\ b^{12}& b^{22}& b^{32}\end{array}\right)\in {\mathbb{R}}^{d_l\times 3}$$给定参数矩阵$W^O\in {\mathbb{R}}^{d_l\times d_l}$，则输出矩阵为$$O=W^O\cdot B\in {\mathbb{R}}^{d_l\times 3}$$综上所述则有$$O=\mathrm{MultiHead}(Q,K,V)=W^O\cdot \mathrm{Concat}\left(\begin{array}{l}{V}^1\cdot \mathrm{softmax}\left(\frac{(K^1{)}^{\top }\cdot Q^1}{\sqrt{d_m}}\right)\\ \\ V^2\cdot \mathrm{softmax}\left(\frac{(K^2{)}^{\top }\cdot Q^2}{\sqrt{d_m}}\right)\end{array}\right)$$

Mask Self-Attention

  如下图左半部分所示，$\mathrm{Self}$-$\mathrm{Attention}$的输出向量 b i , i ∈ { 1 , 2 , 3 , 4 } b^i, i \in \{1,2,3,4\} bi,i∈{1,2,3,4}综合了输入向量 a i , i ∈ { 1 , 2 , 3 , 4 } a^i, i \in \{1,2,3,4\} ai,i∈{1,2,3,4}的全部信息，由此可见，$\mathrm{Self}$-$\mathrm{Attention}$在实际编程中支持并行运算。如下图右半部分所示，$\mathrm{Mask}\mathrm{Self}$-$\mathrm{Attention}$的输出向量$b^i$只利用了已知部分输入的向量$a^i$的信息。例如，$b1$只是与$a^1$有关；$b^2$与$a^1$和$a^2$有关；$b^3$与$a^1$，$a^2$和$a^3$有关；$b^4$与$a^1$，$a^2$，$a^3$和$a^4$有关。$\mathrm{Mask}\mathrm{Self}$-$\mathrm{Attention}$在$\mathrm{Transformer}$中被用到过两次。

• $\mathrm{Transformer}$的$\mathrm{Encoder}$中如果输入一句话的$\mathrm{word}$长度小于指定的长度，为了能够让长度一致往往会用$0$进行填充，此时则需要用$\mathrm{Mask}\mathrm{Self}$-$\mathrm{Attention}$来计算注意力分布。
• $\mathrm{Transformer}$的$\mathrm{Decoder}$的输出是有时序关系的，当前的输出只与之前的输入有关，所以此时算注意力分布时需要用到$\mathrm{Mask}\mathrm{Self}$-$\mathrm{Attention}$。

Transformer模型

 以上对$\mathrm{Transformer}$中的核心内容即自注意力机制进行了详细解剖，接下来会对$\mathrm{Transformer}$模型架构进行介绍。$\mathrm{Transformer}$模型是由$\mathrm{Encoder}$和$\mathrm{Decoder}$两个模块组成，具体的示意图如下所示，为了能够对$\mathrm{Transformer}$内部的操作细节进行更清晰的展示，下图以矩阵运算的视角对$\mathrm{Transformer}$的原理进行讲解。
 $\mathrm{Encoder}$模块操作的具体流程如下所示：

• $\mathrm{Encoder}$的输入由两部分组成分别是词编码矩阵$I\in {\mathbb{R}}^{n\times l\times d}$和位置编码矩阵$P\in {\mathbb{R}}^{n\times l\times d}$，其中$n$表示句子数目，$l$表示一句话单词的最大数目，$d$表示的是词向量的维度。位置编码矩阵$P$表示的是每个单词在一句里的所有位置信息，因为$\mathrm{Self}$-$\mathrm{Attention}$计算注意力分布的时候只能给出输出向量和输入向量之间的权重关系，但是不能给出词在一句话里的位置信息，所以需要在输入里引入位置编码矩阵$P$。位置编码向量生成方法有很多。一种比较简单粗暴的方式就是根据单词在句子中的位置生成一个$\mathrm{one}$-$\mathrm{hot}$的位置编码；还有的方法是将位置编码当成参数进行训练学习；在该论文里是利用三角函数对位置进行编码，具体的公式如下所示$$\mathrm{PE}(pos,2i)=\sin (\frac{pos}{1000^{2i/d}}),\mathrm{PE}(pos,2i+1)=\cos (\frac{pos}{1000^{2i/d}})$$其中$\mathrm{PE}$表示的是位置编码向量，$pos$表示词在句子中的位置，$i$表示编码向量的位置索引。
• 输入矩阵$I+P$通过线性变换生成矩阵$Q$，$K$，$V$。在实际编程中是将输入$I+P$直接赋值给$Q$，$K$，$V$。如果输入单词长度小于最大长度并$0$来填充的时候，还要相应引入$\mathrm{Mask}$矩阵。
• 将矩阵$Q$，$K$，$V$输入到$\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$模块中进行注意分布的计算得到矩阵$I^{\prime }\in {\mathbb{R}}^{n\times l\times d}$，计算公式为$$I^{\prime }=\mathrm{MultiHead}(Q,K,V)$$具体的计算细节参考上文关于$\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$原理的讲解不在这里赘述。然后将原始输入$I+P$与注意力分布$I^{\prime }$进行残差计算得到输出矩阵$I+P+I^{\prime }\in {\mathbb{R}}^{n\times l\times d}$。
• 对矩阵 I + P + I ′ = { x i j k } n l d I+P+I^{\prime}=\{x_{ijk}\}^{nld} I+P+I′={xijk​}nld进行层归一化操作得到$I^{\prime \prime }\in {\mathbb{R}}^{n\times l\times d}$，具体的计算公式为 { μ i j = ∑ k = 1 d x i j k σ i j = ∑ k = 1 d ( x i j k − μ i j ) 2 ⟹ x ^ i j k = x i j k − u i j σ i j , i ∈ { 1 , ⋯   , n } , j ∈ { 1 , ⋯   , l } , k ∈ { 1 , ⋯   , d } \left\{\begin{aligned}\mu^{ij}&=\sum\limits_{k=1}^d x_{ijk}\\\sigma^{ij}&=\sqrt{\sum\limits_{k=1}^d\left(x_{ijk}-\mu^{ij}\right)^2}\end{aligned}\right. \Longrightarrow \hat{x}_{ijk}=\frac{x_{ijk}-u^{ij}}{\sigma^{ij}},\quad i\in\{1,\cdots,n\},j\in\{1,\cdots,l\},k\in\{1,\cdots,d\} ⎩ ⎨ ⎧​μijσij​=k=1∑d​xijk​=k=1∑d​(xijk​−μij)2 ​​⟹x^ijk​=σijxijk​−uij​,i∈{1,⋯,n},j∈{1,⋯,l},k∈{1,⋯,d}
• 将$I^{\prime \prime }$输入到全连接神经网络中得到$I^{\prime \prime \prime }\in {\mathbb{R}}^{n\times l\times d}$ ，然后再让全连接神经网络的输入$I^{\prime \prime }$与输出$I^{\prime \prime \prime }$进行残差计算得到$I^{\prime \prime }+I^{\prime \prime \prime }$，接着对$I^{\prime \prime }+I^{\prime \prime \prime }$进行层归一化操作。
• 以上是一个$\mathrm{Block}$的操作原理，将$N$个$\mathrm{Block}$进行堆叠就组成了$\mathrm{Encoder}$的模块，得到的最后输出为$I^N\in {\mathbb{R}}^{n\times l\times d}$。这里需要注意的是$\mathrm{Encoder}$模块中的各个组件的操作顺序并不是固定的，也可以先进行归一化操作，然后再计算注意力分布，再归一化，再预测等。

 $\mathrm{Decoder}$模块操作的具体流程如下所示：

• $\mathrm{Decoder}$的输入也由两部分组成分别是词编码矩阵$O\in {\mathbb{R}}^{n_1\times l_1\times d}$和位置编码矩阵$P^O\in {\mathbb{R}}^{n_1\times l_1\times d}$。因为$\mathrm{Decoder}$的输入是具有时顺序关系的（即上一步的输出为当前步输入）所以还需要输入$\mathrm{Mask}$矩阵$M$以便计算注意力分布。
• 输入矩阵$O+P^O$通过线性变换生成矩阵$\hat{Q}$，$\hat{K}$，$\hat{V}$。在实际编程中是将输入$O+P^O$直接赋值给$\hat{Q}$，$\hat{K}$，$\hat{V}$。如果输入单词长度小于最大长度并$0$来填充的时候，还要相应引入$\mathrm{Mask}$矩阵。
• 将矩阵$\hat{Q}$，$\hat{K}$，$\hat{V}$以及$\mathrm{Mask}$矩阵$M$输入到$\mathrm{Mask}\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$模块中进行注意分布的计算得到矩阵$O^{\prime }\in {\mathbb{R}}^{n_1\times l_1\times d}$，计算公式为$$O^{\prime }=\mathrm{MaskMultiHead}(\hat{Q},\hat{K},\hat{V},M)$$具体的计算细节参考上文关于$\mathrm{Mask}\mathrm{Self}$-$\mathrm{Attention}$的讲解不在这里赘述。然后将原始输入$O+P^O$与注意力分布$O^{\prime }$进行残差计算得到输出矩阵$O+P^O+O^{\prime }\in {\mathbb{R}}^{n_1\times l_1\times d}$。接着再对矩阵$O+P^O+O^{\prime }$进行层归一化操作得到$O^{\prime \prime }\in {\mathbb{R}}^{n_1\times l_1\times d}$。
• $\mathrm{Encoder}$的输出$I^N$通过线性变换得到$Q^N$和$K^N$，$O^{\prime }$进行线性变换得到${\hat{V}}^{\prime }$，利用矩阵$Q^N$和$K^N$和${\hat{V}}^{\prime }$进行交叉注意力分布的计算得到$O^{\prime \prime \prime }$，计算公式为$$O^{\prime \prime \prime }=\mathrm{MultiHead}(Q^N,K^N,{\hat{V}}^{\prime })$$这里的交叉注意力分布综合$\mathrm{Encoder}$输出结果和$\mathrm{Decoder}$中间结果的信息。实际编程编程中将$I^N$直接赋值给$\hat{Q}$和$\hat{K}$，$O^{\prime }$直接赋值给${\hat{V}}^{\prime }$。然后将$O^{\prime \prime }$与注意力分布$O^{\prime \prime \prime }$进行残差计算得到输出矩阵$O^{\prime \prime }+O^{\prime \prime \prime }$。
• 接着对$O^{\prime \prime }+O^{\prime \prime \prime }$进行层归一操作得到$O^{\prime \prime \prime \prime }$，再将$O^{\prime \prime \prime \prime }$输入到全连接神经网络中得到$O^{\prime \prime \prime \prime \prime }$，接着再做一步残差操作得到$O^{\prime \prime \prime \prime }+O^{\prime \prime \prime \prime \prime }$，最后再进行一层归一化操作。
• 以上是一个$\mathrm{Block}$的操作原理，将$N$个$\mathrm{Block}$进行堆叠就组成了$\mathrm{Decoder}$的模块，得到的输出为$O^N\in {\mathbb{R}}^{n_1\times l_1\times d}$。然后在词汇字典中找到当前预测最大概率的单词，并将该单词词向量作为下一阶段的输入，重复以上步骤，直到输出“$\mathrm{end}$”字符为止。

代码示例

 $\mathrm{Transformer}$具体的代码示例如下所示为一个国外博主视频里的代码，并根据上文对代码的一些细节进行了探讨。根据上文中$\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$原理示例图可知，严格来看$\mathrm{Multi}$-$\mathrm{Head}\mathrm{Attention}$在求注意分布的时候中间其实是有两步线性变换。给定输入向量$x\in {\mathbb{R}}^{256\times 1}$ 第一步线性变换直接让向量$x$赋值给$q$，$k$，$v$，这一过程以下程序中有所体现，在这里并不会产生歧义。第二步线性变换产生多$\mathrm{Head}$，假设$\mathrm{Head}=8$的时候，按理说$q$要与$8$个矩阵$W^{q1},\cdots ,W^{q8}$进行线性变换得到$8$个$q^1,\cdots ,q^8$，同理$k$要与$8$个矩阵$W^{k1},\cdots ,W^{k8}$进行线性变换得到$8$个$k^1,\cdots ,k^8$，$v$要与$8$个矩阵$W^{v1},\cdots ,W^{v8}$进行线性变换得到$8$个$v^1,\cdots ,v^8$，如果按照这个方式在程序实现则需要定义24个权重矩阵，非常的麻烦。以下程序中有一个简单的权重定义方法，通过该方法也可以实现以上多$\mathrm{Head}$的线性变换，以向量$q=(q_1,\cdots ,q_{256}{)}^{\top }\in {\mathbb{R}}^{256\times 1}$为例：

• 首先将向量$q$进行截断分成$\mathrm{Head}=8$个向量，即为$$\{\begin{array}{rl}{q}^{(1)}& =(\mathbf{E},0,0,0,0,0,0,0)\cdot q\\ q^{(2)}& =(0,\mathbf{E},0,0,0,0,0,0)\cdot q\\ q^{(3)}& =(0,0,\mathbf{E},0,0,0,0,0)\cdot q\\ q^{(4)}& =(0,0,0,\mathbf{E},0,0,0,0)\cdot q\\ q^{(5)}& =(0,0,0,0,\mathbf{E},0,0,0)\cdot q\\ q^{(6)}& =(0,0,0,0,0,\mathbf{E},0,0)\cdot q\\ q^{(7)}& =(0,0,0,0,0,0,\mathbf{E},0)\cdot q\\ q^{(8)}& =(0,0,0,0,0,0,0,\mathbf{E})\cdot q\end{array}$$其中$q^{(i)}\in {\mathbb{R}}^{32\times 1}$是$q$的第$i$个截断向量，$\mathbf{E}\in {\mathbb{R}}^{32\times 32}$是单位矩阵，$0\in {\mathbb{R}}^{32\times 32}$是零矩阵。
• 然后对 q ( i ) , i ∈ { 1 , ⋯   , 8 } q^{(i)},i\in \{1,\cdots,8\} q(i),i∈{1,⋯,8}用相同的权重矩阵$W\in {\mathbb{R}}^{32\times 32}$进行线性变换，此时可以发现，训练过程的时候只需要更新权重矩阵$W$即可，而且可以进行多$\mathrm{Head}$线性变换，$8$个权重矩阵可以表示为：$$\{\begin{array}{rl}{W}^{q1}& =W\cdot (\mathbf{E},0,0,0,0,0,0,0)=(W,0,0,0,0,0,0,0)\\ W^{q2}& =W\cdot (0,\mathbf{E},0,0,0,0,0,0)=(0,W,0,0,0,0,0,0)\\ W^{q3}& =W\cdot (0,0,\mathbf{E},0,0,0,0,0)=(0,0,W,0,0,0,0,0)\\ W^{q4}& =W\cdot (0,0,0,\mathbf{E},0,0,0,0)=(0,0,0,W,0,0,0,0)\\ W^{q5}& =W\cdot (0,0,0,0,\mathbf{E},0,0,0)=(0,0,0,0,W,0,0,0)\\ W^{q6}& =W\cdot (0,0,0,0,0,\mathbf{E},0,0)=(0,0,0,0,0,W,0,0)\\ W^{q7}& =W\cdot (0,0,0,0,0,0,\mathbf{E},0)=(0,0,0,0,0,0,W,0)\\ W^{q8}& =W\cdot (0,0,0,0,0,0,0,\mathbf{E})=(0,0,0,0,0,0,0,W)\end{array}$$其中权重矩阵 W q i ∈ R 32 × 256 , i ∈ { 1 , ⋯   , 8 } W^{qi}\in\mathbb{R}^{32 \times 256},i\in\{1,\cdots,8\} Wqi∈R32×256,i∈{1,⋯,8}。

import torch
import torch.nn as nn
import os

class SelfAttention(nn.Module):
	def __init__(self, embed_size, heads):
		super(SelfAttention, self).__init__()
		self.embed_size = embed_size
		self.heads = heads
		self.head_dim = embed_size // heads

		assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads"

		self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

	def forward(self, values, keys, query, mask):
		N =query.shape[0]
		value_len , key_len , query_len = values.shape[1], keys.shape[1], query.shape[1]

		# split embedding into self.heads pieces
		values = values.reshape(N, value_len, self.heads, self.head_dim)
		keys = keys.reshape(N, key_len, self.heads, self.head_dim)
		queries = query.reshape(N, query_len, self.heads, self.head_dim)
		
		values = self.values(values)
		keys = self.keys(keys)
		queries = self.queries(queries)

		energy = torch.einsum("nqhd,nkhd->nhqk", queries, keys)
		# queries shape: (N, query_len, heads, heads_dim)
		# keys shape : (N, key_len, heads, heads_dim)
		# energy shape: (N, heads, query_len, key_len)

		if mask is not None:
			energy = energy.masked_fill(mask == 0, float("-1e20"))

		attention = torch.softmax(energy/ (self.embed_size ** (1/2)), dim=3)

		out = torch.einsum("nhql, nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads*self.head_dim)
		# attention shape: (N, heads, query_len, key_len)
		# values shape: (N, value_len, heads, heads_dim)
		# (N, query_len, heads, head_dim)

		out = self.fc_out(out)
		return out


class TransformerBlock(nn.Module):
	def __init__(self, embed_size, heads, dropout, forward_expansion):
		super(TransformerBlock, self).__init__()
		self.attention = SelfAttention(embed_size, heads)
		self.norm1 = nn.LayerNorm(embed_size)
		self.norm2 = nn.LayerNorm(embed_size)

		self.feed_forward = nn.Sequential(
			nn.Linear(embed_size, forward_expansion*embed_size),
			nn.ReLU(),
			nn.Linear(forward_expansion*embed_size, embed_size)
		)
		self.dropout = nn.Dropout(dropout)

	def forward(self, value, key, query, mask):
		attention = self.attention(value, key, query, mask)

		x = self.dropout(self.norm1(attention + query))
		forward = self.feed_forward(x)
		out = self.dropout(self.norm2(forward + x))
		return out


class Encoder(nn.Module):
	def __init__(
			self,
			src_vocab_size,
			embed_size,
			num_layers,
			heads,
			device,
			forward_expansion,
			dropout,
			max_length,
		):
		super(Encoder, self).__init__()
		self.embed_size = embed_size
		self.device = device
		self.word_embedding = nn.Embedding(src_vocab_size, embed_size)
		self.position_embedding = nn.Embedding(max_length, embed_size)

		self.layers = nn.ModuleList(
			[
				TransformerBlock(
					embed_size,
					heads,
					dropout=dropout,
					forward_expansion=forward_expansion,
					)
				for _ in range(num_layers)]
		)
		self.dropout = nn.Dropout(dropout)


	def forward(self, x, mask):
		N, seq_length = x.shape
		positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)
		out = self.dropout(self.word_embedding(x) + self.position_embedding(positions))
		for layer in self.layers:
			out = layer(out, out, out, mask)

		return out


class DecoderBlock(nn.Module):
	def __init__(self, embed_size, heads, forward_expansion, dropout, device):
		super(DecoderBlock, self).__init__()
		self.attention = SelfAttention(embed_size, heads)
		self.norm = nn.LayerNorm(embed_size)
		self.transformer_block = TransformerBlock(
			embed_size, heads, dropout, forward_expansion
		)

		self.dropout = nn.Dropout(dropout)

	def forward(self, x, value, key, src_mask, trg_mask):
		attention = self.attention(x, x, x, trg_mask)
		query = self.dropout(self.norm(attention + x))
		out = self.transformer_block(value, key, query, src_mask)
		return out

class Decoder(nn.Module):
	def __init__(
			self,
			trg_vocab_size,
			embed_size,
			num_layers,
			heads,
			forward_expansion,
			dropout,
			device,
			max_length,
	):
		super(Decoder, self).__init__()
		self.device = device
		self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)
		self.position_embedding = nn.Embedding(max_length, embed_size)
		self.layers = nn.ModuleList(
			[DecoderBlock(embed_size, heads, forward_expansion, dropout, device)
			for _ in range(num_layers)]
			)
		self.fc_out = nn.Linear(embed_size, trg_vocab_size)
		self.dropout = nn.Dropout(dropout)

	def forward(self, x ,enc_out , src_mask, trg_mask):
		N, seq_length = x.shape
		positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)
		x = self.dropout((self.word_embedding(x) + self.position_embedding(positions)))

		for layer in self.layers:
			x = layer(x, enc_out, enc_out, src_mask, trg_mask)

		out =self.fc_out(x)
		return out


class Transformer(nn.Module):
	def __init__(
			self,
			src_vocab_size,
			trg_vocab_size,
			src_pad_idx,
			trg_pad_idx,
			embed_size = 256,
			num_layers = 6,
			forward_expansion = 4,
			heads = 8,
			dropout = 0,
			device="cuda",
			max_length=100
		):
		super(Transformer, self).__init__()
		self.encoder = Encoder(
			src_vocab_size,
			embed_size,
			num_layers,
			heads,
			device,
			forward_expansion,
			dropout,
			max_length
			)
		self.decoder = Decoder(
			trg_vocab_size,
			embed_size,
			num_layers,
			heads,
			forward_expansion,
			dropout,
			device,
			max_length
			)


		self.src_pad_idx = src_pad_idx
		self.trg_pad_idx = trg_pad_idx
		self.device = device


	def make_src_mask(self, src):
		src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
		# (N, 1, 1, src_len)
		return src_mask.to(self.device)

	def make_trg_mask(self, trg):
		N, trg_len = trg.shape
		trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
			N, 1, trg_len, trg_len
		)
		return trg_mask.to(self.device)

	def forward(self, src, trg):
		src_mask = self.make_src_mask(src)
		trg_mask = self.make_trg_mask(trg)
		enc_src = self.encoder(src, src_mask)
		out = self.decoder(trg, enc_src, src_mask, trg_mask)
		return out


if __name__ == '__main__':
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(device)
	x = torch.tensor([[1,5,6,4,3,9,5,2,0],[1,8,7,3,4,5,6,7,2]]).to(device)
	trg = torch.tensor([[1,7,4,3,5,9,2,0],[1,5,6,2,4,7,6,2]]).to(device)

	src_pad_idx = 0
	trg_pad_idx = 0
	src_vocab_size = 10
	trg_vocab_size = 10
	model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, device=device).to(device)
	out = model(x, trg[:, : -1])
	print(out.shape)