Transformer 代码笔记

之前写过一篇关于 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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# From https://github.com/OpenNMT/OpenNMT-py
class MultiHeadedAttention(nn.Module):
def __init__(self, head_count, model_dim, dropout=0.1):
assert model_dim % head_count == 0
self.dim_per_head = model_dim // head_count
self.model_dim = model_dim

super(MultiHeadedAttention, self).__init__()
self.head_count = head_count

self.linear_keys = nn.Linear(model_dim, head_count * self.dim_per_head)
self.linear_values = nn.Linear(model_dim, head_count * self.dim_per_head)
self.linear_query = nn.Linear(model_dim, head_count * self.dim_per_head)
self.softmax = nn.Softmax(dim=-1)
self.dropout = nn.Dropout(dropout)
self.final_linear = nn.Linear(model_dim, model_dim)

可以看出这其实就是对 model_dim 做了拆分,比如你的输入是一个 (batch_size, seq_length, embed_dim) 的 Tensor,这里其实就是对 embed_dim 按指定的 head_count 做了拆分,拆成了三个部分,也就是所谓 Self-attention(自注意力)。

接下来是重要的 forward:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def forward(self, key, value, query, mask=None):
batch_size = key.size(0)
dim_per_head = self.dim_per_head
head_count = self.head_count
key_len = key.size(1)
query_len = query.size(1)

def shape(x):
"""Projection."""
return x.view(
batch_size, -1, head_count, dim_per_head).transpose(1, 2)

def unshape(x):
"""Compute context."""
return x.transpose(1, 2).contiguous().view(
batch_size, -1, head_count * dim_per_head)

# 1) Project key, value, and query.
# (batch_size, seq_len, head_count * dim_per_head)
key = self.linear_keys(key)
value = self.linear_values(value)
query = self.linear_query(query)
# (batch_size, head_count, seq_len, dim_per_head)
key = shape(key)
value = shape(value)
query = shape(query)

key_len = key.size(2)
query_len = query.size(2)

# 2) Calculate and scale scores. (q·k^t)/sqrt(d_k)
# **Scaled Dot-Product Attention**
query = query / math.sqrt(dim_per_head)
# (batch_size, head_count, seq_len, dim_per_head) * (batch_size, head_count, dim_per_head, seq_len)
# => (batch_size, head_count, query_len, key_len)
query_key = torch.matmul(query, key.transpose(2, 3))
scores = query_key
scores = scores.float()

if mask is not None:
mask = mask.unsqueeze(1) # (batch_size, 1, 1, seq_len)
scores = scores.masked_fill(mask, -1e18)

# 3) Apply attention dropout and compute context vectors.
attn = self.softmax(scores).to(query.dtype)
drop_attn = self.dropout(attn)
# (batch_size, head_count, seq_len, dim_per_head)
context_original = torch.matmul(drop_attn, value)
# **concat**: (batch_size, seq_len, head_count * dim_per_head)
context = unshape(context_original)
# (batch_size, seq_len, model_dim)
output = self.final_linear(context)

# Return multi-head attn
# (batch_size, head_count, query_len, key_len)
attns = attn.view(batch_size, head_count, query_len, key_len)
return output, attns

