浅谈Transfomer代码原理解读

前言

最近transfomer太火了，虽然火但是tansformer也限制了我们这些穷人去用它来炼丹（太费内存了）。
这里我推荐一篇文献综述可以看到tansformer在各方面的应用Learning Texture Transformer Network for Image Super-Resolution

Transformer架构

来张用烂的图吧，transformer模型本质上是一个Encoder-Decoder的结构。输入序列先进行Embedding，经过Encoder之后结合上一次output再输入Decoder，最后用softmax计算序列下一个单词的概率。就是传说中的自注意力机制。

Input Embedding and Output Embedding

Embedding

因为我们的输入是英文单词和中文单词，所以要对它们进行编码Word Embedding + Position Embedding。

Word Embedding 可以采用word2vec也可以采用pytorch中的onehot 利用nn.Embedding 。这里我感觉两种没有什么差别，可能在一些竞赛中有点差别。
　这里我采用 nn.Embedding 实现，使用Pre-trained，用 Pre-trained的Embeddings并固化，设为 Trainable。这样在 training 过程中不断地对 Embeddings 进行改进。
其中d_model表示embedding的维度，即词向量的维度；vocab表示词汇表的数量。

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model  #表示embedding的维度
 
    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

Positional Embedding

在了解Positional Embedding我们先了解RNN 和Lstm
RNN相关算法只能从左向右依次计算或者从右向左依次计算，这种机制带来了两个问题：　
(1) 时间片 t 的计算依赖 t−1 时刻的计算结果，这样限制了模型的并行能力；
　(2) 顺序计算的过程中信息会丢失，尽管LSTM等门机制的结构一定程度上缓解了长期依赖的问题，但是对于特别长期的依赖现象，LSTM依旧无能为力。
　　Transformer的提出解决了上面两个问题：

(1) 首先它使用了Attention机制，将序列中的任意两个位置之间的距离是缩小为一个常量；
　　(2) 其次它不是类似RNN的顺序结构，因此具有更好的并行性，符合现有的GPU框架。
因为Transformer没有像RNN 和Lstm那样依次计算所以没有了位置相关性。所以采用了位置编码Positional Embedding解决了这个问题同时，能够更好地表达的词与词之间的距离。使得 Transformer 可以衡量 word 位置有关的信息。使用公式来计算 positional encoding向量。

Positional Encoding的公式如下：

其中，pos指的是这个 word 在这个句子中的位置；2i指的是 embedding 词向量的偶数维度，2i+1指的是embedding 词向量的奇数维度。
"x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)" 这行代码表示；输入模型的整个Embedding是Word Embedding与Positional Embedding直接相加之后的结果

# Positional Encoding
class PositionalEncoding(nn.Module):
    "实现PE功能"
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
         
        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)           # [1, max_len, d_model]
        self.register_buffer('pe', pe)
         
    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)

在下图中，每一行对应一个词向量的位置编码，20字(行)的位置编码实例，词嵌入大小为512(列)。你可以看到它从中间分裂成两半。这是因为左半部分的值由一个函数(使用正弦)生成，而右半部分由另一个函数(使用余弦)生成。然后将它们拼在一起而得到每一个位置编码向量。而且可以看出相邻之间单词编码相同。

Encoder and decoder

编码组件部分由一堆编码器（encoder）构成（论文中是将6个编码器）。解码组件部分也是由相同数量（与编码器对应）的解码器（decoder）组成的。所有的编码器和解码器在结构上都是相同的。
从编码器输入的句子首先会经过一个自注意力（self-attention）层，这层帮助编码器在对每个单词编码时关注输入句子的其他单词。


自注意力层的输出会传递到前馈（feed-forward）神经网络中。每个位置的单词对应的前馈神经网络都完全一样。解码器中也有编码器的自注意力（self-attention）层和前馈（feed-forward）层。除此之外，这两个层之间还有一个注意力层，用来关注输入句子的相关部分（和seq2seq模型的注意力作用相似）。

Muti-Head-Attention

