transformer如何实现并行

 

RNN 无法并行

 

我们先看一个典型的基于RNN的Encoder-Decoder结构

 

输入是：“机器学习“，输出是“machine learning”。模型的大概工作时序是：Encoder部分，输入序列逐个送进RNN，计算出最后时刻的隐藏状态c，作为上下文信息传给Decoder。Decoder部分，将c和t-1时刻的输出作为t时刻的输入，逐步计算预测得到t时刻的输出。这个结构中，Encoder和Decoder部分都是无法并行化的，这是由RNN结构本身决定的（t时刻计算依赖t-1时刻的输出）。值得注意的一点是，这里我们讨论的并行化，不是指多个“输入-输出”样本批量处理，而是指对于单个的“输入-输出”样本，模型中Encoder和Decoder都需要历经t次迭代。

RNN之所以不支持并行化是因为它天生是个时序结构，t时刻依赖t-1时刻的输出，而t-1时刻又依赖t-2时刻，如此循环往前，我们可以说t时刻依赖了前t时刻所有的信息。RNN的实质目的是学习时序数据中的局部依赖关系，实际上深度学习很多模型的目的都是为了学习样本之间的局部关系，如CNN也是为了学习空间数据中局部依赖关系。为了完成对长距离样本之间依赖关系的学习，CNN采用多次卷积、池化等操作，这样需要加大网络的深度，才能更多的学习到更大范围的样本局部依赖关系。RNN虽然理论上可以建立长距离依赖关系，但由于存在梯度消失、网络容量等问题，实际上效果也要打折扣。

Transformer并行化

Transformer中的Encoder和Decoder分别如何支持并行化训练。

Encoder支持并行化

理解Encoder部分能支持并行化训练的关键就在于理解自注意力机制：

 不严谨的说，自注意力机制就是利用 xi之间两两相关性作为权重的一种加权平均，将每一个xi 映射到 zi。

这样说可能还有点抽象，如上图所示，   x1 和x2 表示“Thinking”和“Machines”的embedding。

此时可以说  x1 和x2 都只表示自己的word，没有考虑局部依赖关系。

自注意力机制首先计算了 x1 和x1 之间的相关性， x1 和x2之间的相关性，记为 a11和a12 ，

 ，通过对x1,x2 这样加权平均方式得到的z1 来表示“Thinking”，此时我们说z1

已经考虑了和其他元素之间的依赖关系。

这里我们举例是一个简单的长度为2的句子，可以推广到长度为N的句子处理上，

z1 是所有X1:N 的加权平均，哪个元素的权重大，我们说它获得了更高的“注意力”。

同理，我们可以可以计算Z2：N ，每个 Zi 都是所有X1:N 的不同加权平均，因为每个Zi 都可以看作是对应的单词在考虑全局依赖后的一种embedding，不同的 zi 对不同的xi

注意力分布不同。

换句话说，在计算每一个zi

时，需要的是全体的x1：N ，而并不依赖 zi-1 ，所以Encoder可以并行化计算所有的zi ，而无需像RNN那样严格按照时序依次迭代计算。

 

 

 

Decoder支持并行化

假设目标语句为：“I love China”这样一个长度为3的句子。首先我们给句子加入一个起始符号：“<start> I love China”。对于解码器模块，我们先来看一个正常的处理思路。

第一轮：给解码器模块输入“<start>” 和 编码器的输出结果，解码器输出“I”

第二轮：给解码器模块输入“<start> I” 和 编码器的输出结果，解码器输出“Iove”

第三轮：给解码器模块输入“<start> I love” 和 编码器的输出结果，解码器输出“China”

第四轮：给解码器模块输入“<start> I love China” 和 编码器的输出结果，解码器输出“<end>”，至此完成。

这就是一个“完美”解码器按时序处理的基本流程。之所以说完美，是因为解码器每次都成功的预测出了正确的单词，如果编码器在某一轮预测错了，那么给下一轮解码器的输入中就包含了错误的信息，然后继续解码。

Transformer之所以能支持Decoder部分并行化训练，是基于以下两个关键点：

1. teacher force
2. masked self attention

对于teacher force，在其他seq2seq模型中也有应用。它是指在每一轮预测时，不使用上一轮预测的输出，而强制使用正确的单词。还以上面这个例子来说，第二轮时，给解码器模块输入“<start> I” 和 编码器的输出结果，解码器没有正确预测出“Iove”，而是得到了“want”。如果没有采用teacher force，在第三轮时，解码器模块输入的就是“<start> I want”。如果采用了 teacher force，第三轮时，解码器模块输入的仍然是“<start> I love”。通过这样的方法可以有效的避免因中间预测错误而对后续序列的预测，从而加快训练速度。而Transformer采用这个方法，为并行化训练提供了可能，因为每个时刻的输入不再依赖上一时刻的输出，而是依赖正确的样本，而正确的样本在训练集中已经全量提供了。值得注意的一点是：Decoder的并行化仅在训练阶段，在测试阶段，因为我们没有正确的目标语句，t时刻的输入必然依赖t-1时刻的输出，这时跟之前的seq2seq就没什么区别了。

对于masked self attention，它本质就是一种自注意力，只是加了一个mask，也正是这个mask为Decoder提供了并行化训练的可能。为了更好的说明mask的作用，我自己手画了两个简图，画的比较丑，请读者凑合着看。

 

 

上图是一个普通的自注意力机制示意图，实际上Encoder部分的自注意力机制大概就是这个关系。图中为了不太凌乱，我没有画出z2 、z3、z4 与x1：4 的关联，实际上Z1:4 和x1:4

是两两关联的。

有的读者可能发现了，这是不是类似于一个的全连接网络？

其实自注意力机制本质就是一种全连接网络，只是普通的全连接网络每层的节点数量在建立模型时就已经确定了，对于一个模型这些连接的权重也通过训练得到。

而自注意力机制可以适应不同长度的输入，这些权重是由K、Q动态生成的，我们要训练得到的是产生K、Q的参数矩阵。这一点和我们要说的并行化训练关系不大，只是做一点补充。

以上这个自注意力模型在训练Decoder无法直接使用，因为Decoder是在做解码，解码时要求第t时刻的输出，是看不到t时刻以后的样本的。换句话说，我们需要把

z1和x2:4 之间的连接去掉，z1 只与x1 相关。同理 只与 相关， 只与 相关 ， 与 全相关。如下图所示：

 

 

REF：

https://zhuanlan.zhihu.com/p/368592551

https://zhuanlan.zhihu.com/p/368592551