论文阅读与实现 Attention Is All You Need

论文阅读笔记
字数统计: 4.4k   |   阅读时长≈ 19 分钟

　　Transformer 是我从入门学习 NLP 开始就早有耳闻的内容，也是我之后的研究生生涯的最重要的基础框架，通过这篇论文再结合 Pytorch 版本的简单代码实现来了解 Transformer 内部的实现原理，包括位置编码，mask，attention的实现，encoder和decoder的构筑以及最终测试时贪心编码的运用。不过在阅读完这份简单代码并跟着实现之后还会有一些问题，由此会进一步改进代码。

1. 论文介绍

1.1 论文背景

　　这篇论文的目的是为了减少在以循环神经网络为基础的网络架构的计算步长，在循环神经网络中，由于对于一个长度为 seq_length 的句子，在一次迭代中数据就会经过seq_length个循环神经单元，由于每一个循环神经元的输入需要前一个循环神经元的输出，因此这种架构的模型的时间开销较大。而 Transformer 则是采用了一种全新的基于 attention 的架构替换掉原有的循环神经网络，并另使用一个位置向量来保存输入序列中处于不同位置的词，这样的构筑使得模型能够并行计算序列中的每一个序列元素对应的输出。

1.2 模型优势

图(1-2-1)

　　如图(1-2-1)展示了 Transformer 模型与循环神经网络、卷积神经网络在 Complexity per Layer , Sequential Operations , Maximum Path Length 这三个层次进行分析。参数 n 代表输入序列的长度，d代表词表示的维度，k表示在CNN网络中的核大小，r代表对 self-attention 进行限制的最长范围。

　　Complexity per Layer 是对每一层的计算复杂度进行考量，对于 Self-Attention 而言，序列中每一个输入词维度为 d ，对于一个词的输出，需要有每一个词的参与，求得每一个词对这个词的权重向量，因此就有了 n * d 的复杂度，而一共需要计算 n 个词的输出，因此总共的复杂度为 n * n * d 。

　　Sequential Operations是与并串行相关的指标，在设备并行能力足够的情况下，经过一层的步数，对于 Recurrent 而言，一层需要经过 n 个循环神经元，Self-Attention 则近似于只需要常数大小的全连接层。

　　MPL(Maximum Path Length) 的指标是针对于这类序列类型的输入，它所衡量的是在一个序列中最大间隔的两个词之间的关联性即论文中所提到的 long-range dependencies，在 Recurrent 中，输入序列的第一个词到最后一个词之间经过了 n 层，因此它的 MPL 为n，而对于 Self-Attention 而言，它的第一个和最后一个之间仅仅只通过了一个常数级别的全连接层 ，即经过全连接层得到对应的 K 和 Q，DotProductAttention 来得到结果。

　　Self-Attention (restricted) 是在输入序列较长的情况下，计算复杂度是一个平方增长，因此通过添加一个限制，即在计算一个词的输出的时候，仅考量该词附近的 r 个词来进行 self-attention 操作而非所有的词，因而将复杂度降为一个线性复杂度。

　　从 Complexity per Layer 上考量，当输入序列较小的时候 Self-Attention 相较于 Recurrent 来说每层的计算复杂度更低，而输入的嵌入向量较小时则相反，而通常处理的数据中基本都是嵌入维度远大于输入维度，并且 Self-Attention 还可以加入范围限制来进一步缩小计算复杂度，因此从这个层面上来看 Self-Attention 的优势更大。从 Sequential Operations 和 MPL 上考量，利用 self-attention 机制能够将这两个指标降为一个常数级别，相较于 Recurrent 依赖于序列长度，更加具有优势。

2. 模型架构

2.1 整体模型架构

　　如图(2-1-1)所示，输入的是一个 [batch_size, sequence_lenght] 的单词 id 序列，左半部分有 N 个相同的 Encoder(编码器)，右半部分有 N 个相同的 Decoder(解码器)，这里 N 设置为 6。初始输入序列经过一层 Embedding 后得到 [batch_size, sequence_length, d_model] 序列向量。Positional Encoding 用于为输入序列添加位置信息被添加进输入向量中。Encoder 中的每一层是由一个 Multi-Head Attention 和一个 FeedForwardNet 构成，并且这两个结构并入了残差网络结构，在其后还会有一个 LayerNorm 层。