Self Attention

例如，下列句子是我们想要翻译的输入句子：
The animal didn’t cross the street because it was too tired
这个“it”在这个句子是指什么呢？它指的是street还是这个animal呢？这对于人类来说是一个简单的问题，但是对于算法则不是。
当模型处理这个单词“it”的时候，自注意力机制会允许“it”与“animal”建立联系。
随着模型处理输入序列的每个单词，自注意力会关注整个输入序列的所有单词，帮助模型对本单词更好地进行编码。
如果你熟悉RNN（循环神经网络），回忆一下它是如何维持隐藏层的。RNN会将它已经处理过的前面的所有单词/向量的表示与它正在处理的当前单词/向量结合起来。而自注意力机制会将所有相关单词的理解融入到我们正在处理的单词中
当我们在编码器#5（栈中最上层编码器）中编码“it”这个单词的时，注意力机制的部分会去关注“The Animal”，将它的表示的一部分编入“it”的编码中。

它实怎样用矩阵来实现呐？
计算自注意力的第一步就是从每个编码器的输入向量（每个单词的词向量）中生成三个向量。也就是说对于每个单词，我们创造一个查询向量、一个键向量和一个值向量。这三个向量是通过词嵌入与三个权重矩阵后相乘创建的。key = linear_k(x)，query = linear_q(x)，value = linear_v(x)。如下图，X1与WQ权重矩阵相乘得到q1, 就是与这个单词相关的查询向量。最终使得输入序列的每个单词的创建一个查询向量、一个键向量和一个值向量。
可以发现这些新向量在维度上比词嵌入向量更低。他们的维度是64，而词嵌入和编码器的输入/输出向量的维度是512. 但实际上不强求维度更小，这只是一种基于架构上的选择，它可以使多头注意力（multiheaded attention）的大部分计算保持不变。

什么是查询向量(query )、键向量(key )和值向量(value )？

计算自注意力的第二步是计算得分。假设我们在为这个例子中的第一个词“Thinking”计算自注意力向量，我们需要拿输入句子中的每个单词对“Thinking”打分。这些分数决定了在编码单词“Thinking”的过程中有多重视句子的其它部分。这些分数是通过打分单词（所有输入句子的单词）的键向量与“Thinking”的查询向量相点积来计算的。所以如果我们是处理位置最靠前的词的自注意力的话，第一个分数是q1和k1的点积，第二个分数是q1和k2的点积。

第三步和第四步是将分数除以8(8是论文中使用的键向量的维数64的平方根，这会让梯度更稳定。这里也可以使用其它值，8只是默认值)，然后通过softmax传递结果。softmax的作用是使所有单词的分数归一化，得到的分数都是正值且和为1。

这个softmax分数决定了每个单词对编码当下位置（“Thinking”）的贡献。显然，已经在这个位置上的单词将获得最高的softmax分数，但有时关注另一个与当前单词相关的单词也会有帮助。
第五步是将每个值向量乘以softmax分数(这是为了准备之后将它们求和)。这里的直觉是希望关注语义上相关的单词，并弱化不相关的单词(例如，让它们乘以0.001这样的小数)。
第六步是对加权值向量求和（译注：自注意力的另一种解释就是在编码某个单词时，就是将所有单词的表示（值向量）进行加权求和，而权重是通过该词的表示（键向量）与被编码词表示（查询向量）的点积并通过softmax得到。），然后即得到自注意力层在该位置的输出(在我们的例子中是对于第一个单词)。

这样自自注意力的计算就完成了。得到的向量就可以传给前馈神经网络。然而实际中，这些计算是以矩阵形式完成的，以便算得更快。
得到query(Q)，key(K)，value(V)之后按照下面的公式计算attention(Q, K, V)：

Muti-Head-Attention

