The Annotated Transformer（一）

原文地址：http://nlp.seas.harvard.edu/2018/04/03/attention.html
摘要：基于复杂循环卷积神经网络的主要序列转换模型包含了编码器（Encoder）和解码器。表现最好的模型在解码器和编码器之间通过注意力机制进行连接。我们提出了一个新的简单网络结构，Transformer，其仅基于注意力机制，完全和循环卷积分开。两种机器翻译的实验表明这些模型不仅翻译质量更高，而且平行计算能力更好，需要的训练时间显著减少。我们的模型在VMT2014 English-to-German翻译任务上达到了28.4BLEU，超过包括集成学习的已知的最好成绩，提高了超过2BLEU。在VMT2014 English-to-German翻译任务上，我们的模型建立的简单的模型在8个GPU上训练3.5天后，得到了最高BLEU得分41.8。从文献上看训练成本是非常小的。我们将Transformer应用在大训练集以及小训练集的英语成分句法分析（Constituency parsing）上表明Transformer可以很好概括其他的任务。

Prelims

# !pip install http://download.pytorch.org/whl/cu80/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl numpy matplotlib spacy torchtext seaborn 

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
from torch.autograd import Variable
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")
%matplotlib inline

Background

模型结构：减少序列计算的目标也构成了扩展神经 GPU, ByteNet 及 ConvS2S的基础，这些模型也是采用卷积神经网络作为基本构建块，并行计算所有输入输出位置的隐藏表示。在这些模型中，将两个任意输入或输出位置的信号关联起来的所需计算量随着位置的距离增长，在ConvS2S中呈线性增长，在ByteNet中呈对数增长。在Transformer中，这被减少到常数值，尽管采用平均注意力加权位置而导致有效分辨率降低，这是我们与多头注意力相抵消的效果。

自注意力，有时称为内部注意力，是将单个序列的不同位置关联起来计算出一个序列的表示。自注意力已经成功用于各种任务，包括阅读理解，抽象概括，文本蕴涵和学习任务独立的句子表示。端到端存储器网络是基于循环注意机制而不是序列对齐重复，并且已经证明在简单语言问答和语言建模任务上表现良好。

然而，据我们所知，Transformer是第一个完全依靠自注意力而不使用序列对齐的RNN或卷积的转换模型来计算其输入和输出的表征。

Model Architecture

大多数有竞争力神经序列转换模型具有编码器 - 解码器框架。编码器将输入信号序列（x1，…，xn）映射到连续序列z=(z1,…,zn)。输入Z，解码器一次一个元素生成输出序列 (y1,…,ym)。在每个步骤中，模型都是自回归的，在生成下一个元素时，之前生成的元素作为附加输入被消耗掉。

class EncoderDecoder(nn.Module):
    """
    标准Encoder-Decoder架构。本模型和其他许多模型的基础。
    """
    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)

class Generator(nn.Module):
    "定义标准 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)

Transformer遵循如下整体架构，使用堆叠的自注意力和逐点、完全连接层用于编码器和解码器，分别如图1的左半部分和右半部分所示。

Encoder and Decoder Stacks
Encoder
编码器由N=6完全相同的层堆栈组成。

def clones(module, N):
    "生成N个相同层。"
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Encoder(nn.Module):
    "核心编码器是N层的堆栈。"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

在每两个子层周围采用残差连接，然后是在归一化层。

class LayerNorm(nn.Module):
    "建立一个归一化层模块。"
    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

也就是说，每个子层的输出是layerNorm（x+Sublayer(x)），其中Sublayer(x)是由子层本身实现的函数。我们将dropout应用于每个子层的输出，然后将其添加到子层输入并归一化。

为了方便进行残差连接，模型中的所有子层以及嵌入层都会生成维度dmodel=512的输出。

