深度学习论文: Attention is All You Need及其PyTorch实现
深度学习论文：注意力机制及其PyTorch实现
主流高端的神经序列转换模型主要运用基于编码器-解码器架构的设计方案，在编码阶段负责将输入符号序列转化为连续空间表达，在解码阶段则通过自回归机制完成输出符号序列的逐步生成。每个阶段内，模型通过自回归机制将已生成的部分输出作为后续输入来推导下一个输出单元的内容。
1 Encoder and Decoder Stacks
1-1 Encoder 编码器
编码模块采用N=6个叠接结构进行堆叠而成。每个编码器块包含两个子块：第一个是多路注意力机制模块，第二个是前馈神经网络模块。为了优化性能，在各子块间引入了残差连接机制，并实施了归一化处理。具体而言，在每个子块的输出端通过LayerNorm(x + Sublayer(x))进行计算得出（其中Sublayer(x)代表该子块的功能实现）。为了确保残差连接的有效性，在整个编码模块中所有子块以及嵌入层均设置输出维度为dmodel=512的特征空间。
def clones(module, N):
"Produce N identical layers."
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
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
class Encoder(nn.Module):
"Core encoder is a stack of N layers"
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 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)))
class EncoderLayer(nn.Module):
"Encoder is made up of self-attn and feed forward (defined below)"
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)
代码解读1-2 Decoder解码器
解码器也由6个结构一致的层进行叠放排列。与编码器不同，在解码器中除了包含两个子层之外还新增了一个专门用于对编码器输出执行多头注意力机制的第三子层。在架构上与编码器相似，在解码器中每个子层也都采用了残差连接并配合使用了标准化处理。此外，在优化自注意力机制方面还对解码器堆叠中的相关组件进行了改进以确保生成过程中的位置预测仅依赖于已知信息这通过在目标位置处施加掩码以及将输出嵌入向量后移一位的方式来实现。
class Decoder(nn.Module):
"Generic N layer decoder with masking."
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):
"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):
"Follow Figure 1 (right) for connections."
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)
代码解读2 Attention
Transformer巧妙地运用了三种不同的多头注意力机制：
在"encoder-decoder attention"模块中, 查询由上一层的解码器输出构成, 而其对应的键和值则来源于编码器输出的信息. 这种结构设计确保了该层中的任何位置都能够聚焦于输入序列的所有细节信息, 完美模仿了传统序列到序列模型中典型的encoder-decoder注意机制.
在编码器内部内置于一种自注意力机制中，在这一机制中使用的是来自上一层输出的键、值与查询以实现信息传递路径的有效连接从而保证了编码器能够全面地捕捉到上一层的信息细节
解码器同样采用了基于自注意力机制的设计，在这种架构下，每个解码器位置能够聚焦并处理包含自身在内的整个解码域内的相关信息。为了确保信息只能单向流动以维持自回归特性这一核心要求，在缩放点积注意力机制中我们特意设计了一种巧妙的方法：通过在缩放点积注意力中将所有非法连接的状态初始化为负无穷大（−∞），从而使得无法向左传播信息。
2-1 Scaled Dot-Product Attention
其包含维度为d_{k}的查询向量和键向量；此外还包括维度为d_{v}的值向量。具体而言，在计算过程中会先通过所有查询向量与相应键向量之间的点积来生成中间结果；随后将这些中间结果均除以\sqrt{d_{k}}进行归一化处理；最后通过应用Softmax函数来确定各值向量的重要性权重。
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-2 Multi-Head Attention
相比仅通过dmodel维度的一次性注意力计算, 该方法采用查询、键和值分别经过独立 learnable 线性变换的方式, 经过h次变换后分别映射至d_k, d_k, d_v维空间以显著提升效果。随后我们同时对各子空间执行注意力计算以获取对应的中间结果, 最后将各子空间的输出结果合并后经 learnable 线性变换得到最终输出。
在本研究中设置了h=8个并行的注意力层，并将其称作多个注意头。每个注意头的维度被设置为d_k = d_v = d_{\text{model}}/h = 64。鉴于每个注意头的尺寸有所减少，在计算开销方面与完整尺寸单个注意头相当。
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)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
"Implements Figure 2"
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)
代码解读3 Position-wise Feed-Forward Networks
除了注意力机制模块外（或注意），编码器与解码器的每一层都包含一个全连接前馈网络段（或网架），它们分别独立地并行应用于每一个位置（或位点），经过两个线性变换层后再接入一个ReLU激活函数模块（或激活组件）。
这里的输入和输出的维度是d_{model}=512，而内层的维度是d_{ff}=2048。
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))))
代码解读4 Embeddings and Softmax
与其他序列转换模型具有相似特征，在本研究中我们采用学习嵌入技术将输入标记和输出标记映射到d_{model}维向量空间进行处理。随后，解码器输出经过线性变换并结合softmax函数计算得到预测概率分布。值得注意的是，在本模型设计中，嵌入层与位于softmax之前的线性变换共用相同的权重矩阵；此外，在嵌入层的计算过程中，默认会将权重值乘以\sqrt{d_{model}}进行缩放处理。
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)
代码解读5 Positional Encoding
为了提高模型性能，在编码器与解码器堆栈底部的输入嵌入中添加了“位置信息编码”。其中，所使用的定位信息编码具有与嵌入层相同维数的特点：这是因为只有这样才能够实现定位信息与原始嵌入数据的有效融合。具体来说，定位信息编码的方式多样：既可以采用预定义固定的模式进行生成（Fixed），也可以通过可学习的方式动态调整（Learnable）。
在这里，选用不同频率的正弦和余弦函数：
在这一过程中,我们设定了变量pos代表位置,并将维度i作为变量.这表明每个位置编码的维度都对应一个正弦波.这些正弦波的波长从2π开始,以等比数列的方式延伸至10000·2π.选择这种函数的原因在于其能够使模型更容易地关注相对位置关系,具体而言,即对于任意固定的偏移量k,在这一情况下PE_{pos+k}可被表示为PE_{pos}的线性函数.
在编码器与解码器堆栈中的嵌入层和位置编码层上采用了Dropout技术。其中针对基础模型，采用的比例为P_{drop}=0.1。
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 自注意力机制几乎就是万能的 https://arxiv.org/pdf/1706.03762.pdf
2 图解Transformer https://jalammar.github.io/illustrated-transformer/
3 注解Transformer http://nlp.seas.harvard.edu/2018/04/03/attention.html