图(2-1-1)

2.2 Embedding

　　如图(2-1-1)的inputs，是正常的词级别的 Embedding，将预处理后(代码2-2-1)得到的序数序列通过 Embedding 层转变为一个 d_model 大小的词向量。

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
sentences = [
    ['ich mochte ein bier P', 'S i want a beer .', 'i want a beer . E'],  # enc_inputs, dec_inputs, dec_outputs(label)
    ['ich mochte ein cola P', 'S i want a coke .', 'i want a coke . E']
]

src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4, 'cola': 5}
src_idx2word = {i: w for i, w in enumerate(src_vocab)}  # 构造输入词典（德语）
src_vocab_size = len(src_vocab)

tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'coke': 5, 'S': 6, 'E': 7, '.': 8}
tgt_idx2word = {i: w for i, w in enumerate(tgt_vocab)}  # 构造输出词典（英语）
tgt_vocab_size = len(tgt_vocab) 

def make_data(sentences):  # 将单词转换为字典中的序数
    enc_inputs, dec_inputs, dec_outputs = [], [], []
    for i in range(len(sentences)):
        enc_input = [src_vocab[n] for n in sentences[i][0].split()]
        dec_input = [tgt_vocab[n] for n in sentences[i][1].split()]
        dec_output = [tgt_vocab[n] for n in sentences[i][2].split()]
        enc_inputs.append(enc_input)
        dec_inputs.append(dec_input)
        dec_outputs.append(dec_output)
    return torch.tensor(enc_inputs), torch.tensor(dec_inputs), torch.tensor(dec_outputs)
enc_inputs, dec_inputs, dec_outputs = make_data(sentences)
print(enc_inputs)  # 其他同理
"""
tensor([[1, 2, 3, 4, 0],
        [1, 2, 3, 5, 0]])
"""
代码(2-2-1)

　　位置编码是图(2-1-1)中的 Position Encoding 部分，套用的论文中给出的公式(图2-2-2)，原论文中有提到位置编码也可以通过训练的方式得到，但是最终的结果和采用这种 sin 和 cos 编码的结果相似，并且这个是生成一个固定的位置编码向量，因此也更加节省资源。采用这样的编码的好处是对于任意一个位置PE_i，它与它 k 个之后的位置PE_i+k 的 内积是一个只与 k 有关的式子，也就是说即便是处于不同位置的两对词，只要它们的相对位置是相同的，那么它们的内积也是相同的，并且这个式子也保证了每一个不同位置的位置编码的唯一性。(代码2-2-3)给出了位置编码的具体实现。

图(2-2-2)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class PositionEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        PE = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # max_len, 1
        div = 10000 ** (torch.arange(0, d_model, 2) / d_model)
        PE[:, 0::2] = torch.sin(pos / div)
        PE[:, 1::2] = torch.cos(pos / div)
        PE = PE.unsqueeze(0).transpose(0, 1)
        self.register_buffer('PE', PE)

    def forward(self, x):
        """

        :param x: length, batch, d_model
        :return:
        """
        position_embeddings = self.PE[:x.shape[0], :, :]
        output = x + position_embeddings
        return self.dropout(output)
代码(2-2-3)

　　实现位置编码主要的难点在于对 python 中的广播特性不是特别熟悉，通过实践总结出当我们需要生成由多个变量决定的 tensor，在这个 tensor 中一个变量的变换将代表一个 shape（对于 PositionEncoding 而言就是变量 pos 和变量 i），按照顺序首先创建第一个变量对应的 tensor：[n1, 1]，然后再将这个变量与第二个变量的向量形式 [n2] 相乘（也不一定是相乘，可以是任何的一种函数）得到 tensor：[n1, n2]，如果有第三个变量那么就将这个相乘得到的变量扩展最后一维：[n1, n2, 1]，然后再乘以第三个变量的向量形式 [n1, n2, n3]，以此类推。

