Transformer（论文 + PyTorch源码解读）

Transformer模型早在2017年就出现了，当时实验室的分享也有关于这个的。但我当时没有意识到这篇论文的厉害之处，听名字感觉像是那种昙花一现的论文，也没有关注它。直到最近出现了BERT这一神物之后，方才后知后觉此时Transformer已然这么有用！因此，这才仔仔细细地撸了这篇“古老”的论文和源码，这里将主要对照论文和相应的PyTorch源码进行逐一对照解读。因笔者能力有限，如有不详实之处，读者可移步至文末的传送门去看更多细节，并欢迎指出~

文章目录

前言

2017年6月，Google发布了一篇论文《Attention is All You Need》，提出了Transformer模型。正如论文的名称所说，其旨在全部利用Attention方式来替代掉RNN的循环机制，从而能并行化计算并实现提速。同时，在特定的任务上，这个模型也超过了当时Google神经机器翻译模型。笔者主要阅读了论文及两篇博客（链接见文末的传送门），这里主要是对这些内容做一个整合和提炼~

一. 背景

在Transformer出现之前，LSTM、GRU等RNN系列网络以及encoder-decoder+attention架构基本上铸就了所有NLP任务的铁桶江山。但RNN的一个缺陷在于是自回归的模型，只能串行的一步一步进行计算，无法并行化。因此有一些网络如ByteNet和ConvS2S都是以此为切入点，使用CNN作为基本构建模块，这样可以并行计算所有输入和输出位置的隐层表示。但在这些模型中，关联来自两个任意输入或输出位置的信号所需的操作数量会随着位置之间的距离而增长，如ConvS2S呈线性增长、ByteNet呈对数增长， 这使得学习较远位置之间的依赖变得更加困难。而在Transformer中，两个输入之间的距离对其计算来说没有影响，都是一样的，它没有使用RNN和卷积，可以进行并行计算。

二. Transformer整体架构

下面是从论文中截出的Transformer整体结构图：
Transformer（论文 + PyTorch源码解读）

这个图乍一看非常唬人，但实际上仔细看的话仍旧是熟悉的Encoder-Decoder架构，左边的是Encoder，右边的是Decoder。下面将一一进行剖析。

三. Transformer编码器

首先来看Encoder部分（左半部分），它是由N层方框里面的内容堆叠起来的。对于每一层来说，都由两部分构成：一部分是multi-head self-attention机制，另一部分是一个简单的全连接前馈网络。在每一部分上，都使用残差+layer normalization来进行处理。论文中，这样的方框有6个，即N=6N=6，模型的隐层单元数dmodel=512d_{model} = 512。

1. 自注意力机制

Encoder内部没有使用RNN，取而代之的是一种self-attention（自注意力）机制。

一般我们用的attention机制，可以抽象为输入一个查询（query），去查询键值对（key-value pair）中的key，然后得到一个概率分布，再据此对value进行加权相加，获取当前query下的注意力表征。而我们的query，往往是Decoder中某一个step的输出，key-value pair往往是encoder的输出。

论文里面使用的也是这种attention机制，只不过其query、key、value都是由encoder的输出经过不同的变换而来，也即self-attention，所有的东西都是自己。他们定义了一种叫“Scaled Dot-Product Attention”的计算方式，用于计算给定query、key和value下的注意力表征，如下图（左）所示：
Transformer（论文 + PyTorch源码解读）
这里的QQ、KK和VV分别表示query、key和value矩阵，它们的维度分别为Lq∗dkL_q * d_k、Lk∗dkL_k * d_k、Lk∗dvL_k * d_v。计算公式为：
Transformer（论文 + PyTorch源码解读）
一般我们经常使用的attention计算方式有两种：一种是乘性attention，即使用内积的方式；另一种是加性attention，即使用额外一层隐藏层来计算。这两种计算方式理论上复杂度是差不多的，但乘性attention因为可以用矩阵运算，会更节省时间和空间。对照着上图（左）和公式来看，这个公式与乘性attention计算方式的唯一不同就在于使用了一个缩放因子1dk\frac{1}{\sqrt{d_k}}。这里为何要进行缩放呢？论文中给出了解释：在dkd_k比较小的时候，不加缩放的效果和加性attention的效果差不多，但当dkd_k比较大的时候，不加缩放的效果就明显比加性attention的效果要差，怀疑是当dkd_k增长的时候，内积的量级也会增长，导致softmax函数会被推向梯度较小的区域，为了缓解这个问题，加上了这个缩放项进行量级缩小。

