之前写过一篇关于 Attention Is All You Need 的 论文笔记,不过那时候写的笔记都没有深入 Code 环节,再加上其实已经有了一篇 The Annotated Transformer,也没必要做重复工作。不过现在 Transformer 已经大放异彩到几乎成为了标准配件,所以觉得有必要单独拿出来就组件角度再次学习一遍,于是就有了这篇文章。
本文代码主要基于 OpenNMT,另外也参考了一点 fairseq,这俩都是 PyTorch 实现的。Tensorflow 实现的版本相对更多一些,详见 Appendix 部分。
Transformer 中,无论 Encoder 还是 Decoder,Self-Attention 都是一个更加基础的组件,除此之外还有一个 PositionEncoder。我们就先看一下这两个更基础的组件,然后再分别学习 Encoder 和 Decoder。
Multi-Head Attention
以下代码来自 OpenMNT,先看模型的定义(为了便于理解,做了部分删减和调整):
1 | # From https://github.com/OpenNMT/OpenNMT-py |
可以看出这其实就是对 model_dim 做了拆分,比如你的输入是一个 (batch_size, seq_length, embed_dim)
的 Tensor,这里其实就是对 embed_dim 按指定的 head_count 做了拆分,拆成了三个部分,也就是所谓 Self-attention(自注意力)。
接下来是重要的 forward:
1 | def forward(self, key, value, query, mask=None): |
归纳一下步骤:
线性变换后把输入的维度(
head_count * dim_per_head
)展开。根据以下公式计算 Attention,Softmax 是返回的 Attention,它的维度为:
(batch_size, head_count, seq_len, dim_per_head)
。Attention 经过 dropout 后与 V 相乘得到 context,它的维度为:
(batch_size, head_count, seq_len, dim_per_head)
,还原后经过线性变换作为 output 输出,其维度与输入的 qkv 维度一致。
这么一看它的思想并不复杂,其实就是把输入的 Query,Key 和 Value 变为 Value 的加权(Attention)结果,权重来自 Q 和 K。为了进一步理解,我们可以假设 head_count = 1,其实它的本质就是 K 和 Q 矩阵乘法 Softmax 后得到 Attention,然后与 V 相乘后得到加权后的 context,再线性变换得到 output。
这个思想和最一开始的 Attention 是类似的,不同的是之前的 Attention 是根据 Encoder 的 output 和 Decoder 每个 time step 的 output 去计算的;而 Self-attention 是根据 K 和 Q 计算的,K 和 Q 有几种不同的用法,其实是包括了之前的 Attention:
- query 来自上一个 Decoder layer,memory keys 和 values 来自 Encoder 的 output,可以参考 Luong Attention。
- Encoder 包含 Self-attention,key value 和 query 来自相同的位置,即前一层的输出。Encoder 的每个位置都可以注意到前一层的所有位置。
- Decoder 与 Encoder 类似,不同的是需要将所有不合法连接 mask 以防止左边信息溢出,这里应该是因为 Decoder 时后面的 Token 还没有生成,自然不能用来计算 Attention。
这让 Attention 这种机制更加普遍,有没有更加理解那个标题《Attention Is All You Need》。作者当时的想法可能是,既然能够通过 Encoder 的 Output 来计算 Attention 权重,那是不是也可以根据 Encoder(或 Decoder)的前一层来计算下一层的 “Attention” 权重呢,于是就有了这篇文章。
再抽象一点去思考,其实它就是利用已有信息去计算一个更加 “好” 的 Context 表示。这也给了我们一点启示,一些好的机制是不是也可以迁移到更泛的领域呢?Transformer 的这种机制可以说是里程碑式的创新和进步,它采用了和 RNN 完全不同的建模机制,相比 RNN 的按 token 粒度切分,它从 Hidden 的粒度去切分,达到了相当的效果,训练过程却更加高效。自然这种建模方式也可以同时按 token 粒度切分去构建语言模型(比如 GPT-2),其实也就是上面的第三种用法。
Position Encoder
位置编码主要用于对位置不敏感的模型,所以 Transformer 是需要的,最常见的就是绝对位置编码,可以把索引作为编码,或者随机初始化让模型自己学习。一般位置编码会加在 WordEmbedding 上。
Transformer 中使用了相对位置编码中一种称为 Sinusoidal Positional Encoding 的技术:
其中,pos 是位置,i 是维度。
1 | # From https://github.com/OpenNMT/OpenNMT-py |
有些代码 emb 的 batch_size 在前,此时只要 pe.unsqueeze(0)
即可。此外,论文里提到的 这篇文章 有其他一些位置编码技术可以参考。
Transformer Encoder
Encoder 是由 N 个 Encoder Layer 组成,它的实现也比较简单:
1 | # From https://github.com/OpenNMT/OpenNMT-py |
Encoder 就是几个 EncoderLayer,在看 Layer 之前,提一下这个 mask,这里的 mask 其实是根据你输入 batch 中 sequence 的长度对 pad 部分进行 mask,因为 padding 部分其实是无意义的,所以这里会被设置为一个很小的负数。还是举个例子:
1 | # 样例数据 (seq_len, batch_size) |
上面的 out
和 mask
都会传入 TransformerEncoderLayer:
1 | # From https://github.com/OpenNMT/OpenNMT-py |
可以看出这部分就是 论文 中左边那部分,包含两个子模块:Multi-Head Attention 和 Feed Forward。都是归一化后传入对应的模块,然后结果和输入来一个残差连接。这里 PositionwiseFeedForward
的归一化没有体现出来,是在它的内部完成的,这个组件在 Decoder 也会用到,它的定义如下:
1 | # From https://github.com/OpenNMT/OpenNMT-py |
所以,mask 最终是在 Multi-Head Attention 里面把 input 中长度填充的部分给 mask 掉。如果不 mask,那些位置的值就是 pad 的元素值(一般为 0),Encoder 时这里可以选择 mask(提供 lengths)或者不 mask。现在大部分的操作都是会 mask 的。再举个例子:
1 | # lengths |
以上就是 Encoder 部分了,简单总结一下:
- 它由 N 个 EncoderLayer 构成
- 每个 EncoderLayer 包含两个组件:MultiHeadedAttention 和 PositionwiseFeedForward,两个组件都是以归一化的 input 作为输入,输出和 input 做残差连接。
- mask 主要对输入中长度补足的 Token 做处理(设置为一个很小的负数)。
这里的 Attention 也就是 Self-Attention,key query 和 value 来自相同的位置,Attention 的是前一层的位置。
Transformer Decoder
Decoder 要复杂一些,我们先以这篇 文章 为例,和 OpenNMT 的大同小异,只不过后者考虑了更多的情况,其中和 Encoder 中重复或类似的部分就一笔带过了。
1 | # From: http://nlp.seas.harvard.edu/2018/04/03/attention.html |
Decoder Layer 由三个部分组成,对应了三个核心组件:Masked Multi-Head Attention,Multi-Head Attention 和 FeedForward,FeedForward 和 Encoder 中的一样,即 PositionwiseFeedForward,前两个 Attention 分别对应第一部分(Multi-Head Attention)中的另外两种用法(第一种就是 Encoder 中的 Self-Attention)。说了是不同的用法,自然组件其实是一样的(上面的 self_attn 和 src_attn 其实是一个),只是怎么使用的问题。前者其实是使用了 masked Self-Attention,和 Encoder 中的情况非常类似,只不过这里 mask 的是后面未生成的位置;后者其实就是类似 Luong Attention 的机制,key 和 value 来自 Encoder 的 outputs,query 来自 Decoder,这也是最开始的 Attention。所以我们只需重点关注一下这个 masked Self-Attention 即可,具体而言就是这个 mask 如何使用的问题。
1 | # From: http://nlp.seas.harvard.edu/2018/04/03/attention.html |
先看一个例子:
1 | # 假设输入的 sequence 为 (batch_size, seq_len) |
其实这里是做了两个层面的 mask,第一个 mask 掉 padding 的位置,第二个 mask 掉未来的位置。再举一个真实的例子:
1 | # (batch_size, seq_len) |
这里如果你的输入输出都是单句,第一个 mask 其实是没意义的,因为每句都是完整的,不需要截断或者 padding,自然也就不需要 mask padding 的位置了。
以上就是 Decoder 部分了, 简单总结一下:
- 它由 N 个 DecoderLayer 构成
- 每个 DecoderLayer 包含三个组件:Masked Multi-Head Attention,Multi-Head Attention 和 PositionwiseFeedForward,均以归一化的 input 作为输入,输出和 input 做残差连接(注意看 SublayerConnection)。
- Masked Multi-Head Attention 主要对未来的 Token 位置进行 mask,Multi-Head Attention 其实是一个 Context Attention,它的 key value 均为 Encoder 的 outputs。
补充一下残差连接那里的实现说明:
1 | # 注意这里 self_attn 的参数 x 不是前面的那个 x |
再举个小例子:
1 | x = "X" |
这里的 sublayer_connection_forward
其实就相当于上面的 sublayer[i]
,sublayer 其实就是那个 lambda 函数,所以 sublayer 是先调用了 norm,再调用 attn,然后和 x 做了残差连接。这个和 Encoder 中是一样的,作者这里抽象了一个 SublayerConnection,个人觉得是非常优雅的。