从零开始复现GPT2（二）：模型实现和掩码机制

---

源码地址：https://gitee.com/guojialiang2023/gpt2

---

模型

掩码机制

定义了两个类：PadMasking 和 FutureMasking。被设计用于在GPT-2中处理序列数据。让我们一步步详细解释每个类的功能和工作原理。

PadMasking 类

功能

PadMasking 类用于创建一个掩码张量（mask tensor），该张量标识输入序列中的填充（padding）位置。在处理不等长的序列时，通常需要用特定的填充符号（如0）将它们填充到相同的长度。这个掩码有助于模型识别哪些部分是填充，从而在训练和推理过程中忽略这些部分。

构造函数 (init)

• 参数 pad_idx 是用于表示填充值的索引。
• super().__init__() 是调用父类 nn.Module 的构造函数。

forward 方法

• 输入 x 是一个张量，表示输入序列。
• offset 是一个整数，默认为0，用于在掩码张量中添加额外的列。
• 方法首先识别出哪些元素是填充元素（is_pad），然后创建一个零张量（shifted），其大小与输入的最后一个维度（除去序列长度）相同，并且在该维度上添加 offset 数量的额外列。
• 然后，将 shifted 和 is_pad 在最后一个维度上连接起来，形成最终的掩码张量。
• 这个掩码张量会扩展到与输入张量 x 的形状相匹配。

FutureMasking 类

功能

FutureMasking 类用于创建一个未来掩码（future mask），在自回归模型GPT-2中，此掩码用于确保在预测每个位置的输出时，模型只能使用该位置之前的输入（即避免使用未来的信息）。

forward 方法

• 输入 x 是一个张量，表示输入序列。
• offset 是一个整数，默认为0，它决定了未来掩码的偏移量。
• 方法首先创建一个大小为 (seq_len, seq_len + offset) 的全1张量，然后使用 triu（上三角）函数来保留偏移量之后的上三角部分，并将其余部分设置为0。这样确保了每个位置只能访问其之前和偏移量定义的未来位置。
• 接着，将未来掩码张量的形状调整为适应输入张量 x 的维度。
• 最后，扩展未来掩码张量以匹配输入张量 x 的形状。

import torch
import torch.nn as nn

class PadMasking(nn.Module):
    """
    Tensor          Type            Shape
    ===========================================================================
    input           long            (..., seq_len)
    ---------------------------------------------------------------------------
    output          float           (..., seq_len, seq_len + offset)
    ===========================================================================
    """
    def __init__(self, pad_idx: int):
        super.__init__()
        self.pad_idx = pad_idx

    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        is_pad = (x == self.pad_idx).unsqueeze(-2)
        shifted = torch.zeros(x.size()[:-1] + (1, offset,),
                              dtype=torch.bool, device=x.device)

        mask = torch.cat((shifted, is_pad), dim=-1)
        return mask.expand(x.shape + mask.shape[-1:])

class FutureMasking(nn.Module):
    """
    Tensor          Type            Shape
    ===========================================================================
    input           long            (..., seq_len)
    ---------------------------------------------------------------------------
    output          float           (..., seq_len, seq_len + offset)
    ===========================================================================
    """
    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        seq_len = x.size(-1)

        # Create shifted upper triangular matrix.
        future = torch.ones((seq_len, seq_len + offset),
                            dtype=torch.bool, device=x.device)
        future = future.triu(offset + 1)

        mask = future.view((1,) * (x.ndim - 1) + future.size())
        return mask.expand(x.shape + mask.shape[-1:])

模型实现

代码定义了 TransformerLayer ，它表示 Transformer 模型中的一个层级。以下是对这段代码的详细解释：

1. 输入和输出张量说明：

   • x：输入张量，表示当前时间步的输入特征序列，其形状为 (batch_size, seq_len, dims)。这个张量包含了当前时间步的信息，其中 seq_len 是序列的长度，dims 是特征的维度。
   • past（可选）：一个包含先前时间步注意力信息的数据结构，通常用于自回归生成任务。如果存在，它表示先前时间步的注意力信息，其形状为 (batch_size, past_len, dims)，其中 past_len 是先前时间步的序列长度。
   • mask（可选）：一个掩码张量，用于屏蔽某些位置的注意力计算。它的形状为 (batch_size, seq_len, past_len + seq_len)，通常用于处理注意力的掩码，以限制模型关注的位置。

2. 初始化函数：

   • __init__ 函数用于初始化 Transformer 层。它接受以下参数：
     • heads：注意力头的数量。
     • dims：隐藏单元的维度。
     • rate：前馈神经网络（Feed-Forward）的增长率。
     • dropout（可选）：用于添加 dropout 的概率，默认为 0.1。
   • 在初始化中，创建了以下子模块：
     • self.attn：多头注意力层，用于处理注意力计算。
     • self.ff：位置前馈神经网络层，用于处理前馈传递。
     • self.ln_attn 和 self.ln_ff：Layer Normalization 层，用于规范化层输入。