归纳一下步骤:

  • 线性变换后把输入的维度(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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# From https://github.com/OpenNMT/OpenNMT-py
class PositionalEncoding(nn.Module):
def __init__(self, dropout, dim, max_len=5000):
# (max_len, dim)
pe = torch.zeros(max_len, dim)
# (max_len, 1)
position = torch.arange(0, max_len).unsqueeze(1)
# (dim/2),式子括号里分母
div_term = torch.exp(
(torch.arange(0, dim, 2, dtype=torch.float) * -(math.log(10000.0) / dim)))
# (max_len, dim/2),0,2,4... 列
pe[:, 0::2] = torch.sin(position.float() * div_term)
# (max_len, dim/2),1,3,5... 列
pe[:, 1::2] = torch.cos(position.float() * div_term)
# (max_len, 1, dim)
pe = pe.unsqueeze(1)
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
self.dim = dim
def forward(self, emb, step=None):
"""Embed inputs.
Args:
emb (FloatTensor): Sequence of word vectors
``(seq_len, batch_size, self.dim)``
step (int or NoneType): If stepwise (``seq_len = 1``), use
the encoding for this position.
"""
emb = emb * math.sqrt(self.dim)
if step is None:
# 给 emd 加 position encoding
emb = emb + self.pe[:emb.size(0)]
else:
emb = emb + self.pe[step]
emb = self.dropout(emb)
return emb

有些代码 emb 的 batch_size 在前,此时只要 pe.unsqueeze(0) 即可。此外,论文里提到的 这篇文章 有其他一些位置编码技术可以参考。

Transformer Encoder

Encoder 是由 N 个 Encoder Layer 组成,它的实现也比较简单:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# From https://github.com/OpenNMT/OpenNMT-py
class TransformerEncoder(EncoderBase):
"""
Args:
num_layers (int): number of encoder layers
d_model (int): size of the model
heads (int): number of heads
d_ff (int): size of the inner FF layer
dropout (float): dropout parameters
embeddings (onmt.modules.Embeddings):
embeddings to use, should have positional encodings
Returns:
(torch.FloatTensor, torch.FloatTensor):
* embeddings ``(src_len, batch_size, model_dim)``
* memory_bank ``(src_len, batch_size, model_dim)``
"""
def __init__(self, num_layers, d_model, heads, d_ff, dropout,
attention_dropout, embeddings):
super(TransformerEncoder, self).__init__()

self.embeddings = embeddings
self.transformer = nn.ModuleList(
[TransformerEncoderLayer(d_model, heads, d_ff, dropout, attention_dropout)
for i in range(num_layers)])
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
def forward(self, src, lengths):
# (seq_len, batch_size, embedding_size)
emb = self.embeddings(src)
# (batch_size, seq_len, embedding_size)
out = emb.transpose(0, 1).contiguous()
# lengths: (batch_size, )
# (batch_size, 1, seq_len)
mask = ~sequence_mask(lengths).unsqueeze(1)
for layer in self.transformer:
out = layer(out, mask)
out = self.layer_norm(out)
return emb, out.transpose(0, 1).contiguous(), lengths

Encoder 就是几个 EncoderLayer,在看 Layer 之前,提一下这个 mask,这里的 mask 其实是根据你输入 batch 中 sequence 的长度对 pad 部分进行 mask,因为 padding 部分其实是无意义的,所以这里会被设置为一个很小的负数。还是举个例子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 样例数据 (seq_len, batch_size)
{'source_var': tensor([
[ 11, 48, 4, 32],
[ 49, 2573, 54, 17],
[1186, 637, 3453, 1159],
[ 52, 747, 7, 7],
[ 164, 7, 3, 3],
[1187, 3, 0, 0],
[ 128, 0, 0, 0],
[1188, 0, 0, 0],
[ 7, 0, 0, 0],
[ 3, 0, 0, 0]]),
'source_lengths': [10, 6, 5, 5]
}
# 返回的 mask 为 (batch_size, 1, 1, seq_len)
tensor([[[False, False, False, False, False, False, False, False, False, False]],

[[False, False, False, False, False, False, True, True, True, True]],

[[False, False, False, False, False, True, True, True, True, True]],

[[False, False, False, False, False, True, True, True, True, True]]])

上面的 outmask 都会传入 TransformerEncoderLayer:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
# From https://github.com/OpenNMT/OpenNMT-py
class TransformerEncoderLayer(nn.Module):
"""
A single layer of the transformer encoder.
Args:
d_model (int): the dimension of keys/values/queries in
MultiHeadedAttention, also the input size of
the first-layer of the PositionwiseFeedForward.
heads (int): the number of head for MultiHeadedAttention.
d_ff (int): the second-layer of the PositionwiseFeedForward.
dropout (float): dropout probability(0-1.0).
"""

def __init__(self, d_model, heads, d_ff, dropout, attention_dropout,
max_relative_positions=0):
super(TransformerEncoderLayer, self).__init__()

self.self_attn = MultiHeadedAttention(heads, d_model, dropout=attention_dropout)
self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
self.dropout = nn.Dropout(dropout)

def forward(self, inputs, mask):
"""
Args:
inputs (FloatTensor): ``(batch_size, src_len, model_dim)``
mask (LongTensor): ``(batch_size, 1, src_len)``
Returns:
(FloatTensor):
* outputs ``(batch_size, src_len, model_dim)``
"""
input_norm = self.layer_norm(inputs)
context, _ = self.self_attn(
input_norm, input_norm, input_norm, mask=mask, attn_type="self")
out = self.dropout(context) + inputs
return self.feed_forward(out)

可以看出这部分就是 论文 中左边那部分,包含两个子模块:Multi-Head Attention 和 Feed Forward。都是归一化后传入对应的模块,然后结果和输入来一个残差连接。这里 PositionwiseFeedForward 的归一化没有体现出来,是在它的内部完成的,这个组件在 Decoder 也会用到,它的定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# From https://github.com/OpenNMT/OpenNMT-py
class PositionwiseFeedForward(nn.Module):
""" A two-layer Feed-Forward-Network with residual layer norm.

Args:
d_model (int): the size of input for the first-layer of the FFN.
d_ff (int): the hidden layer size of the second-layer
of the FNN.
dropout (float): dropout probability in :math:`[0, 1)`.
"""
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.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
self.dropout_1 = nn.Dropout(dropout)
self.relu = nn.ReLU()
self.dropout_2 = nn.Dropout(dropout)

def forward(self, x):
"""Layer definition.
Args:
x: ``(batch_size, input_len, model_dim)``
Returns:
(FloatTensor): Output ``(batch_size, input_len, model_dim)``.
"""
inter = self.dropout_1(self.relu(self.w_1(self.layer_norm(x))))
output = self.dropout_2(self.w_2(inter))
return output + x

所以,mask 最终是在 Multi-Head Attention 里面把 input 中长度填充的部分给 mask 掉。如果不 mask,那些位置的值就是 pad 的元素值(一般为 0),Encoder 时这里可以选择 mask(提供 lengths)或者不 mask。现在大部分的操作都是会 mask 的。再举个例子:

1
2
3
4
5
6
7
8
9
# lengths
# (batch_size, ),每个元素为序列的长度
lens = torch.Tensor([5, 4, 3, 2])
# (batch_size, 1, 1, seq_len), (4, 1, 1, 5)
mask = ~sequence_mask(lens).unsqueeze(1)
# (batch_size, head_count, query_len, key_len)
score = torch.randint(1, 10, (4, 8, 5, 5))
# (batch_size, head_count, query_len, key_len)
score = score.masked_fill(mask, -1e18)

以上就是 Encoder 部分了,简单总结一下:

  • 它由 N 个 EncoderLayer 构成
  • 每个 EncoderLayer 包含两个组件:MultiHeadedAttention 和 PositionwiseFeedForward,两个组件都是以归一化的 input 作为输入,输出和 input 做残差连接。
  • mask 主要对输入中长度补足的 Token 做处理(设置为一个很小的负数)。

这里的 Attention 也就是 Self-Attention,key query 和 value 来自相同的位置,Attention 的是前一层的位置。

Transformer Decoder

Decoder 要复杂一些,我们先以这篇 文章 为例,和 OpenNMT 的大同小异,只不过后者考虑了更多的情况,其中和 Encoder 中重复或类似的部分就一笔带过了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# From: http://nlp.seas.harvard.edu/2018/04/03/attention.html
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(m, m, x, src_mask))
return self.sublayer[2](x, self.feed_forward)

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

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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# From: http://nlp.seas.harvard.edu/2018/04/03/attention.html
class Batch:
"Object for holding a batch of data with mask during training."
def __init__(self, src, trg=None, pad=0):
self.src = src
self.src_mask = (src != pad).unsqueeze(-2)
if trg is not None:
self.trg = trg[:, :-1]
self.trg_y = trg[:, 1:]
self.trg_mask = self.make_std_mask(self.trg, pad)
self.ntokens = (self.trg_y != pad).data.sum()
@staticmethod
def make_std_mask(tgt, pad):
"Create a mask to hide padding and future words."
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)))
return tgt_mask

