Transformer - 基础分析与实现

Transformer

我们知道，自注意力同时具有并行计算和最短的最大路径长度这两个优势。因此，使用自注意力来设计深度架构是很有吸引力的。对比之前仍然依赖循环神经网络实现输入表示的自注意力模型，transformer 模型完全基于注意力机制，没有任何卷积层或循环神经网络层。尽管 transformer 最初是应用于在文本数据上的序列到序列学习，但现在已经推广到各种现代的深度学习中，例如语言、视觉、语音和强化学习领域。

模型

Transformer 作为编码器－解码器架构的一个实例，其整体架构图如下图展示。正如所见到的，transformer是由编码器和解码器组成的。与之前基于 Bahdanau 注意力实现的序列到序列的学习相比，transformer的编码器和解码器是基于自注意力的模块叠加而成的，源（输入）序列和目标（输出）序列的嵌入（embedding）表示将加上位置编码（positional encoding），再分别输入到编码器和解码器中。

上图中概述了 transformer 的架构。从宏观角度来看，transformer 的编码器是由多个相同的层叠加而成的，每个层都有两个子层（子层表示为 $\mathrm{sublayer}$）。第一个子层是多头自注意力（multi-head self-attention）汇聚；第二个子层是基于位置的前馈网络（positionwise feed-forward network）。具体来说，在计算编码器的自注意力时，查询、键和值都来自前一个编码器层的输出。受中残差网络的启发，每个子层都采用了残差连接（residual connection）。在transformer中，对于序列中任何位置的任何输入 $\mathbf{x}\in {\mathbb{R}}^d$，都要求满足 $\mathrm{sublayer}(\mathbf{x})\in {\mathbb{R}}^d$，以便残差连接满足 $\mathbf{x}+\mathrm{sublayer}(\mathbf{x})\in {\mathbb{R}}^d$。在残差连接的加法计算之后，紧接着应用层规范化（layer normalization）。因此，输入序列对应的每个位置，transformer编码器都将输出一个 $d$ 维表示向量。

Transformer解码器也是由多个相同的层叠加而成的，并且层中使用了残差连接和层规范化。除了编码器中描述的两个子层之外，解码器还在这两个子层之间插入了第三个子层，称为编码器－解码器注意力（encoder-decoder attention）层。在编码器－解码器注意力中，查询来自前一个解码器层的输出，而键和值来自整个编码器的输出。在解码器自注意力中，查询、键和值都来自上一个解码器层的输出。但是，解码器中的每个位置只能考虑该位置之前的所有位置。这种掩蔽（masked）注意力保留了自回归（auto-regressive）属性，确保预测仅依赖于已生成的输出词元。

我们已经描述并实现了基于缩放点积多头注意力和位置编码。接下来，我们将实现transformer模型的剩余部分。

import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l
1
2
3
4
5

基于位置的前馈网络

基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的（positionwise）的原因。在下面的实现中，输入X的形状（批量大小，时间步数或序列长度，隐单元数或特征维度）将被一个两层的感知机转换成形状为（批量大小，时间步数，ffn_num_outputs）的输出张量。

# 基于位置的前馈网络(由relu函数组成的两层感知机)
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    def __init__(self, ffn_num_inputs, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        
        self.dense1 = nn.Linear(ffn_num_inputs, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)
        
    # 前向传播函数，计算两层全连接层的结果
    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))
1
2
3
4
5
6
7
8
9
10
11
12
13

下面的例子显示，改变张量的最里层维度的尺寸，会改变成基于位置的前馈网络的输出尺寸。因为用同一个多层感知机对所有位置上的输入进行变换，所以当所有这些位置的输入相同时，它们的输出也是相同的。

ffn = PositionWiseFFN(4, 4, 8)            # 定义逐位前馈网络
ffn.eval()                               
# 逐位前馈网络的输出形状(batch_size, num_steps, ffn_num_outputs(num_hiddens))
ffn(torch.ones((2, 3, 4)))[0]
1
2
3
4

tensor([[-0.1728, -0.5686, -0.0589,  0.6131, -0.3308, -0.2334, -0.2196,  0.3750],
        [-0.1728, -0.5686, -0.0589,  0.6131, -0.3308, -0.2334, -0.2196,  0.3750],
        [-0.1728, -0.5686, -0.0589,  0.6131, -0.3308, -0.2334, -0.2196,  0.3750]],
       grad_fn=)
