从零实现 nanoGPT（三）：实现因果自注意力，让模型读懂上下文

前言

在上一篇《从零实现 nanoGPT（二）：用 Bigram 跑通第一个语言模型》中，我们实现了一个最简单的字符级语言模型 BigramLanguageModel。

通过这个模型，我们已经完成了语言模型的基本流程：

输入一段字符对应的 Token ID。
模型预测每个位置的下一个字符。
使用交叉熵损失衡量预测结果。
通过反向传播更新模型参数。
根据预测概率逐个生成新字符。

但是，Bigram 模型存在一个明显的问题：它在预测下一个字符时，只能看到当前位置的字符。

例如：

css 复制代码

输入一：c a t
输入二：h a t

对于最后一个位置，Bigram 模型看到的都是字符 t，因此得到的下一字符预测分布完全相同。

无论前面出现的是 ca 还是 ha，都不会影响它对 t 后面字符的预测。

显然，这与我们使用语言的方式不同。

一个字符接下来应该出现什么，往往需要结合前面的内容进行判断。

那么，如何让模型在预测当前位置的下一个字符时，利用之前出现过的其他字符呢？

这就是本文要解决的问题。

我们将从零实现一个单头因果自注意力（Causal Self-Attention） ，并将它接入上一篇的语言模型。

本文完成后，我们将拥有：

字符嵌入（Token Embedding）。
位置嵌入（Position Embedding）。
单头因果自注意力（Single-Head Causal Self-Attention）。
一个能够利用上下文预测下一个字符的语言模型。

需要说明的是，这还不是完整的 GPT。后续我们会继续加入多头注意力、前馈神经网络、残差连接和 LayerNorm。

一、回顾上一篇：Bigram 为什么无法理解上下文？

上一篇中，我们使用 nn.Embedding 实现了 Bigram 模型的核心结构：

ini 复制代码

self.token_embedding_table = nn.Embedding(
    vocab_size,
    vocab_size
)

这里的 vocab_size 表示词表大小。

由于我们使用的是字符级语言模型，因此每个不同的字符都对应一个 Token ID。

假设词表中有 65 个字符，那么：

ini 复制代码

vocab_size = 65

嵌入表的大小就是：

scss 复制代码

(65, 65)

每个字符 ID 对应一行长度为 65 的向量，这个向量直接表示模型预测下一个字符时，对词表中各个字符给出的分数。

例如：

lua 复制代码

idx = torch.tensor([[1, 2, 3, 4]])

logits = self.token_embedding_table(idx)

得到：

ini 复制代码

idx.shape     = (1, 4)
logits.shape  = (1, 4, 65)

每个位置都会产生一个长度为 65 的预测向量。

但这里有一个重要的问题。

当前位置的输出完全由当前位置的 Token ID 决定。

假设：

lua 复制代码

a = torch.tensor([[1, 2, 3]])
b = torch.tensor([[4, 5, 3]])

虽然两个序列前面的字符不同，但最后一个 Token ID 都是 3。

因此：

ini 复制代码

logits_a = model.token_embedding_table(a)
logits_b = model.token_embedding_table(b)

torch.equal(logits_a[:, -1, :], logits_b[:, -1, :])

结果为：

python 复制代码

True

原因很简单：两个位置都查询了嵌入表的第 3 行。

这就是 Bigram 的局限。

它能够学习相邻字符之间的统计关系，却无法根据更长的上下文调整预测。

为了突破这个限制，我们需要让不同位置之间发生信息交互。

二、第一步：把字符 ID 转换成特征向量

2.1 从预测分数到字符特征

在 Bigram 模型中，我们直接将字符映射为词表大小的预测向量：

scss 复制代码

nn.Embedding(vocab_size, vocab_size)

现在，我们不再直接生成预测分数，而是先把每个字符映射为一个特征向量：

scss 复制代码

nn.Embedding(vocab_size, n_embd)

其中：

vocab_size：词表大小。
n_embd：字符特征向量的维度。

例如：

ini 复制代码

vocab_size = 65
n_embd = 32

embedding = nn.Embedding(vocab_size, n_embd)

idx = torch.tensor([[1, 2, 3, 4]])
x = embedding(idx)

print(x.shape)

输出：

css 复制代码

torch.Size([1, 4, 32])

这意味着，我们把 4 个字符 ID 转换成了 4 个长度为 32 的向量。

可以表示为：

csharp 复制代码