先看一个例子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# 假设输入的 sequence 为 (batch_size, seq_len)
ts = torch.Tensor([[2, 3, 1],
[2, 3, 0]]).type(torch.LongTensor)

# (batch_size, 1, seq_len)
(ts != 0).unsqueeze(-2)
"""
tensor([[[ True, True, True]],

[[ True, True, False]]])
"""

# (1, seq_len, seq_len)
subsequent_mask(ts.size(-1))
"""
tensor([[[ True, False, False],
[ True, True, False],
[ True, True, True]]])
"""

# (batch_size, seq_len, seq_len)
(ts != 0).unsqueeze(-2) & subsequent_mask(ts.size(-1))
"""
tensor([[[ True, False, False],
[ True, True, False],
[ True, True, True]],

[[ True, False, False],
[ True, True, False],
[ True, True, False]]])
"""

其实这里是做了两个层面的 mask,第一个 mask 掉 padding 的位置,第二个 mask 掉未来的位置。再举一个真实的例子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# (batch_size, seq_len)
batch.src # (1, 8)
# tensor([[ 61, 89, 560, 129, 11, 159, 57977, 3]])

# (batch_size, seq_len)
batch.src_mask # (1, 1, 8)
# tensor([[[True, True, True, True, True, True, True, True]]])