2.3 Encoding

2.3.1 Multi-Head Attention

　　在了解 Multi-Head Attention 模块之前需要先了解一下自注意力机制，注意力机制有 Additive Attention 和 Dot-Product Attention，如图(2-3-1)所示，Additive Attention 是将 k 和 q 进行拼接，再通过一个全连接层得到一个权重系数，由此得到一个 q 对应的所有的权重系数，再进行 softmax 得到权重向量。而 Dot-Product Attention 则是通过 K，Q，V 三个权重矩阵来进行运算，图(2-3-2)的方程展示了具体的运算规则，注意这里的公式是通过 Scaled 之后的公式，如果是通常的 Dot-Product Attention 则不需要除以根号dk，这里进行 Scaled 操作是由于 Dot-Product Attention 本身随着 dk 的增大导致点积之后得到的值会很大，从而使得 softmax 的梯度会非常小。图(2-3-3)我这两个 Attention 方式的理解图，代码(2-3-4)展示了 Dot-Product Attention 的详细过程。

图(2-3-1)

图(2-3-2)

图(2-3-3)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super(ScaledDotProductAttention, self).__init__()
        self.d_k = d_k

    def forward(self, Q, K, V, attn_mask):
        """
        :param Q: batch_size, n_heads, len_q, d_k
        :param K: batch_size, n_heads, len_k, d_k
        :param V: batch_size, n_heads, len_v, d_v
        :param attn_mask: batch_size, n_heads, len, len
        :return:
        """
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(self.d_k)  # batch_size, n_heads, len_q, len_k
        scores.masked_fill_(attn_mask, -1e9)  # 将 pad 的部分 mask 掉
        attn = nn.Softmax(dim=-1)(scores)
        context = torch.matmul(attn, V)
        return context, attn
代码(2-3-4)

　　代码(2-3-4)中传入的 attn_mask 矩阵中 shape 的第一个 len 由原来的 [batch_size, n_heads, len] 复制扩展而得到，而第二个 len 则是该批量中每个序列的长度对应的 mask 矩阵。通过调用 masked_fill_ 方法将点积过后的矩阵对应应该被 mask 的内容改为一个极小的负数，这样在后续的 softmax 操作时则可以将它们都置为逼近于0的值。而至于传入的 attb_mask 矩阵和 Q，K，V 如何得到则可以参照接下来的 Multi-Head-Attention 层的代码(2-3-5)内容。

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
def get_attn_pad_mask(seq_q, seq_k):
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)
    return pad_attn_mask.expand(batch_size, len_q, len_k)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, d_k, d_v, n_heads, device):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)
        self.n_heads = n_heads
        self.d_model = d_model
        self.d_k = d_k
        self.d_v = d_v
        self.device = device

    def forward(self, input_Q, input_K, input_V, attn_mask):
        """

        :param input_Q:
        :param input_K:
        :param input_V:
        :param attn_mask:
        :return:
        """
        residual = input_Q
        batch_size = input_Q.shape[0]
        Q = self.W_Q(input_Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(input_K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(input_V).view(batch_size, -1, self.n_heads, self.d_v).transpose(1, 2)  # batch_size, n_heads, seq, d_v
        attn_mask = attn_mask.unsqueeze(1).repeat(1, self.n_heads, 1, 1)
        context, attn = ScaledDotProductAttention(self.d_k)(Q, K, V, attn_mask)
        context = context.transpose(1, 2).reshape(batch_size, -1, self.n_heads * self.d_v)
        output = self.fc(context)
        return nn.LayerNorm(self.d_model).to(self.device)(output + residual), attn
代码(2-3-5)

　　在 Encoder 层中 Q，K，V 都是由 encoder 的输入序列产生，这里制造多头的方式用了一个全连接层，将输出的维度设置为 n_heads * d_model ，并将输出进行 reshape，通过一步全连接操作实现获取多头的 Q，K，V。多头注意力机制能够让模型拥有不同的注意力目标，每一个注意力的注意内容有所不同，从而使得模型能够获取更多值得注意的内容。然后将多头注意力得到的每一个输出进行拼接，通过一个全连接层，最后经过残差网络和层标准化得到最终的输出结果。而get_attn_pad_mask 函数中的 eq 能获取 seq_k 中为0的部分（也就是pad的部分）和不为0的部分的信息，并将它们分别置为 True 和 False。

2.3.2 FeedForward

　　FeedForward 的实现很简单，就是使用了两个全连接层，并在中间套用了一个 ReLU 函数，代码(2-3-5)展示了 FFN 的具体过程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class PosWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, device):
        super(PosWiseFeedForward, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False)
        )
        self.d_model = d_model
        self.device = device

    def forward(self, x):
        residual = x
        output = self.fc(x)
        return nn.LayerNorm(self.d_model).to(self.device)(output + residual)