Token ID
   ↓
[1, 2, 3, 4]
   ↓
Embedding
   ↓
[
  [0.12, 0.34, ..., 0.56],
  [0.21, 0.15, ..., 0.72],
  [0.45, 0.11, ..., 0.38],
  [0.33, 0.29, ..., 0.81]
]

上面的数字只是用于说明结构，并非实际运行结果。

每一个向量都包含 32 个可学习的数值。

这些数值在初始化时没有我们预先规定的语义，而是随着模型训练逐渐调整。

2.2 为什么不直接预测下一个字符？

因为我们希望字符特征能够与其他位置的特征交互。

Bigram 的计算过程是：

复制代码

Token ID
   ↓
Embedding
   ↓
logits

现在，我们希望变成：

复制代码

Token ID
   ↓
Embedding
   ↓
与前面字符的信息交互
   ↓
logits

其中，负责信息交互的模块就是自注意力。

三、什么是自注意力？

自注意力的核心思想并不复杂：

序列中的每个位置，都可以根据自身的需求，从其他允许访问的位置获取信息。

假设输入：

复制代码

h e l l o

当模型处理最后一个字符 o 时，如果能够结合前面的 h e l l，就不必只依赖 o 本身进行预测。

但是，前面不同位置的信息，未必具有相同的重要程度。

我们希望模型自己学习：

应该关注哪些位置？
各个位置应该占多大的权重？
最终应该从这些位置提取什么信息？

注意力机制使用三个向量完成这一过程：

名称	含义	作用
Query（Q）	查询向量	当前的位置想寻找什么信息
Key（K）	键向量	每个位置提供什么匹配条件
Value（V）	值向量	每个位置实际提供什么信息

这里的"寻找""匹配"和"提供"都是帮助理解的比喻。

实际上，Q、K、V 都是由输入特征经过可训练的线性层得到的。

ini 复制代码

q = self.query(x)
k = self.key(x)
v = self.value(x)

3.1 Q、K、V 从哪里来？

假设输入特征：

ini 复制代码

x.shape = (B, T, C)

其中：

B：Batch Size，一批中包含多少条序列。
T：Sequence Length，序列长度。
C：Embedding Dimension，特征维度。

我们定义三个线性层：

ini 复制代码

self.query = nn.Linear(n_embd, head_size, bias=False)
self.key = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)

这里的 head_size 表示单个注意力头处理的特征维度。

假设：

ini 复制代码

n_embd = 32
head_size = 32

那么：

ini 复制代码

q = self.query(x)
k = self.key(x)
v = self.value(x)

三个张量的形状均为：

css 复制代码

(B, T, 32)

注意，Q、K、V 虽然来自同一个输入 x，但它们经过的是不同的线性变换，因此通常具有不同的数值。

四、如何计算注意力分数？

我们已经得到了 Q、K、V。

接下来需要解决的问题是：

当前位置应该对其他位置分配多大的注意力？

4.1 使用点积衡量匹配程度

注意力使用 Query 和 Key 的点积计算匹配分数：
$sij =qi⊤kj$ $s_{ij}=q_i\^\\top k_j$ $sij=qi⊤kj$

其中：

$(qi) (q_i)$ (qi)：第 (i) 个位置的 Query。
$(kj) (k_j)$ (kj)：第 (j) 个位置的 Key。
$( sij ) (s_{ij})$ (sij)：第 (i) 个位置对第 (j) 个位置的匹配分数。

用 PyTorch 实现：

ini 复制代码

scores = q @ k.transpose(-2, -1)

为什么需要转置？

因为：

ini 复制代码

q.shape = (B, T, head_size)
k.shape = (B, T, head_size)

转置 K 的最后两个维度：

scss 复制代码

k.transpose(-2, -1)

(B, T, head_size)
        ↓
(B, head_size, T)

因此矩阵乘法为：

r 复制代码

(B, T, head_size)
         @
(B, head_size, T)
         ↓
     (B, T, T)

得到一个注意力分数矩阵。

其中第 i 行、第 j 列表示：第 i 个位置对第 j 个位置的匹配分数。

4.2 一个简单的数值例子

假设一条序列只有三个位置，Query 和 Key 的维度都是 2：
$Q=$ $1 0 0 1 1 1$ ,K= $1 0 0 1 1 1$ Q= \begin{bmatrix} 1&0\\ 0&1\\ 1&1\end{bmatrix} , K= \begin{bmatrix} 1&0\\ 0&1\\ 1&1 \end{bmatrix} Q= 101011 ,K= 101011

