理解大语言模型（二）——从零开始实现GPT-2

一、概述

大语言模型这个商业术语正如其名，强调了这类模型的一个共同特点，那就是“大”。这主要体现在三个方面：首先，这类模型拥有大规模的模型参数，其数量级通常在数十亿到数千亿之间；其次，为了训练这些模型，需要大规模的数据集，语料库的总长度常常达到万亿级别；最后，由于前两个因素的影响，训练这些模型的成本也相当巨大。2023年，从零开始训练一个最先进的大语言模型需要数千台专业服务器，花费高达数百万美元。

从技术角度看，大语言模型并没有一个明确的定义。通常，它指的是包含注意力机制且用于自然语言处理的神经网络模型。尽管不同的大语言模型在结构上存在较大差异，但从发展历史来看，它们都有一个共同的祖先：Transformer。图1左侧展示了模型的详细结构，这是从Transformer模型的原始论文中摘取的，因此在相关文献中被广泛引用。这个图示中包含大量细节，可能会让读者迷失方向。因此，本书更倾向于使用图1右侧的简化示意图，以便更清晰地理解模型的整体架构。

图1

Transformer模型具备完整的编码器和解码器结构，因此通常应用于序列到序列模式²。从注意力的角度来看，它包含3种不同类型的注意力机制，分别是双向注意力，用于编码器；单向注意力，用于解码器；以及交叉注意力，用于编码器和解码器的协同工作。

复杂的结构提高了模型在翻译等任务中的性能，也使它的应用范围受到限制。为了更广泛地应用这一架构，出现了两种不同的改进和简化方式：一种是仅使用图1中的编码器部分（只包含双向注意力），通常用于自编码模式，最著名的代表是BERT；另一种是只包含图1中的解码器部分（只包含单向注意力），通常用于自回归模式，其中最著名的是GPT³。

就结构而言，以GPT为代表的单向注意力模型是最简单的，在工程处理和训练数据准备方面也最为便捷。也许正因如此，这类模型取得了最引人瞩目的成就。因此，本章的讨论重点是这类模型的经典代表：GPT-2。从实用角度来看，尽管存在更卓越的单向注意力模型，但它们通常规模巨大，难以在普通的家用计算机上运行，更不用说训练了。相比之下，GPT-2的规模适中，适合在家用计算机上运行，我们可以下载、使用或修改该模型，以便更好地理解其原理。（但要注意，最好在配备GPU的服务器上进行模型训练，在家用计算机上训练模型可能需要非常长的时间）。

二、模型结构

总体来说，GPT-2的结构可以分为3个主要部分，自下而上分别是嵌入层、多次重复的解码块，以及语言建模头，如图11-7右侧所示。其中，解码块（Decoder Block）是至关重要的组成部分，包含4个核心元素：多头单向自注意力（Masked Multi-Head Attention）、残差连接、层归一化和多层感知器⁴。多层感知器是一个相对被人熟知的概念，残差连接和层归一化是提高模型训练效果的关键技术，具体的细节可以参考其他文章[TODO]，这里不再详述。或许会让读者感到困惑的是多头单向自注意力，它只是自注意力机制的改进版本，在初步理解时，可以将其等同于普通的注意力。基于上述内容，在深入细节之前，再从宏观上讨论一下这个模型的独特之处。

GPT-2的图示与神经网络图示有一些显著不同，特别是解码块。在神经网络发展的早期，研究人员通常从仿生学的角度来构建模型，因此引入了神经元、全连接和隐藏层等概念。随着研究的深入，学术界发现神经网络的核心本质是线性计算和非线性变换的多层叠加。因此，研究人员突破神经元连接方式的限制，设计了卷积神经网络和循环神经网络。不仅如此，他们还改进了神经元的内部结构，设计了长短期记忆网络。

GPT-2的解码块延续了这一创新思路。尽管它的内部结构难以用传统的图示表示，但这并不重要，因为它的设计承载了神经网络的核心理念。如这篇文章所述，注意力机制只涉及线性计算，层归一化和残差连接也同样如此。因此，整个解码块实际上是线性计算和非线性变换的叠加。其中，非线性变换来自多层感知器，这也是解码块中包含多层感知器的原因之一。

