【论文学习】Transformer：模型搭建

Transformer由论文Attention Is All You Need提出，该模型舍弃了RNN和CNN，完全基于注意力机制实现了高效、并行化的训练，成为了NLP领域的里程碑式模型。本篇文章介绍Transformer的模型结构并基于PyTorch完成模型的搭建。

0 引言

  Transformer基于seq2seq架构实现NLP领域的典型任务，如机器翻译、文本生成等。

  在论文中，作者以机器翻译任务为主，详细介绍了Transformer的结构，其总体框架如图1(来源于论文)所示，一共包含四个部分：

• 输入部分：输入词向量预处理

  • 源文本嵌入层及位置编码器
  • 目标文本嵌入层及位置编码器

• 编码器部分

  • 多头自注意力机制子层和规范层以及一个残差连接
  • 前馈前连接子层和规范层以及一个残差连接

• 解码器部分

  • 多头自注意力机制子层和规范层以及一个残差连接
  • 多头注意力机制子层和规范层以及一个残差连接
  • 前馈前连接子层和规范层以及一个残差连接

• 输出部分：输出向量预处理

  • 线性层

  • softmax层

图1 Transformer模型框架

1 输入部分

1.1 文本嵌入层

  文本嵌入层将文本中词汇数字表示转变为高维向量表示，更便于捕捉词汇间的关系。其代码实现如下：

class Embedding(nn.Module):
    def __init__(self, d_model, vocab):
        '''
        文本嵌入层
        :param d_model: 词嵌入的维度
        :param vocab: 词表的大小
        '''
        super().__init__()
        # 嵌入层
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model
    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

1.2 位置编码器

  位置编码器在词向量中添加每个词在句子中的位置信息，使模型具有学习词序信息的能力。论文中采用正余弦交替的位置编码，公式如下：
$$P{E}_t^{(i)}=\{\begin{array}{l}sin(w_{k}t)，i=2k\\ cos(w_{k}t)，i=2k+1\end{array}$$

$$w_k=\frac{1}{10000^{2k/d\_model}},k=0,1,2,...,\frac{d\_model}{2}-1$$

其中$i$表示词向量$t$所处的位置，$d\_model$表示词向量$t$的维度。

  其代码实现如下：

class PositionEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.2, max_len=5000):
        '''
        位置编码层:正余弦交替编码
        :param d_model: 词向量维度
        :param dropout: 置零比例,正则化,防止模型过拟合
        :param max_len: 每个句子的最大长度
        '''
        super().__init__()
        # dropout正则化层
        self.dropout = nn.Dropout(p=dropout)
        # 初始化一个位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        # 初始化一个位置矩阵(句子中每个词向量的位置) [max_len, 1]
        position = torch.arange(0, max_len).unsqueeze(1)
        # 构造位置编码中正余弦频率(共用频率) [d_model // 2,]
        w = 1 / torch.pow(10000, torch.arange(0, d_model, 2) / d_model)
        # 构造最终的位置编码:偶数位置使用正弦编码,奇数位置使用余弦编码 [max_len, d_model]
        pe[:, 0::2] = torch.sin(position * w)
        pe[:, 1::2] = torch.cos(position * w)
        # [max_len, d_model] -> [1, max_len, d_model]
        pe = pe.unsqueeze(0)
        # 位置编码注册为模型buffer,其在优化过程中不会更新,在模型保存后重加载时和模型一起加载
        self.register_buffer('pe', pe)
    def forward(self, x):
        '''
        :param x: [bs(句子数量), length(句子长度), d_model] 序列的嵌入(embedding)表示
        :return x_pos:[bs, length, d_model] 加上词序信息的词向量
        '''
        length = x.size(1)  # 句子的长度
        x = x + Variable(self.pe[:, :length, :], requires_grad=False)
        return self.dropout(x)

2 编码器部分

  编码器部分结构如图2所示，由N个编码器层组成，每个编码器层由两个子层连接而成：

