【大语言模型基础】Transformer模型Torch代码详解和训练实战

一、Transformer概述

Transformer是由谷歌在17年提出并应用于神经机器翻译的seq2seq模型，其结构完全通过自注意力机制完成对源语言序列和目标语言序列的全局依赖建模。

Transformer由编码器和解码器构成。下图展示了它的结构，其左侧和右侧分别对应着编码器（Encoder）和解码器（Decoder）结构，它们均由若干个基本的 Transformer Encoder/Decoder Block（N×表示N次堆叠）。

二、Transformer结构与实现

2.1、嵌入表示层

对于输入文本序列，首先通过**输入嵌入层（Input Embedding）**将每个单词转换为其相对应的向量表示。通常直接对每个单词创建一个向量表示。

注意：在翻译问题中，有两个词汇表，分别对应源语言和目标语言。

由于Transfomer中没有任何信息能表示单词间的相对位置关系，故需在词嵌入中加入位置编码（Positional Encoding）。

具体来说，序列中每一个单词所在的位置都对应一个向量。这一向量会与单词表示对应相加并送入到后续模块中做进一步处理。在训练的过程当中，模型会自动地学习到如何利用这部分位置信息。

2.1.1、词元嵌入层

初始化词汇表（对原始词汇表用**BPE(Byte Pair Encoding)**进行压缩分词，得到最终的词元list）

self.embedding = nn.Embedding(vocab_size, num_hiddens)

2.1.2、位置编码

为了使用序列的顺序信息，通过在输入表示中添加**位置编码（positional encoding）**来注入绝对的或相对的位置信息。

位置编码可以通过学习得到也可以直接固定得到。接下将介绍基于正弦函数和余弦函数的固定位置编码。

假设输入\(\mathbf{X} \in \mathbb{R}^{n \times d}\)表示包含一个序列中\(n\)个词元的\(d\)维嵌入表示。 位置编码使用相同形状的位置嵌入矩阵\(\mathbf{P} \in \mathbb{R}^{n \times d}\) 输出 \(\mathbf{X} +\mathbf{P}\)， 矩阵第行\(pos\)、第列\(2i\)和列上\(2i+1\)的元素为：

\[\begin{split}\begin{aligned} p_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d}}\right),\\p_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d}}\right).\end{aligned}\end{split} \]

其中，\(pos\)表示单词所在的位置，\(2i\)和\(2i+ 1\)表示位置编码向量中的对应维度，\(d\) 则对应位置编码的总维度。

通过上面这种方式计算位置编码有这样几个好处：

• 首先，正余弦函数的范围是在 [-1,+1]，导出的位置编码与原词嵌入相加不会使得结果偏离过远而破坏原有单词的语义信息。

• 其次，依据三角函数的基本性质，可以得知第\(pos + k\)个位置的编码是第\(pos\)个位置的编码的线性组合，这就意味着位置编码中蕴含着单词之间的距离信息。

class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)

2.1、多头自注意力(Multi-Head-self-Attention)

2.2.1、自注意力机制

1） 缩放点积注意力（scaled dot-product attention）

考虑到在\(d\)过大时，点积值较大会使得后续Softmax操作溢出导致梯度爆炸，不利于模型优化。故将注意力得分除以\(\sqrt{d}\)进行缩放。

注：当\(m=1\)时，就是传统的注意力机制（1个\(q\), 多个\(k\),\(v\)）。

import math
import torch
from torch import nn

class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 设置transpose_b=True为了交换keys的最后两个维度
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)

为批量处理数据或在自回归处理时避免信息泄露等情况，在Token序列中填充[mask]Token，从而使一些值不纳入注意力汇聚计算。这里可指定一个有效序列长度（即Token个数）， 以便在计算softmax时过滤掉超出指定范围的位置。

注：该缩放点积注意力的实现使用了dropout进行正则化。

masked_softmax函数实现了掩码\(softmax\)操作（masked softmax operation）， 其中任何超出有效长度的位置都被掩蔽并置为\(0\)（将掩码位置的注意力系数变为无穷小\(-inf\)，\(Softmax\)后的值为一个接近\(0\)的值）

def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e9)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X

2）自注意力

当\(n=m\)时，且\(\mathbf{Q}\)、\(\mathbf{K}\)、\(\mathbf{V}\)均源于输入\(\mathbf{X} \in\mathbb R^{n\times d}\)经过不同的线性变换时，缩放点积注意力即推广为自注意力。

这时，每个查询都会关注所有的键值对并生成一个注意力输出。 由于查询、键和值来自同一组输，故称为Self-Attention。

2.2.2、多头自注意力

class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = DotProductAttention(<