# (batch_size, seq_len)
batch.trg # (1, 9)
# tensor([[ 2, 59, 1348, 12, 90, 32926, 70, 16, 5]])

# (batch_size, 1, seq_len)
batch.trg_mask # (1, 9, 9)
"""
tensor([[[ True, False, False, False, False, False, False, False, False],
[ True, True, False, False, False, False, False, False, False],
[ True, True, True, False, False, False, False, False, False],
[ True, True, True, True, False, False, False, False, False],
[ True, True, True, True, True, False, False, False, False],
[ True, True, True, True, True, True, False, False, False],
[ True, True, True, True, True, True, True, False, False],
[ True, True, True, True, True, True, True, True, False],
[ True, True, True, True, True, True, True, True, True]]])
"""

这里如果你的输入输出都是单句,第一个 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
2
3
4
# 注意这里 self_attn 的参数 x 不是前面的那个 x
sublayer[0](x, lambda x: self_attn(x, x, x, tgt_mask))
# 实际上是这样的,后面的 x 是 self_attn 的参数
sublayer[0](x, lambda y: self_attn(y, y, y, tgt_mask))

再举个小例子:

1
2
3
4
5
6
7
8
9
x = "X"
def attn(x):
return x + " attned"
def norm(x):
return x + " normed"
def sublayer_connection_forward(x, sublayer):
return x + sublayer(norm(x))
sublayer_connection_forward(x, lambda x: attn(x))
# 'XX normed attned'

这里的 sublayer_connection_forward 其实就相当于上面的 sublayer[i],sublayer 其实就是那个 lambda 函数,所以 sublayer 是先调用了 norm,再调用 attn,然后和 x 做了残差连接。这个和 Encoder 中是一样的,作者这里抽象了一个 SublayerConnection,个人觉得是非常优雅的。

Appendix