最近在看《从零构建大语言模型》这本书，跟着思路自己也动手码了一个基础款的 LLM，代码不多，拢共 400 来行。

所以，这篇读书笔记不打算空谈理论，就想换个方式：试着把这 400 行代码的来龙去脉讲清楚。通过解说代码，把从零构建一个大模型需要经历哪些环节、每个环节的目标和大概原理给串起来。

笔者并非该方向的专业人士，很多东西不会讲得太深（我自己也不懂），所以很多时候都是点到即止。这篇文章更适合那些和我一样的开发者，希望能从一个接地气的工程角度，对大模型的构建原理有个宏观的了解。

400 行代码从零构建一个大语言模型

我们将采用一种贴近软件开发者思维的代码执行流视角，从程序的入口 main 函数开始，顺着调用栈逐层深入，探究数据处理、模型构建、训练循环的每一个细节。当一个模块被调用时，我们便深入其中，直至最底层的实现。

这趟旅程将遵循以下路线图：

  • 从 main 函数出发：探寻程序的入口与总调度中心。
  • 深入数据流水线：看原始文本如何被加工成模型能够消化的食粮。
  • 解构核心引擎：层层剖析 GPTModel、TransformerBlock 直至最深处的 MultiHeadAttention 机制。
  • 启动训练循环：见证模型如何通过损失计算和权重更新，从随机变得智能。
  • 见证文本生成：观察训练好的模型如何像我们一样，逐字逐句地"思考"和"创作"。

让我们即刻出发，揭开大语言模型背后的代码之谜。

1. 数据准备与采样

prepare_train_and_val_data 的具体实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def prepare_train_and_val_data(file_path, tokenizer) -> tuple[DataLoader, DataLoader]:
    """准备训练和验证数据加载器"""
    
    # 步骤 1: 读取原始文本
    with open(file_path, "r", encoding="utf-8") as file:
        text_data = file.read()

    # 步骤 2: 按比例分割成训练和验证两部分：90%训练，10%验证
    train_ratio = 0.9
    split_idx = int(train_ratio * len(text_data))
    train_data = text_data[:split_idx]
    val_data = text_data[split_idx:]

    # 步骤 3: 分别为两部分文本创建 DataLoader
    train_loader = create_dataloader(
        tokenizer,
        train_data,
        batch_size=2,
        max_length=GPT_CONFIG_124M["context_length"],
        stride=GPT_CONFIG_124M["context_length"],
        drop_last=True,
        shuffle=True,
        num_workers=0,
    )

    val_loader = create_dataloader(
        tokenizer,
        val_data,
        batch_size=2,
        max_length=GPT_CONFIG_124M["context_length"],
        stride=GPT_CONFIG_124M["context_length"],
        drop_last=False,
        shuffle=False,
        num_workers=0,
    )
    return train_loader, val_loader

1.1 划分数据集 prepare_train_and_val_data

prepare_train_and_val_data 的代码并不复杂，就是将数据集划分为训练集和验证集，那么第一个问题就来了：模型为何需要训练集和验证集？

和人类一样，模型通过看例子来学习。我们给它一本"教科书"和配套的"练习册"（训练集），让它反复练习，寻找规律。但我们如何知道它是真的学会了，还是仅仅背下了答案（过拟合）呢？

答案是，我们需要一场它从未见过的模拟考试（验证集）。

  • 训练集 (Training Set)：这是模型学习的唯一资料。模型会尽全力去拟合训练集中的数据，目标是在这个数据集上获得尽可能低的出错率（损失）。这对应了代码中 90% 的文本数据 。
  • 验证集 (Validation Set)：这是一份被隔离的数据，模型在训练过程中绝对不能用它来更新自己的权重。我们只在训练的特定阶段用它来“考”一下模型，看看模型在“新题型”上的表现如何。这对应了代码中 10% 的文本数据 。

如果模型在训练集上表现优异（比如损失很低），但在验证集上表现糟糕，这就亮起了过拟合的红灯。这说明模型只是死记硬背了训练题，而没有学到普适的规律。prepare_train_and_val_data 函数的核心使命，就是为模型准备好这两份至关重要的数据集。

它调用了 create_dataloader 函数。我们必须注意 train_loader 和 val_loader 在配置上的一个关键区别：

  • train_loader 的 shuffle 设置为 True 。

    这是为了保证训练的有效性。在每一轮 (epoch) 训练开始时，DataLoader 都会将训练样本的顺序完全打乱。这就像我们学习时会打乱单词卡片的顺序一样，可以防止模型学到样本的出场顺序这种无关信息，从而迫使它学习更具泛化性的语言规律。

  • val_loader 的 shuffle 设置为 False 。

    这是为了保证评估的客观性和一致性。验证集是我们的模拟考试，我们希望每次考试的卷子（题目顺序）都是一样的，这样才能客观地比较模型在不同训练阶段的得分，判断它是否真的在进步。

1.2 构建数据集 create_dataloader

prepare_train_and_val_data 只是一个调度者，真正的数据加工发生在它调用的 create_dataloader 和 GPTDataset 中。

1
2
3
4
5
6
7
8
9
10
11
def create_dataloader(tokenizer, txt, batch_size=4, max_length=256, stride=128, shuffle=True, drop_last=True, num_workers=0):
    """创建数据加载器"""
    dataset = GPTDataset(txt, tokenizer, max_length, stride)
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers,
    )
    return dataloader

create_dataloader 是一个简单的封装，它的核心是做了两件事：

  1. dataset = GPTDataset(txt, tokenizer, max_length, stride)：实例化一个 GPTDataset 对象。
  2. dataloader = DataLoader(dataset, ...)：将这个 dataset 对象包装成一个 PyTorch 的 DataLoader。

1.2.1 Dataset & DataLoader

在深入 GPTDataset 结构之前，这里我们先对 PyTorch 中的 2 个关键数据类型 Dataset 和 DataLoader 进行简要介绍，对于 Pytorch 更详细的介绍可参考笔者这篇 告别死记硬背：一份真正理解 PyTorch 核心设计的指南。

简而言之，Dataset 和 DataLoader 是为了解决"数据集是什么"和"如何使用数据集"这 2 个核心问题，更具体的来说，在数据准备阶段，我们可能会面临以下几个问题：

  1. 原始数据格式各异，如何统一读取？
  2. 数据集可能非常大，无法一次性载入内存，怎么办？
  3. 训练时需要对数据进行批量 (batching)、打乱 (shuffling) 和预处理 (preprocessing)，如何高效实现？
  4. 如何利用多核 CPU 来加速数据加载，避免 GPU 等待？

PyTorch 的解决方案就是 Dataset 和 DataLoader ：

  • Dataset：它定义了"数据集"是什么。这是一个抽象类，你只需要继承它并实现两个方法：__len__(返回数据集大小) 和 __getitem__ (根据索引 idx 返回一条数据)。它解决了如何获取单条数据的问题，将数据访问的逻辑封装起来。
  • DataLoader：它定义了"如何使用数据集"。它接收一个 Dataset对象，并在此基础上，优雅地解决了所有工程问题：
    • batch_size：自动将单条数据打包成一个 batch。
    • shuffle=True：在每个 epoch 开始时自动打乱数据顺序。
    • num_workers：启动多个子进程并行加载数据，极大地提高了数据供给效率。
    • collate_fn：自定义如何将多条样本合并成一个 batch，对于处理非标准数据（如不同长度的句子）非常有用。

它们之间的关系如图所示：

Dataset 与 DataLoader

1.2.2 GPTDataset

现在我们可以来看 GPTDataset 的具体逻辑了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class GPTDataset(Dataset):
    """GPT训练数据集"""
    def __init__(self, txt, tokenizer, max_length, stride) -> None:
        self.input_ids = []
        self.target_ids = []

        token_ids = tokenizer.encode(txt)

        # 使用滑动窗口创建训练样本
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i+max_length]
            target_chunk = token_ids[i+1:i+max_length+1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))
    
    def __len__(self):
        return len(self.input_ids)
    
    def __getitem__(self, idx):
        return self.input_ids[idx], self.target_ids[idx]

GPTDataset 继承了 PyTorch 的 Dataset 类型，我们重点来看它的构造函数 __init__()，它分为 2 个步骤：

  1. 将文本数据转为词元 ID 列表；
  2. 使用滑动窗口逐个构建输入-目标对，构建整个数据集，分别置于 input_ids 和 target_ids 这 2 个字段中。

1.2.2.1 文本词元化

现在到了本篇的第一个真正意义上的理论环节，我们需要先搞清楚词元化（即下面这一行代码）到底是在做什么？为什么要这样？有哪些具体的方式方法？

1
token_ids = tokenizer.encode(txt)

包括大语言模型在内的深度神经网络模型是无法直接处理原始文本的。由于文本数据是离散的，因此我们无法直接用它来执行神经网络训练所需的数学运算。我们需要一种将单词表示为连续值的向量格式的方法（通常是张量 Tensor）。

将数据转换为向量格式的过程通常称为嵌入（embedding），如下图所示：

要理解 Tensor，我们需要先建立一个最重要的心智模型：Tensor 的每一个维度 (dimension) 都有其特定的语义含义。

一个典型的 4D Tensor (B, C, H, W) 在计算机视觉中，其形状 (16, 3, 224, 224) 并不是一串孤立的数字，它的意思是：

  • B (Batch size) = 16: 这个 Tensor 里有 16 张独立的图像。
  • C (Channels) = 3: 每张图像有 3 个通道（R, G, B）。
  • H (Height) = 224: 每张图像的高度是 224 像素。
  • W (Width) = 224: 每张图像的宽度是 224 像素。

在多个维度综合起来语义含义越接近的词，它们的词嵌入向量在空间表示中就越相近，也就越"相似"，如下图所示：

当把文本转为词嵌入向量之后，我们的训练模型就可以识别这些数据并利用它们进行学习了。