Multi-Head Self Attention 实际上是由h个Self Attention 层并行组成，原文中h=8.
1.它扩展了模型专注于不同位置的能力。在上面的例子中，虽然每个编码都在z1中有或多或少的体现，但是它可能被实际的单词本身所支配。如果我们翻译一个句子，比如“The animal didn’t cross the street because it was too tired”，我们会想知道“it”指的是哪个词，这时模型的“多头”注意机制会起到作用。
2.它给出了注意力层的多个“表示子空间”（representation subspaces）。接下来我们将看到，对于“多头”注意机制，我们有多个查询/键/值权重矩阵集(Transformer使用八个注意力头(我们采用6个，可以达到效果)，因此我们对于每个编码器/解码器有6个矩阵集合)。这些集合中的每一个都是随机初始化的，在训练之后，每个集合都被用来将输入词嵌入(或来自较低编码器/解码器的向量)投影到不同的表示子空间中。

在“多头”注意机制下，我们为每个头保持独立的查询/键/值权重矩阵，从而产生不同的查询/键/值矩阵。和之前一样，我们拿X乘以WQ/WK/WV矩阵来产生查询/键/值矩阵。

如果我们做与上述相同的自注意力计算，只需八次不同的权重矩阵运算，我们就会得到八个不同的Z矩阵。

这给我们带来了一点挑战。前馈层不需要8个矩阵，它只需要一个矩阵(由每一个单词的表示向量组成)。所以我们需要一种方法把这八个矩阵压缩成一个矩阵。那该怎么做？其实可以直接把这些矩阵拼接在一起，然后用一个附加的权重矩阵WO与它们相乘。

看我们在例句中编码“it”一词时，不同的注意力“头”集中在哪里？
当我们编码“it”一词时，一个注意力头集中在“animal”上，而另一个则集中在“tired”上，从某种意义上说，模型对“it”一词的表达在某种程度上是“animal”和“tired”的代表。

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
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
         
    def forward(self, query, key, value, mask=None):
        """
        实现MultiHeadedAttention。
           输入的q，k，v是形状 [batch, L, d_model]。
           输出的x 的形状同上。
        """
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
         
        # 1) 这一步qkv变化:[batch, L, d_model] ->[batch, h, L, d_model/h]
        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) 计算注意力attn 得到attn*v 与attn
        # qkv :[batch, h, L, d_model/h] -->x:[b, h, L, d_model/h], attn[b, h, L, L]
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        # 3) 上一步的结果合并在一起还原成原始输入序列的形状
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        # 最后再过一个线性层
        return self.linears[-1](x)

Add & Norm

每个编码器中的每个子层（自注意力，ffnn）在其周围都有残差连接，后续再进行层归一化（layer-normalization）步骤。如果我们要对向量和与自注意力相关的层规范操作进行可视化，则看起来应该像这样：

# Position-wise Feed-Forward Networks
class PositionwiseFeedForward(nn.Module):
    "实现FFN函数"
    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))))

Feed-Forward Network

Feed-Forward Network可以细分为有两层，第一层是一个线性激活函数，第二层是激活函数是ReLU。可以表示为：

# Position-wise Feed-Forward Networks
class PositionwiseFeedForward(nn.Module):
    "实现FFN函数"
    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))))

总结

如果我们设想由2个编码器解码器堆栈组成的Transformer，它看起来像这样：

Encoder

def clones(module, N):
    "产生N个相同的层"
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
 
class Encoder(nn.Module):
    """N层堆叠的Encoder"""
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
         
    def forward(self, x, mask):
        "每层layer依次通过输入序列与mask"
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

Decoder

# Decoder部分
class Decoder(nn.Module):
    """带mask功能的通用Decoder结构"""
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
         
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

Mask-Multi-Head-Attention