class SublayerConnection(nn.Module):
    """
    残差层连接后跟着归一层。
    注意：为了代码简单性，归一化是第一个而不是最后一个。
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "将残差连接应用于具有相同大小的任何子层。"
        return x + self.dropout(sublayer(self.norm(x)))

每层有两个子层。 第一层是多头自注意力机制，第二层是简单的，位置完全连接的前反馈网络。

class EncoderLayer(nn.Module):
    "编码器由自注意和前反馈组成（定义如下）"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

Decoder
解码器也由N = 6的相同的层堆栈组成。

class Decoder(nn.Module):
    "具有masking的通用N层解码器。"
    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)

除了每个编码器层中的两个子层之外，解码器还插入第三子层，其对编码器堆栈的输出执行多头注意力。 与编码器类似，我们在每个子层周围使用残差连接，然后进行归一化。

class DecoderLayer(nn.Module):
    "解码器由self_attn，src-attn和feed_forward组成（定义如下）"
    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):
        "根据图1 (右) 进行连接"
        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)

我们还修改解码器堆栈中的自注意子层以避免当前位置关注后续的位置信息。 掩码层结合偏移了一个位置的输出嵌入层，确保了位置i的预测仅依赖于小于i的位置处的已知输出。

def subsequent_mask(size):
    "屏蔽后续位置。"
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0

以下注意力掩码表明下方显示允许每个目标字（行）可以查看的位置（列）。 在训练期间，单词被阻止用于处理未来的单词。

Attention
Attention可以被描述为将查询(query)和一组键值对(key-value)映射到输出，其中查询，键，值和输出都是向量。输出被计算为值的加权和，其中分配给每个值的权重由查询（query）与对应键（key）的兼容性函数计算。
我们特别关注“Scaled Dot-Product Attention”。输入包括查询（query)、维度dk的键(key)以及维度dv的值(value)。我们使用所有key计算query的点积，除以根号dk，并应用softmax函数来获得值的权重。

实际上，我们同时计算一组查询的注意函数，将它们打包成矩阵Q。键和值也一起打包成矩阵K和V.我们计算的输出矩阵为：

def attention(query, key, value, mask=None, dropout=None):
    "计算 '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/√dk 的比例因子外，点积（乘法）注意和我们的算法是一样的。附加注意使用具有单个隐藏层的前馈网络来计算兼容性函数。虽然两者在理论上的复杂性相似，但在实践中，点积注意更快，更节省空间，因为它可以使用高度优化的矩阵乘法代码来实现。对于较小的dk值，两种机制的表现相似，但是对于更大的dk值在不缩放时附加注意力优于点积注意。
我们猜想是对于较大的dk值，点积值大幅增长，使得softmax函数进入到梯度极其小的区域（为了说明点积变大的原因，假设分量q和k是平均值为0且方差为1的独立随机变量。那么他们的点积q ⋅k=∑qiki的平均值为0且方差为dk）。为了抵消这种影响，我们将点积乘上1/√dk 进行缩小。

多头注意允许模型共同关注来自不同位置的不同表征的子空间信息。 只有单独的注意头的话，平均计算会抑制这种情况。

其中分量是参数矩阵

在这项工作中，我们采用h = 8个平行注意力层或头部。 对其中每一个我们使用dk = dv = dmodel / h = 64。 由于减少了每个头部的尺寸，总计算成本与全维度的单头注意力相似。

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "考虑模型尺寸和头数。"
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # 我们假设d_v 一直等于d_k
        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):
        "实现图2"
        if mask is not None:
            # 相同的掩码应用于所有h 头.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        
        # 1) 从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)注意批量中所有投影向量。
        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) “Concat”使用视图并应用最终线性。 
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

Applications of Attention in our Model
Transformer以三种不同的方式使用多头注意：
1）在“编码器 - 解码器注意”层中，查询（queries）来自先前的解码器层，并且存储的键（keys）和值(values)来自编码器的输出。 这允许解码器中的每个位置都可以注意输入序列中的所有位置。这模拟了序列到序列模型中典型的编码器 - 解码器注意机制。
2）编码器包含自注意层。 在自注意层中，所有键，值和查询来自相同的位置，是编码器中前一层的输出。 编码器中的每个位置都可以注意编码器前一层中的所有位置。
3）类似地，解码器中的自注意层允许解码器中的每个位置注意解码器中的所有位置直到并包括该位置。 我们需要防止解码器中的向左信息流以保持自回归属性。 通过屏蔽所有softmax输入（设置为-∞），则其对应于非法连接， 实现缩放的dot-product注意。

Position-wise Feed-Forward Networks
除了注意子层之外，编码器和解码器中每层都包含一个完全连接的前馈网络，该网络分别相同地应用于每个位置。 其包括采用ReLU激活的两个线性变换。

虽然线性变换在不同位置上是相同的，但它们在层与层之间使用不同的参数。 另一种描述这种情况的方法是两个内核大小为1的卷积。输入和输出的维数是dmodel = 512，内部层的维度dff = 2048。

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))))

Embeddings and Softmax

与其他序列转换模型类似，我们使用学习嵌入将输入词和输出词转换为维度dmodel的向量。 我们还使用通常的学习线性变换和softmax函数将解码器输出转换为预测的下一个词的概率。 在我们的模型中，两个嵌入层和pre-softmax线性变换之间采用相同的权重矩阵。 在嵌入层中，我们将这些权重乘以√dmodel 。

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

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

Positional Encoding
由于我们的模型不包含递归和卷积，为了使模型能够利用序列的顺序，我们必须增加序列中词的相对或绝对位置的一些信息。 为此，我们将“位置编码”添加到编码器和解码器堆栈底部的输入嵌入。 位置编码具有与嵌入相同的维度dmodel，因此可以将两者相加。 位置编码的选择、学习和固定有许多选择。

在这项工作中，我们使用不同频率的正弦和余弦函数：

pos是位置，i是维度。
也就是说，位置编码的每个维度对应的正弦曲线。波长形成了一个从2π到10000⋅2π的几何级数。
我们选择这个函数是因为我们假设它允许模型很容易地学习相对位置，因为对于任何固定偏移k，PEpos + k可以表示为PEpos的线性函数。此外，在编码器和解码器堆栈中我们将dropout应用于嵌入和位置编码之和。 对于基本模型，我们使用Pdrop = 0.1的速率。

class PositionalEncoding(nn.Module):
    "实现PE功能。"
    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)

下图是位置编码根据位置添加正弦波。 波的频率和偏移对于每个维度是不同的。

plt.figure(figsize=(15, 5))
pe = PositionalEncoding(20, 0)
y = pe.forward(Variable(torch.zeros(1, 100, 20)))
plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
plt.legend(["dim %d"%p for p in [4,5,6,7]])
None

我们还尝试使用学习的位置嵌入，并发现这两个版本生成的结果几乎相同。 我们选择了正弦曲线版本，因为它允许模型外推到比训练期间遇到的序列更长的序列。

未完待续