3. 前向传播函数：

   • forward 函数用于执行 Transformer 层的前向传播。
   • 首先，通过 Layer Normalization (self.ln_attn) 规范化输入 x，得到 a。
   • 接下来，调用多头注意力层 (self.attn) 来计算注意力信息。如果存在 past，则将其传递给注意力层以处理自回归生成任务。计算结果包括注意力输出 a 和更新后的 past。
   • 将输入 x 与注意力输出 a 相加，以获得注意力层的输出。
   • 通过 Layer Normalization (self.ln_ff) 规范化输出，然后将其传递给位置前馈神经网络层 (self.ff) 进行前馈传递。
   • 最后，将前馈传递的输出与注意力层的输出相加，得到 Transformer 层的最终输出 x。

4. 输出：

   • 如果模型处于训练模式 (self.training=True)，则返回单个张量 x 作为输出。
   • 如果模型处于生成模式，通常会返回一个元组 (x, past)，其中 x 是输出张量，past 包含了先前时间步的注意力信息，用于下一个时间步的生成。

表示了 Transformer 模型中的一个层级，包括多头注意力层和前馈神经网络层，以及 Layer Normalization 层，用于处理序列数据和生成下一个时间步的输出。这个层级在整个 Transformer 模型中可以多次堆叠以构建更深的模型。

class TransformerLayer(nn.Module):
    """
    Tensor          Type            Shape
    ===========================================================================
    x               float           (..., seq_len, dims)
    past (*)        float           (..., past_len, dims)
    mask            bool            (..., seq_len, past_len + seq_len)
    ---------------------------------------------------------------------------
    output 1        float           (..., seq_len, dims)
    output 2 (*)    float           (..., past_len + seq_len, dims)
    ===========================================================================
    """
    def __init__(self, heads: int, dims: int, rate: int, dropout: float = 0.1):
        super().__init__()
        self.attn = AttentionLayer(heads, dims, dropout)
        self.ff = PositionwiseFeedForward(dims, rate, dropout)
        self.ln_attn = LayerNorm(dims)
        self.ln_ff = LayerNorm(dims)

    def forward(self,
                x: torch.Tensor,
                past: Optional[Past] = None,
                mask: Optional[torch.Tensor] = None,
                ) -> Union[torch.Tensor, Tuple[torch.Tensor, Past]]:
        # Layer normalizations are performed before the layers respectively.
        a = self.ln_attn(x)
        a, past = self.attn(a, a, a, past, mask)

        x = x + a
        x = x + self.ff(self.ln_ff(x))

        return x if self.training else (x, past)

在这段代码中，past 是一个可选参数，用于在注意力层中存储先前时间步的注意力键（key）和值（value）。这个参数的主要目的是用于处理自回归模型，如 Transformer 中的解码器部分。在上述代码中，past 用于存储先前时间步的注意力键和值，并在下一个时间步中与当前时间步的键和值进行拼接，以供注意力计算和线性变换使用。让我解释一下 past 的作用和用途：

1. 处理自回归生成：自回归生成是一种生成序列的方式，其中模型按顺序生成输出，每个时间步都依赖于先前的生成结果。在这种情况下，为了生成当前时间步的输出，需要使用先前时间步的信息。past 就是用来存储这些先前时间步的信息的数据结构。
2. 减少计算复杂度：存储先前时间步的注意力键和值可以减少计算复杂度。在每个时间步，模型可以重复使用先前的键和值，而不必重新计算。这对于长序列和大模型尤其重要，因为重新计算可能非常昂贵。
3. 支持增量解码：使用 past 可以实现增量解码，即逐步生成序列而不是一次性生成整个序列。这在一些应用中很有用，比如文本生成和语音合成。