一个完整的文本处理步骤，大概如下图所示：

  1. 输入文件 This is an example.：这是所有处理的起点，是我们希望模型去理解和回应的原始、非结构化的人类语言。

  2. 词元化：原始文本被切分成独立的单元：This, is, an, example, .。

    计算机模型无法一次性理解一整个句子。它需要将句子分解成更小的、标准化的单元，这些单元被称为词元 (Token)。如图中的文字描述，词元既可以是单词，也可以是标点符号之类的特殊字符。

  3. 转换为词元 ID：每个词元被映射到一个唯一的整数：This -> 40134, is -> 2052, an -> 133, example -> 389, . -> 12。

    计算机不认识字符串 This，但它能高效地处理数字 40134。这一步是将语言世界映射到数字世界的关键。每一个 ID 都对应着模型词汇表（一个巨大的“字典”）中的一个条目。

  4. 生成词元嵌入：一串数字 ID 变成了多个向量（关于词嵌入的具体细节，我们会在后续进行展开）。

    即我们前面提到的，单个数字 ID（如 40134）本身是孤立的，不包含任何语义信息，所以我们需要将其转为词嵌入向量，在训练过程中，模型会不断调整这些向量，使得意思相近的词元，其向量在空间中的位置也相互靠近。

  5. 模型处理与输出：这些嵌入向量组成的序列，最终被送入类 GPT 的纯解码器 Transformer 。这是模型的核心大脑。Transformer 模型会分析这些向量之间的关系，理解整个句子的上下文，然后进行计算。经过后续处理步骤（如选择概率最高的词元）后，模型会生成一个输出文本。

[!IMPORTANT]

小结一下，从人类语言 (字符串) -> 语言单元 (词元) -> 机器语言 (数字 ID) -> 数学对象 (嵌入向量) -> 模型输入，每一步都是为了让原始的、非结构化的文本，变得结构化、数值化，并富含语义信息，最终成为能够被神经网络高效处理的原料。

一个简单的分词器实现如下所示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class SimpleTokenizer:
    def __init__(self, vocab) -> None:
        self.str_to_int = vocab
        self.int_to_str = {i:s for s, i in vocab.items()}
    
    def encode(self, text):
        preprocessed = re.split(r'([,.:;?_!"()\']|--|\s)', text)
        preprocessed = [item.strip() for item in preprocessed if item.strip()]
        preprocessed = [item if item in self.str_to_int else "<|unk|>" for item in preprocessed]
        ids = [self.str_to_int[s] for s in preprocessed]
        return ids
    
    def decode(self, ids):
        text = " ".join([self.int_to_str[i] for i in ids])
        # Remove the spaces before specific punctuation marks.
        text = re.sub(r'\s+([,.?!"()\'])', r'\1', text)
        return text
  1. __init__ 初始化词典，里面每一个词元都唯一对应一个 ID；
  2. encode 原始文本转为一系列词元 ID，对于不识别的词元，会使用 <|unk|> 特殊标识进行占位，一般来说，还会使用诸如 <|endoftext|> 等特殊标识符来表示文本结束等特殊语义。
  3. decode 将词元 ID 列表转回原始文本。

回到本篇的代码实现：

1
token_ids = tokenizer.encode(txt)

这里我们使用的是现有的 Python 开源库 tiktoken，它基于 Rust 的源代码非常高效地实现了 BPE（Byte Pair Encoder） 算法。

1
2
3
4
5
6
7
8
9
import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
text1 = "Hello, do you like tea?"
text2 = "In the sunlit terraces of the palace."
text = " <|endoftext|> ".join((text1, text2))
ids = tokenizer.encode(text, allowed_special={"<|endoftext|>"})
print(ids)
print(tokenizer.decode(ids))

输出：

1
2
[15496, 11, 466, 345, 588, 8887, 30, 220, 50256, 554, 262, 4252, 18250, 8812, 2114, 286, 262, 20562, 13]
Hello, do you like tea? <|endoftext|> In the sunlit terraces of the palace.

通过输出，我们可以看到 <|endoftext|> 词元被分配了一个较大的词元 ID，即 50256。事实上，用于训练 GPT-2、GPT-3 和 ChatGPT 中使用的原始模型的 BPE 分词器的词汇总量为 50257，这意味着 <|endoftext|> 被分配了最大的词元 ID。

另外，BPE 分词器可以正确地编码和解码未知单词，比如 someunknownPlace。BPE 分词器是如何做到在不使用<|unk|>词元的前提下处理任何未知词汇的呢？

BPE 算法的原理是将不在预定义词汇表中的单词分解为更小的子词单元甚至单个字符， 从而能够处理词汇表之外的单词。因此，得益于 BPE 算法，如果分词器在分词过程中遇到不熟悉的单词，它可以将其表示为子词词元或字符序列，如下图所示。

将未知单词分解为单个字符的能力确保了分词器以及用其训练的大语言模型能够处理任何文本，即使文本中包含训练数据中不存在的单词。

1.2.2.2 滑动窗口进行数据采样

分析完了词元化的背后底层逻辑后，我们来看这一部分的代码：

1
2
3
4
5
6
7
8
9
10
class GPTDataset(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride) -> None:
				# ...

        # 使用滑动窗口创建训练样本
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i+max_length]
            target_chunk = token_ids[i+1:i+max_length+1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

要理解这段代码，我们需要回归到大语言模型（文本模型）是唯一任务：根据你给出的上文，猜出下一个词应该是什么。

例如，对于句子 Time is an illusion，我们可以为模型制作如下一系列的练习题：

  • 问题：Time -> 答案：is
  • 问题：Time is -> 答案：an
  • 问题：Time is an -> 答案：illusion

模型需要通过海量的这类"问答对"进行练习，才能逐渐掌握语言的规律。如果手动去制作上亿个这样的问答对，显然是不现实的。代码中的"滑动窗口"机制，就是为了解决这个问题。我们用一个具体的例子来解释这个 for 循环：

  • 假设 max_length = 5
  • 假设一段文本分词后的 token_ids 是 [10, 20, 30, 40, 50, 60]

当 for 循环第一次执行时 (i=0)：

  • input_chunk = token_ids[0:5] 会切出 [10, 20, 30, 40, 50]
    • 这就是提供给模型的上下文，也就是问题。
  • target_chunk = token_ids[1:6] 会切出 [20, 30, 40, 50, 60]
    • 这就是模型需要预测的正确答案。

滑动窗口示意图

所以我们通过这样一个 for 循环，就可以根据传入的文本 txt 快速生成大量的输入-目标对供给模型进行训练和检验。

[!IMPORTANT]

回到 prepare_train_and_val_data 函数，现在我们可以用一句话概括它的全部工作：它是一个数据准备总管，负责将一本原始小说，严格划分为用于学习的训练集和用于考试的验证集，并最终将它们都加工成模型可以直接使用的、一批一批的、包含(输入-目标)对的标准化数据传送带。

2. 初始化模型与优化器

  1. GPTModel() 初始化一个 GPT 模型实例；
  2. model.to(device) 是 PyTorch 中用于将模型移动到指定设备（CPU 或 GPU）的方法，深度学习模型在 GPU 上训练速度比 CPU 快很多，通过这种方式可以确保模型和输入数据在同一个设备上。
  3. torch.optim.AdamW() 创建一个优化器（optimizer），用于训练神经网络模型。AdamW 是一种优化算法，在训练过程中，优化器会接收损失函数计算出的梯度、使用 AdamW 算法更新模型参数和帮助模型逐步收敛到最优解。在本篇中，我们不对这个进行过多的解释，因为这并不在我们的核心学习目标上。

2.1 模型配置

在深入代码细节之前，我们先看 GPT_CONFIG_124M 这个配置字典。它就像是建造 GPT 模型大厦的设计蓝图，定义了模型的规模和所有关键参数：

1
2
3
4
5
6
7
8
9
GPT_CONFIG_124M = {
    "vocab_size": 50257,    # 词汇表大小
    "context_length": 256,  # 上下文长度
    "emb_dim": 768,         # 嵌入维度
    "n_heads": 12,          # 注意力头数量
    "n_layers": 12,         # Transformer层数
    "drop_rate": 0.1,       # dropout率
    "qkv_bias": False       # QKV线性层是否使用偏置
}

  • vocab_size: 词汇表里有多少个不同的词元 (Token)。50257 是 GPT-2 使用的标准词汇表大小，即我们前面讨论的 BPE 分词器的词汇表大小。

  • context_length: 模型一次能处理的最长文本长度（以词元计）。这里是 256，意味着模型一次最多能看 256 个词元。

  • emb_dim: 嵌入维度。这是模型内部表示每个词元的向量长度。768 维意味着每个词都会被转换成一个包含 768 个数字的向量，这是模型理解语言的基础。

  • n_heads 和 n_layers: 这两个参数共同决定了模型的深度和宽度。n_layers=12 表示我们的模型会堆叠 12 个 TransformerBlock，而 n_heads=12 表示在每个 Block 内部的注意力机制都有 12 个"头"，让模型能从多个角度分析文本。

  • drop_rate: Dropout 比率。这是防止模型过拟合的重要技术。0.1 表示在训练时，每个神经元有 10% 的概率被临时"关闭"，迫使模型学习更鲁棒的特征表示。

  • qkv_bias: 查询-键-值偏置。这个参数控制是否在注意力计算中添加偏置项。False 表示不使用偏置，这是 GPT-2 的设计选择，可能有助于模型的稳定性。

有些概念你可能还不认识，没关系，我们继续往下看，待会就懂了！· ·

2.2 模型结构总览

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class GPTModel(nn.Module):
    """完整的GPT模型实现"""
    def __init__(self, cfg) -> None:
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])

        # 堆叠多个Transformer块
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])],
        )

        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        # 词嵌入 + 位置嵌入
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

整个 GPT 模型的架构如下图所示：

GPTModel 类是整个语言模型的顶层封装，其设计目标是构建一个端到端的、具备自回归（auto-regressive）生成能力的序列处理架构。从根本上说，任何一个此类模型都必须解决三个核心问题：

  1. 输入表示 (Input Representation)：如何将离散的、一维的词元 ID 序列，转化为模型能够处理的、富含信息的连续多维向量？
  2. 上下文编码 (Contextual Encoding)：如何对输入序列中的每个元素进行深度处理，使其向量表示能够充分融合整个序列（尤其是其上文）的上下文信息？
  3. 输出投影 (Output Projection)：如何将模型内部经过深度处理的上下文向量，重新映射回词汇表空间，以生成对下一个词元的概率预测？