1
2
3
4

残差连接和层规范化

现在让我们关注上图中的“加法和规范化（add&norm）”组件。正如在本节开头所述，这是由残差连接和紧随其后的层规范化组成的。两者都是构建有效的深度架构的关键。

在多头注意力机制中，我们解释了在一个小批量的样本内基于批量规范化对数据进行重新中心化和重新缩放的调整。层规范化和批量规范化的目标相同，但层规范化是基于特征维度进行规范化。尽管批量规范化在计算机视觉中被广泛应用，但在自然语言处理任务中（输入通常是变长序列）批量规范化通常不如层规范化的效果好。

以下代码对比不同维度的层规范化和批量规范化的效果。

ln = nn.LayerNorm(2)                   # 层规范化
bn = nn.BatchNorm1d(2)                 # 批量规范化

# 数据X
X = torch.tensor([[1, 2], [2, 3]], dtype=torch.float32)
# 在训练模式下计算X的均值和方差
print('layer norm:', ln(X), '\nbatch norm:', bn(X))
1
2
3
4
5
6
7

layer norm: tensor([[-1.0000,  1.0000],
        [-1.0000,  1.0000]], grad_fn=) 
batch norm: tensor([[-1.0000, -1.0000],
        [ 1.0000,  1.0000]], grad_fn=)
1
2
3
4

现在我们可以使用残差连接和层规范化来实现AddNorm类。暂退法也被作为正则化方法使用。

class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        # 定义暂退法
        self.dropout = nn.Dropout(dropout)
        # nn.LayerNorm()的初始化参数为批量化的形状normalized_shape
        self.ln = nn.LayerNorm(normalized_shape)
    
    # 前向传播函数，首先进行残差连接操作，再进行层规范化
    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)
1
2
3
4
5
6
7
8
9
10
11
12

残差连接要求两个输入的形状相同，以便加法操作后输出张量的形状相同。

add_norm = AddNorm([3, 4], 0.5)
add_norm.eval()
# 输出结果仍为(batch_size, num_steps, dimension)
add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape
1
2
3
4

torch.Size([2, 3, 4])
1

编码器

有了组成transformer编码器的基础组件，现在可以先实现编码器中的一个块。下面的EncoderBlock类包含两个子层：多头自注意力和基于位置的前馈网络，这两个子层都使用了残差连接和紧随的层规范化。