class Transformer(nn.Module):
    """
    Tensor          Type            Shape
    ===========================================================================
    x               long            (..., seq_len)
    past (**)       float           (..., past_len, dims)
    ---------------------------------------------------------------------------
    output 1        float           (..., seq_len, dims)
    output 2 (**)   float           (..., past_len + seq_len, dims)
    ===========================================================================
    """

    def __init__(self,
                 layers: int,
                 pad_idx: int,
                 words: int,
                 seq_len: int,
                 heads: int,
                 dims: int,
                 rate: int = 4,
                 dropout: float = 0.1,
                 bidirectional: bool = True):
        super().__init__()
        self.bidirectional = bidirectional
        self.pad_masking = PadMasking(pad_idx)
        self.future_masking = FutureMasking()

        self.positional_embedding = PositionalEmbedding(seq_len, dims)
        self.token_embedding = TokenEmbedding(words, dims)
        self.dropout_embedding = nn.Dropout(dropout)

        self.transformers = nn.ModuleList([
            TransformerLayer(heads, dims, rate, dropout)
            for _ in range(layers)])
        self.ln_head = LayerNorm(dims)

    def forward(self,
                x: torch.Tensor,
                past: Optional[List[Past]] = None,
                use_grad_ckpt: bool = False
                ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Past]]]:
        # 获取过去信息长度，用于对齐模型形状，因为k，v是逐步增大的，所以mask的形状会发生改变，offset代表会变的那个维度改变后的长度
        offset = past[0][0].size(-2) if past is not None else 0

        # Create masking tensor.
        mask = self.pad_masking(x, offset)
        if not self.bidirectional:
            mask = mask + self.future_masking(x, offset)

        # Use token embedding and positional embedding layers.
        x = self.token_embedding(x) + self.positional_embedding(x, offset)
        x = self.dropout_embedding(x)

        # Apply transformer layers sequentially.
        present = []
        for i, transformer in enumerate(self.transformers):
            if self.training and use_grad_ckpt:
                transformer = partial(torch.utils.checkpoint.checkpoint,
                                      transformer)

            x = transformer(x, past[i] if past is not None else None, mask)

            if not self.training:
                present.append(x[1])
                x = x[0]

        x = self.ln_head(x)
        x = self.token_embedding(x, transposed=True)

        return x if self.training else (x, present)

offset的作用：

offset用于获取过去信息长度，用于对齐模型形状，因为$k，v$是逐步增大的，所以mask的形状会发生改变，offset代表会变的那个维度改变后的长度，拿$k$举例，$v$同理，$k=(batch\_size,seq\_len,dim)$，在第一步中，$k_1=(batch\_size,seq\_le{n}_1,dim)$，第二步中，由于要进行拼接\增量（详见上一节注意力机制部分），$k_2=(batch\_size,seq\_le{n}_1+seq\_le{n}_2,dim)$，offset获取的就是这个$seq\_le{n}_1+seq\_le{n}_2$

PadMasking中

让我们通过一个简单的例子来说明这个张量是如何表示对输入序列进行偏移的，其中填充标记被移动到序列的开始位置。
假设我们有以下输入序列 x：python x = tensor([1, 2, 0, 3, 4, 0, 5])
其中，0 表示填充标记。现在，假设 offset 为 2，即我们希望将填充标记移到序列的开始位置。
原始输入序列 x
x = tensor([1, 2, 0, 3, 4, 0, 5])
对应的 is_pad 张量
is_pad = tensor([False, False, True, False, False, True, False])
对应的 shifted 张量（偏移后）

shifted = tensor([[False, False],
                 [False, False],
                 [False, False],
                 [False, False],
                 [False, False],
                 [False, False],
                 [False, False]])

对应的 mask 张量

mask = tensor([[False, False, False],
            [False, False, False],
            [True, False, False],
            [False, False, False],
            [False, False, False],
            [True, False, False],
            [False, False, False]])

在这个例子中，is_pad 表示了原始输入序列中的填充标记位置。shifted 张量是一个大小为 (7, 2) 的零张量，表示对输入序列进行了偏移，并且填充标记被移动到序列的开始位置。最后，mask 张量通过将 shifted 和 is_pad 拼接在一起而生成，用于在注意力计算中屏蔽填充标记和偏移后的部分。mask第一列即为当前x，之后的部分为之前的信息，不需要屏蔽(之前已经做过屏蔽)，所以全为False

FutureMasking中

通过在序列的未来位置添加偏移，创建了一个上三角的掩码，以避免模型在训练时看到未来的信息。
让我用一个简单的数列来说明如何使用上三角矩阵来屏蔽未来信息。
考虑以下数列：
$\text{数列}=[2,4,6,8,10]$
我们的任务是预测下一个数字。在模型训练的过程中，我们希望模型在生成每个数字时只能依赖于当前数字及其之前的数字，而不能依赖于之后的数字。
现在，我们将使用上三角矩阵作为掩码，以便在模型的自注意力机制中屏蔽未来信息。矩阵的形式如下：
$$\begin{array}{ccccc}1& 1& 1& 1& 1\\ 0& 1& 1& 1& 1\\ 0& 0& 1& 1& 1\\ 0& 0& 0& 1& 1\\ 0& 0& 0& 0& 1\end{array}$$
其中，矩阵的每个元素 $M_{ij}$表示模型在生成位置$i$ 的时候是否能够依赖位置$j$。(1代表被遮盖，表示不能依赖此项，0代表可依赖)
按行看，生成第一个元素只能依赖1，生成第二个元素能依赖1和2
这样，模型在训练时能够更好地捕捉数列中的因果关系，而不受到未来数字的干扰。这就是使用上三角矩阵来屏蔽未来信息的基本思想。