代码(2-3-5)

　　这一部分的作用我个人没有想得太明白，参考网上的资料大致是说：由于在 Multi-Head Attention 中基本上是采用的点积的方式来得到输出，而点积方式相比于全连接层而言表达能力更差，因此需要在后边增加一个 FFN 以增强表达能力。

2.3.3 Encoder代码

　　Encoder 的输入是原始的词 id 序列，首先通过一层 Embedding 层和 Position-Encoding 层得到输入的序列向量表示，然后再进入 N=6 个 Encoding Layer 即 Multi-Head Attention 和 FeedForward 的组合。代码(2-3-6)展示了 Encoding 层的内容。

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
class EncoderLayer(nn.Module):
    def __init__(self, d_model, d_k, d_v, n_heads, d_ff, device):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention(d_model=d_model, d_k=d_k, d_v=d_v, n_heads=n_heads, device=device)
        self.pos_ffn = PosWiseFeedForward(d_model=d_model, d_ff=d_ff, device=device)

    def forward(self, enc_inputs, enc_self_attn_mask):
        """
        :param enc_inputs: batch_size, src_len, d_model
        :param enc_self_attn_mask: batch_size, src_len, src_len
        :return:
        """
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)
        enc_outputs = self.pos_ffn(enc_outputs)
        return enc_outputs, attn

class Encoder(nn.Module):
    def __init__(self, src_vocab_size, d_model, d_k, d_v, n_heads, d_ff, n_layers, device):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)
        self.pos_emb = PositionEncoding(d_model=d_model)
        self.layers = nn.ModuleList(EncoderLayer(d_model=d_model,
                                                 d_k=d_k,
                                                 d_v=d_v,
                                                 n_heads=n_heads,
                                                 d_ff=d_ff,
                                                 device=device)
                                    for _ in range(n_layers))

    def forward(self, enc_inputs):
        """

        :param enc_inputs: batch_size, src_len
        :return:
        """
        enc_outputs = self.src_emb(enc_inputs)  # batch_size, src_len, d_model
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # batch_size, src_len, d_model
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)  # batch_size, src_len, src_len
        enc_self_attns = []
        for layer in self.layers:
            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)
            enc_self_attns.append(enc_self_attn)
        return enc_outputs, enc_self_attns
代码(2-3-6)

2.4 Decoder