计算：
$QK⊤=$ $1 0 1 0 1 1 1 1 2$ QK^\top= \begin{bmatrix} 1&0&1\\ 0&1&1\\ 1&1&2 \end{bmatrix} QK⊤= 101011112

观察第三行：
$1,1,2$ $1,1,2$ $1,1,2$

它表示第三个位置对三个位置的原始匹配分数。

不过，这些数值还不能直接作为注意力权重。

我们还需要进行缩放和 Softmax。

4.3 为什么要除以根号 head_size？

注意力使用缩放点积：
$S= QK⊤ dk S=\frac{QK^\top}{\sqrt{d_k}}$ S=dk QK⊤

其中 $(dk) (d_k)$ (dk) 就是 Key 的最后一维大小，也就是这里的 head_size。

对应代码：

ini 复制代码

scores = scores * (head_size ** -0.5)

等价于：

ini 复制代码

scores = scores / math.sqrt(head_size)

为什么需要缩放？

在一定的随机分布假设下，向量点积的方差会随着维度增大而增大。

如果直接将较大的分数输入 Softmax，就可能得到特别集中的概率分布。

例如：

scss 复制代码

softmax([1, 2, 3])
≈ [0.090, 0.245, 0.665]

而：

scss 复制代码

softmax([10, 20, 30])
≈ [0.000, 0.000, 1.000]

第二种情况中，概率几乎全部集中在最后一个位置。

这种过度集中的分布可能导致 Softmax 的梯度变得很小，不利于训练。

除以 $( dk ) (\sqrt{d_k})$ (dk )可以控制点积分数的尺度，使训练更加稳定。

五、为什么需要因果遮罩？

到这里，我们已经能够计算序列中各个位置之间的注意力分数。

但还有一个非常关键的问题。

我们不能让模型看到未来。

5.1 回顾第一篇的数据构造

第一篇中，我们构造过错位的输入与目标：

复制代码

原始字符：h e l l o

输入 x：  h e l l
目标 y：  e l l o

对于输入的第一个位置：

复制代码

输入：h
目标：e

模型需要根据 h 预测 e。

如果这个位置可以通过注意力读取输入中的第二个字符 e，就相当于提前看到了答案。

因此，我们必须限制每个位置能够访问的范围：

markdown 复制代码

                    可访问的位置

                  h   e   l   l

当前位置 h        ✓   ×   ×   ×
当前位置 e        ✓   ✓   ×   ×
当前位置 l        ✓   ✓   ✓   ×
当前位置 l        ✓   ✓   ✓   ✓

也就是说：

第 (i) 个位置，只能关注第 (0) 到第 (i) 个位置，不能访问后面的字符。

这就是因果注意力中的"因果"约束。

5.2 使用下三角矩阵构造 Mask

PyTorch 提供了：

scss 复制代码

torch.tril()

用于提取矩阵的下三角部分。

例如：

scss 复制代码

import torch

mask = torch.tril(torch.ones(4, 4))

print(mask)

输出：

csharp 复制代码

tensor([
    [1., 0., 0., 0.],
    [1., 1., 0., 0.],
    [1., 1., 1., 0.],
    [1., 1., 1., 1.]
])

其中：

1 表示允许关注。
0 表示禁止关注。

我们将这个矩阵作为注意力遮罩。

5.3 为什么使用 masked_fill？

假设某个位置计算出的注意力分数为：

csharp 复制代码

scores = torch.tensor([
    [1.0, 2.0, 3.0],
    [2.0, 1.0, 3.0],
    [1.0, 2.0, 3.0]
])

对应的因果遮罩：

ini 复制代码

mask = torch.tril(torch.ones(3, 3))

执行：

ini 复制代码

scores = scores.masked_fill(
    mask == 0,
    float("-inf")
)

得到：

csharp 复制代码

tensor([
    [1., -inf, -inf],
    [2., 1., -inf],
    [1., 2., 3.]
])

为什么要填充负无穷？

因为接下来要计算 Softmax：
$esi ∑jesj \frac{e^{s_i}}{\sum_j e^{s_j}}$ ∑jesjesi

而：
$e−∞=0 e^{-\infty}=0$ e−∞=0

因此，被遮住的位置对应的注意力权重就会变为 0。

ini 复制代码