GPTModel 的结构正是围绕这三个核心问题，划分成了三个逻辑清晰的功能区块。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class GPTModel(nn.Module):
    def __init__(self, cfg):
        # 区块一：输入表示层
        self.tok_emb = nn.Embedding(...)
        self.pos_emb = nn.Embedding(...)
        self.drop_emb = nn.Dropout(...)

        # 区块二：上下文编码器堆栈
        self.trf_blocks = nn.Sequential(...)

        # 区块三：输出投影层
        self.final_norm = LayerNorm(...)
        self.out_head = nn.Linear(...)

    def forward(self, in_idx):
        # 执行区块一的功能
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(...)
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        
        # 执行区块二的功能
        x = self.trf_blocks(x)
        
        # 执行区块三的功能
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

2.3 输入表示层：从离散符号到情境化向量

数据流的第一步是将输入的词元索引 in_idx 转换为包含位置信息的向量表示。

1
2
3
4
5
# GPTModel forward 方法的起始部分
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds
x = self.drop_emb(x)

2.3.1 词元嵌入

1
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])

正如前文所说的，计算机无法直接处理"单词"这样的符号。为了进行数学运算，必须将每个离散的词元映射到一个高维的连续向量空间中。这个过程被称为嵌入(Embedding)。nn.Embedding 是一个简单的查找表。它本质上是一个权重矩阵，维度为 (vocab_size, emb_dim)。输入一个词元的索引，它会返回该索引对应的行向量。这个向量是可训练的，模型在训练过程中会不断调整这些向量，使得在向量空间中语义相近的词元彼此靠近。

2.3.2 位置嵌入

理论上，词元嵌入非常适合作为大语言模型的输入。然而，大语言模型存在一个小缺陷——它们的自注意力机制（见后文）无法感知词元在序列中的位置或顺序。嵌入层的工作机制是，无论词元 ID 在输入序列中的位置如何，相同的词元 ID 始终被映射到相同的向量表示，如下图所示。

举个最简单的例子："人咬狗"和"狗咬人"在上述机制看来，包含的词元集合是相同的。然而，顺序在自然语言中至关重要。

为了实现这一点，可以将位置信息进行嵌入，一般有以下 3 种位置嵌入方式：

  • 绝对位置嵌入 (absolute positional embedding)：直接与序列中的特定位置相关联。对于输入序列的每个位置，该方法都会向对应词元的嵌入向量中添加一个独特的位置嵌入，以明确指示其在序列中的确切位置。例如，序列中的第一个词元会有一个特定的位置嵌入，第二个词元则会有另一个不同的位置嵌入，以此类推。这种方式可以是可学习的，也可以是通过固定的数学函数（如正弦/余弦函数）生成的。
  • 相对位置嵌入 (relative positional embedding)：关注的是词元之间的相对位置或距离，而非它们的绝对位置。该方法通常在计算注意力分数时，引入一个与词元间距离相关的偏置项，从而让模型学习的是词元之间的"间隔"关系，而不是它们在序列中的"具体坐标"。这种方法使得模型能够更好地适应不同长度（包括在训练过程中从未见过的长度）的序列。
  • 旋转位置嵌入 (Rotary Positional Embedding, RoPE)：通过一种创新的方式将位置信息融入自注意力机制中。它并非将位置向量直接添加到词元嵌入上，而是根据词元的绝对位置，对其在注意力计算中使用的查询（Query）和键（Key）向量进行旋转。这种精妙的旋转操作使得任意两个词元之间的注意力分数，能够自然地表示出它们的相对位置关系，从而让模型在处理位置信息时既高效又具备强大的长度泛化能力。

OpenAI 的 GPT 模型和本书使用的都是绝对位置嵌入：

1
2
3
4
5
6
7
8
9
10
11
12
class GPTModel(nn.Module):
    """完整的GPT模型实现"""
    def __init__(self, cfg) -> None:
				# ...
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        # ...

    def forward(self, in_idx):
        # ...
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        # ...
  1. 参数初始化：self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"]) 这行代码会初始化一个权重矩阵，也称为查找表。该矩阵的维度是 (context_length, emb_dim)。矩阵的每一行都是一个向量，且每一行都唯一对应一个从 0 到 context_length - 1 的绝对位置索引。这些行向量是模型的可训练参数，其初始值通常是随机设定的。
  2. 向量查找：当模型处理一个具体输入时，torch.arange(seq_len) 会首先生成一个包含该输入序列所有位置索引的张量 (tensor)，例如 [0, 1, 2, ..., seq_len-1]。随后，这个位置索引张量被传递给 self.pos_emb 层。该层会根据索引值，从第一步初始化的权重矩阵中，精确地查找并提取出每一行对应的位置向量，最终构成一个维度为 (seq_len, emb_dim) 的位置嵌入张量 pos_embeds。
  3. 信息融合：x = tok_embeds + pos_embeds 执行向量的逐元素加法操作。此操作将代表词元语义信息的 tok_embeds 张量与上一步生成的位置信息 pos_embeds 张量合并。

如下图所示：

2.3.3 dropout 掩码

至此，我们已经通过词元嵌入和位置嵌入的结合，得到了一个信息完备的输入向量。这个向量既包含了词元的语义信息，也明确了其在序列中的顺序，可以说是为模型准备了一份完美的"学习材料"。

然而，在将这份完美的材料送入 Transformer 的核心进行深度加工之前，我们还需要进行一个看似矛盾，却至关重要的操作——故意引入一些不确定性。为什么要这么做呢？这是为了防止模型在训练中变得过于依赖输入的每一个细节，从而陷入"死记硬背"的陷阱，也就是我们常说的"过拟合"。为了让模型学会从不完美的信息中也能提取核心规律，我们需要引入一种正则化技术。这正是我们接下来要讨论的 Dropout。

Dropout 技术示意图

它的主要目的是防止模型过拟合，提升模型的泛化能力。在训练期间，它会随机地将一部分输入数据置为零，迫使模型不能过度依赖于任何少数的特征，从而学习到更加鲁棒的模式。在评估和预测时，Dropout 会自动失效，不会对数据做任何改动。

1
2
3
4
5
6
7
8
9
10
11
class GPTModel(nn.Module):
    """完整的GPT模型实现"""
    def __init__(self, cfg) -> None:
        # ...
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
				# ...

    def forward(self, in_idx):
        # ...
        x = self.drop_emb(x)
        # ...

我们本次的实现总共会有 2 个地方应用到 dropout 技术：

  1. 在词元嵌入和位置嵌入相加之后，进入第一个 Transformer Block 之前，即上面的代码所做的事情。这是模型遇到的第一层正则化。它直接作用于融合了语义和位置信息的初始输入向量 x。通过随机将输入向量中的某些特征置为零，它迫使后续所有的 Transformer Block 都不能过度依赖输入向量中的任何单一维度。这相当于从源头上增加了训练难度，要求整个模型学习到对输入特征扰动不敏感的、更本质的规律。
  2. 在多头注意力模块内部，计算出注意力权重 attn_weights 并经过 softmax 归一化之后，在用它去加权 Value 向量之前。这种 dropout 不作用于输入向量本身，而是作用于注意力权重。注意力权重决定了在生成一个词的表示时，应该关注上下文中其他词的程度。在这里应用 dropout，会随机地将某些词与词之间的注意力连接切断（权重置为0）。这可以防止模型在学习时走捷径，比如过度依赖于某个特定的前文词汇。它鼓励模型去考虑更广泛的上下文信息，而不是仅仅依赖几个最强的信号。（关于注意力模块，我们后面会详细讨论）

[!IMPORTANT] 到目前为止，我们已经完整地剖析了 GPTModel 的输入表示层。我们从第一性原理出发，理解了为什么需要将离散的词元 ID，通过词元嵌入（Token Embedding）和位置嵌入（Positional Embedding），转化为一个融合了语义与顺序信息的、信息完备的高维向量 x。

这个过程的本质，是将人类的符号语言，翻译成了神经网络能够进行数学运算的、结构化的内部语言。

我们还探讨了 Dropout 技术。它像一个严格的教练，通过在训练中随机遮盖部分信息，强迫模型不能死记硬背，必须学会从不完整的信息中提炼出更本质、更鲁棒的规律，从而提升其泛化能力。

现在，我们有了一批准备就绪、信息丰富且经过初步正则化处理的训练材料。然而，此时此刻，序列中的每一个向量虽然知道了自己是谁以及在哪，但它仍然是一个独立的、上下文无关的个体。它并不知道自己与其他词元之间存在着怎样复杂的句法和语义关联。

那么，模型是如何让这些孤立的向量开始"交流"，理解彼此之间的关系，并最终形成对整个序列的深度理解呢？

答案，就藏在 Transformer 架构的革命性核心——自注意力机制 (Self-Attention Mechanism) 之中。下面我们就来深入 LLM 中最关键的部分，探究 Transformer 架构的层层细节！

2.4 核心处理层：Transformer 块的堆叠

在 GPTModel 中，我们共使用了 n_layers 个 TransformerBlock：

1
2
3
4
5
6
7
class GPTModel(nn.Module):
    def __init__(self, cfg) -> None:
       	# ...
        # 堆叠多个Transformer块
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])],
        )

我们先来看 TransformerBlock 的结构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class TransformerBlock(nn.Module):
    """Transformer块：多头注意力 + 前馈网络 + 残差连接"""
    def __init__(self, cfg) -> None:
        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

它的结构示意图如下所示：

2.4.1 注意力机制

深入探讨大语言模型核心的自注意力机制之前，让我们考虑一下在大语言模型出现之前的没有注意力机制的架构中所存在的问题。假设我们想要开发一个将文本从一种语言翻译成另一种语言的语言翻译模型。如下图所示，由于源语言和目标语言的语法结构不同，我们无法简单地逐个单词进行翻译。

这正是传统的序列处理模型（如 RNN）一个根本缺陷的体现：信息瓶颈。它们通过一个循环结构顺序处理文本，导致序列末端的信息很难直接关联到序列开头的遥远信息。

[!IMPORTANT]

自注意力机制 (Self-Attention) 的提出，正是为了打破这种信息瓶颈。其根本思想是：为序列中的每个元素，建立与其他所有元素的直接连接，并动态计算这些连接的强度（即注意力权重）。这样，模型在处理任何一个词元时，都能拥有一个全局视野，直接审视并借鉴整个上下文。