　　相较于 Encoder，Decoder 在 self-Multi-Head Attention 和 FeedForward 层之间新加入了一个 Multi-Head Attention，用于结合来自于 Decoder 自身的输入和来自于 Encoder 的输出，并且由于对于输出序列中的某一位置而言，它的后置位置的内容应该是不可见的，因此需要对该序列做一个 mask 操作，将后置序列变得无法影响到这一个位置的输出。代码(2-4-1)给出了 Decoder 的具体实现。与 Encoder 不同的是， Decoder 中的 mask 需要加入额外的get_attn_subsequence_mask 用于掩盖后置内容的 mask 矩阵。假设对于某一个序列的结果是['I', 'have', 'a', 'dream', 'E']，那么训练过程中的 decoder（注意不是 encoder 的输入） 的输入序列应该是['S', 'I', 'have', 'a', 'dream'] 在 decoder 中就能够通过上述的 mask 矩阵来实现这样的操作：对于模型如果想要预测 'dream' 这个位置的内容，那么在 decoder 中对于这一个预测目标上的输入相当于是 ['S', 'I', 'have', 'a']（因为后面的 ['dream'] 已经被 mask 掉了）。这里运用了一种 teaching force 的机制，在训练过程中的 decoder 的输入始终使用的是最正确的输入词，而并未使用通过decoder 生成出的结果来作为新的词并入到输入序列，而在测试阶段或者说是模型的实际使用阶段，通常是没有给定好的 decoder 的输入序列的，这个时候就只能通过初始输入词 'S' 开始一步一步地预测出每一个词，并将新的词并入到 decoder 的输入序列中，直到预测出一个表示结束的 'E' ，这一过程又称为贪心编码，代码(2-4-2)展示了具体的实现过程。

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
def get_attn_subsequence_mask(seq):
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    subsequence_mask = np.triu(np.ones(attn_shape), k=1)
    subsequence_mask = torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask

class DecoderLayer(nn.Module):
    def __init__(self, d_model, d_k, d_v, n_heads, d_ff, device):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention(d_model, d_k, d_v, n_heads, device)
        self.dec_enc_attn = MultiHeadAttention(d_model, d_k, d_v, n_heads, device)
        self.pos_ffn = PosWiseFeedForward(d_model, d_ff, device)

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        """

        :param dec_inputs: batch, tgt_len, d_model
        :param enc_outputs: batch, src_len, d_model
        :param dec_self_attn_mask: batch, tgt_len, tgt_len
        :param dec_enc_attn_mask: batch, tgt_len, src_len
        :return:
        """
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)
        dec_outputs = self.pos_ffn(dec_outputs)
        return dec_outputs, dec_self_attn, dec_enc_attn

class Decoder(nn.Module):
    def __init__(self, tgt_vocab_size, d_model, d_k, d_v, n_heads, d_ff, n_layers, device):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = PositionEncoding(d_model=d_model)
        self.layers = nn.ModuleList(DecoderLayer(d_model=d_model,
                                                 d_k=d_k,
                                                 d_v=d_v,
                                                 n_heads=n_heads,
                                                 d_ff=d_ff,
                                                 device=device)
                                    for _ in range(n_layers))
        self.device = device

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        dec_outputs = self.tgt_emb(dec_inputs)
        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).to(self.device)
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).to(self.device)
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).to(self.device)
        dec_self_attn_mask = torch.gt((dec_self_attn_subsequence_mask + dec_self_attn_pad_mask), 0).to(self.device)
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)
        dec_self_attns, dec_enc_attns = [], []
        for layer in self.layers:
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs,
                                                             dec_self_attn_mask, dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        return dec_outputs, dec_self_attns, dec_enc_attns

代码(2-4-1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def greedy_decoder(model, enc_input, start_symbol):
    """
    贪心编码
    :param model:
    :param enc_input:
    :param start_symbol:
    :return:
    """
    enc_outputs, enc_self_attns = model.encoder(enc_input)
    dec_input = torch.zeros(1, 0).type_as(enc_input.data)
    terminal = False
    next_symbol = start_symbol
    while not terminal:
        dec_input = torch.cat([dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
                              -1)
        dec_outputs, _, _ = model.decoder(dec_input, enc_input, enc_outputs)
        projected = model.projection(dec_outputs)
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        next_word = prob.data[-1]
        next_symbol = next_word
        if next_symbol == tgt_vocab["E"]:
            terminal = True
    greedy_dec_predict = dec_input[:, 1:]
代码(2-4-2)

小节

　　通过一篇 Transformer 就学习或者说是巩固了许多编程上的小技巧，比如 tensor 的广播机制，如何写一些常用的 mask 矩阵；结合代码也完全理解了 transformer 中 Encoder 和 Decoder 架构的实现原理，Position Encoding，Dot-Product Attention的具体实现，还有在seq2seq中常见的生成序列是如何通过贪心编码在测试中生成结果等等。

  • Copyrights © 2022 Tanthen

请我喝杯咖啡吧~

支付宝
微信