论文里面还提到，只使用一个attention的计算方式未免太过单薄，所以他们提出了multi-head（多头）注意力机制。将注意力的计算分散到不同的子空间进行，以期能从多方面进行注意力的学习，具体做法如上图（右）所示。并行地将QQ、KK和VV通过不同的映射矩阵映射到不同的空间（每个空间是一个头），再分别在这些空间中对应着进行单个“Scaled Dot-Product Attention”的学习，最后将得到的多头注意力表征进行拼接，经过一个额外的映射层映射到原来的空间。其公式如下：
Transformer（论文 + PyTorch源码解读）
这里的WiQ∈Rdmodel∗dkW_i^Q \in R^{d_{model} * d_k}，WiK∈Rdmodel∗dkW_i^K \in R^{d_{model} * d_k}，WiV∈Rdmodel∗dvW_i^V \in R^{d_{model} * d_v}，WO∈Rhdv∗dmodelW^O \in R^{hd_v * d_{model}}。表示第ii个头的变换矩阵，hh表示头的个数。

在论文里面，h=8h = 8，并且dk=dv=dmodel/h=64d_k = d_v = d_{model} / h = 64。可见这里虽然分了很多头去计算，但整体的维度还是不变的，因此计算量还是一样的。

2. 前馈网络

这部分是整体架构图中的Feed Forward模块，其实就是一个简单的全连接前馈网络。它由两层全连接及ReLU**函数构成，计算公式如下：
Transformer（论文 + PyTorch源码解读）

这里的全连接是Position-wise逐位置的，即设前面的attention输出的维度为B∗Length∗dmodelB * Length * d_{model}，则变换时，实际上是只针对dmodeld_{model}进行变换，对于每个位置（Length维度）上，都使用同样的变换矩阵。

在论文中，这里的dmodeld_{model}仍然是512，两层全连接的中间隐层单元数为dff=2048d_{ff} = 2048。

3. add & norm

在整体架构图中，还有一个部分是add&norm，这其实是借鉴了图像中的残差思想。在self-attention和feed forward计算之后都会加上一个残差变换，同时也会加上Layer Normalization（参见： https://arxiv.org/pdf/1607.06450.pdf ，用在有循环机制的网络里面效果较好）。设输入为xx，则输出为LayerNorm(x+SubLayer(x))LayerNorm(x+SubLayer(x))，这里的SubLayerSubLayer即是self-attention或feed forward层。

四. Transformer解码器

接着来看Decoder部分（右半部分），它同样也是由N层（在论文中，仍取N=6N = 6）堆叠起来，对于其中的每一层，除了与Encoder中相同的self-attention及Feed Forward之外，还在中间插入了一层传统encoder-decoder框架中的attention层，即将decoder的输出作为query去查询encoder的输出，同样用的是multi-head attention，使得在decode的时候能看到encoder的所有输出。

同时，作为decoder，在预测当前步的时候，是不能知道后面的内容的，即attention需要加上mask，将当前步之后的分数全部置为−∞-\infty，然后再计算softmax，以防止发生数据泄露。

五. Embedding和Softmax

最后来看Transformer整体架构图中的两头，即Embedding和Softmax部分。

1. Embedding和Softmax层

论文中将embedding层的参数与最后的Softmax层之前的变换层参数进行了共享，并且在embedding层，将嵌入的结果乘上dmodel\sqrt{d_{model}}。

2. 位置编码层

细心的读者可能发现了，在整体架构图中，还有一个叫Positional Encoding的东西，这是个啥？

Transformer虽然摒弃了RNN的循环结构和CNN的局部相关性，但对于序列来说，最重要的其实还是先后顺序。看前面self-attention的处理方式，实际上与“词袋”模型没什么区别，这样忽略了位置信息的缺陷肯定是要通过一定的手段来弥补。

论文中提出了一个非常“smart”的方式来加入位置信息，就是这里的Positional Encoding，它对于每个位置pospos进行编码，然后与相应位置的word embedding进行相加，构成当前位置的新word embedding。它采用如下的公式为每个pospos进行编码：
Transformer（论文 + PyTorch源码解读）

其中，ii表示embedding向量中的位置，即dmodeld_{model}中的每一维。选择这种sin函数有两种好处：1）可以不用训练，直接编码即可，而且不管什么长度，都能直接得到编码结果；2）能表示相对位置，根据sin(α+β)=sinαcosβ+cosαsinβsin(\alpha+\beta)=sin\alpha cos\beta + cos\alpha sin\beta，PEpos+kPE_{pos+k}可以表示为PEposPE_{pos}的线性变换，这为表达相对位置信息提供了可能性。

六. PyTorch实现

对于PyTorch实现部分，主要参考的是 http://nlp.seas.harvard.edu/2018/04/03/attention.html 。这里将针对核心部分进行剖析和解读：

1. 多头注意力机制

multi-head attention可用于三个地方，分别是Encoder和Decoder中各自的self-attention部分，还有Encoder-Decoder之间的attention部分。但其实这三个地方的不同仅仅在于query、key、value和mask的不同，因此当将这4部分作为参数传入时，模型的计算方式便可抽象为如下的形式：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)  # (3 + 1)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        
        # 1) Do all the linear projections in batch from d_model => h x d_k 
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]
        
        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

