Transformer解读（论文 + PyTorch源码）_pytorch的transformer论文

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整体结构图：

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

三. Transformer细节剖析

1. 编码器

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

2. 自注意力机制

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下的注意力表征，如下图（左）所示：

这里的$Q$、$K$和$V$分别表示query、key和value矩阵，它们的维度分别为$L_q\ast d_k$、$L_k\ast d_k$、$L_k\ast d_v$。计算公式为：

一般我们经常使用的attention计算方式有两种：一种是乘性attention，即使用内积的方式；另一种是加性attention，即使用额外一层隐藏层来计算。这两种计算方式理论上复杂度是差不多的，但乘性attention因为可以用矩阵运算，会更节省时间和空间。对照着上图（左）和公式来看，这个公式与乘性attention计算方式的唯一不同就在于使用了一个缩放因子$\frac{1}{\sqrt{d_k}}$。这里为何要进行缩放呢？论文中给出了解释：在$d_k$比较小的时候，不加缩放的效果和加性attention的效果差不多，但当$d_k$比较大的时候，不加缩放的效果就明显比加性attention的效果要差，怀疑是当$d_k$增长的时候，内积的量级也会增长，导致softmax函数会被推向梯度较小的区域，为了缓解这个问题，加上了这个缩放项进行量级缩小。

论文里面还提到，只使用一个attention的计算方式未免太过单薄，所以他们提出了multi-head（多头）注意力机制。将注意力的计算分散到不同的子空间进行，以期能从多方面进行注意力的学习，具体做法如上图（右）所示。并行地将$Q$、$K$和$V$通过不同的映射矩阵映射到不同的空间（每个空间是一个头），再分别在这些空间中对应着进行单个“Scaled Dot-Product Attention”的学习，最后将得到的多头注意力表征进行拼接，经过一个额外的映射层映射到原来的空间。其公式如下：

这里的$W_i^Q\in R^{d_{model}\ast d_k}$，$W_i^K\in R^{d_{model}\ast d_k}$，$W_i^V\in R^{d_{model}\ast d_v}$，$W^O\in R^{h{d}_v\ast d_{model}}$。表示第$i$个头的变换矩阵，$h$表示头的个数。

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

3. 前馈网络

这部分是整体架构图中的Feed Forward模块，其实就是一个简单的全连接前馈网络。它由两层全连接及ReLU激活函数构成，计算公式如下：

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

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

4. add & norm

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

5. 解码器

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

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

6. 位置编码层

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

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

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

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

7. 一些Tricks和技巧

• Embedding和Softmax：论文中将embedding层的参数与最后的Softmax层之前的变换层参数进行了共享（参见：https://arxiv.org/pdf/1608.05859.pdf ），并且在embedding层，将嵌入的结果乘上$\sqrt{d_{model}}$。

• 初始化：看代码里面的初始化方式采用的是PyTorch里面的nn.xavier_uniform，不知道是不是必须的，这个还是要具体问题具体尝试？

• 优化器：论文里面提到了他们用的优化器，是以${\beta }_1=0.9$，${\beta }_2=0.98$和$ϵ=10^{-9}$的Adam为基础，而后使用一种warmup的学习率调整方式来进行调节。具体公式如下：基本上就是先用一个固定的warmup_steps进行学习率的线性增长（热身），而后到达warmup_steps之后会随着step_num的增长而逐渐减小，他们用的$warmup\_steps=4000$，这个可以针对不同的问题自己尝试。
  $$l_{rate}=d_{model}^{-0.5}\cdot min(step\_nu{m}^{-0.5},step\_num\cdot warmup\_step{s}^{-1.5})$$

• 正则化：论文在训练的时候采用了两种正则化的方式。1）dropout：主要用在每个SubLayer计算结束之后，比如self-attention或feed forward，然后再与输入进行add & norm，同时也作用在经过了位置编码后的embedding上，他们取的$P_{drop}=0.1$；2）标签平滑：即Label Smoothing（参见： https://arxiv.org/pdf/1512.00567.pdf ），其实还是从图像上搬过来的，具体操作可以看下一节的代码实现。这里论文取的${ϵ}_{ls}=0.1$，他们发现会损失困惑度，但能提升准确率和BLEU值！

四. PyTorch实现

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

1. 编解码器

这是一个通用的Encoder-Decoder架构：

class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many 
    other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask,
                            tgt, tgt_mask)
    
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
    
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

要注意的是，src和tgt都需要传入mask进行计算。

2. 多头注意力机制

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

参数里面的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
1
2
3
4
5
6
7
8
9
10
11

3. 前馈网络

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

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))))
1
2
3
4
5
6
7
8
9
10

4. 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)))
1
2
3
4
5
6
7
8
9
10
11
12
13

在forward函数里面，彷佛与前面的$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
1
2
3
4
5
6
7
8
9
10
11
12

5. 位置编码

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

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

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

6. 关于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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

对于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
1
2
3
4
5

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

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

7. 优化器

这里其实是手动实现了上一节提到的带warmup的学习率调节公式，代码比较简单：

class NoamOpt:
    "Optim wrapper that implements rate."
    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0
        
    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()
        
    def rate(self, step = None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * \
            (self.model_size ** (-0.5) *
            min(step ** (-0.5), step * self.warmup ** (-1.5)))
        
def get_std_opt(model):
    return NoamOpt(model.src_embed[0].d_model, 2, 4000,
            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

主要调节是在rate这个函数中，model_size即为$d_{model}$，warmup即为$warmup\_steps$，factor可以理解为初始的学习率。

8. 标签平滑

Label Smoothing这里，我看到的PyTorch版本是用KL散度损失，对于输出的分布，从原始的one-hot分布转为在groundtruth上使用一个confidence值，而后其他的所有非groudtruth标签上采用$\frac{1-confidence}{odim-1}$作为概率值进行平滑。具体代码如下：

class LabelSmoothing(nn.Module):
    "Implement label smoothing."
    def __init__(self, size, padding_idx, smoothing=0.0):
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(size_average=False)
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None
        
    def forward(self, x, target):
        assert x.size(1) == self.size
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (self.size - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, Variable(true_dist, requires_grad=False))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

这里的size是输出词表的大小，smoothing是用于分摊在非groundtruth上面的概率值。

五. 为啥要用Self-Attention？

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

论文从计算复杂度、序列操作数以及最大路径长度三个角度比较了不同的层，包括Self-Attention、Recurrent、Convolutional等，如下表所示：

表里面的$n$代表序列长度，$d$代表向量维度，$k$表示kernel的大小，$r$表示受限的memory的长度（主要是针对过长序列，直接使用self-attention未免太大）。从表中的数据看起来，好像Self-Attention确实比较优良。

六. 总结

1. 模型特点：采用全attention的方式，完全摒弃了RNN和CNN的做法。
2. 优势：训练速度更快，在两个翻译任务上取得了SoTA。
3. 不足：在decode阶段还是自回归的，即还是不能并行，而且对于每个step的计算，都是要重新算一遍，没有前面的记忆。

传送门

论文：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实现的核心源码博客，有原理，也有对应的代码段，非常适合对照学习！