weights = F.softmax(scores, dim=-1)

结果近似为：

csharp 复制代码

tensor([
    [1.0000, 0.0000, 0.0000],
    [0.7311, 0.2689, 0.0000],
    [0.0900, 0.2447, 0.6652]
])

注意第一行：

csharp 复制代码

[1, 0, 0]

第一个位置只能使用自己的信息。

第二行：

csharp 复制代码

[0.7311, 0.2689, 0]

第二个位置可以使用前两个位置的信息，但不能读取第三个位置。

第三行没有未来位置，因此可以关注全部三个位置。

这就实现了因果约束。

六、使用注意力权重汇总信息

经过前面的计算，我们已经得到注意力权重矩阵：

ini 复制代码

weights.shape = (B, T, T)

其中：

css 复制代码

weights[b, i, j]

表示第 b 条序列中的第 i 个位置，对第 j 个位置分配的注意力权重。

接下来，我们使用这些权重对 Value 向量进行加权求和。

ini 复制代码

out = weights @ v

形状变化：

r 复制代码

weights: (B, T, T)

v:       (B, T, head_size)

out:     (B, T, head_size)

6.1 具体例子

假设三个位置的 Value 分别为：
$V=$ $1 0 0 2 2 1$ V= \begin{bmatrix} 1&0\ 0&2\ 2&1 \end{bmatrix} V= $10 02 21$

假设第三个位置的注意力权重为：
$0.2,0.3,0.5$ $0.2,0.3,0.5$ $0.2,0.3,0.5$

那么第三个位置输出的特征向量就是：
$o3 =0.2$ $1,0$ +0.3 $0,2$ +0.5 $2,1$ = $1.2,1.1$ \begin{aligned} o_3 &=0.2 $1,0$ +0.3 $0,2$ +0.5 $2,1$ \ &= $1.2,1.1$ \end{aligned} o3=0.2 $1,0$ +0.3 $0,2$ +0.5 $2,1$ = $1.2,1.1$

这个新的向量不再仅仅来自第三个位置，而是结合了前面多个位置的信息。

这正是自注意力与 Bigram 最大的区别：

Bigram 直接查表预测；自注意力则先让各个位置的信息交互，再进行预测。

将整个过程写成一个公式：
$softmax⁡ ( QK⊤ dk +M) V \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}}+M \right)V$ softmax(dk QK⊤+M)V

