构建完整的 GPT 结构

从单个部件到完整的模型

上一篇我们实现了注意力头，使模型能够看见上下文。但单个注意力头还不构成 GPT，正如单个气缸还不构成一台发动机。

本篇要把其余的部件补齐——多头注意力、前馈网络、残差连接、层归一化、位置编码——再把它们组装成一个完整的 GPT 模型类。本篇代码偏多，但每个部件都相当短小，我们逐一来看。先把本篇会用到的几个超参数列出来，它们决定了这个模型的各项规格：

n_embd = 64        # 每个字向量的长度
n_head = 4         # 注意力头的数量
n_layer = 4        # Transformer 层数
block_size = 32    # 上下文长度
dropout = 0.1      # 一个用于抑制过拟合的部件，第七篇详细说明

多头注意力：多个头各自关注不同侧面

上一篇实现的是单个注意力头。但语言中的关系是多种多样的——有的关系属于语法层面（主语与谓语的对应），有的属于语义层面（某个代词指代前文的哪个名词）。指望单个注意力头同时兼顾所有这些关系，要求过高。

解决办法并不复杂：设置多个注意力头并行运作，让它们各自关注不同的方面。每个头独立地执行一遍上一篇那套查询、键、值的运算，关注前文的不同侧面;运算完成后，把各个头的输出拼接起来，再经过一个线性层做一次融合。这就是多头注意力multi-head attentionmulti-head attentionMulti-head attention runs several attention heads side by side, each free to focus on a different kind of relationship — one head may track grammar, another may track which noun a pronoun refers to. Their outputs are concatenated and then blended by a linear layer. A single head would be stretched too thin trying to capture every kind of dependency at once. 。

先把上一篇的注意力头封装成一个规范的模块：

import torch
import torch.nn as nn
from torch.nn import functional as F


class Head(nn.Module):
    """单个注意力头，即上一篇那套运算"""
    def __init__(self, head_size: int):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # 下三角矩阵，用于因果掩码；register_buffer 表示它不是参数、不参与训练
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5     # 计算关注程度并缩放
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # 屏蔽未来
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        v = self.value(x)
        return wei @ v


class MultiHeadAttention(nn.Module):
    """多个注意力头并行运作，再拼接、融合"""
    def __init__(self, num_heads: int, head_size: int):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)        # 拼接之后的融合层
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)   # 拼接各头的输出
        return self.dropout(self.proj(out))

Head 就是上一篇代码的模块化版本。MultiHeadAttention 把若干个 Head 装入 nn.ModuleList，前向计算时让它们各自运算一遍，用 torch.cat 拼接起来，再经过一个线性层 proj 把各头的信息融合在一起。

前馈网络：对收集到的信息做加工

注意力机制所做的，是让每个字"收集"前文的信息。但收集回来之后，还需要一个环节让模型对这些信息做进一步的加工。这个环节就是前馈网络feed-forward networkfeed-forward networkThe feed-forward network is a small two-layer network applied to each position independently. Where attention gathers information from other tokens, the feed-forward network is where each token digests what it has gathered. It widens the vector, passes it through a non-linearity, then narrows it back, giving the model room to compute on the collected information. 。

它的结构相当朴素：对每个位置的向量，独立地做两次线性变换，中间夹一个非线性激活函数。具体而言，先把向量放大到 4 倍宽度，经过一道激活，再压缩回原来的宽度。

class FeedForward(nn.Module):
    """逐位置的小型网络，负责加工注意力收集到的信息"""
    def __init__(self, n_embd: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),    # 放大到 4 倍宽度
            nn.ReLU(),                         # 非线性激活
            nn.Linear(4 * n_embd, n_embd),    # 压缩回原宽度
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

中间的 nn.ReLU() 是激活函数激活函数激活函数是夹在两次线性变换之间的一道非线性运算，本篇用的是 ReLU。它的必要性在于：如果一个网络从头到尾全是线性变换，那么再多层叠加，效果也始终等价于单独一层，表达能力被彻底锁死。插入一道非线性之后，多层叠加才真正产生更强的表达力，模型才能拟合语言中那些曲折的、非线性的规律。 ，作用是为模型引入"非线性"。如果模型全部由线性变换构成，即便叠加多层，其效果也等同于单层，无法表达复杂的规律;夹入一道非线性，模型才能够拟合真实语言中那些曲折的关系。这两个部件的分工可以这样记忆：注意力负责"字与字之间相互参考",前馈网络负责"每个字对收集到的信息独立加工"。

残差连接：为梯度的回传留出通路

现在已经有了两个主要部件——多头注意力与前馈网络。我们希望把它们叠加许多层，使模型更强。但深层网络存在一个长期存在的问题：层数一多，第四篇所讲的反向传播在向前回传梯度时会逐层衰减，传到最前面几层时已经微弱到几乎不起作用，那些层便难以得到训练。

残差连接residual connectionresidual connectionA residual connection adds a component output back onto its input, written as x = x + f(x). It opens a direct path through the network so that, during backpropagation, gradients can reach the early layers without being weakened layer by layer. It is the single trick that makes networks dozens or hundreds of layers deep actually trainable. 是解决这一问题的经典办法，其形式简单到出乎意料：把一个部件的输入，直接加到它的输出之上。

写成代码就是 x = x + 部件(x)，而不是 x = 部件(x)。

它为何有效。因为这个"把输入加回去"的操作，相当于在网络中开辟了一条直通的通路。反向传播的梯度，除了走部件内部那条逐层衰减的路径之外，还能够沿着这条加法通路畅通地回传到前面的层。有了它，几十层乃至上百层的网络才能够得到训练。这是使深层网络得以成立的关键一步。