在 TransformerBlock 的内部，MultiHeadAttention 模块是其第一个、也是最为关键的子层。它是整个 GPT 模型"智能"的根本来源。要理解它，我们不能一蹴而就。很幸运的是，《从零构建大语言模型》的作者 Sebastian Raschka，为我们提供了一条从简单到复杂的演进路径，如下图所示，这部分的内容非常精华且重要，所以笔者将尽可能将这部分的内容进行完整记录，以帮助读者们更好的理解自注意力机制。

2.4.1.1 没有可训练权重的简单自注意力机制

在深入研究包含可训练权重的复杂版本之前，书中首先实现了一个不含任何可训练权重的简化自注意力机制，以便阐明其核心概念 。

如上图所示，这个机制的目标是为输入序列中的每一个词元（Token），计算出一个上下文向量（Context Vector） 。这个上下文向量是一种增强版的嵌入，它不仅包含了当前词元自身的信息，还融合了序列中所有其他词元的信息 。这对于理解句子中单词间的关系至关重要 。

计算这个上下文向量的过程分为三步：

  1. 计算注意力分数：衡量每个词对其他词的"相关性"或"相似度"。计算每对词之间的点积，得到相似度分数。点积在这里可以被看作是一种衡量相似度的方式：两个向量的点积越大，代表它们之间的对齐程度或相似度越高，注意力分数也越高 。
  2. 归一化获取注意力权重：得到的注意力分数是一些原始的数值，它们的尺度不一。为了使其规范化并易于解释，我们使用 Softmax 函数对这些分数进行处理 。Softmax 函数能将一组任意实数转换为一个概率分布，确保所有输出值的和为 1，并且每个值都是正数。这样得到的数值就是"注意力权重"，代表了在当前查询下，序列中每个词元的重要性 。
  3. 计算上下文向量：最后一步，将序列中的每一个词元嵌入向量与其对应的注意力权重相乘，然后将所有结果向量相加 。最终得到的向量就是我们想要的上下文向量，它是整个输入序列的加权和，权重由刚刚计算出的注意力权重决定 。

这3个步骤实现了自注意力机制的核心思想：

  • 看：计算每个词对其他词的关注度
  • 权衡：将关注度转换为权重
  • 融合：根据权重融合所有词的信息

最终效果： 每个词的向量表示都包含了整个序列的上下文信息，而不仅仅是自己的信息。这样模型就能理解词与词之间的关系，比如 "journey starts" 中的 "starts" 会更多地关注 "journey" 的信息。

现在我们用代码来演示一下，假设我们有以下输入：

1
2
3
4
5
6
7
8
9
10
inputs = torch.tensor(
    [
        [0.43, 0.15, 0.89],   # Your
        [0.55, 0.87, 0.66],   # journey
        [0.57, 0.85, 0.64],   # starts
        [0.22, 0.58, 0.33],   # with
        [0.77, 0.25, 0.10],   # one
        [0.05, 0.80, 0.55],   # step
    ]
)

笔者将尝试从第一性原理出发，对代码进行一步步拆解，说明白这个过程为什么能够实现我们期望的目标：为每个词元（Token）生成一个包含了上下文信息的向量。

这个问题的核心在于，我们要证明最终的上下文向量 (Context Vector) 的确融合了其他词元的信息，并且是根据"相关性"来融合的。

我们将以你例子中的词 starts (第三个词元) 为例，来全程追踪它的变化。

starts 的原始输入向量是: [0.57, 0.85, 0.64]。这个向量只代表 starts 本身，它对句子中的其他词一无所知，是孤立的。我们的目标是生成一个新的向量，让这个新的向量知道它前面有 Your journey，后面有 with one step。

第一步我们计算注意力分数，是为了发现"谁与我最相关"：

1
2
3
4
5
6
7
8
9
10
11
12
attn_scores = torch.empty(6, 6)
for i, x_i in enumerate(inputs):
    for j, x_j in enumerate(inputs):
        attn_scores[i, j] = torch.dot(x_i, x_j)

# 结果如下：
# tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
#         [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
#         [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
#         [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
#         [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
#         [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

attn_scores 的第 3 行是：[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605]。这行数字告诉我们，starts 这个词与句子中每个词的原始相关性得分分别是：

  • 与 Your 的相关性: 0.9422
  • 与 journey 的相关性: 1.4754 <-- 非常高
  • 与 starts (自身) 的相关性: 1.4570 <-- 非常高
  • 与 with 的相关性: 0.8296
  • 与 one 的相关性: 0.7154
  • 与 step 的相关性: 1.0605

仅从这一步看，这个机制已经成功地从数学上发现了 starts 与 journey 之间的紧密关系，因为它们的点积分数是最高的之一。这完全符合我们对语言的直觉（"旅程"和"开始"在语义上强相关）。

第一步得到的分数是原始的、未经缩放的数值，不易于作为权重使用。比如 1.4754 究竟代表多大的重要性？我们无法直接判断。所以第二步就是进行归一化获取注意力权重：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 2. 归一化获取注意力权重
attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)

# -1 表示在最后一个维度进行归一化，因为 attn_scores 是一个 [行, 列]，所以这里是在列上进行归一化，使得每行的值（在列维度的总和）为 1
# "沿着列的方向" = 在每一行内部，从左到右（列 0 到列 5）进行归一化
# 每一行都独立进行这个过程，结果每一行的 6 个数字加起来都等于 1
row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
print("Row 2 sum:", row_2_sum)
print("All row sums:", attn_weights.sum(dim=-1))


# 结果如下：
# tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
#         [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
#         [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
#         [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
#         [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
#         [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])
# Row 2 sum: 1.0
# All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

现在，这些数字的意义变得非常清晰了： 当模型在处理 starts 这个词时，它应该将它的注意力这样分配：

  • 13.90% 的注意力给 Your
  • 23.69% 的注意力给 journey <-- 权重最高
  • 23.26% 的注意力给 starts (自身)
  • 12.42% 的注意力给 with
  • 11.08% 的注意力给 one
  • 15.65% 的注意力给 step

这一步将第一步发现的相关性，转化为了具体的、可操作的重要性权重。它明确地告诉我们，为了理解 starts，我们需要重点参考 journey 和 starts 自身的信息。

接下来，这是最关键的一步，我们终于要创造那个"增强版"的向量（上下文向量）了。我们的目标是为词 i (例如 starts) 创建一个新的表示 \(C_i\)。这个新的表示 \(C_i\) 必须满足两个条件：

  1. 它必须包含所有其他词 j (从 Your 到 step) 的信息。
  2. 每个词 j 贡献的信息量，应该由我们刚刚算出的注意力权重 \({w_ij}\) 来决定。权重越高的词，影响越大。

现在，让我们思考一下，在数学上，特别是向量空间中，有什么运算可以同时满足这两个条件？答案就是加权平均 (Weighted Average) 或 加权和 (Weighted Sum)。

1
2
3
4
5
6
7
8
9
10
# 3. 计算上下文向量
all_contexts_vec = attn_weights @ inputs

# 结果如下：
# tensor([[0.4421, 0.5931, 0.5790],
#         [0.4419, 0.6515, 0.5683],
#         [0.4431, 0.6496, 0.5671],
#         [0.4304, 0.6298, 0.5510],
#         [0.4671, 0.5910, 0.5266],
#         [0.4177, 0.6503, 0.5645]])

矩阵乘法 attn_weights @ inputs 是一种非常高效的、一次性为所有词计算加权求和的方式。本质上是通过以下方式计算的： \[ C_i = \sum_{j=1}^N w_{ij} \cdot V_j \] 其中，\(w_{ij}\) 是词 i 对词 j 的注意力权重，\(V_j\) 是词 j 的原始向量。

现在我们来看这个公式的两个部分：

第一部分：\(w{ij}\cdot V{j}\) （缩放）：通过这一步，我们为句子中的每一个词，都生成了一个待贡献的向量。这个向量的意义和原始词一样，但它的影响力已经被其对应的注意力权重精确地调整好了。

第二部分：∑j=1N （求和）：这一步是把所有这些待贡献的向量全部加起来：\(C_{starts}=(w_{s,y}⋅V_y)+(w_{s,j}⋅V_j)+(w_{s,s}⋅V_s)+…\)。这一步的几何意义是：在那个高维的意义空间里，我们从原点出发，先沿着加强版 journey 向量走一段，再接着走削弱版 one向量的方向，再走加强版 starts 自身的方向...... 把所有词的贡献都走完，最终到达的那个新的位置，就是我们的上下文向量 \(C_{starts}\)。

这个三步过程之所以能起到我们想要的作用，是因为它完美地模拟了人类理解语言的一个核心逻辑：

  1. 关注焦点：当我们读到一个词时，我们会本能地寻找与它最相关的词。这个机制通过计算点积，找到了这些相关词。
  2. 分配精力：我们不会对所有相关的词都投入相同的精力。这个机制通过 Softmax，将"相关性"转化为"重要性"权重，量化了应该投入多少精力。
  3. 综合理解：我们基于这些焦点和精力分配，在大脑中形成对当前词的综合理解。这个机制通过加权求和，将所有词的信息根据重要性权重融合在一起，生成了最终的上下文向量。

通过这个过程，输出的每一个向量都从"我是谁"的孤立状态，变成了"在这样一个句子里，我是谁"的上下文感知状态，从而实现了我们的最终目标。

2.4.1.2 实现带可训练权重的自注意力机制

在上个阶段，注意力机制本身没有自己独立的、可以在训练中被优化的参数。模型在训练时，虽然可以学习和调整输入的 x 向量（即词嵌入本身），但它无法学习如何更好地计算注意力。无论输入的向量如何变化，计算注意力的公式始终不变。而且这种方式过于僵化。一个词元的向量表示 x 需要同时承载多种信息，它既要代表自身的语义，又要能很好地跟其他词元的向量进行点积来判断相关性。这就像要求一个人同时扮演运动员和裁判员，角色发生了混淆，难以做到最优。

所以第二个版本的根本问题就是：如何让模型学会去关注什么？如何让相似度的计算方式本身变得灵活和强大？

