Self-Attention与谷歌的Transformer架构
本文翻译自一篇非常赞的解释Transformer的文章，原文链接。Transformer是在"Attention is All You Need"中提出的，其中的TF应用是Tensor2Tensor的子模块。哈佛的NLP团队专门制作了对应的PyTorch的指南说明。
1. A High-Level Look
我们先将整个模型视为黑盒，比如在机器翻译中，接收一种语言的句子作为输入，然后将其翻译成其他语言输出。
细看下，其中由编码组件、解码组件和它们之间的连接层组成。
编码组件是六层编码器首位相连堆砌而成，解码组件也是六层解码器堆成的。
编码器是完全结构相同的，但是并不共享参数，每一个编码器都可以拆解成以下两个子部分。
编码器的输入首先流过一个self-attention层，该层的作用是当它编码某个词时能够帮助编码器看到输入序列中的其他单词。后面，我们会细看self-attention的内部结构。
self-attention的输出流向一个前向网络，每个位置对应的前向网络是完全独立。
解码器同样也有这些子层，但是在两个子层间增加了attention层，该层有助于解码器能够关注到输入句子的相关部分，这与 seq2seq model 的Attention作用相似。
2. Bringing The Tensors Into the Picture
以上讲述了模型最主要的组件，现在我们看看向量/Tensor是如何在各个组件之间流动并完成输入到输出的。
正如NLP应用的常见例子，先看看如何将输入单词使用embedding algorithm转成向量。
每个词被嵌入到512维向量上，此处用格子表示向量
词的嵌入向量化只发生在编码器的最底层输入时，所有编码器接收一个长度为512的向量作为输入，除了底层编码器直接接收词向量化作为输入，其他编码器的输入是前个编码器的输出。向量的维度作为超参数可以手动设置，通常设置为训练集的最长句子的长度。
在对输入序列做词的向量化之后，它们流经编码器的如下两个子层。
这里能看到Transformer的一个关键特性，每个位置的词仅仅流过它自己的编码器路径。在self-attention层中，这些路径两两之间相互依赖， 而前向网络层这些路径相互独立，因此这些路径在流经前向网络时可以并行执行。
3. Now We’re Encoding !
正如之前所提，编码器接收向量的列表作为输入将其送入self-attention处理，再送入前向网络，最后将输入传入下一个编码器。
每个位置的词向量被送入self-attention模块，然后是前向网络(对每个向量都是完全相同的网络结构)
4. Self-Attention at a High Level
下面我们逐步分解下Self-Attention是如何工作的。
以下面这句话为例，作为我们想要翻译的输入语句
“The animal didn’t cross the street because it was too tired".
句子中"it"指的是什么呢？“it"指的是"street” 还是“animal”？对人来说很简单的问题，但是对算法而言并不简单。
当模型处理单词“it”时，self-attention允许将“it”和“animal”联系起来。当模型处理每个位置的词时，self-attention允许模型看到句子的其他位置信息作为辅助线索来更好地编码当前词。如果你对RNN熟悉，应该能够理解RNN的隐藏状态是如何将之前的词向量表示融合进入当前词向量表示。Transformer使用self-attention来将相关词的理解编码到当前词中。
当编码"it"时（编码器的最后层输出），部分attention集中到"the animal"，并将其表示合并到“it”的编码中
上图是Tensor2Tensor notebook的可视化例子
5. Self-Attention in Detail
我们从矩阵计算角度看看如何用向量计算self-attention。
第一步，对编码器的每个词嵌入向量，创建一个query向量, 一个key向量, 一个value向量，生成方法为用词嵌入向量分别乘以三个在训练过程中学习的矩阵WQ，WK，WV。【注意：不是每个词向量独享3个matrix，而是所有输入共享3个转换矩阵；权重矩阵是基于输入位置的转换矩阵；有个可以尝试的点，如果每个词独享一个转换矩阵，会不会效果更厉害呢？】生成的q,k,v新向量的维度往往比输入词向量的维度要小，目的是为了让多头attention的计算更稳定，计算流程如下图所示：
query/key/value向量提取计算对于attention是有益的，下面将详细讲解。
第二步，计算attention就是计算一个分值。对“Thinking Matchines”这句话，当计算“Thinking”这个词的attention 分值时，我们需要计算每个关联“Thinking”的词评估分，这个分决定着编码“Thinking”时（某个固定位置时），每个输入词需要给予多少关注度。
这个分数值通过“Thinking”对应query向量与所有词的key向量依次做点积得到。所以当我们处理位置#1时，第一个分值是q1和k1的点积，第二个分值是q1和k2的点积。这就像将当前词作为搜索的query，去和句子中所有词（包含该词本身）的key去匹配，看看相关度有多高。我们用q1代表Thinking对应的query vector，k1及 k2分别代表Thinking以及Machines对应的key vector，则计算Thinking的attention score的时候我们需要计算 q1与k1 ,k2的点乘，同理，我们计算Machines的attention score的时候需要计算 q2与 k1 ,k2的点乘。如上图中所示我们分别得到了q1与k1 ,k2的点乘积，然后我们进行尺度缩放与softmax归一化，如下图所示
第三步和第四步，除以8（k向量维度的均方根，论文中为64），这样梯度会更稳定。然后对输出分值向量作softmax操作，归一化分值使得全为正数且加和为1。
softmax分值决定着在这个位置，每个词的表达程度（关注度）。很明显，处在这个位置的词应该有最高的归一化分数，但大部分时候其它位置的词对当前位置总是有表达意义。
第五步，将softmax分值与value向量按位相乘。保留关注词的value值，削弱非相关词的value值。
第六步，将所有加权向量加和，产生该位置的self-attention的输出结果。
上述就是self-attention的计算过程，生成的向量流入前向网络。在实际应用中，上述计算是以速度更快的矩阵形式进行的。下面我们看下在词级别的矩阵计算。
6. Matrix Calculation of Self-Attention
第一步，计算query/key/value 矩阵，将所有输入词向量合并成输入矩阵X，并且将其分别乘以权重矩阵WQ , WK , WV。
输入矩阵X的每一行表示输入句子的一个词嵌入向量
最后，鉴于我们使用矩阵处理，将步骤2~6合并成一个计算self-attention层输出的公式。
矩阵形式的self-attention计算
7. The Beast With Many Heads
论文进一步增加了multi-headed的机制到self-attention上，在如下两个方面提高了attention层的效果：
多头机制扩展了模型集中于不同位置的能力。在上面的例子中，z1包含了其他词的少许信息，但是却由该词本身主导。在其他情况下，比如翻译 “The animal didn’t cross the street because it was too tired”时，我们想知道单词"it"指的是什么。
多头机制赋予attention多种子空间表达。像下面的例子所示，在多头下有多组query/key/value矩阵，而非仅仅一组（论文中使用8-heads）。每一组都是随机初始化，经过训练之后，输入向量可以被映射到不同的子表达空间中。
每个head都有一组Q/K/V矩阵
采用相同的计算方式，我们采用八组不同的Q/K/V矩阵，可以得到八个不同的Z矩阵。
然而前向网络并不能接收八个矩阵， 所以要将八个矩阵合并成一个。我们将矩阵级联并乘以一个加权矩阵WO, 加权矩阵同样参与参数训练，最后得到的矩阵送入前向网络。
上述就是多头自注意机制的内容，下面尝试着将它们放到一个图上可视化如下。
现在加入attention heads之后，重新看下当编码“it”时，不同的attention head会集中到何处。
编码"it"时，一个attention head集中于"the animal"，另一个head集中于“tired”，某种意义上讲，模型对“it”的表达合成了“animal”和“tired”两者
如果我们将所有的attention heads都放入到图中，理解起来就更复杂一些。
8. Representing The Order of The Sequence Using Positional Encoding
截止到目前为止，我们还没有讨论如何理解输入语句中词的顺序。
为解决词序的利用问题，Transformer对每个词嵌入向量新增一个描述向量，这些向量遵循模型学习的指定模式，来决定词的位置，或者序列中不同词的距离关系。当其映射到Q/K/V向量以及点乘的attention时，增加的这些值能提供词嵌入向量间的距离的有用描述。
为了能够给模型提供词序的信息，新增位置编码向量，每个向量值都遵循指定模式
如果假设位置向量有4维，实际的位置向量编码将如下所示：
一个只有4维的位置向量表示例子
所谓的指定模式是什么样的呢？
在下图中，每一行表示一个词向量的位置编码，所以第一行是我们将要加到句子第一个词向量上的位置编码。每个行有512值，每个值范围在[-1,1]，涂色是方便将模式可视化。
一个20个词(rows)的真实例句，每个词512维编码(cols)。可以观察中间显著的分隔，原因是左半部分是用sine函数生成，右侧是用cosine生成
位置向量编码公式在论文的3.5节有提到，也可以看代码get_timing_signal_ld()，对位置编码而言并不只有一种方法。需要注意的是，编码方法必须能够处理未知长度的序列。
上面的位置编码方法来自于Transformer的Transformer2Transformer实现，然而论文中的方法并不是对sine和cosine的简单级联，而是两个向量相互交叠，如下图所示，生成方法详见生成代码。
9. The Residuals
编码器结构中值得提出注意的一个细节是，在每个子层中（self-attention, ffnn），都有残差连接，并且紧跟着layer-normalization。
如果我们可视化向量和layer-norm操作，将如下所示：
在解码器中也是如此，假设两层编码器+两层解码器组成Transformer，其结构如下：
10. The Decoder Side
现在我们已经了解了编码器侧的大部分概念，也基本了解了解码器的工作方式，下面看下他们是如何共同工作的。
编码器从输入序列的处理开始，最后的编码器的输出被转换为K和V的注意力向量，它们将被送入解码器的"encoder-decoder atttention"层，帮助解码器获取到输入序列的不同位置的注意力。
在编码之后，是解码过程；解码的每一步输出输出序列的一个元素
下面的步骤一直重复直到特殊符号出现表示解码器完成了翻译输出。每一步的输出被喂到下一个解码器中。正如编码器的输入所做的处理，对解码器的输入增加位置向量。
在解码器中的self attention 层与编码器中的稍有不同，在解码器中，self-attention 层仅仅关注早于当前输出的位置。实现方式是在softmax之前，遮挡未来位置（将它们设置为-inf）。
"Encoder-Decoder Attention "层工作方式跟multi-headed self-attention是一样的，除了一点，它从前层获取输出创建Query矩阵，而Key和Value则来自于编码器输出。
总结一下，对于decoder来讲，我们注意到有两个与encoder不同的地方，一个是第一级的Masked Multi-head，另一个是第二级的Multi-Head Attention不仅接受来自前一级的输出，还要接收encoder的输出。第一级decoder的key, query, value均来自前一层decoder的输出，但加入了Mask操作，即我们只能attend到前面已经翻译过的输出的词语，因为翻译过程我们当前还并不知道下一个输出词语，这是我们之后才会推测到的。而第二级decoder也被称作encoder-decoder attention layer，即它的query来自于之前一级的decoder层的输出，但其key和value来自于encoder的输出，这使得decoder的每一个位置都可以attend到输入序列的每一个位置。
11. The Final Linear and Softmax Layer
解码器最后输出浮点向量，如何将它转成词？这是最后的线性层和softmax层的主要工作。
线性层是个简单的全连接层，将解码器的最后输出映射到一个非常大的logits数值向量上。假设模型是从有1万个单词（输出的词表）的数据集中训练得到。那么，logits向量就有1万维，每个值表示是某个词的概率分值。softmax层将这些数值转换成概率值（都是正值，且加和为1），最高值对应的维度上的词就是这一步的输出单词。
12. Recap of Training
现在我们已经了解了一个训练完毕的Transformer的前向过程，现在需要看看训练过程。
在训练时，模型将经历上述的前向过程，当我们在标记好的训练集上训练时，可以对比预测输出与实际输出。
为了可视化，假设输出一共只有6个单词（“a”, “am”, “i”, “thanks”, “student”, “”）
模型的词表是在训练之前的预处理中生成的
一旦定义了词表，我们就能够构造一个同维度的向量来表示每个单词，比如one-hot编码，下面举例编码“am”。
举例采用one-hot编码输出词表
下面让我们讨论下模型的损失函数，优化损失函数的过程将导向一个愈加准确的模型。
13. The Loss Function
我们用一个简单的例子来示范训练，比如翻译“merci”为“thanks”。那意味着输出的概率分布指向单词“thanks”，但是由于模型是随机初始化的，不太可能就是期望的输出。
由于模型参数是随机初始化的，因此只能输出随机值。对比真实输出，然后利用误差反向传播调整模型权重，从而让输出更接近真实输出
如何对比两个概率分布呢？简单采用 cross-entropy或者Kullback-Leibler divergence中的一种。
鉴于这是个极其简单的例子，更真实的情况是，使用一个句子作为输入。比如，输入是“je suis étudiant”，期望输出是“i am a student”。在这个例子下，我们期望模型输出连续的概率分布满足如下条件：
1 每个概率分布都与词表同维度。
2 第一个概率分布对“i”具有最高的预测概率值。
3 第二个概率分布对“am”具有最高的预测概率值。
4 一直到第五个输出指向同样来自词表的"<end of sentense>"标记。
单个句子的目标概率分布
在足够大的训练集上训练足够时间之后，我们期望产生的概率分布如下所示：
训练好之后，模型的输出是我们期望的翻译。当然，这并不意味着这一过程是来自训练集。注意，每个位置都能有值，即便与输出近乎无关，这也是softmax对训练有帮助的地方。
现在，模型每步只产生一个输出并选择最高概率，扔掉其他的部分，这种产生预测结果的方法叫做贪婪解码。另外一种方法是beam search，每一步仅保留最头部高概率的两个输出，根据这俩输出再预测下一步，再保留头部高概率的两个输出，重复直到预测结束。保留输出个数作为超参可试验调整。
14. Some details about source code
我们可以通过阅读源码加深理解，主要的部分是common_attention library以及调用该library的Transformer模型。
我们先来看看common_attention library 中的multihead-attention method（为简化省去很多argument及logic，集中于主要的logic，而且以单一head为例并加入了一些comment，感兴趣可以详细查阅源码）:
def multihead_attention(query_antecedent,
memory_antecedent,
...):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
...
Returns:
The result of the attention transformation. The output shape is
[batch_size, length_q, hidden_dim]
"""
#计算q, k, v矩阵
q, k, v = compute_qkv(query_antecedent, memory_antecedent， ...)
#计算dot_product的attention
x = dot_product_attention(q, k, v, ...)
x = common_layers.dense(x, ...)
return x其中compute_qkv定义为：
def compute_qkv(query_antecedent,
memory_antecedent,
...):
"""Computes query, key and value.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels]
...
Returns:
q, k, v : [batch, length, depth] tensors
"""
# 注意这里如果memory_antecedent是None，它就会设置成和query_antecedent一样，encoder的
# self-attention调用时memory_antecedent 传进去的就是None。
if memory_antecedent is None:
memory_antecedent = query_antecedent
q = compute_attention_component(
query_antecedent,
...)
# 注意这里k,v均来自于memory_antecedent。
k = compute_attention_component(
memory_antecedent,
...)
v = compute_attention_component(
memory_antecedent,
...)
return q, k, v
def compute_attention_component(antecedent,
...):
"""Computes attention compoenent (query, key or value).
Args:
antecedent: a Tensor with shape [batch, length, channels]
name: a string specifying scope name.
...
Returns:
c : [batch, length, depth] tensor
"""
return common_layers.dense(antecedent, ...)其中dot_product_attention定义为：
def dot_product_attention(q,
k,
v,
...):
"""Dot-product attention.
Args:
q: Tensor with shape [..., length_q, depth_k].
k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
match with q.
v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
match with q.
Returns:
Tensor with shape [..., length_q, depth_v].
"""
# 计算Q, K的矩阵乘积。
logits = tf.matmul(q, k, transpose_b=True)
# 利用softmax将结果归一化。
weights = tf.nn.softmax(logits, name="attention_weights")
# 与V相乘得到加权表示。
return tf.matmul(weights, v)我们再来看看Transformer模型中是如何调用的，对于encoder：
def transformer_encoder(encoder_input,
hparams,
...):
"""A stack of transformer layers.
Args:
encoder_input: a Tensor
hparams: hyperparameters for model
...
Returns:
y: a Tensors
"""
x = encoder_input
with tf.variable_scope(name):
for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
with tf.variable_scope("layer_%d" % layer):
with tf.variable_scope("self_attention"):
# layer_preprocess及layer_postprocess包含了一些layer normalization
# 及residual connection, dropout等操作。
y = common_attention.multihead_attention(
common_layers.layer_preprocess(x, hparams),
#这里注意encoder memory_antecedent设置为None
None,
...)
x = common_layers.layer_postprocess(x, y, hparams)
with tf.variable_scope("ffn"):
# 前馈神经网络部分。
y = transformer_ffn_layer(
common_layers.layer_preprocess(x, hparams),
hparams,
...)
x = common_layers.layer_postprocess(x, y, hparams)
return common_layers.layer_preprocess(x, hparams)对于decoder:
def transformer_decoder(decoder_input,
encoder_output,
hparams,
...):
"""A stack of transformer layers.
Args:
decoder_input: a Tensor
encoder_output: a Tensor
hparams: hyperparameters for model
...
Returns:
y: a Tensors
"""
x = decoder_input
with tf.variable_scope(name):
for layer in range(hparams.num_decoder_layers or hparams.num_hidden_layers):
layer_name = "layer_%d" % layer
with tf.variable_scope(layer_name):
with tf.variable_scope("self_attention"):
# decoder一级memory_antecedent设置为None
y = common_attention.multihead_attention(
common_layers.layer_preprocess(x, hparams),
None,
...)
x = common_layers.layer_postprocess(x, y, hparams)
if encoder_output is not None:
with tf.variable_scope("encdec_attention"):
# decoder二级memory_antecedent设置为encoder_output
y = common_attention.multihead_attention(
common_layers.layer_preprocess(x, hparams),
encoder_output,
...)
x = common_layers.layer_postprocess(x, y, hparams)
with tf.variable_scope("ffn"):
y = transformer_ffn_layer(
common_layers.layer_preprocess(x, hparams),
hparams,
...)
x = common_layers.layer_postprocess(x, y, hparams)
return common_layers.layer_preprocess(x, hparams)15. Go Forth And Transform
希望本文能够帮助读者对Transformer的主要概念理解有个破冰效果，如果想更深入了解，建议参考下面链接：
a. 阅读 Attention Is All You Need，Transformer的博客文章Transformer: A Novel Neural Network Architecture for Language Understanding，Tensor2Tensor使用说明。
b. 观看"Łukasz Kaiser’s talk"，梳理整个模型及其细节。
c. 尝试一下项目Jupyter Notebook provided as part of the Tensor2Tensor repo
d.尝试下项目Tensor2Tensor
相关工作