大模型原理：实现TransformerBlock块类（连接注意力层和线性层等）

作者：马育民 • • 阅读：10029

介绍

实现 Transformer块，这是GPT和其他大语言模型架构的基本构建块。在参数量为 1.24亿 的GPT-2架构中，这个块被 重复了许多次
它结合了我们之前提及的多个概念：

  • 掩码多头注意力
  • 层归一化
  • dropout
  • 前馈层
  • GELU激活函数

Transformer块核心思想

自注意力机制 在 多头注意力块 中用于 识别 和 分析 输入序列中 元素之间的关系。

相比之下，前馈神经网络则在每个位置上对数据进行单独的修改。

这种组合不仅提供了对输入更细致的理解和处理，而且提升了模型处理复杂数据模式的整体能力

Transformer块示意图。输入的词元被嵌入到 768 维的向量中。每一行对应一个词元的向量表示。Transformer块的输出是与输入具有相同维度的向量，这些向量可以传递到大语言模型的后续层中

实现

准备

准备

从 大模型原理：全部代码 复制以下类或函数：

  • MultiHeadAttention
  • GELU
  • FeedForward
  • LayerNorm

创建 TransformerBlock块类

这段代码定义 TransformerBlock 类，该类实现功能如下：

  1. 层归一化(LayerNorm)应用于这两个组件之前。这种方法也被称为 前层归一化(Pre-LayerNorm)。较早的架构（如最初的Transformer模型）在自注意力和前馈神经网络之后才应用层归一化，这种方法被称为 后层归一化(Post-LayerNorm)，这通常会导致较差的训练效果。

  2. 多头注意力机制(MultiHeadAttention)。配置字典（cfg，比如GPT_CONFIG_124M）进行配置。

  3. 前馈神经网络(FeedForward)。配置字典（cfg，比如GPT_CONFIG_124M）进行配置。

  4. dropout 应用于这两个组件之后，以便对模型进行正则化并防止过拟合。

  5. TransformerBlock 类还实现了前向传播，其中每个组件后面都跟着一个 快捷连接，将块的输入加到其输出上。这个关键特性有助于在训练过程中使梯度在网络中流动，并改善深度模型的学习效果

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # 在注意力块中添加快捷连接
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        # 将原始输入添加回来
        x = x + shortcut

        # 在前馈层中添加快捷链接
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        # 将原始输入添加回来
        x = x + shortcut
        return x

执行

使用之前定义的 GPT_CONFIG_124M 配置字典来实例化一个 Transformer 块，并输入一些测试数据：

GPT_CONFIG_124M = {
    "vocab_size": 50257,     # 词汇表大小
    "context_length": 1024,  # 上下文长度
    "emb_dim": 768,          # 嵌入维度
    "n_heads": 12,           # 注意力头的数量
    "n_layers": 12,          # 层数
    "drop_rate": 0.1,        # dropout率
    "qkv_bias": False        # 查询-键-值偏置
}

torch.manual_seed(123)
x = torch.rand(2, 4, 768)
block = TransformerBlock(GPT_CONFIG_124M)
output = block(x)
print("输入的 shape:", x.shape)
print("输出的 shape:", output.shape)

执行结果：

输入的 shape: torch.Size([2, 4, 768])
输出的 shape: torch.Size([2, 4, 768])

由打印结果可知：Transformer块在 输出中维持了输入的维度，这表明Transformer架构在处理数据序列时不会改变它们在网络中的形状。

提示：保持形状不变并非偶然，而是设计的，每个输出向量直接对应一个输入向量，保持一一对应的关系。

总结

实现了Transformer块（结合了层归一化、前馈神经网络、GELU激活函数和快捷连接），具备了构建GPT架构所需的所有构建块。如图

提示：其中黑色勾选标记表示我们已构建完成的块

原文出处：/show_1GW2elrInYWr.html

吉ICP备15003043号