解决方案是引入角色分工，让专业的角色做专业的事。第二个版本中，我们不再直接使用原始的输入向量 x，而是引入三个独立的可训练线性变换层 (nn.Linear)：Query (Q)、Key (K) 和 Value (V)。

  • 查询向量 (Query)：代表当前这个词，主动去查询句子中其他词与自己的关系。可以理解为：我 (starts) 是谁？
  • 键向量 (Key)：代表句子中的每个词，用来被其他词查询的。可以理解为：我是 (journey)，你可以通过这个‘键’来了解我。
  • 值向量 (Value)：代表句子中每个词所携带的真正信息。一旦查询完毕，确定了关系密切度，我们就从这个值中提取信息。

至关重要的是，这三个权重矩阵 \(W_q\)、\(W_k\)、\(W_v\) 是可训练的。这意味着在训练过程中，模型会不断优化它们，学会如何将原始输入 x 转换成最有效的 Q、K 和 V，从而学会如何更好地去关注，让注意力的计算本身变得灵活而强大 。（所以才称这个版本是带可训练权重的自注意力机制）

Linear transformation of the word embedding to obtain Query, Key, and Value vectors — From(https://epichka.com/blog/2023/qkv-transformer/)

代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import torch
import torch.nn as nn

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        """
        d_in (int): 输入向量的维度。
        d_out (int): 查询(Query)、键(Key)和值(Value)向量的输出维度。
        qkv_bias (bool): 是否为Q, K, V线性层添加偏置项。
        """
        super().__init__()
        # 使用 nn.Linear 层来定义 Q, K, V 的线性变换。
        # 相比手动创建 nn.Parameter，nn.Linear 提供了更优的权重初始化，并可选择性地包含偏置项，
        # 是 PyTorch 中实现线性变换的标准做法。
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        """
        x (torch.Tensor): 输入张量，形状为 [批量大小, 序列长度, d_in]。
        """
        
        # 1. 线性投影：将输入 x 转换为 Q, K, V 表示
        # 每个输入词元的向量都会通过独立的线性层，生成其在查询、键、值三个空间中的新表示。
        keys = self.W_key(x)      # 形状: [批量大小, 序列长度, d_out]
        queries = self.W_query(x)  # 形状: [批量大小, 序列长度, d_out]
        values = self.W_value(x)    # 形状: [批量大小, 序列长度, d_out]

        # 2. 计算注意力分数：通过点积衡量 Query 和 Key 的相似度
        # queries 与 keys 的转置 (.T) 进行矩阵相乘，得到一个注意力分数矩阵。
        # 矩阵中的每个元素 attn_scores[i, j] 代表第 i 个查询与第 j 个键之间的原始相关性。
        attn_scores = queries @ keys.T

        # 3. 缩放与归一化：将分数转换为最终的注意力权重
        # a. 缩放(Scaling): 将分数除以 key 向量维度的平方根。这一步对于稳定训练至关重要，
        #    可以防止在维度过高时，点积结果过大导致 softmax 梯度消失。
        # b. 归一化(Normalization): 使用 softmax 函数将缩放后的分数转换为概率分布，
        #    确保每一行（代表每个查询）的注意力权重总和为1。
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1] ** 0.5, dim=-1,
        )

        # 4. 计算上下文向量：对 Value 向量进行加权求和
        # 将上一步得到的注意力权重矩阵与 values 矩阵相乘。
        # 这一步是根据注意力权重，对所有词元的 Value 信息进行加权聚合，
        # 最终为每个词元生成一个融合了全局上下文信息的新向量。
        context_vec = attn_weights @ values
        return context_vec

大体流程可参考下图理解：

Dot-product attention procedure —From(https://epichka.com/blog/2023/qkv-transformer/)

[!NOTE]

缩放点积注意力的原理：对嵌入维度进行归一化是为了避免梯度过小，从而提升训练性能。例如，在类 GPT 大语言模型中，嵌入维度通常大于 1000，这可能导致点积非常大，从而在反向传播时由于 softmax 函数的作用导致梯度非常小。当点积增大时，softmax 函数会表现得更像阶跃函数，导致梯度接近零。这些小梯度可能会显著减慢学习速度或使训练停滞。

因此，通过嵌入维度的平方根进行缩放解释了为什么这种自注意力机制也被称为缩放点积注意力机制。

2.4.1.3 利用因果注意力隐藏未来词汇

对于像 GPT 这样用于文本生成的模型，有一个核心要求：在预测序列中的下一个词元时，模型只能看到当前位置及之前的信息，绝不能偷看未来的词元。标准的自注意力机制会一次性访问整个输入序列，这显然不符合要求。为了解决这个问题，我们引入了因果注意力（Causal Attention），也称为掩码注意力（Masked Attention）。

实现因果注意力的关键在于掩码（Masking）操作。具体做法是在计算出注意力分数之后、应用 Softmax 函数之前，对注意力分数矩阵进行修改 。我们会创建一个"上三角"掩码矩阵，其中主对角线及以下的元素为 0，而主对角线以上的元素为负无穷大 (-inf) 。

当这个掩码矩阵被加到注意力分数矩阵上时，所有代表"未来"位置的分数都会变成负无穷大 。经过 Softmax 函数处理后，这些负无穷大的值对应的概率会变为 0 。这样一来，任何词元在计算其上下文向量时，其注意力权重都只会分布在它自身及之前的位置上，从而有效地隐藏了未来的词汇 。

代码实现逻辑如下：

1
2
3
4
5
6
7
8
9
# diagonal=1 表示不包含主对角线，只包含上三角部分
mask = torch.triu(torch.ones(context_length, context_length), diagonal=1) 
print(mask)
# masked_fill() 用指定值填充掩码为True的位置
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
print(masked)
# masked / keys.shape[-1]**0.5 进行缩放，torch.softmax 进行归一化
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=1)
print(attn_weights)

输出：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
tensor([[0., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 0.]])
tensor([[-0.0763,    -inf,    -inf,    -inf,    -inf,    -inf],
        [-0.0408,  0.0038,    -inf,    -inf,    -inf,    -inf],
        [-0.0423,  0.0025,  0.0074,    -inf,    -inf,    -inf],
        [-0.0090,  0.0404,  0.0416,  0.0233,    -inf,    -inf],
        [-0.0586, -0.0207, -0.0141, -0.0227,  0.1097,    -inf],
        [ 0.0040,  0.0504,  0.0502,  0.0317,  0.0320,  0.0354]],
       grad_fn=<MaskedFillBackward0>)
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.4921, 0.5079, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3259, 0.3364, 0.3376, 0.0000, 0.0000, 0.0000],
        [0.2442, 0.2529, 0.2531, 0.2498, 0.0000, 0.0000],
        [0.1919, 0.1971, 0.1980, 0.1968, 0.2161, 0.0000],
        [0.1632, 0.1686, 0.1686, 0.1664, 0.1664, 0.1668]],
       grad_fn=<SoftmaxBackward0>)

此外，为了防止模型在训练中过拟合，还可以在注意力权重矩阵上应用我们前面提到的 Dropout 技术。即在训练过程中，随机地将一部分注意力权重置为零，这有助于增强模型的泛化能力 。

代码实现逻辑如下：

1
2
3
4
5
6
7
dropout = torch.nn.Dropout(0.5) # 使用 50% 的 dropout 率
example = torch.ones(6, 6)
print(example)
print(dropout(example))

# 对权重矩阵进行 dropout
print(dropout(attn_weights))

输出：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
tensor([[1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1.]])
tensor([[2., 2., 0., 2., 2., 0.],
        [0., 0., 0., 2., 0., 2.],
        [2., 2., 2., 2., 0., 2.],
        [0., 2., 2., 0., 0., 2.],
        [0., 2., 0., 2., 0., 2.],
        [0., 2., 2., 2., 2., 0.]])
tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.6790, 0.6818, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.5059, 0.5040, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.4336, 0.0000],
        [0.3242, 0.3345, 0.0000, 0.3339, 0.3433, 0.3292]],
       grad_fn=<MulBackward0>)

可以看到大约一半的值被置为 0，且原来的值被放大了，用于位置权重的整体平衡。

一个完整的简单因果注意力类的参考实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import torch
import torch.nn as nn

class CausalAttention(nn.Module):
    """
    一个实现了因果自注意力（Causal Self-Attention）的模块。
    
    因果特性确保在处理序列中的任何一个词元（token）时，
    注意力机制只能关注到当前位置及之前位置的词元，而不能看到未来的词元。
    这对于自回归（auto-regressive）的文本生成任务至关重要。
    """
    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
        """
        初始化方法。
        
        参数:
        d_in (int): 输入嵌入向量的维度。
        d_out (int): 查询(Query)、键(Key)和值(Value)向量的输出维度。
        context_length (int): 模型的最大序列长度，用于创建因果掩码。
        dropout (float): 应用于注意力权重矩阵的 dropout 比率。
        qkv_bias (bool): 是否为 Q, K, V 的线性层添加偏置项。
        """
        super().__init__()
        self.d_out = d_out
        # 定义用于生成 Query, Key, Value 的线性变换层
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        # 定义 Dropout 层，用于正则化，防止过拟合
        self.dropout = nn.Dropout(dropout)
        # 创建并注册因果掩码（causal mask）
        # register_buffer 将一个张量注册为模块的缓冲区，它不会被视为模型参数（即不会在训练中被更新），
        # 但会随着模型移动（例如，.to(device)）。
        self.register_buffer(
            'mask',
            # torch.triu 创建一个上三角矩阵。diagonal=1 表示主对角线（及以下）的元素都为0，
            # 只有主对角线上方的元素为1。这个矩阵用于屏蔽未来的位置。
            torch.triu(torch.ones(context_length, context_length), diagonal=1)
        )

    def forward(self, x):
        """
        执行前向传播。
        
        参数:
        x (torch.Tensor): 输入张量，形状为 [批量大小, 序列长度, 输入嵌入维度 d_in]。
        """
        # 获取输入的维度信息
        b, num_tokens, d_in = x.shape

        # 1. 将输入 x 投影到 Query, Key, Value 空间
        keys = self.W_key(x)      # 输出形状: [b, num_tokens, d_out]
        queries = self.W_query(x) # 输出形状: [b, num_tokens, d_out]
        values = self.W_value(x)  # 输出形状: [b, num_tokens, d_out]

        # 2. 计算注意力分数
        # 将 keys 的最后两个维度转置，以便进行矩阵乘法
        # 形状从 [b, num_tokens, d_out] 变为 [b, d_out, num_tokens]
        # queries @ keys.transpose(...) 计算每个查询与所有键的点积
        attn_scores = queries @ keys.transpose(1, 2)  # 输出形状: [b, num_tokens, num_tokens]
        
        # 3. 应用因果掩码
        # masked_fill_ 是一个原地操作（in-place），它会直接修改 attn_scores 张量
        # self.mask.bool()[:num_tokens, :num_tokens] 会选取与当前输入序列长度匹配的掩码部分
        # 并将所有需要屏蔽的“未来”位置的分数填充为负无穷大
        attn_scores.masked_fill_(
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf
        )
        
        # 4. 缩放分数并应用 softmax 得到注意力权重
        # a. 缩放 (Scaling): 除以 key 维度的平方根，稳定梯度
        # b. Softmax: 将分数转换为概率分布。由于未来位置的分数是负无穷，
        #    经过 softmax 后它们的权重将变为0。
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )

        # 5. 应用 Dropout
        # 在训练阶段，随机将一些注意力权重置为0，以防止过拟合
        attn_weights = self.dropout(attn_weights)

        # 6. 计算上下文向量
        # 将注意力权重与 value 向量相乘，得到加权和
        context_vec = attn_weights @ values # 输出形状: [b, num_tokens, d_out]
        return context_vec