从类比的角度来看，注意力机制是对循环神经网络的改进。尽管在图示上无法准确表示，但解码块实际上是循环神经网络和多层感知器的组合。在理解复杂神经网络时，读者不应固守传统的图示，而应从计算意义的角度理解各个运算步骤的作用，这有助于更深刻地理解模型。

图2

在其他文献中，GPT-2的结构经常被表示为图2左侧的形式，它与Transformer的图示相似，更易于理解，然而它并不是模型的精确表示。对比图2左右两侧的图示可以发现，在左侧的图示中，层归一化分别放置在多层感知器（对应Feed Forward）和多头单向自注意力（对应Masked Multi-Head Attention）之后，与模型的实际设计不相符。然而这无伤大雅，这种差异并不会对整体效果产生重大影响。这个例子再次强调了在理解神经网络时，应重点关注关键组件，而对非关键部分的理解应具有一定的灵活性，完全没必要死记特定模型的结构。

三、多头单向注意力

多头单向注意力是多个单向注意力的组合。为了更清晰地表述，可以将这篇文章中讨论的注意力机制称为单头注意力。程序清单1是单头单向注意力组件的代码实现，其中有两个关键点需要注意。

注意力的计算需要3个关键参数，分别是query、key和value。在模型中，采用3个独立的线性回归模型生成这些向量，具体实现请参考第8—10行和第17—20行。由于模型中使用了层归一化，这3个线性模型都不需要截距项。
单向注意力的实现需要依赖一个上三角矩阵，也就是掩码（mask）。具体的实现细节见第12、13和21行。由于GPT-2对模型能够处理的文本长度有限制（关于这一点的详细原因请参考后文），为了提高计算效率，在创建模型时使用参数sequence_len提前生成能够覆盖最长文本的矩阵tril。值得注意的是，上三角矩阵的作用是辅助注意力计算，它并不需要参与模型的训练。为了实现这一点，使用register_buffer来记录生成的矩阵。

程序清单1 单头单向注意力

 1 |  def attention(query, key, value, dropout, mask=None):
 2 |      ......
 3 |  
 4 |  class MaskedAttention(nn.Module):
 5 |  
 6 |      def __init__(self, emb_size, head_size):
 7 |          super().__init__()
 8 |          self.key = nn.Linear(emb_size, head_size, bias=False)
 9 |          self.query = nn.Linear(emb_size, head_size, bias=False)
10 |          self.value = nn.Linear(emb_size, head_size, bias=False)
11 |          # 这个上三角矩阵不参与模型训练
12 |          self.register_buffer(
13 |              'tril', torch.tril(torch.ones(sequence_len, sequence_len)))
14 |          self.dropout = nn.Dropout(0.4)
15 |  
16 |      def forward(self, x):
17 |          B, T, C = x.shape  # C = emb_size
18 |          q = self.query(x)  # (B, T, H)
19 |          k = self.key(x)    # (B, T, H)
20 |          v = self.value(x)  # (B, T, H)
21 |          mask = self.tril[:T, :T]
22 |          out, _ = attention(q, k, v, self.dropout, mask)
23 |          return out         # (B, T, H)

从组件的功能角度来看，单头单向注意力的主要任务是进行特征提取。为了尽可能全面地提取特征信息，多头单向注意力采用了反复提取的策略。简而言之，对于相同的输入，它会使用多个结构相同但具有不同模型参数的单头单向注意力组件来提取特征⁵，并对得到的多个特征进行张量拼接⁶和映射。

从张量形状的角度来看，由于在模型中使用了残差连接，因此注意力组件的输入形状和输出形状最好相同。然而，单头注意力组件并没有这种保证，通常情况下，单头注意力的输出形状会小于输入形状。为了确保张量形状的一致性，可以使用多头注意力，通过对多个张量进行拼接的方式来实现这一目标。