• 第一个子层包括一个多头自注意力机制和规范层以及一个残差连接
• 第二个子层包括一个前馈全连接子层和规划层以及一个残差连接
  

图2 编码器结构 >

2.1 掩码张量

  在Transformer中，生成的attention信息中，可能会包含未来的信息，为了避免未来的信息被提前使用，会使用掩码张量对部分信息进行遮掩。其构建代码如下：

def subsequent_mask(size):
    attn_shape = (1, size, size)

    # 首先构建一个全1矩阵, 再将其形成上三角矩阵(对角线元素也舍弃)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 构造一个下三角矩阵(包含对角线元素)
    return torch.from_numpy(1 - subsequent_mask)

2.2 注意力机制

我们观察事物时，大脑能够快速的把注意力放在事物最具辨识度的部分从而作出判断。基于这种形式，在模型中引入注意力机制，帮助模型聚焦于对当前任务更为关键的信息，降低对其他信息的关注，提高模型对特定任务处理的效率和准确率。

  论文中使用的注意力计算方法如图3所示，计算公式如下：
$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
其中$Q(query),K(key),V(value)$由输入词向量$X$通过线性变化得到，当$Q=K=V$时称为自注意力机制，$d_k$为词嵌入维度。

图3 注意力机制结构

  注意力计算的实现代码如下：

def attention(query, key, value, mask=None, dorpout=None):
    '''
    :param query: [bs, head, length, d_k] Q
    :param key: [bs, head, length, d_k] K
    :param value: [bs, head, length, d_k] V
    :return:
    '''
    d_k = query.size(-1)  # d_model/d_k 词嵌入维度
    # 计算Qk^T,得到注意力张量scores  [bs, head, length, length]
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    # 判断是否使用掩码张量
    if mask is not None:
        scores = scores.mask_fill(mask == 0, -1e9)
    # 对scores的最后一维度进行softmax操作 [bs, head, length, length]
    p_attn = F.softmax(scores, dim=-1)
    # 判断是否使用dropout
    if dorpout is not None:
        p_attn = dorpout(p_attn)
    # 得到最终的注意力表达式 [bs, head, length, d_k]
    attn = torch.matmul(p_attn, value)
    return attn, p_attn

2.3 多头注意力机制

  多头注意力机制结构如图4所示，其将输入词向量在词嵌入维度上d\_model进行切分，分为$h$个头，每个头的词嵌入维度为$d\_model/h$，分别进行自注意力机制运算后再对结果在通道上进行拼接。

  多头注意力机制让每个注意力去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义信息具有更丰富的表达。

图4 多头注意力机制结构

  多头注意力机制实现代码如下：

def clones(module, N):
    '''用于生成相同网络层的克隆函数,其参数不共享'''
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        '''
        多头注意力机制
        :param h: 词向量在词嵌入维度上划分的头数
        :param d_model: 词嵌入维度
        :param dropout: 
        '''
        super().__init__()
        # 每个头的分割词向量维度d_k
        assert d_model % h == 0
        self.d_k = d_model // h
        # 分割头数
        self.h = h
        # 得到线性变化层,分别用于Q,K,V矩阵以及对拼接矩阵的线性操作
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        # 注意力张量
        self.attn = None
        # dropout正则化
        self.dropout = nn.Dropout(p=dropout)
    def forward(self, query, key, value, mask=None):

        # 样本数
        batch_size = query.size(0)

        # 多头处理环节：首先对QKV进行线性操作, 然后按照头数进行划分
        # q,k,v:[bs, length, head, d_k] -> [bs, head, length, d_k]
        query, key, value = \
            [model(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
             for model, x in zip(self.linears, (query, key, value))]
        # 对每个头进行注意力计算
        # x: 注意力张量[bs, head, length, d_k]
        # attn: 注意力张量权重[bs, head, length, length]
        x, self.attn = attention(query, key, value, mask=mask, dorpout=self.dropout)
        # 将多头的注意力张量进行合并 [bs, head, length, d_k] -> [bs, length, head, d_k] -> [bs, length, head*d_k(d_model)]
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
        # 多头注意力机制返回结果 [bs, length, d_model]
        return self.linears[-1](x)

2.4 前馈全连接层

  前馈全连接层由两个线性操作和ReLU激活函数组成，其公式如下：
$$FFN(x)=W_{2}max(0,W_{1}x+b_1)+b_2$$
  其代码实现如下：

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_hidden, dropout=0.1):
        '''
        前馈全连接层
        :param d_model: 词嵌入维度
        :param d_hidden: 隐藏层维度
        :param dropout:
        '''
        super().__init__()
        self.w1 = nn.Linear(d_model, d_hidden)
        self.w2 = nn.Linear(d_hidden, d_model)
        self.act = F.relu
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):
        '''
        :param x: 词向量 (bs, length, d_model)
        :return:
        '''
        y = self.dropout(self.act(self.w1(x)))
        return self.w2(y)