2.4.1.4 将单头注意力扩展到多头注意力

虽然带有可训练权重的因果自注意力机制已经非常强大，但它仍然有局限性：模型在某个位置只能学习到一种注意力模式。为了让模型能够从不同角度、不同表示子空间共同关注信息，原始 Transformer 论文引入了多头注意力（Multi-Head Attention）机制 。

"多头"的核心思想是并行地运行多次注意力计算，而不是只进行一次 。具体实现如下：

  1. 分割成多个头：我们不再只有一组 \(W_q\)、\(W_k\)、\(W_v\) 权重矩阵，而是为每个头都创建一组独立的权重矩阵 。例如，如果我们有 12 个头，那我们就有 12 组这样的矩阵。
  2. 并行计算注意力：每个头都独立地对输入执行缩放点积注意力计算（包含因果掩码）。由于每个头拥有不同的权重矩阵，它们会将输入投影到不同的表示子空间，从而学习到输入序列的不同方面特征 。例如，一个头可能关注语法结构，另一个头可能关注语义关联。
  3. 拼接与投影：在所有头都完成计算后，我们会得到多个输出上下文向量。我们将这些向量拼接（concatenate）在一起，形成一个更长的向量 。
  4. 最终线性投影：最后，这个拼接后的长向量会通过一个额外的线性层（out_proj）进行投影，将其维度恢复到模型期望的维度，并融合所有头学习到的信息 。

在代码中，可以通过实现一个简单的 MultiHeadAttentionWrapper 类来达到这一目标，MultiHeadAttentionWrapper 类堆叠了多个之前实现的 CausalAttention 模块实例。

1
2
3
4
5
6
7
8
9
10
11
12
13
class MultiHeadAttentionWrapper(nn.Module):
    """一个实现多头注意力的封装类"""
    def __init__(self, d_in, d_out, context_length,
                dropout, num_heads, qkv_bias=False) -> None:
        super().__init__()
        self.heads = nn.ModuleList(  # 堆叠多个 CausalAttention
            [CausalAttention(
                d_in, d_out, context_length, dropout, qkv_bias,
            ) for _ in range (num_heads)]
        )

    def forward(self, x):
        return torch.cat([head(x) for head in self.heads], dim=-1)

书中还提到了一种更高效的实现方式：与其创建多组独立的权重矩阵，不如创建一个更大的权重矩阵，一次性完成对所有头的查询、键、值向量的计算，然后通过重塑（reshape）和转置（transpose）操作将结果分割成多个头 。这种方法在数学上是等价的，但利用了现代硬件进行大规模矩阵运算的优势，计算效率更高。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
"""
输入: [2, 6, 3]
    ↓
Q,K,V: [2, 6, 2]
    ↓
重塑: [2, 6, 2, 1] (2个头，每个头1维)
    ↓
转置: [2, 2, 6, 1] (2个头并行计算)
    ↓
注意力: [2, 2, 6, 6] (每个头有自己的注意力矩阵)
    ↓
上下文: [2, 2, 6, 1] (每个头的结果)
    ↓
合并: [2, 6, 2] (所有头的结果合并)
    ↓
输出投影: [2, 6, 2]
"""
class MultiHeadAttention(nn.Module):
    """一个高效的多头注意力类"""
    def __init__(self, d_in: int, d_out: int,
                context_length: int, dropout: float, num_heads: int, qkv_bias=False):
        super().__init__()
        assert (d_out % num_heads == 0), "d_out must be divisible by num_heads"

        self.d_out = d_out  # 输出维度
        self.num_heads = num_heads # 头数量
        self.head_dim = d_out // num_heads # 每个头的维度
        
        # 初始化可训练的权重矩阵，分别代表查询向量、键向量、值向量
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

        # 使用一个线性层来组合头的输出
        self.out_proj = nn.Linear(d_out, d_out)

        # 掩码 + dropout
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length), diagonal=1)
        )

    def forward(self, x):
        # [batch_size, sequence_length, embedding_dim]
        b, num_tokens, d_in = x.shape

        # 计算权重矩阵 Q/K/V
        keys: Tensor = self.W_key(x)
        queries: Tensor = self.W_query(x)
        values: Tensor = self.W_value(x)

        # 重塑为多头格式
        # 将 [batch, seq_len, d_out] 重塑为 [batch, seq_len, num_heads, head_dim]
        # [2, 6, 4] -> [2, 6, 2, 2]
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

        # 调整维度顺序，让每个头独立计算注意力，便于批量处理所有头
        # 从形状 (b, num_tokens, num_heads, head_dim)
        # 转换到 (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)
        queries = queries.transpose(1, 2)

        # 计算注意力分数，这样每个批次(2)的每个头(2)都有了一个 6×6 的注意力分数矩阵
        attn_scores = queries @ keys.transpose(2, 3) # [2, 2, 6, 2] @ [2, 2, 2, 6] = [2, 2, 6, 6]
        mask_bool: Tensor = self.mask.bool()[:num_tokens, :num_tokens]

        # 应用因果掩码
        attn_scores.masked_fill_(mask_bool, -torch.inf)

        # 归一化权重
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1) # [2, 2, 6, 6]

        # 使用 dropout 掩码减少过拟合
        attn_weights = self.dropout(attn_weights) # [2, 2, 6, 6]

        # 每个头计算自己的上下文向量
        context_vec: Tensor = (attn_weights @ values).transpose(1, 2) # [2, 2, 6, 6] @ [2, 2, 6, 2] = [2, 2, 6, 2] -> [2, 6, 2, 2]
        # 重塑回原始格式 [2, 6, 2, 2] -> [2, 6, 4]
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        # 通过输出投影层 [2, 6, 4]
        context_vec = self.out_proj(context_vec)
        return context_vec

第二个版本 MultiHeadAttention 之所以更好，根本原因在于它将多次小规模的独立计算，整合为一次大规模的并行计算，从而最大化地利用了现代硬件（尤其是 GPU）的并行处理能力。

第一个版本 MultiHeadAttentionWrapper 的根本问题是计算被拆散了。对于一个有 12 个头的模型，这意味着要执行 12 组独立的 Q, K, V 矩阵乘法。在 GPU 上，每次独立的矩阵乘法都需要一次内核启动（kernel launch），这个启动本身是有开销的。执行 12 次小规模的计算，其总开销远大于执行 1 次等效的大规模计算。这就像让一个工人去搬 12 次箱子，每次只搬一个，远不如让他用推车一次性搬完 12 个箱子来得快。

这里面的向量变化可能有一些复杂，感兴趣的读者可以参考下图进行辅助理解。

多头注意力完整流程可视化

[!IMPORTANT]

到此为止，我们已经走完了一条从最简陋到最完备的注意力机制演进之路。我们从一个不带任何可训练参数的简单点积模型出发，理解了"看-权衡-融合"的核心思想；接着，通过引入可学习的 QKV 矩阵，赋予了模型学会如何去关注的能力；随后，我们用因果掩码为模型戴上了眼罩，强制它遵守时间顺序，只能回顾过去；最后，通过多头机制和高效的并行化实现，我们构建出了 GPT 模型真正的认知核心——MultiHeadAttention 模块。

这个模块是 Transformer 架构的灵魂。它为模型提供了一个动态的、可学习的机制，使其能够在处理每一个词元时，都能审视全局（或全局的过去），并精确地计算出上下文中每一个其他词元对当前词元的重要性，最终生成一个富含深度上下文信息的新表示。

然而，一个强大的引擎（MultiHeadAttention）本身还不足以构成一辆性能优越的赛车（TransformerBlock）。我们还需要稳定系统、传动装置和进一步的加工环节。这就引出了我们接下来的问题：

  • 模型在通过注意力机制融合了上下文信息之后，如何对这些新信息进行进一步的加工和思考？
  • 当我们把 12 个这样强大的计算层堆叠在一起时，如何保证训练过程的稳定，防止梯度消失或爆炸？
  • 在经过如此复杂的变换后，如何确保原始的、未经处理的信息不会在层层传递中丢失？

让我们先来回顾一下 TransformerBlock 的结构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class TransformerBlock(nn.Module):
    """Transformer块：多头注意力 + 前馈网络 + 残差连接"""
    def __init__(self, cfg) -> None:
        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"])