其中 (M) 是加性因果遮罩：
$Mij = { 0, j≤i −∞, j>i M_{ij}= \begin{cases} 0,&j\le i\ -\infty,&j>i \end{cases}$ Mij={0,j≤i −∞,j>i

我们目前实现的，就是这个公式。

七、使用 PyTorch 实现单头因果自注意力

前面已经分别介绍了 Q、K、V、注意力分数、Mask 和加权求和。

现在将这些计算封装成一个完整的 PyTorch 模块。

在上一篇的 model.py 中新增一个类：

python 复制代码

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


class Head(nn.Module):
    """单头因果自注意力"""

    def __init__(self, n_embd, head_size, block_size):
        super().__init__()

        # Q、K、V 线性投影
        self.query = nn.Linear(
            n_embd, head_size, bias=False
        )
        self.key = nn.Linear(
            n_embd, head_size, bias=False
        )
        self.value = nn.Linear(
            n_embd, head_size, bias=False
        )

        # 因果遮罩
        self.register_buffer(
            "tril",
            torch.tril(torch.ones(block_size, block_size))
        )

    def forward(self, x):
        B, T, C = x.shape

        # 生成 Q、K、V
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        # 计算注意力分数
        scores = q @ k.transpose(-2, -1)

        # 缩放
        scores = scores * (k.shape[-1] ** -0.5)

        # 因果遮罩
        scores = scores.masked_fill(
            self.tril[:T, :T] == 0,
            float("-inf")
        )

        # Softmax
        weights = F.softmax(scores, dim=-1)

        # 加权求和
        out = weights @ v

        return out

这段代码虽然不长，但已经完成了一个单头因果自注意力的完整前向计算。

7.1 为什么使用 register_buffer？

注意这里：

less 复制代码

self.register_buffer(
    "tril",
    torch.tril(torch.ones(block_size, block_size))
)

我们没有直接使用：

matlab 复制代码

self.tril = torch.tril(
    torch.ones(block_size, block_size)
)

而是使用 register_buffer()。

这是因为因果遮罩具有两个特点：

它是模型需要保存和使用的张量。
它不需要通过梯度下降进行更新。

register_buffer() 可以让这个张量成为模块的 Buffer，而不是可训练参数。

这样，当执行：

vbnet 复制代码

model.to("mps")

时，遮罩也会随着模型移动到 MPS 设备。

默认情况下，Buffer 还会被包含在模型的 state_dict() 中。

7.2 为什么是 self.tril $,$ ？

我们在初始化时创建的是：

scss 复制代码

(block_size, block_size)

大小的下三角矩阵。

但实际输入序列长度 T 不一定等于 block_size。

例如：

ini 复制代码

block_size = 32
T = 10

此时注意力分数矩阵的大小是：

scss 复制代码

(10, 10)

因此，只需要使用遮罩左上角的：

ruby 复制代码

self.tril[:T, :T]

得到：

scss 复制代码

(10, 10)

大小的遮罩即可。

这也要求输入长度不能超过 block_size。

八、为什么还需要位置嵌入？

现在我们已经实现了自注意力，但还有一个容易忽略的问题。

注意力机制本身没有显式提供字符在序列中的绝对位置。

假设输入：

css 复制代码

a b c

和：

css 复制代码

c b a

两条序列的字符不同顺序，语义也可能完全不同。

如果我们只使用字符 Embedding，并且不考虑因果遮罩的差别，那么标准自注意力并不知道某个字符原本处于第几个位置。

为了给模型提供位置信息，我们引入位置嵌入。

8.1 创建位置嵌入表

ini 复制代码

self.position_embedding_table = nn.Embedding(
    block_size,
    n_embd
)

假设：

ini 复制代码

block_size = 32
n_embd = 32

那么位置嵌入表的大小就是：

scss 复制代码

(32, 32)

它为位置 0 到位置 31 分别分配一个可训练的位置向量。

注意，这里的位置指的是当前输入窗口内部的位置。

8.2 如何得到位置编号？

假设当前序列长度为：

ini 复制代码

T = 4

那么可以执行：

ini 复制代码

positions = torch.arange(T)

得到：

scss 复制代码

tensor([0, 1, 2, 3])

接着：

ini 复制代码

position_emb = self.position_embedding_table(positions)

得到：

scss 复制代码

(4, 32)

大小的位置特征。

8.3 将位置特征与字符特征相加

ini 复制代码

token_emb = self.token_embedding_table(idx)

position_emb = self.position_embedding_table(
    torch.arange(T, device=idx.device)
)

x = token_emb + position_emb

形状为：

makefile 复制代码

token_emb:     (B, T, n_embd)
position_emb:  (T, n_embd)

x:             (B, T, n_embd)

这里使用了 PyTorch 的广播机制。

位置嵌入会沿 Batch 维度广播，添加到每条序列的对应位置上。

因此：

x_i=E_{\\text{token}}(t_i)+E_{\\text{pos}}(i)

其中：

(E_{\text{token}})：字符嵌入。
(E_{\text{pos}})：位置嵌入。
(i)：位置编号。

这样，即使相同字符出现在不同位置，输入注意力模块的向量也会不同。

九、构建 AttentionLanguageModel

到这里，我们已经有了：

Token Embedding。
Position Embedding。
Causal Self-Attention。

接下来将它们组合为一个可以训练和生成文本的语言模型。

为了与第二篇保持一致，我们继续使用：

ini 复制代码

forward(idx, targets=None)

并返回：

复制代码

logits, loss

生成函数也继续使用自回归方式，每次只生成一个新 Token。

在 model.py 的 Head 类后面添加：

ini 复制代码

class AttentionLanguageModel(nn.Module):
    def __init__(
        self,
        vocab_size,
        n_embd=32,
        block_size=32
    ):
        super().__init__()

        self.block_size = block_size

        # 字符嵌入
        self.token_embedding_table = nn.Embedding(
            vocab_size, n_embd
        )

        # 位置嵌入
        self.position_embedding_table = nn.Embedding(
            block_size, n_embd
        )

        # 单头因果自注意力
        self.attention = Head(
            n_embd=n_embd,
            head_size=n_embd,
            block_size=block_size
        )

        # 预测下一个字符
        self.lm_head = nn.Linear(
            n_embd, vocab_size
        )

    def forward(self, idx, targets=None):
        B, T = idx.shape

        if T > self.block_size:
            raise ValueError(
                "输入长度不能超过 block_size"
            )

        # 字符嵌入
        token_emb = self.token_embedding_table(idx)

        # 位置嵌入
        positions = torch.arange(
            T, device=idx.device
        )
        position_emb = self.position_embedding_table(
            positions
        )

        # 合并字符特征与位置特征
        x = token_emb + position_emb

        # 因果自注意力
        x = self.attention(x)

        # 输出预测分数
        logits = self.lm_head(x)

        # 计算损失
        loss = None

        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(B * T, -1),
                targets.reshape(B * T)
            )

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens):

        for _ in range(max_new_tokens):

            # 只保留最近 block_size 个 Token
            idx_cond = idx[:, -self.block_size:]

            # 模型预测
            logits, _ = self(idx_cond)

            # 取最后一个位置
            logits = logits[:, -1, :]

            # 转换为概率
            probs = F.softmax(logits, dim=-1)

            # 采样下一个 Token
            idx_next = torch.multinomial(
                probs, num_samples=1
            )

            # 拼接到原序列
            idx = torch.cat(
                (idx, idx_next), dim=1
            )

        return idx