2.5 规范化层

  规范化层实现对神经网络数值的标准化，使其特征数值满足均值为0，方差为1的分布，有利于模型的收敛。其代码实现如下：

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        '''
        规范化层
        :param features: 输入特征的维度
        :param eps:防止分母为0
        '''

        self.a2 = nn.Parameter(torch.ones(features))  # 随着模型训练
        self.b2 = nn.Parameter(torch.zeros(features))  # 随着模型训练
        self.eps = eps

    def forward(self, x):
        '''
        :param x: 词向量 [bs, length, d_model]
        :return:
        '''
        mean = x.mean(-1, keepdim=True)  # 均值
        std = x.std(-1, keepdim=True)  # 方差
        # 标准化 [bs, length, d_model]
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

2.6 子层连接结构(残差连接)

  对子层结构进行封装，实现代码如下：

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        '''
        :param size: 词嵌入维度
        :param dropout:
        '''
        self.norm = LayerNorm(size)  # 标准化
        self.dropout = nn.Dropout(p=dropout)
    def forward(self, x, sublayer):
        '''
        :param x: 词向量 [bs, length, d_model]
        :sublayer: MultiHeadedAttention / PositionwiseFeedForward
        :return:
        '''
        return x + self.dropout(sublayer(self.norm(x)))

2.7 编码器层

  结合以上函数，实现编码器层：

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        '''
        :param size: 词嵌入维度
        :param self_attn: 多头注意力机制
        :param feed_forward: 前馈全连接层
        :param dropout:
        '''
        self.self_attn = self_attn
        self.feed_forward = feed_forward

        # 构造子层连接
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size
    def forward(self, x, mask):
        y = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))  # 第一个子层(多头注意力+标准化+残差连接)
        y = self.sublayer[1](y, self.feed_forward)  # 第二个子层(前馈全连接+标准化+残差连接)
        return y

2.8 编码器

  基于编码器层，构造编码器：

class Encoder(nn.Module):
    def __init__(self, layer, N):
        '''
        :param layer: 编码器层
        :param N: 编码器层数
        '''
        super().__init__()
        self.layers = clones(layer, N)  # 构造N层编码器
        self.norm = LayerNorm(layer.size)  # 构造标准化层
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

3 解码器部分

  解码器部分结构如图5所示，由N个编码器层组成，每个编码器层由三个子层连接而成：

• 多头自注意力机制子层和规范层以及一个残差连接
• 多头注意力机制子层和规范层以及一个残差连接
• 前馈前连接子层和规范层以及一个残差连接
  

图5 解码器结构