答案就藏在构成 TransformerBlock 的另外几个关键组件中。接下来，我们将把目光从注意力机制本身移开，去探索环绕在它周围的左膀右臂——前馈网络 (FeedForward Network)、层归一化 (Layer Normalization) 和 残差连接 (Shortcut/Residual Connections)，看看它们是如何协同工作，共同构成 Transformer 架构坚实可靠的核心处理单元的。

  • 层归一化 (LayerNorm): 它的根本作用是稳定训练过程。在数据经过复杂的注意力计算或前馈网络变换后，其数值分布可能会变得非常不稳定。层归一化就像一个调节器，在每个子层处理之前，都将数据拉回到一个标准的、易于处理的分布上，确保信息流的稳定。
  • 前馈神经网络 (FeedForward Network): 如果说注意力机制负责融合来自上下文的信息，那么前馈网络则负责对这些融合后的信息进行加工和思考。它是一个小型的、独立处理每个词元位置的神经网络，用于提取更高级、更抽象的特征，增加模型的非线性表达能力。
  • 快捷连接 (Shortcut/Residual Connection): 这是训练深度网络的关键技巧。它允许信息绕过某个处理层（如注意力或前馈网络），直接传递到下一层。这确保了即使在经过多达 12 层甚至更多的深度变换后，最原始的输入信息也不会完全丢失，同时极大地缓解了深度学习中的梯度消失问题，让深度堆叠成为可能。

2.4.2 使用层归一化进行归一化激活

一个深度神经网络就像一个多级信息加工流水线。数据（信号）在每一层都会被权重矩阵进行复杂的数学变换。当层数很深时，每一层微小的变化都可能被逐层放大。这会导致两个极端问题 ：

  1. 信号爆炸：某些层的输出值变得非常大，导致后续计算溢出，训练过程崩溃。
  2. 信号消失：某些层的输出值变得非常小，接近于零，导致信息无法有效传递到更深层，模型学不到东西。 这两种情况统称为内部协变量偏移 (Internal Covariate Shift)，它使得训练过程极其不稳定，就像在一条崎岖不平的山路上开车，油门（学习率）稍有不慎就会冲出赛道。

第一性原理解决方案：强制信号标准化

最直接的解决方案，就是在信息进入每个核心处理单元（如注意力和前馈网络）之前，强制进行一次校准或标准化。层归一化正是扮演了这个角色。 它的核心思想是，不管上一层传来的数据分布如何，它都强行将这批数据的均值调整为 0，方差调整为 1 。这相当于在流水线的每个关键工序前都安装了一个稳压器，确保无论输入信号如何波动，进入工序的信号始终是稳定、标准化的。

在 TransformerBlock 中，层归一化被放置在多头注意力和前馈网络之前 (self.norm1 和 self.norm2) 。这确保了这两个进行核心计算的模块接收到的输入始终处于一个稳定且易于处理的范围内，从而极大地稳定了整个深度模型的训练过程。

LayerNorm 的代码实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
class LayerNorm(nn.Module):
    """层归一化实现"""
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x-mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

让我们从第一性原理出发，来根本性地解释 LayerNorm 的这份代码实现。它的每一行都服务于一个核心目的：在保持模型表达能力的同时，稳定深度网络的训练过程。我们可以将这个实现拆解为两个核心部分来理解：强制标准化 和 可学习的自适应调整。

2.4.2.1 第一部分：强制标准化 - 解决信号失控问题

这是代码的核心计算部分：

1
2
3
4
# forward 方法中的核心计算
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x-mean) / torch.sqrt(var + self.eps)
  • mean = x.mean(dim=-1, keepdim=True) 和 var = x.var(dim=-1, ...)：这两行代码计算了每一个输入样本在其特征维度（emb_dim）上的均值和方差。dim=-1 是关键，它指定了归一化是沿着特征维度进行的，而不是像批归一化（Batch Norm）那样跨批次进行。这使得 LayerNorm 的效果与批次大小无关，在处理可变长度序列时尤其稳定。
  • norm_x = (x-mean) / torch.sqrt(var + self.eps)：这是标准的标准化公式（减去均值，再除以标准差）。它将原始输入 x 转换为了一个均值为0、方差为1的新向量 norm_x。
    • self.eps = 1e-5：eps (epsilon) 是一个极小的常数，它的唯一作用是防止分母为零。如果某个样本的方差恰好为0，没有 eps 就会导致除零错误，eps 保证了计算的数值稳定性 1。
    • unbiased=False：这是一个实现细节，表示在计算方差时分母是 N 而不是 N-1。选择 False 是为了与原始 GPT-2 模型的实现保持兼容，因为其最初是使用 TensorFlow 实现的，而这是 TensorFlow 的默认行为 2。

至此，我们已经强制将输入信号稳定在一个 N(0, 1) 的标准正态分布上。但这又带来了新的问题。

2.4.2.2 第二部分：可学习的自适应调整 - 恢复模型的表达能力

这是在 __init__ 中定义并在 forward 最后使用的部分：

1
2
3
4
5
6
# __init__ 中
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

# forward 的最后一步
return self.scale * norm_x + self.shift

根本问题：将每一层的输入都强制变为均值为0、方差为1的分布，这种做法可能过于暴力和死板。它虽然稳定了训练，但也可能限制了模型的表达能力。也许对于某个特定层来说，一个均值为 10、方差为 5 的输入分布才是最优的。我们不希望因为追求稳定而扼杀了模型学习这种分布的可能性。

解决方案：在强制标准化之后，再赋予模型撤销或重新调整这种标准化的能力。这是通过两个可学习的参数 scale 和 shift 来实现的。

  • self.scale (增益)：这是一个与特征维度相同大小的可学习向量。它与标准化后的 norm_x 进行逐元素相乘。它被初始化为全1，所以在训练刚开始时，它不起任何作用（乘以 1 等于不变）。
  • self.shift (偏置)：这也是一个可学习的向量。它被加到缩放后的结果上。它被初始化为全 0，所以在训练开始时，它也不起作用（加上 0 等于不变）。

这步的精髓在于：模型在训练过程中，可以通过反向传播自由地学习 scale 和 shift 的最佳值。

  • 如果模型发现强制标准化 N(0, 1) 对当前层来说是最好的，它就会让 scale 保持接近 1，shift 保持接近 0。
  • 如果模型发现一个不同的分布更好，它就可以学会相应的 scale 和 shift 值，将 norm_x 线性变换到任何它认为最优的均值和方差。

总的来说，LayerNorm 的实现是一个精妙的两步过程：

  1. 先稳定：通过强制的标准化，将可能失控的输入信号拉回到一个稳定的 N(0, 1) 分布，解决了深度网络训练不稳定的根本问题。
  2. 后放开：通过引入可学习的 scale 和 shift 参数，赋予模型恢复甚至创造全新分布的自由度，解决了强制标准化可能带来的表达能力受限的问题。

最终，这个实现既保证了训练的稳定性，又保留了模型的灵活性和表达能力。

2.4.3 实现具有 GELU 激活函数的前馈神经网络

自注意力机制的核心是加权求和 (attn_weights @ values)。虽然计算权重时有 softmax 引入了非线性，但信息融合的最后一步本质上是一个线性组合。如果整个 TransformerBlock 只依赖于注意力机制来处理信息，那么模型的表达能力将受到限制。它擅长融合信息，但在对融合后的信息进行深度加工方面能力不足。

第一性原理解决方案：为每个词元提供独立的非线性处理空间

为了弥补这一不足，我们需要一个专门的组件来对注意力机制输出的上下文向量进行进一步的、更复杂的非线性变换。前馈网络 (FFN) 就是这个组件。 它通常由两个线性层和一个非线性激活函数（如 GELU）组成 。它会对序列中的每一个词元向量独立地进行一次"升维-非线性激活-降维"的操作 。

  1. 升维：第一个线性层将向量维度扩大（例如从 768 维扩展到 3072 维），这为模型提供了更广阔的特征空间来表示和加工信息。
  2. 非线性激活 (GELU)：这是关键一步，打破了线性变换的局限，允许模型学习输入和输出之间更复杂、更抽象的关系。
  3. 降维：第二个线性层将维度恢复到原始大小，以便于下一层 TransformerBlock 处理。

前馈神经网络 FeedForward 的实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class FeedForward(nn.Module):
    """前馈神经网络"""
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"])
        )

    def forward(self, x):
        return self.layers(x)
      
class GELU(nn.Module):
    """GELU激活函数"""
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))
  1. 引入非线性：通过 GELU 激活函数，让模型有能力学习复杂的数据模式。
  2. 深度加工信息：通过"升维-降维"的结构，为模型提供一个更广阔的计算空间来提取和转换特征，同时保持整个 TransformerBlock 输入输出维度的一致性，使其能够被方便地深度堆叠。

2.4.4 添加快捷连接

当网络非常深时（例如堆叠 12 层 Transformer 块），会遇到两个致命问题：

  1. 梯度消失 (Vanishing Gradients)：在训练时，用于更新权重的梯度信号需要从最后一层反向传播到第一层。每经过一层，梯度都会被乘以该层的权重。在深层网络中，这些连乘操作很可能导致梯度信号迅速衰减，等传到浅层网络时已经微乎其微，导致浅层参数几乎不更新，模型无法有效训练 。
  2. 信息退化 (Information Degradation)：输入向量 x 每经过一个 TransformerBlock ，都会被复杂的注意力机制和前馈网络完全重构。在经过多层变换后，最原始、最直接的语义和位置信息可能会被冲淡甚至丢失。

第一性原理解决方案：建立信息/梯度的"高速公路"

解决方案出奇地简单而有效：在每个复杂处理单元（如注意力和前馈网络）旁边，建立一条直连通道，让输入可以直接跳过这个单元，与该单元的输出相加。这就是快捷连接或残差连接。

在 TransformerBlock 中，残差连接的实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class TransformerBlock(nn.Module):
    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

残差连接在这段代码中体现在以下两个关键操作上：

  1. shortcut = x: 这是分叉路口，将原始信息备份到 shortcut 变量中，开辟了直连通道。
  2. x = x + shortcut: 这是十字路口汇合，将主干道上经过复杂处理的信息与旁路上的原始信息重新组合。

具体如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# 第一个子层：多头注意力 + 残差连接

# ------------------- 残差连接的起点 -------------------
# 1. 保存原始输入：在进行任何变换之前，我们先把原始的输入 x 保存到一个名为 shortcut 的变量中。
#    这就相当于开辟了一条“快捷通道”或“旁路”，让原始信息可以绕过复杂的处理。
shortcut = x