整个模型的结构如下：

css 复制代码

输入 Token ID
      │
      ▼
Token Embedding ──────┐
                      ├── 相加
Position Embedding ───┘
                      │
                      ▼
              Causal Attention
                      │
                      ▼
                  Linear
                      │
                      ▼
                   logits
                      │
                      ▼
               CrossEntropy

这里有两个地方需要重点理解。

9.1 lm_head 是什么？

ini 复制代码

self.lm_head = nn.Linear(
    n_embd, vocab_size
)

在 Bigram 中：

scss 复制代码

nn.Embedding(vocab_size, vocab_size)

直接给出了下一字符的预测分数。

现在，经过注意力之后，我们得到的是长度为 n_embd 的特征向量。

因此需要一个线性层，将其映射到词表大小：

css 复制代码

(B, T, n_embd)
       ↓
     lm_head
       ↓
(B, T, vocab_size)

输出的 logits 仍然表示每个位置预测下一个字符的分数。

所以，我们不需要修改第一篇的数据标签，也不需要修改第二篇的交叉熵训练目标。

9.2 为什么 generate 要截取上下文？

注意这一行：

ini 复制代码

idx_cond = idx[:, -self.block_size:]

我们在模型初始化时使用了：

scss 复制代码

nn.Embedding(block_size, n_embd)

以及：

matlab 复制代码

torch.tril(
    torch.ones(block_size, block_size)
)

因此，模型最多只能处理 block_size 个位置。

但生成文本时，序列会不断增长。

例如：

erlang 复制代码

初始长度：1
生成之后：2
继续生成：3
...
生成之后：100

当序列长度超过 block_size，不能直接将整个序列输入模型。

所以我们每次只取最近的 block_size 个 Token：

css 复制代码

idx[:, -self.block_size:]

例如：

复制代码

完整序列长度：100
block_size：32

当前预测只使用最后 32 个 Token

这样就可以在有限上下文窗口内持续进行自回归生成。

需要注意：窗口滑动后，位置嵌入的编号会从 0 重新开始。这是本篇采用的简化方案。

十、接入上一篇的训练循环

现在模型已经实现，我们可以继续使用前两篇的数据集和训练代码。

这一步不需要重新下载 Shakespeare 数据，也不需要重新构建 Tokenizer。

我们只需要替换模型，并继续使用原来的训练流程。

10.1 初始化新模型

在 train.py 中导入：

javascript 复制代码

from model import AttentionLanguageModel

然后初始化：

ini 复制代码

block_size = 32
n_embd = 32

model = AttentionLanguageModel(
    vocab_size=meta["vocab_size"],
    n_embd=n_embd,
    block_size=block_size
).to(device)

这里沿用第一篇得到的词表信息。

注意，block_size 需要与训练批次使用的窗口长度保持一致。

10.2 创建优化器

ini 复制代码

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3
)

特别注意：替换模型后，必须为新模型创建优化器。

因为优化器保存的是创建时传入的参数引用。如果继续使用旧的优化器，它可能不会更新新模型的参数。

10.3 训练过程

继续使用上一篇已经实现的 get_batch() 获取训练数据：

ini 复制代码

xb, yb = get_batch("train")

这里假设仍然沿用上一篇的 get_batch(split) 接口。如果你的实现需要额外传入数据集或设备等参数，保持原来的调用方式即可。