参数里面的h和d_model分别表示注意力头的个数，以及模型的隐层单元数。注意到在__ini__函数中，定义了self.linears = clones(nn.Linear(d_model, d_model), 4)，clone(x, N)即为深拷贝N份，这里定义了4个全连接函数，实际上是3+1，其中的3个分别是Q、K和V的变换矩阵，最后一个是用于最后将多头concat之后进行变换的矩阵。

在forward函数中，是首先将query、key和value进行相应的变换，然后需要经过attention这个函数的计算，这个函数实际上就是“Scaled Dot Product Attention”这个模块的计算，如下所示：（注意这里面的mask方式）

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

2. 前馈网络

前面说了，前馈网络实际上就是两层全连接，其代码如下：

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

3. add & norm

这是残差模块+LayerNormalization的实现方式：

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))

在forward函数里面，彷佛与前面的LayerNorm(x+SubLayer(x))LayerNorm(x+SubLayer(x))不太一样，其实这里都可以的，主要是看任务，自己实验。

下面是LayerNormalization的实现，其实PyTorch里面已经集成好了nn.LayerNorm，这里列出来只是方便读者看清其原理，为了代码简洁，可以直接使用PyTorch里面实现好的函数。

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

4. 位置编码

位置编码相关的代码如下所示：

class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], 
                         requires_grad=False)
        return self.dropout(x)

可见，这里首先是按照最大长度max_len生成一个位置，而后根据公式计算出所有的向量，在forward函数中根据长度取用即可，非常方便。注意要设置requires_grad=False，因其不参与训练。

5. 关于mask

在Transformer里面，Encoder和Decoder的attention计算都需要相应的mask处理，但功能却不同。在Encoder中，mask主要是为了让那些在一个batch中长度较短的序列的padding不参与attention的计算，而在Decoder中，还要考虑不能发生数据泄露。那这些具体是怎么实现的呢？看下面的代码：

class Batch:
    "Object for holding a batch of data with mask during training."
    def __init__(self, src, trg=None, pad=0):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if trg is not None:
            self.trg = trg[:, :-1]
            self.trg_y = trg[:, 1:]
            self.trg_mask = \
                self.make_std_mask(self.trg, pad)
            self.ntokens = (self.trg_y != pad).data.sum()
    
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & Variable(
            subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
        return tgt_mask

对于src的mask方式就比较简单，直接把pad给mask掉即可。对于trg的mask计算略微复杂一些，不仅需要把pad给mask掉，还需要进行一个subsequent_mask的操作，其代码如下：

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0

这里是给定一个序列长度size，生成一个下三角矩阵，在主对角线右上的都是false，其示意图如下：

Transformer（论文 + PyTorch源码解读）

经过&得到的mask即为最终trg需要的mask。

七. 为啥要用Self-Attention？

论文专门开了一个章节来阐释为什么选用self-attention这种方式来代替RNN和CNN，我这也开一个章节专门讲一下吧，以示尊重。

论文从计算复杂度、序列操作数以及最大路径长度三个角度比较了不同的层，包括Self-Attention、Recurrent、Convolutional等，如下表所示：
Transformer（论文 + PyTorch源码解读）
表里面的nn代表序列长度，dd代表向量维度，kk表示kernel的大小，rr表示受限的memory的长度（主要是针对过长序列，直接使用self-attention未免太大）。从表中的数据看起来，好像Self-Attention确实比较优良。

八. 总结

  1. 模型特点：采用全attention的方式，完全摒弃了RNN和CNN的做法。
  2. 优势：训练速度更快，在两个翻译任务上取得了SoTA。
  3. 不足：在decode阶段还是自回归的，即还是不能并行。

传送门

论文：https://arxiv.org/pdf/1706.03762.pdf
源码：https://github.com/tensorflow/tensor2tensor （TensorFlow）
https://github.com/OpenNMT/OpenNMT-py （PyTorch）
https://github.com/awslabs/sockeye (MXNet)
参考：https://jalammar.github.io/illustrated-transformer 一个优质的英文博客，有很好的可视化图例，适合不进行原理深究或只关注实现的入门级博客。其后面还有很多好的资源可以用来参考，一些是googleblog，还有视频等，可以收藏后慢慢研读！
http://nlp.seas.harvard.edu/2018/04/03/attention.html PyTorch实现的核心源码博客，有原理，也有对应的代码段，非常适合对照学习！

相关文章：

猜你喜欢
相关资源
相似解决方案
热门标签
Java Python linux javascript Mysql C# Docker 算法 前端 SpringBoot Redis Vue spring 设计模式 .net core .net kubernetes c++ 数据库 数据结构 大数据 js 机器学习 微服务 Android Go 程序员 面试 JVM ASP.net core 云原生 人工智能 后端 PHP git CSS golang k8s Nginx Django mybatis 深度学习 多线程 React 架构 devops 爬虫 云计算 Spring Boot LeetCode