# ------------------- 主处理路径 (F(x)) -------------------
# 2. 对输入进行复杂变换：
#    - 先进行层归一化
#    - 再通过多头注意力机制
#    - 最后应用 Dropout
#    这一系列操作的结果，更新了变量 x。现在的 x 已经不再是原始输入，
#    而是经过注意力模块深度加工后的“增量信息”或“残差”。
x = self.norm1(x)
x = self.att(x)
x = self.drop_shortcut(x)

# ------------------- 残差连接的终点 -------------------
# 3. 将原始输入与变换后的结果相加：
#    这行代码是残差连接最关键的体现。我们将“快捷通道”中的原始信息 (shortcut)，
#    与“主处理路径”上经过复杂变换后的增量信息 (x) 进行逐元素相加。
#    这样，最终的输出既包含了新学到的上下文关系，又没有丢失最原始的输入信息。
x = x + shortcut

2.5 输出层：从向量到概率分布

2.6 文本生成

至此，我们已经完成了 GPTModel 的全部架构代码实现。

我们从顶层的 GPTModel 容器开始，构建了其核心的可堆叠单元 TransformerBlock。在 TransformerBlock 内部，我们不仅实现了其进行上下文信息融合的核心模块——MultiHeadAttention，还集成了确保其稳定和高效运行的关键辅助组件：LayerNorm、FeedForward 网络和残差连接。修改修改修改修改

当前，我们拥有一个结构上完整、参数可扩展的 GPT 模型蓝图。然而，必须明确的是，这个模型的所有可训练参数（nn.Embedding、nn.Linear、LayerNorm 中的权重和偏置）均由随机值初始化。因此，尽管模型结构已经完备，但它不具备任何语言知识，无法执行任何有意义的任务，其输出将是无意义的随机内容。

不过，在让我们的模型具备输出有意义的内容之前，我们还是先来实现模型文本生成的能力。GPT 模型将输出张量转化为生成文本的过程涉及多个步骤，如下图所示。这些步骤包括解码输出张量、根据概率分布选择词元，以及将这些词元转换为人类可读的文本。

下图更加详细地展示的下一词元生成过程说明了 GPT 模型如何在给定输入的情况下生成下一个词元。在每一步中，模型输出一个矩阵，其中的向量表示有可能的下一个词元。将与下一个词元对应的向量提取出来，并通过 softmax 函数转换为概率分布。在包含这些概率分数的向量中，找到最高值的索引，这个索引对应于词元 ID。然后将这个词元 ID 解码为文本，生成序列中的下一个词元。最后，将这个词元附加到之前的输入中，形成新的输入序列，供下一次迭代使用。这个逐步的过程使得模型能够按顺序生成文本，从最初的输入上下文中构建连贯的短语和句子。

让我们来实现一个文本生成工具，如下代码所示，generate_and_print_sample 是一个用于快速验证和展示的便捷工具，它封装了从编码、生成到解码的全过程，并妥善处理了模型的训练/评估模式切换。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def generate_and_print_sample(model, tokenizer, device, start_context):
    """生成文本样本"""
    model.eval()
    context_size = model.pos_emb.weight.shape[0]
    encoded = text_to_token_ids(start_context, tokenizer).to(device)
    with torch.no_grad():
        token_ids = generate_text_simple(
            model=model, idx=encoded,
            max_new_tokens=50, context_length=context_size,
        )
    decoded_text = token_ids_to_text(token_ids, tokenizer)
    print(decoded_text.replace("\n", " "))
    model.train()
    
def generate_text_simple(model, idx, max_new_tokens, context_length):
    """使用模型生成文本"""
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_length:]
        with torch.no_grad():
            logits = model(idx_cond)

        logits = logits[:, -1, :]
        probas = torch.softmax(logits, dim=-1)
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)
        idx = torch.cat((idx, idx_next), dim=1)
    
    return idx

def text_to_token_ids(text, tokenizer):
    """将文本转换为token ID"""
    encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)
    return encoded_tensor

def token_ids_to_text(token_ids, tokenizer) -> str:
    """将token ID转换回文本"""
    flat = token_ids.squeeze(0)
    return tokenizer.decode(flat.tolist())

我们尝试调用一下：

1
generate_and_print_sample(model, tokenizer, device, "Every effort moves you")

可以看到输出是毫无意义的：

1
Every effort moves you Mexican rarity implementing NouPsychCle..." Contributamong enable lacked complications tendon conclud Nearly oddly insign Champions senseless poopuclear shuts dove aspirinentionrous Miniasions fearsomeRanked adore disadvantages disregkeepvocensed eased museums William glovesople Palace shooters increases felony chops Batteryracuse Advertising cease

那么，如何将这个参数随机化的架构，转变为一个能够理解和生成语言的功能性模型呢？答案是通过模型训练 (Model Training)。

3. 训练模型

本篇我们将进入将架构赋予生命的核心环节：实现训练循环 (Training Loop)。我们将详细介绍模型如何通过处理大量文本数据，在一个反复迭代的过程中，系统性地调整其内部数以亿计的参数。

我们将具体探讨损失函数 (Loss Function) 的计算、优化器 (Optimizer) 的作用以及反向传播 (Backpropagation) 的机制，这些是驱动模型从随机状态向智能状态收敛的根本动力。

3.1 模型训练流程

训练流程的根本目的，就是通过一个系统性的、迭代的优化过程，让初始化的 GPTModel 这个"空壳大脑"通过学习海量的数据样本，逐步调整其内部参数，最终掌握预测下一个词元的规律。

模型学习的数学基础是梯度下降 (Gradient Descent)。其核心思想可以归结为：

  1. 定义目标：我们需要一个损失函数 (Loss Function) 来量化模型当前预测与真实答案之间的差距。差距越大，损失值越高。
  2. 寻找方向：通过微积分计算损失函数对模型中每一个参数的梯度 (Gradient)。梯度指明了在该参数上，能让损失值上升最快的方向。
  3. 进行修正：我们让参数朝着梯度的相反方向迈出一小步。这一小步的步长由学习率 (Learning Rate) 控制。
  4. 反复迭代：不断重复"预测->计算损失->计算梯度->更新参数"的过程，模型的参数就会被逐步优化，使得损失值越来越小，预测越来越准。

train_model_simple 函数就是这一原理的精确代码实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def train_model_simple(model, train_loader, val_loader,
                    optimizer, device, num_epochs,
                    eval_freq, eval_iter, start_context, tokenizer):
    """训练模型的主循环"""
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1

    for epoch in range(num_epochs):
        model.train()
        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()
            optimizer.step()
            tokens_seen += input_batch.numel()
            global_step += 1

            # 定期评估
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device, eval_iter)
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                track_tokens_seen.append(tokens_seen)
                print(f"Ep {epoch+1} (Step {global_step:06d}):"
                    f"Train loss {train_loss:.3f}, "
                    f"Val loss {val_loss:.3f}")

        # 每3个epoch生成样本
        if epoch % 3 == 0:
            generate_and_print_sample(model, tokenizer, device, start_context)
    return train_losses, val_losses, track_tokens_seen

我们可以将这个函数的结构分解为三个层次：外层循环、核心学习循环和监控系统。

3.1.1 外层循环

for epoch in range(num_epochs): 定义了模型需要完整地看几遍整个训练数据集。一个 Epoch 代表对所有训练数据的一次完整遍历。让模型反复看同样的数据，是为了使其有机会从不同的批次组合和随机顺序中，更深入地学习数据中的模式。

3.1.2 核心学习循环

for input_batch, target_batch in train_loader: 是学习发生的真正场所。对于从 train_loader 中取出的每一个数据批次，模型都会严格执行梯度下降的"四步曲"：

  1. 清空旧梯度optimizer.zero_grad()

    PyTorch 的梯度计算默认是累加的。如果不手动清零，当前批次计算出的梯度会和之前所有批次的梯度叠加在一起，导致错误的更新方向。因此，在每次计算新梯度前，必须先“清空缓存”。

  2. 前向传播与计算损失 loss = calc_loss_batch(...)

    我们需要知道模型在当前参数下的表现有多差。calc_loss_batch 函数内部会调用 model(input_batch)，完成一次前向传播，得到预测的 logits。然后，使用交叉熵损失函数 cross_entropy 来计算预测 logits 和真实 target_batch 之间的差距，得到一个量化误差的标量 loss。

  3. 反向传播计算梯度 loss.backward()

    知道了总误差（loss）后，我们需要将这个误差"分摊"到每一个导致误差的参数上，即计算损失对每一个模型参数的偏导数。这是 PyTorch autograd 引擎的核心功能。这一行代码会自动地、高效地完成整个反向传播过程，计算出模型中所有可训练参数的梯度，并存储在它们的 .grad 属性中。

    不熟悉反向传播概念的读者，可参考：大白话解释反向传播算法

  • 更新模型参数 optimizer.step()

    有了修正方向（梯度）后，需要一个执行者来实际地调整参数。优化器（如 AdamW）会根据 loss.backward() 计算出的梯度，以及自身的更新规则（如学习率），去更新模型中的每一个参数，完成一次学习和进化。

3.1.3 监控系统

仅仅闷头学习是不够的，我们还需要知道学得怎么样。这个函数内置了两套监控系统：

  • 定量评估 (if global_step % eval_freq == 0)：每隔 eval_freq 步，就调用 evaluate_model 函数。该函数会暂停训练 (model.eval())，在不计算梯度 (torch.no_grad()) 的模式下，快速计算模型在训练集和验证集上的损失。通过观察这两个损失的变化，我们可以清晰地了解模型的学习状态。
  • 定性观察 (if epoch % 3 == 0)：每隔几个 epoch，就调用 generate_and_print_sample 函数。它会给模型一个固定的开头 (start_context)，让模型在当前的学习状态下续写一段文本。通过观察从最初的“胡言乱语”到逐渐生成通顺句子的过程，我们可以获得最直观的反馈。

3.2 计算文本生成损失

在 train_model_simple 中，有两个关键的函数：

  • calc_loss_batch
  • evaluate_model
  • generate_and_print_sample：这个前面我们已经介绍过了，用于生成文本并打印出来。
参考资料
许可协议

本文采用 署名-非商业性使用-相同方式共享 4.0 国际 许可协议，转载请注明出处。

分享文章

快来参与讨论吧~