class EncoderBlock(nn.Module):
    """Transformer编码器块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                dropout, use_bias=False, **kwargs):
        
        super(EncoderBlock, self).__init__(**kwargs)
        
        # 编码器块的第一个结构：多头自注意力 + 加&规范化层
        self.attention = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout,
            use_bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        
        # 编码器块的第二个结构：逐位前馈网络 + 加&规范化层
        self.ffn = PositionWiseFFN(
            ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(norm_shape, dropout)
    
    # 编码器的前向传播函数
    def forward(self, X, valid_lens):

        # Step1:先进行多头自注意力，再进行规范化操作
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        # Step2:进行逐位前馈网络后，再进行规范化操作
        return self.addnorm2(Y, self.ffn(Y))
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

正如我们所看到的，transformer 编码器中的任何层都不会改变其输入的形状。

X = torch.ones((2, 100, 24))                                               # 测试变量X
valid_lens = torch.tensor([3, 2])                                          # 有效长度[3, 2]
# EncoderBlock的参数(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input，
#    ffn_num_hiddens，num_heads, dropout)
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)      # 定义编码器块
encoder_blk.eval()
# EncoderBlock输出参数形状为(batch_size, num_steps, num_hiddens)
encoder_blk(X, valid_lens).shape
1
2
3
4
5
6
7
8

torch.Size([2, 100, 24])
1

在实现下面的 transformer编码器 的代码中，我们堆叠了 num_layers 个 EncoderBlock类 的实例。由于我们使用的是值范围在 $-1$ 和 $1$ 之间的固定位置编码，因此通过学习得到的输入的嵌入表示的值需要先乘以嵌入维度的平方根进行重新缩放，然后再与位置编码相加。

class TransformerEncoder(d2l.Encoder):
    """transformer编码器"""
    def __init__(self, vocab_size, key_size, query_size, value_size,
                num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                num_heads, num_layers, dropout, use_bias=False, **kwargs):
        
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens                                           # 隐藏层单元数num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)                   # 嵌入层embedding
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)         # 位置编码pos_encoding
        self.blks = nn.Sequential()                                              # 存放多个编码器块
        
        for i in range(num_layers):
            # 不断地向 Sequence 内追加编码器块
            self.blks.add_module('block' + str(i),
                                EncoderBlock(key_size, query_size, value_size, num_hiddens,
                                            norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                                            dropout, use_bias))
            
            
    def forward(self, X, valid_lens, *args):
        # 因为位置编码值在-1和1之间，
        # 因此嵌入值乘以嵌入维度的平方根进行缩放，
        # 然后再与位置编码相加。
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)                                              # 多个编码层块进行编码操作
            # 追加 attention 参数权重
            self.attention_weights[i] = blk.attention.attention.attention_weights
        
        return X
        
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

下面我们指定了超参数来创建一个两层的 transformer 编码器。 Transformer编码器输出的形状是（批量大小，时间步数目，num_hiddens）。

encoder = TransformerEncoder(
    200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
encoder.eval()
# TransformerEncoder输出形状(批量大小、时间步、num_hiddens)
encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape
1
2
3
4
5

torch.Size([2, 100, 24])
1

解码器

如上图所示，transformer解码器也是由多个相同的层组成。在 DecoderBlock 类中实现的每个层包含了三个子层：解码器自注意力、“编码器-解码器”注意力和基于位置的前馈网络。这些子层也都被残差连接和紧随的层规范化围绕。

正如在本节前面所述，在掩蔽多头解码器自注意力层（第一个子层）中，查询、键和值都来自上一个解码器层的输出。关于序列到序列模型（sequence-to-sequence model），在训练阶段，其输出序列的所有位置（时间步）的词元都是已知的；然而，在预测阶段，其输出序列的词元是逐个生成的。因此，在任何解码器时间步中，只有生成的词元才能用于解码器的自注意力计算中。为了在解码器中保留自回归的属性，其掩蔽自注意力设定了参数 dec_valid_lens ，以便任何查询都只会与解码器中所有已经生成词元的位置（即直到该查询位置为止）进行注意力计算。

# 解码器块DecoderBlock
class DecoderBlock(nn.Module):
    """解码器中第 i 个块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        # i表示解码器中的第i个块
        self.i = i
        
        # 解码器块的第一个结构：带掩蔽的多头注意力 + 加&规范化
        self.attention1 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        
        # 解码器块的第二个结构：多头注意力(加载编码器结果) + 加&规范化
        self.attention2 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        
        # 解码器块的第三个结构：逐位前馈网络 + 加&规范化
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    
    def forward(self, X, state):
        # 编码器输出结果state[0] ==> (batch_size, num_steps, num_hiddens)
        # 实际有效长度enc_valid_lens
        enc_outputs, enc_valid_lens = state[0], state[1]
        
        # 训练阶段，输出序列的所有词元都在同一时间处理
        # 因此state[2][self.i]初始化为None
        # 预测阶段，输出序列是通过词元一个接着一个解码的，
        # 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示！！！
        if state[2][self.i] is None:
            key_values = X
        else:
            key_values = torch.cat((state[2][self.i], X), axis=1)
            
        state[2][self.i] = key_values
        if self.training:
            batch_size, num_steps, _ = X.shape
            # dec_valid_lens的开头:(batch_size,num_steps),
            # 其中每一行是[1,2,...,num_steps]
            dec_valid_lens = torch.arange(
                1, num_steps + 1, device=X.device).repeat(batch_size, 1)
        else:
            dec_valid_lens = None

        # 多头自注意力部分 + 残差规范层
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        
        # 编码器－解码器多头注意力部分 + 残差规范层
        # enc_outputs的开头:(batch_size,num_steps,num_hiddens)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        
        #前馈网络部分 + 残差规范层
        return self.addnorm3(Z, self.ffn(Z)), state
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
58
59
60

为了便于在“编码器－解码器”注意力中进行缩放点积计算和残差连接中进行加法计算，编码器和解码器的特征维度都是num_hiddens。

decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, 0)
decoder_blk.eval()
X = torch.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
# 解码器的输出形状仍为(batch_size, num_steps, num_hiddens)
decoder_blk(X, state)[0].shape
1
2
3
4
5
6

torch.Size([2, 100, 24])
1

现在我们构建了由 num_layers 个 DecoderBlock 实例组成的完整的 transformer解码器。最后，通过一个全连接层计算所有 vocab_size 个可能的输出词元的预测值。解码器的自注意力权重和编码器解码器注意力权重都被存储下来，方便日后可视化的需要。

# Transformer解码器
class TransformerDecoder(d2l.AttentionDecoder):
    
    def __init__(self, vocab_size, key_size, query_size, value_size, 
                num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens                                         # 隐单元个数num_hiddens
        self.num_layers = num_layers                                           # 解码器块个数num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)                 # 词嵌入矩阵embedding
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)       # 位置编码
        self.blks = nn.Sequential()
        
        for i in range(num_layers):                                            # 加入多个解码器块
            self.blks.add_module('block' + str(i),
                                DecoderBlock(key_size, query_size, value_size, num_hiddens,
                                            norm_shape, ffn_num_input, ffn_num_hiddens,
                                            num_heads, dropout, i))
            
        self.dense = nn.Linear(num_hiddens, vocab_size)                        # 最终拼接全连接层
    
    
    # 初始化解码器状态信息，获取(编码器输出信息，编码器信息有效长度，第i层解码器块状态)
    def init_state(self, enc_outputs, enc_valid_lens, *args):
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]
    
    
    # 定义前向传播函数
    def forward(self, X, state):
        
        # 根据嵌入层进行位置编码
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        
        # 初始化注意力权重
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
        
        for i, blk in enumerate(self.blks):
            # 多次执行解码器块，更新输出信息X和状态信息state
            X, state = blk(X, state)
            
            # 解码器自注意力权重
            self._attention_weights[0][i] = blk.attention1.attention.attention_weights
            #“编码器-解码器”自注意力权重
            self._attention_weights[1][i] = blk.attention2.attention.attention_weights
        
        #返回全连接层后的输出结果(batch_size, num_steps, vocab_size)和状态信息state
        return self.dense(X), state
    
    @property
    def attention_weights(self):
        return self._attention_weights
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

训练

依照transformer架构来实例化编码器－解码器模型。在这里，指定transformer的编码器和解码器都是2层，都使用4头注意力。与之前类似，为了进行序列到序列(Seq2Seq)的学习，我们在“英语－法语”机器翻译数据集上训练transformer模型。

num_hiddens, num_layers, dropout, batch_size, num_step = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]

# 加载数据集
train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_step)

# Transformer的编码器架构
encoder = TransformerEncoder(
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)

# Transformer的解码器架构
decoder = TransformerDecoder(
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)

net = d2l.EncoderDecoder(encoder, decoder)
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

loss 0.034, 3032.0 tokens/sec on cpu
1

模型训练后，我们用它将几个英语句子翻译成法语并计算它们的BLEU分数。

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']

for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_step, device, True)
    print(f'{eng} => {translation}', f'bleu {d2l.bleu(translation, fra, k=2):.3f}')
1
2
3
4
5
6
7

go . => va ! bleu 1.000
i lost . => j'ai perdu . bleu 1.000
he's calm . => il est calme . bleu 1.000
i'm home . => je suis chez moi . bleu 1.000
1
2
3
4

当进行最后一个英语到法语的句子翻译工作时，让我们可视化 transformer 的注意力权重。编码器自注意力权重的形状为(编码器层数， 注意力头数， num_steps或查询的数目，num_steps或“键-值”对的数目)。

enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((
    num_layers, num_heads, -1, num_step))
enc_attention_weights.shape
1
2
3

torch.Size([2, 4, 10, 10])
1

小结

1、transformer是编码器－解码器架构的一个实践，尽管在实际情况中编码器或解码器可以单独使用。

2、在transformer中，多头自注意力用于表示输入序列和输出序列，不过解码器必须通过掩蔽机制来保留自回归属性。

3、transformer中的残差连接和层规范化是训练非常深度模型的重要工具。

4、transformer模型中基于位置的前馈网络使用同一个多层感知机，作用是对所有序列位置的表示进行转换。