基于上述讨论，多头单向注意力的实现如图3所示，其实现相对简单且易于理解，但它并不是最高效的实现方式。细心的读者可能会注意到红色框内包含循环操作，生成self.heads时也使用了循环操作。这些循环操作不利于并行计算，会影响模型的运算效率。更高效的实现方式是将“多头”操作设计为张量运算的形式，具体细节可以借鉴长短期记忆网络（LSTM）中的程序清单2。

图3

四、解码块

与传统的多层感知器略有不同，解码块中的多层感知器[^ 7]包括两层线性计算和一层非线性变换⁷，如程序清单2所示，而传统的多层感知器通常采用一层线性计算和一层非线性变换的配对结构。实际上，解码块中的多层感知器可以分为两个部分：第一部分是经典的单层感知器，第二部分是一个线性映射层。组件的第二部分不仅可以完成一次线性学习，还保证了组件输入和输出的张量形状相同⁸。

程序清单2 解码块中的多层感知器

 1 |  class FeedForward(nn.Module):
 2 |      
 3 |      def __init__(self, emb_size):
 4 |          super().__init__()
 5 |          self.l1 = nn.Linear(emb_size, 4 * emb_size)
 6 |          self.l2 = nn.Linear(4 * emb_size, emb_size)
 7 |          self.dropout = nn.Dropout(0.4)
 8 |  
 9 |      def forward(self, x):
10 |          x = F.gelu(self.l1(x))
11 |          out = self.dropout(self.l2(x))
12 |          return out

按照模型的图示，将上述的组件组合在一起，就得到了解码块的实现，如图4所示。

图4

五、GPT-2的完整结构与重现

在设计GPT-2的模型结构时，还有最后一个关键细节需要考虑，那就是如何捕捉词元在文本中的位置信息。尽管注意力机制成功地捕捉了词元之间的相关关系，但它却顾此失彼，忽略了词元的位置。回顾一下注意力机制中的内容，可以发现：双向注意力只包含不受位置影响的张量计算。这意味着打乱词元在文本中的位置不会影响双向注意力的计算结果。类似地，对于单向注意力，更改左侧文本中的词元顺序也不会影响计算结果。然而，对于自然语言处理来说，词语在文本中的位置通常至关重要，因此需要想办法让模型能够捕捉词元的位置信息。

有一种非常简单的方法可以实现这一点。在使用循环神经网络进行自然语言处理时，在模型的开头使用了文本嵌入技术。文本嵌入层的输入是词元在字典中的位置，而输出是词元的语义特征，该特征将用于后续的模型计算。对于位置信息，我们完全可以“依葫芦画瓢”，在模型的开头引入一个位置嵌入层。这一层的输入是词元在文本中的位置，输出是与语义特征具有相同形状的位置特征。位置特征和语义特征将被结合在一起，参与后续的模型处理。从人类的角度来看，文本嵌入层学习了词元的语义特征，位置嵌入层学习了词元的位置信息；从模型的角度来看，它们几乎是相同的，都是基于位置信息（字典位置和文本位置）的学习。因此，这个设计虽然简单，却能够有效地捕获词元的位置信息。

将上述内容转化为代码，如程序清单3所示。其中，第6行定义了位置嵌入层。在生成位置嵌入层时，需要确定嵌入层的大小，即最大的文本长度，这也解释了为什么模型只能处理有限长度的文本。除了GPT-2，其他大语言模型也存在类似的限制，这是由注意力机制和位置嵌入导致的，也是这些模型的明显不足。因此，如何克服或放宽这一限制是当前的热门研究方向。

程序清单3 GPT-2

 1 |  class CharGPT(nn.Module):
 2 |  
 3 |      def __init__(self, vs):
 4 |          super().__init__()
 5 |          self.token_embedding = nn.Embedding(vs, emb_size)
 6 |          self.position_embedding = nn.Embedding(sequence_len, emb_size)
 7 |          blocks = [Block(emb_size, head_size) for _ in range(n_layer)]
 8 |          self.blocks = nn.Sequential(*blocks)
 9 |          self.ln = nn.LayerNorm(emb_size)