层归一化：稳定数值的范围

还有一个用于稳定训练的部件：层归一化layer normalizationlayer normalizationLayer normalization rescales each position vector to a standard range — mean zero, variance near one — before it enters a component. As data passes through many operations its scale tends to drift, and that drift makes training unstable. Layer normalization corrects the scale at every step, keeping the numbers well-behaved. 。

数据穿过一层层运算时，数值的大小会发生漂移——某一层输出的数值普遍偏大，下一层又普遍偏小，这种忽大忽小会使训练变得不稳定、难以收敛。层归一化所做的，是在每个部件处理之前，把每个位置的向量重新缩放到一个标准的范围（均值为 0、方差在 1 附近），相当于每经过一道工序就把数值"校正"一次。

PyTorch 直接提供了 nn.LayerNorm，可以直接使用。把多头注意力、前馈网络、残差连接、层归一化组合到一起，就构成了一个 Transformer 块（Block）——GPT 正是由若干个这样的块堆叠而成：

class Block(nn.Module):
    """一个 Transformer 块：注意力与前馈，各自带有残差连接和层归一化"""
    def __init__(self, n_embd: int, n_head: int):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))      # 先层归一化，再注意力，结果经残差加回
        x = x + self.ffwd(self.ln2(x))    # 先层归一化，再前馈，结果经残差加回
        return x

forward 中的那两行，把本篇的部件都用上了：ln1、ln2 是层归一化，sa、ffwd 是两个主要部件，而每行开头的 x = x + 就是残差连接。

位置编码：为模型补充顺序信息

最后一个部件，用于弥补注意力机制的一个盲点。

回顾上一篇，注意力是用查询与键做点积来计算关注程度的。这一运算有一个特点：它只考察两个字向量的内容，并不考察它们的先后位置。也就是说，"猫追狗"与"狗追猫",在注意力看来字向量是同样的一批，它无法区分二者的顺序——而这两句话的含义恰好相反。

因此需要设法把"位置"这一信息也提供给模型。办法相当巧妙：再建立一张表，专门为位置编号。第 0 个位置对应一个向量，第 1 个位置对应一个向量，依此类推。这张位置表称为位置编码positional embeddingpositional embeddingAttention compares token vectors by content alone and is blind to their order, so cat-chases-dog and dog-chases-cat look identical to it. Positional embedding fixes this by keeping a second table, one vector per position, and adding the position vector onto the token vector. Each token then carries both what it is and where it sits. 。然后把"字本身的向量"与"它所在位置的向量"相加，作为这个字真正进入模型的表示。这样每个字便同时携带了"我是哪个字"与"我处在第几个位置"两份信息。

组合成完整的 GPT

各个部件都已齐备，下面组装最终的模型。它的处理流程是：查出字向量、查出位置向量、两者相加、穿过若干个 Transformer 块、最后用一个线性层把每个位置的向量映射回 vocab_size 个评分。

class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)      # 字向量表
        self.position_embedding_table = nn.Embedding(block_size, n_embd)   # 位置向量表
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)               # 最后一道层归一化
        self.lm_head = nn.Linear(n_embd, vocab_size)   # 映射回 vocab_size 个评分

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)                      # (B,T,n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T))       # (T,n_embd)
        x = tok_emb + pos_emb                                          # 字向量 + 位置向量
        x = self.blocks(x)                                             # 穿过所有块
        x = self.ln_f(x)
        logits = self.lm_head(x)                                       # (B,T,vocab_size)
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
        return logits, loss

    def generate(self, idx, max_new_tokens: int):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]      # 只保留最近 block_size 个字
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

与第三篇的 Bigram 模型做一个对比，可以发现两者的骨架高度相似：同样有 forward 返回 logits 与 loss，同样有 generate 做连续预测，同样用交叉熵计算损失。改变的只是中间环节——从"查一张表直接得出评分",变成了"查出字向量与位置向量、穿过若干个 Transformer 块、再得出评分"。

generate 中多了一行 idx[:, -block_size:]，原因在于位置表只有 block_size 个位置，输入长度不能超过这个数值，因此生成时只把最近的 block_size 个字输入模型。

这就是一个完整的 GPT，GPT 这三个字母此时都已对应到位：Generative（能够生成）、Pre-trained（即将进行训练）、Transformer（由 Transformer 块堆叠而成）。它与 OpenAI 的 GPT 采用同一套结构，差别仅在于规模——我们的 n_layer 是 4，大型模型可能是几十乃至上百;我们的 n_embd 是 64，大型模型则是数千。原理上没有任何不同。

本篇要点

• 多头注意力让多个注意力头并行运作，各自关注前文的不同侧面，再拼接并融合。
• 前馈网络对每个位置独立做加工;注意力负责"字与字相互参考",前馈负责"每个字独立加工"。
• 残差连接 x = x + 部件(x) 为反向传播的梯度开辟了一条直通通路，使深层网络得以训练。
• 层归一化在每个部件之前把数值缩放到标准范围，从而稳定训练。
• 注意力本身无法区分字的先后，位置编码用一张位置向量表补充顺序信息，并与字向量相加。
• 完整的 GPT 由字向量与位置向量相加、穿过若干 Transformer 块、再经线性层映射回评分构成;它与真正的 GPT 同构，只是规模较小。

下一篇

模型已经构建完毕，训练循环在第三篇也已具备。下一篇把两者结合起来，完整地训练这个迷你 GPT——设定好超参数，运行训练，观察损失的下降，并查看它生成的文字如何逐轮改善。

参考资料

• Attention Is All You Need（Transformer 原始论文）
• nanoGPT 项目源码
• 本博客 ml-basics 系列《Transformer》《批归一化》《Dropout 详解》