训练部分仍然是：

scss 复制代码

model.train()

logits, loss = model(xb, yb)

optimizer.zero_grad(set_to_none=True)

loss.backward()

optimizer.step()

你会发现：

虽然模型内部结构发生了很大的变化，但训练循环几乎没有改变。

这是因为语言模型的训练目标依然相同：

-\frac{1}{BT}

\sum_{b=1}^{B}

\sum_{t=1}^{T}

\log P(y_{b,t}\mid x_{b,\le t})

]

其中，(x_{b,\le t}) 表示第 (b) 条序列从开始到位置 (t) 的输入上下文。

与 Bigram 不同的是，现在预测可以依赖该位置及其之前的字符，而不只是当前位置的 Token。

10.4 在 Mac 上使用 MPS

如果使用 Apple Silicon Mac，可以继续沿用前面的设备选择方式：

makefile 复制代码

device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)

model = model.to(device)

训练数据也需要移动到相同设备：

ini 复制代码

xb = xb.to(device)
yb = yb.to(device)

只要模型、输入和目标都在相同设备上，就可以通过 MPS 执行支持的张量计算。

十一、验证因果注意力是否正确

代码能正常运行，不代表注意力实现一定正确。

尤其是因果 Mask，如果写错，很可能导致模型在训练过程中读取未来的 Token。

因此，我们需要设计一个测试，检查模型是否真的遵守因果约束。

11.1 测试思路

构造两条输入序列：

css 复制代码

序列 A：1 2 3 4
序列 B：1 2 8 9

它们前两个位置完全相同，但后两个位置不同。

由于模型不能访问未来，因此：

第 0 个位置的输出应该相同。
第 1 个位置的输出应该相同。
后两个位置的输出则可以不同。

如果前两个位置的输出发生变化，就说明模型的前面位置受到了未来输入的影响。

11.2 测试代码

创建：

复制代码

test_attention.py

写入：

ini 复制代码

import torch

from model import AttentionLanguageModel


torch.manual_seed(42)

model = AttentionLanguageModel(
    vocab_size=10,
    n_embd=16,
    block_size=4
)

model.eval()

a = torch.tensor(
    [[1, 2, 3, 4]],
    dtype=torch.long
)

b = torch.tensor(
    [[1, 2, 8, 9]],
    dtype=torch.long
)

with torch.no_grad():
    logits_a, _ = model(a)
    logits_b, _ = model(b)

# 检查输出形状
assert logits_a.shape == (1, 4, 10)

# 检查前两个位置是否受未来输入影响
torch.testing.assert_close(
    logits_a[:, :2, :],
    logits_b[:, :2, :]
)

print("形状检查与因果性检查通过")

执行：

复制代码

python test_attention.py

输出：

复制代码

形状检查与因果性检查通过

这个测试比单纯查看代码中有没有 torch.tril() 更有意义。

它直接验证了模型的行为：

后面的 Token 发生变化，不会影响前面位置的预测结果。

不过，它只是一个针对特定输入的功能测试，并不能代替对所有输入和运行条件的完整验证。

十二、和上一篇的 Bigram 进行对比

到这里，我们已经实现了第一个能够利用上下文的语言模型。

现在回到最开始的问题。

假设两个输入序列分别是：

bash 复制代码

cat
hat

对于 Bigram：

markdown 复制代码

当前位置字符：t
        ↓
查找 t 对应的 Embedding
        ↓
下一字符预测

无论前面是什么，只要当前位置相同，预测分布就相同。

但对于当前的 Attention 模型：

css 复制代码

c a t
  ↓
字符嵌入 + 位置嵌入
  ↓
因果自注意力
  ↓
汇总当前及之前位置的信息
  ↓
预测下一个字符

而：

css 复制代码

h a t
  ↓
字符嵌入 + 位置嵌入
  ↓
因果自注意力
  ↓
汇总当前及之前位置的信息
  ↓
预测下一个字符

因为输入上下文不同，所以 t 所获得的上下文特征也可能不同，最终预测结果也可以不同。

这就是自注意力带来的重要变化。

12.1 对比两个模型的训练表现

我们可以在相同的 Shakespeare 数据集上分别训练：

BigramLanguageModel
AttentionLanguageModel

然后使用上一篇的验证函数评估 Validation Loss。

建议使用相同的：

复制代码

数据集
batch_size
block_size
训练步数
验证样本设置