10 |          self.lm_head = nn.Linear(emb_size, vs)
11 |  
12 |      def forward(self, x):
13 |          B, T = x.shape
14 |          pos = torch.arange(0, T, dtype=torch.long, device=x.device)
15 |          tok_emb = self.token_embedding(x)       # (B, T,  C)
16 |          pos_emb = self.position_embedding(pos)  # (   T,  C)
17 |          x = tok_emb + pos_emb                   # (B, T,  C)
18 |          x = self.blocks(x)                      # (B, T,  C)
19 |          x = self.ln(x)                          # (B, T,  C)
20 |          logits = self.lm_head(x)                # (B, T, vs)
21 |          return logits

与其他模型类似，要在自然语言处理任务中应用该模型，需要完成两个额外的步骤：定义分词器和准备训练数据。

GPT-2采用的分词器是字节级字节对编码分词器。有关此分词器的详细算法和缺陷，请参考利用神经网络学习语言（一）——自然语言处理的基本要素。
GPT-2的训练数据是OpenWebText⁹。在准备训练数据时，采用的方法是在文本的末尾添加一个特殊字符来表示文本结束，然后将所有文本拼接成一个长字符串。在这个长字符串上，截取长度等于sequence_len的训练数据。这种方法有效解决了文本长度不一致的问题，提高了训练效率。

至此，我们终于完成了重现GPT-2所需的一切准备工作。模型的训练过程需要耗费一定的计算资源和时间，根据Andrej Karpathy的实验¹⁰，为了复现最小版本的GPT-2（拥有1.24亿个参数），我们需要一台配备8块A100 40GB显卡的计算机和大约4天的训练时间。

六、Python语言学习任务

尽管没有资源来复现GPT-2，但是可以利用类似的模型来解决较小的自然语言处理任务，比如前文中反复提到的Python语言学习。这将使我们有机会亲身体验该模型的优点。
在Python语言学习任务中，无须改动模型结构，只需调整分词器和训练数据。具体来说，模型将使用字母级别的分词器，训练数据的准备方法与GPT-2非常相似。更多细节可以参考完整代码和这里。

模型的具体结果如图5所示。该模型的规模相对较小，只包含大约240万个参数，训练时间较短，但取得了令人满意的效果。如果进行更长时间的训练或者增加模型规模，能够获得更出色的模型效果。

图5

模型的实现过程参考了OpenAI提供的GPT-2开源版本，Harvard NLP提供的Transformer开源实现，以及Andrej Karpathy的课程“Neural Networks: Zero to Hero”。 ↩︎
经过精心的设计和调整，Transformer模型已经成功应用于自回归和自编码模式。此外，仅包含编码器的模型也能在序列到序列模式和自回归模式下使用（解码器也类似）。 ↩︎
BERT的全称是Bidirectional Encoder Representation from Transformer。GPT的全称是Generative Pre-trained Transformer。 ↩︎
解码块中的多层感知器与传统的多层感知器略有不同，细节请见后文。 ↩︎
这里的设计受到了卷积神经网络中卷积层的启发，多头机制对应卷积层中的通道概念。 ↩︎
为了更好地理解这一方法，可以参考循环神经网络中的隐藏状态。 ↩︎
模型中的非线性变换是GeLU（Gaussian Error Linear Unit），这是对ReLU的一种改进。 ↩︎
在多头单向注意力组件中，最后一个计算步骤也是线性映射。不同的是，多头单向注意力的线性映射并没有改变张量的形状，而这里的线性映射对张量进行了压缩。这个设计使得解码块中的多层感知器呈现出两头细、中间粗的形状，既有助于特征提取（中间越宽，模型可以提取的特征就越多），又能兼顾模型后续的残差连接。 ↩︎
OpenWebText是由OpenAI创建的数据集，用于训练GPT-2模型。尽管它的名字中包含“Open”，但实际上它本身并不是一个开源的数据集。不过，我们可以使用工具datasets来获取由其他研究者创建的开源版本。根据论文“Language Contamination Helps Explain the Cross-lingual Capabilities of English Pretrained Models”的研究，该数据集以英文为主，包含少量中文，这也解释了为什么GPT-2能够理解中文。 ↩︎
具体结果请查阅Andrej Karpathy的GitHub页面。 ↩︎