mask可以理解成遮罩、面具，作用是帮助我们“遮挡”掉我们不需要的东西，即让被遮挡的东西不影响我们的attention过程。
在forward的时候，有两个mask参数可以设置：
key_padding_mask
每一个batch的每一个句子的长度一般是不可能完全相同的，所以我们会使用padding把一些空缺补上。而这里的这个key_padding_mask是用来“遮挡”这些padding的。
这个mask是二元(binary)的，也就是说，它是一个矩阵和我们key的大小是一样的，里面的值是1或0，我们先取得key中有padding的位置，然后把mask里相应位置的数字设置为1，这样attention就会把key相应的部分变为"-inf". (为什么变为-inf我们稍后再说)
attn_mask
这个mask经常是用来遮挡“正确答案”的：
假如你想要用这个模型每次预测下一个单词，我们每一个位置的attention输出是怎么得来的？是不是要看一遍整个序列，然后每一个单词都计算一个attention weight？那也就是说，你在预测第5个词的时候，你其实会看到整个序列，这样的话你在预测之前不就已经知道第5个单词是什么了，这就是作弊了。
我们不想让模型作弊，因为在真实使用这个模型去预测的时候，我们是没有整个序列的信息的。那么怎么办？那就让第5个单词的attention weight=0吧，即声明：我不想看这个单词，我的注意力一点也别分给它。
如何让这个weight=0：
我们先想象一下，我们目前拥有的attention scores是什么样的？(注：attention_score是attention_weight的初始样子，经过softmax之后会变成attention_weight.
attention_score和weight的形状是一样的，毕竟只有一个softmax的差别)

我们之前提到，attention weights的形状是L*L，因为每个单词两两之间都有一个weight。
如下图所示，我用黑笔圈出的部分，就是“我想要预测x2”时，整个sequence的attention score情况。我用叉划掉的地方，是我们希望=0的位置，因为我们想让x2、x3、x4的权值为0，即：预测x2的时候，我们的注意力只能放在x1上。

对于其他行，你可以以此类推，发现我们需要一个三角形区域的attention weight=0（见最底下的图）, 这时候我们的attn_mask这时候就出场了，把这个mask做成三角形即可。
和key_padding_mask不同，我们的attn_mask不是binary的，而是一个“additive mask”。
什么是additive mask呢？就是我们mask上设置的值，会被加到我们原本的attention score上。我们要让三角形区域的weight=0，我们这个三角mask设置什么值好呢？答案是-inf，（这个-inf在key_padding_mask的讲解中也出现了，这里就来说说为什么要用-inf）。我们上面提到了，attention score要经过一个softmax才变成attention_weights.
我们都知道softmax的图像长这样：

当我们attention score的值设置为-inf，通过softmax之后我们的attention weight就会趋近于0了，这就是为什么我们这里的两个mask都要用到-inf。

三角形的样子和softmax前后的情况可见下图：（图源见水印）

def subsequent_mask(size):
    """
    mask后续的位置，返回[size, size]尺寸下三角Tensor
    对角线及其左下角全是1，右上角全是0
    """
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0
    #这里mask是一个下三角矩阵，对角线以及对角线左下都是1，其余都是0。下面是个10维度的下三角矩阵：
tensor([[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)

Encoder-Decoder Multi-head Attention代码，self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) 这行就是Encoder-Decoder Multi-head Attention。

　　query = x，key = m, value = m, mask = src_mask，这里x来自上一个 DecoderLayer，m来自 Encoder的输出。：

```python
class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)
  
    def forward(self, x, memory, src_mask, tgt_mask):
        "将decoder的三个Sublayer串联起来"
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

Softmax

线性层（he Linear layer）是一个简单的完全连接的神经网络，将解码器堆栈产生的向量投射到一个大得多的对数向量中。我们假设自己的模型从训练数据集中共学会了10,000个不同的英语单词（我们模型的“输出词汇表”）。这将使对数向量的宽度变为10,000个单元，每个单元对应各个单词的得分。我们将会通过这样的方式来解释模型的输出。然后，softmax层将会把这些分数转换为概率（全部为正数，各项相加和为1.0）。概率最高的单元被选中，且与该单元相关联的单词将成为该步的输出。该图从底部开始，生成的向量作为解码器堆栈的输出，后续会被转换为文字输出

class Generator(nn.Module):
    """
    Define standard linear + softmax generation step。
    定义标准的linear + softmax 生成步骤。
    """
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)
 
    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)