但需要注意，两个模型的参数数量并不相同，因此这样的对比主要用于观察模型变化，并不是严格控制参数量的实验。

如果想统计参数量：

scss 复制代码

num_params = sum(
    p.numel() for p in model.parameters()
)

print(f"参数量：{num_params:,}")

记录下：

复制代码

Bigram：
训练 Loss
验证 Loss
参数量
生成文本

Attention：
训练 Loss
验证 Loss
参数量
生成文本

观察引入注意力后，模型是否能得到更低的验证损失，以及生成文本是否出现更多上下文相关的局部结构。

这里不预设具体实验结果，因为它们受到训练步数、学习率、初始化和模型规模等因素的影响。

同时还需要认识到：单头注意力并不保证一定优于 Bigram，尤其是在训练不足时。

十三、这一版模型还有哪些不足？

虽然我们已经让模型能够利用上下文，但它距离完整的 GPT 仍然有一定距离。

目前的模型结构是：

markdown 复制代码

Token Embedding
       +
Position Embedding
       ↓
Single Attention Head
       ↓
Linear
       ↓
logits

它仍然存在一些局限。

13.1 只有一个注意力头

目前我们只有一组 Q、K、V。

不同注意力头可以分别学习不同的注意力模式。

例如，有的头可能更关注邻近字符，有的头可能更关注较远的位置。

这里不代表每个头都会自动学到固定的语言规则，而是说多头结构为模型提供了多个不同的特征投影空间。

13.2 缺少前馈神经网络

当前模型在注意力加权求和之后，只使用一个线性层输出 logits。

我们还没有加入 Transformer 中常见的前馈网络（Feed Forward Network）。

FFN 可以对每个位置汇总后的特征进行进一步的非线性变换。

13.3 缺少残差连接和 LayerNorm

完整的 Transformer Block 通常包含：

残差连接。
Layer Normalization。
多头自注意力。
前馈神经网络。

这些结构有助于构建更深的模型，并改善训练时的优化特性。

13.4 上下文长度受到限制

我们使用固定的：

ini 复制代码

block_size = 32

因此模型在一次前向传播中，最多只能处理 32 个 Token。

此外，当前注意力显式构造了：

r 复制代码

(T, T)

大小的注意力分数矩阵。

当序列长度增加时，这部分计算和中间矩阵的规模会快速增长。

后续可以进一步了解 PyTorch 的 scaled_dot_product_attention 等优化实现。

十四、总结

回顾整个 nanoGPT 系列，我们已经完成了三次重要的迭代。

第一篇：让文本变成数据

我们完成了：

markdown 复制代码

Shakespeare 原始文本
         ↓
字符词表
         ↓
Token ID
         ↓
构造输入 x 和标签 y

解决了如何让模型读取文本，以及如何构造下一个字符预测任务的问题。

第二篇：让模型能够学习和生成

我们实现了：

markdown 复制代码

Token ID
    ↓
Bigram Embedding
    ↓
logits
    ↓
CrossEntropy Loss
    ↓
反向传播

跑通了完整的语言模型训练与自回归生成过程。

第三篇：让模型可以利用上下文

我们将模型升级为：

css 复制代码

Token ID
    ↓
字符嵌入 + 位置嵌入
    ↓
Q、K、V
    ↓
计算注意力分数
    ↓
因果 Mask
    ↓
Softmax
    ↓
Value 加权求和
    ↓
lm_head
    ↓
下一个字符预测

至此，模型不再只能依赖当前字符，而是可以结合当前位置之前的内容进行预测。

但它还只有一个注意力头，也没有完整的 Transformer Block。

下一篇：《从零实现 nanoGPT（四）：从单头注意力到完整 Transformer Block》

我们将继续实现：

多头注意力（Multi-Head Attention）。
前馈网络（Feed Forward Network）。
残差连接（Residual Connection）。
层归一化（LayerNorm）。
堆叠多个 Transformer Block。

把这些模块组合起来之后，我们就能得到一个结构更加完整的 GPT 语言模型。

系列文章

《从零实现 nanoGPT（一）：让莎士比亚变成模型能读懂的数据》
《从零实现 nanoGPT（二）：用 Bigram 跑通第一个语言模型》
《从零实现 nanoGPT（三）：实现因果自注意力，让模型读懂上下文》

从零实现 nanoGPT（三）：实现因果自注意力，让模型读懂上下文