3.1 解码器层

  结合编码器部分中实现的模块，解码器层代码如下：

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout=0.1):
        '''
        :param size: 词嵌入维度
        :param self_attn: 多头自注意力机制(Q=K=V)
        :param src_attn: 多头注意力机制(Q!=K=V)
        :param feed_forward: 前馈全连接层
        :param dropout:
        '''
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        # 构造子连接层
        self.sublayers = clones(SublayerConnection(size, dropout), 3)
    def forward(self, x, memory, source_mask, target_mask):
        '''
        :param x: 上一层的输出张量 [bs, length, d_model]
        :param memory: 编码器层的输出 [bs, length, d_model]
        :param source_mask: 多头注意力机制的掩码张量
        :param target_mask: 多头自自注意机制的掩码张量
        '''
        m = memory

        # 第一个子层:多头自注意力机制 + 标准化 + 残差连接
        # 目标数据掩码张量:避免对未来信息的使用
        x = self.sublayers[0](x, lambda x:self.self_attn(x, x, x, target_mask))
        # 第二个子层:多头注意力机制 + 标准化 + 残差连接
        # 源数据掩码张量:遮挡对结果无效的信息
        x = self.sublayers[1](x, lambda x:self.src_attn(x, m, m, source_mask))
        # 第三个子层:前馈全连接层 + 标准化 + 残差连接
        x = self.sublayers[2](x, self.feed_forward)

        return x

3.2 解码器

  基于解码器层，构造解码器：

class Decoder(nn.Module):
    def __init__(self, layer, N):
        '''
        :param layer: 解码器层
        :param N: 数量
        '''
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
    def forward(self, x, memory, source_mask, target_mask):
        '''
        :param x: 数据的嵌入表示 [bs, length, d_model]
        :param memory: 编码器输出 [bs, length, d_model]
        :param source_mask:
        :param target_mask:
        '''
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)

4 输出部分

  输出层包括线性层和softmax层，其代码如下：

class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        '''
        :param d_model: 词嵌入维度
        :param vocab_size: 词表大小
        '''
        super().__init__()
        # 线性层,将词向量转换为词表预测结果
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        '''
        :param x: 解码器的最终输出 [bs, length, d_model]
        :return: 最终预测结果 [bs, length, vocab_size]
        '''
        return F.log_softmax(self.project(x), dim=-1)

5 模型搭建

class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        '''
        编码器 - 解码器结构
        :param encoder: 编码器
        :param decoder: 解码器
        :param source_embed: 源数据嵌入函数
        :param target_embed: 目标数据嵌入函数
        :param generator: 类别生成器
        '''
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator
    def forward(self, source, target, source_mask, target_mask):
        '''
        :param source: 源数据 [bs, length]
        :param target: 目标数据 [bs, length]
        :param source_mask: 源数据掩码张量
        :param target_mask: 目标数据掩码张量
        '''
        return self.decode(self.encode(source, source_mask), source_mask,
                          target, target_mask)
    def encode(self, source, source_mask):
        '''编码器'''
        return self.encoder(self.src_embed(source), source_mask)
    def decode(self, memory, source_mask, target, target_mask):
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)


def make_model(source_vocab, target_vocab, N=6,
               d_model=512, d_hidden=2048, head=8, dropout=0.1):
    '''
    构建模型
    :param source_vocab: 源数据特征(词汇)总数
    :param target_vocab: 目标数据特征(词汇)总数
    :param N: 解码器和编码器个数
    :param d_model: 词嵌入维度
    :param d_hidden: 前馈全连接层隐藏层层数
    :param head: 多头注意力机制头数
    :param dropout:
    '''
    # 深拷贝
    c = copy.deepcopy
    # 实例化多头注意力机制类
    attn = MultiHeadedAttention(head, d_model)
    # 实例化前馈全连接层类
    ff = PositionwiseFeedForward(d_model, dropout)
    # 实例化位置编码类
    position = PositionEncoding(d_model, dropout=dropout)

    # 构造模型
    model = EncoderDecoder(
        encoder=Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        decoder=Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        source_embed=nn.Sequential(Embedding(d_model, source_vocab), c(position)),
        target_embed=nn.Sequential(Embedding(d_model, target_vocab), c(position)),
        generator=Generator(d_model, target_vocab)
    )
    # 初始化参数
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model