前置知识：前面的 micrograd、makemore 系列课程（可选），熟悉 Python，微积分和统计学的基本概念
目标：理解和欣赏 GPT 的工作原理
你可能需要的资料：
Colab Notebook 地址
Twitter 上看到的一份很细致的笔记，比我写得好
ChatGPT#
在 2022 年底问世的 ChatGPT 到如今的 GPT4、Claude3，这些 LLM（大语言模型）已经融入了很多人的日常生活。它们都是概率系统，对于同样的 prompt，它们的答案是多样的。相较于我们之前实现的语言模型，GPT 等可以能够模拟单词、字符或更一般的符号序列，并知道英语中某些单词是如何相互跟随的。从这些模型的视角来看，我们的 prompt 就是一段序列的开始，模型要做的就是补全这个序列。
那么，为这些单词序列建模的神经网络是什么呢？
Transformer#
2017 年，AI 领域里程碑式的论文《Attention is All You Need》提出了 Transformer 架构。
我们所熟知的 GPT，全称是 Generative Pre-trained Transformer （生成式预训练 Transformer）。尽管原文章对应的领域是机器翻译，但其深远影响了整个 AI 领域，稍微改动这个架构就可以应用到大量 AI 应用程序中，也是 ChatGPT 的核心。
当然，本节的目标并不是训练一个 ChatGPT，毕竟那可是一个超级工业级项目，涉及大量数据的训练、预训练和微调过程，我们要做的只是训练一个基于 Transformer 的语言模型，跟前面一样，这也会是一个字符级的语言模型。
搭建模型#
数据集#
使用 toy 级别的小规模数据集 “Tiny Shakespeare”，它深得 Andrej 喜爱。这个数据集基本就是莎士比亚所有走品的大杂烩，文件大小约 1MB。和 ChatGPT 的一个区别是，ChatGPT 的输出单位是 token，类似 “单词块” 的概念，我们后面也会提到。
# 我们总是从一个数据集开始训练, 下载小莎士比亚数据集
!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
# 阅读以检查
with open('input.txt', 'r', encoding='utf-8') as f:
text = f.read()
Tokenize#
chars = sorted(list(set(text))) # set得到序列中不重复的字符的无序序列，转为list来得到排序功能
vocab_size = len(chars)
print(''.join(chars)) # 合并为一个字符串
print(vocab_size)
# 输出(按ASCII码排序)：
# !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz
# 65
这里和我们在前面几节中做的字符表功能是一样的，然后我们需要开发一个能 tokenize (标记化) 输入序列的功能，这个名字的意义是将原始文本作为一个字符串转换为一些整数序列。对于我们要做的字符级模型来说的话，那就只是简单的把单个字符映射为数字。
如果看过前面几节内容的话，这部分代码给你的感觉应该相当熟悉，如 Bigram 中 “创建 Lookup Table 和字符组合的映射关系” 就和这里很相像。
# 创建一个从字符到整数的映射
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: 接收一个字符串，输出一个整数列表（编码器）
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: 接收一个整数列表，输出一个字符串（解码器）
print(encode("hii there"))
print(decode(encode("hii there")))
这里我们同时构建了编码器和解码器，这里的作用就是在字符级别，将字符串和整数相互转换。这只是一个非常简单的 tokenize 算法，人们设想了很多方法，如谷歌的SentencePiece ，它可以将文本分割成子词（subwords），这也是实践中常采用的方法；OpenAI 也有TikToken这个用字节对（byte pairs）来 tokenize。
使用 tiktoken 来编码：gpt2 的词汇表中包含了 50,257 个 token，面对同样的字符串，相较于我们简单的算法，只用了 3 个整数就完成了编码。
# 现在对整个文本数据集进行编码，并将其存储到一个 torch.tensor 中。
import torch
data = torch.tensor(encode(text), dtype=torch.long)
print(data.shape, data.dtype)
print(data[:1000])
之前看到的 1000 个字符在 GPT 中会呈现为这样：
可以看到 0 就是空格，1 就是换行符
目前为止，整个数据集就被重新表示为了一个庞大的整数序列。
训练、验证集划分#
# 将数据分成训练集和验证集，检查模型的过拟合程度
n = int(0.9*len(data)) # 前90%将用于训练，其余部分为验证
train_data = data[:n]
val_data = data[n:]
我们不希望模型完美记住莎翁的作品，而是让它创造模仿莎翁风格的文本。
chunks & batches#
要注意的是，我们不会把整个文本一次输入 Transformer 中，而是使用数据集的 chunks，也就是从训练集中随机抽取小块样本。
分块处理#
Block Size 被用来指定模型训练时每个输入数据块（如文本片段）的固定长度。
block_size = 8
train_data[:block_size+1]
x = train_data[:block_size]
y = train_data[1:block_size+1]
for t in range(block_size):
context = x[:t+1]
target = y[t]
print(f"when input is {context} the target: {target}")
这实际上是一种逐步揭露上下文信息给模型的策略
这个方法能强迫模型学会基于先前的字符（或 token）来预测序列中的下一个字符（或 token），提高推理能力。
分批处理#
为了提高 GPU 擅长的并行运算的效率，我们还要考虑分批训练，将多个批次的文本块堆叠在一个 tensor 中，同时处理多个之间相互独立的数据块。
batch size 的含义就是我们的 Transformer 每一次 forward & backward pass 需要处理多少个独立序列。
torch.manual_seed(1337) # 提供采样和可复现能力
batch_size = 4 # 并行处理独立序列的个数
block_size = 8 # 预测的最大上下文长度
# 这里的作用类似torch中的dataloader
def get_batch(split):
# 生成一小批输入数据 x 和目标数据 y
data = train_data if split == 'train' else val_data
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1:i+block_size+1] for i in ix])
return x, y
xb, yb = get_batch('train')
torch.stack用于沿着新的维度对一系列张量进行堆叠（stack），所有的张量需要具有相同的形状。
可以看到 inputs 的形状是 4x8，每一列都是训练集中的一部分；而 targets 的作用是在模型最后计算损失函数。
for b in range(batch_size): # 批次维度
for t in range(block_size): # 时间维度
context = xb[b, :t+1]
target = yb[b,t]
这样可以更清楚了解 inputs & outpus 两个数组的关系
Bigram#
在 Makemore 系列中，我们深入了解并实现了 bigram 语言模型，现在改用 PyTorch Module 来快速重新实现。
模型搭建#
import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)
class BigramLanguageModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
# 每个token直接从lookup table中读取下一个token的对数
self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
这个 embedding layer 想必也很熟悉，比如输入 24 那就是到 embedding table 中取出第 24 行。
def forward(self, idx, targets=None):
# idx和targets都是整数的（B，T）张量
logits = self.token_embedding_table(idx) # (Batch=4,Time=8,Channel=65)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
在 Makemore 系列中，我们知道了衡量损失的一个好方法：负对数似然损失，在 PyTorch 中对应的实现是 “cross-entropy”（交叉熵）。直观来讲，就是模型对于 logit 对应的正确分类应该有一个很高的概率（信心高），同时其他所有维度都是很低的概率（信心极低）。此时的损失是可以估计的，大概是 -log (1/65)，约等于 4.17，但因为一些熵的存在实际结果会更大一些。
# 从模型中生成
def generate(self, idx, max_new_tokens):
# idx 是当前上下文中索引的 (B, T) 数组
for _ in range(max_new_tokens):
# 获取预测结果
logits, loss = self(idx)
# 仅关注最后一个time step
logits = logits[:, -1, :] # becomes (B, C)
# 应用softmax以获取概率
probs = F.softmax(logits, dim=-1) # (B, C)
# 从分布中抽样
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# 将采样索引添加到运行序列中
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
print(loss)
print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist())) # 序列的第一个字符为换行符(0)
generate 的任务就是把大小为 BxT 的代表上下文信息的 idx 扩展为 $B\times T + 1, + 2 ,+\ldots$，也就是在时间维度中的所有批次维度上继续生成 。
模型未训练时的生成结果，完全是随机的
模型训练#
下面开始训练这个模型。相较于 Makemore 系列使用随机梯度下降 (SGD)，这里我们使用的是更先进和流行的 AdamW 优化器，
# 创建一个 PyTorch 优化器
optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)
优化器基本作用就是获取梯度，并根据梯度更新参数。
batch_size = 32 # 选择更大的 batch-size
for steps in range(100): # 增加步数以获得良好结果
# 采样一个 batch 的数据
xb, yb = get_batch('train')
# 评估损失
logits, loss = m(xb, yb)
optimizer.zero_grad(set_to_none=True) # 将上一步的梯度清零
loss.backward() # 反向传播
optimizer.step() # 相当于“新参数=旧参数-学习率*梯度”，跟我们以前手动实现的梯度下降循环基本一样
print(loss.item())
可以看到我们的优化起作用了，损失在减小
增加训练轮数，最终达到了约 2.48，我们复制之前的采样代码片段再次生成，应该能得到有进步的结果。
初具人形，但又没那么具。
再怎么训练也很难得到理想的结构，因为这个模型假设十分简单（只根据前一个预测后一个 token），各个 token 之间没有联系，这就是使用 transformer 的原因。
Transformer#
如果有 Nvidia 显卡的话可以加速训练：
device = 'cuda' if torch.cuda.is_available() else 'cpu'
这么设定后需要对代码进行一些改动，大概就是要让数据加载、计算和采样生成都在 device (GPU) 上进行，具体细节请见 Andrej 的 lecture 仓库，其中的 bigram.py 是我们的起点。
除此之外，我们的模型分为了训练阶段和评估阶段，不过目前模型中只有一个nn.Embedding 层而两个阶段表现一样，并没有引入 dropout layer、batch norm layer 等。这样做是训练模型中的最佳实践 ( best practice )，因为一些层在训练和推理阶段的行为不同。
Self-Attention#
自注意力机制
在上手 Transformer 之前，我们要做的第一件事是通过简单例子来习惯一个实现 Transformer 内部自注意力核心的数学技巧。
torch.manual_seed(1337)
B,T,C = 4,8,2 # batch, time, channels
x = torch.randn(B,T,C)
# x.shape = torch.Size([4, 8, 2])
我们希望通过一个特定的方式将这些本来独立的 token 结合起来，比如说，第五个 token 不应该能与第 6、7 和 8 号 token 交互通信，因为在这个序列中，这三个属于未来的future tokens。因此，5 号 token 只能与第 4、3、2 和 1 号 token 交互通信，也就是说信息只能从之前的上下文流向当前的 time step，根据这些信息来预测未来信息。
那 token 之间相互通信的最简单的方式是什么呢？
答案很令人意外：对前面的 token 取平均，变为当前背景下的历史特征向量。当然，也很容易猜到这种交互方式太弱了，丢失了大量关于这些 token 空间排列的信息。
v1. 循环#
# 我们希望有 x[b,t] = mean_{i<=t} x[b,i]
xbow = torch.zeros((B,T,C))
for b in range(B):
for t in range(T):
xprev = x[b,:t+1] # (t,c)
xbow[b,t] = torch.mean(xprev, 0)
打印出来就能理解这段代码的作用了
这个方法运算效率较低，我们可以使用矩阵乘法来更高效地完成这个工作。
v2. 矩阵乘法#
# 一个简化示例，说明矩阵乘法如何用于“加权聚合”。
torch.manual_seed(42)
a = torch.ones(3, 3)
b = torch.randint(0,10,(3,2)).float()
c = a @ b
这里就是基本的矩阵乘法运算，c 中的每个数是 a 和 b 中分别对应的行与列的点积。如 c 的 (1, 1) 元素的就是 。
为了实现一样的效果，可以将 a 换为下三角矩阵：
这样可以实现 “分别抽取第一、二行加和” 的效果，通过torch.trill实现
a = torch.tril(torch.ones(3, 3))
现在的作用是加和，因为 a 中的元素都为 1，为了实现加权聚合，可以对 A 的每行进行归一化，使每行元素和为 1：
a = torch.tril(torch.ones(3, 3))
a = a / torch.sum(a, 1, keepdim=True) # keepdim保证广播语义可行
现在 a 的每行和为 1，c 就是 b 中的对应前几行的平均了
回到前面来应用这个更有效的方法：
# 版本2：使用矩阵乘法进行加权聚合
wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)
weight 权重矩阵对应上面的 a 矩阵
# 这里torch会给weight创造一个batch维度
xbow2 = wei @ x # (T, T) @ (B, T, C) ----> (B, T, C)
torch.allclose(xbow, xbow2) # 比较两个张量是否在一定的数值容差范围内相等
# 输出 True,两种方法效果一样
总结一下这个窍门：我们可以用批矩阵乘法来实现加权聚合，权重在这个 T✖️T 矩阵中指定，而加权和是根据维度和权重呈倒三角分布，使得 t 维上的 token 只能从先前的 token 中获取信息。
v3. Softmax#
除此之外，可以使用 Softmax 来实现第三个版本。
其中一个重要的 api 是 torch.masked_fill()，根据指定的掩码 ( mask ) tensor 对输入 tensor 进行填充，如下图所示：
那对于每一行都取 Softmax 会发生什么呢？在前面的章节中提过，Softmax 是一个归一化操作，这里的作用就是对每行中 “过去的” 元素使用下三角矩阵乘法进行加权聚合：
# 版本3：使用Softmax
tril = torch.tril(torch.ones(T, T))
wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
xbow3 = wei @ x
除了编码 token 的身份，还会对 token 的位置进行编码：
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每个token直接从查找表中读取下一个token的logits
self.token_embedding_table = nn.Embedding(vocab_size, n_embd) # token编码
self.position_embedding_table = nn.Embedding(block_size, n_embd) # 位置编码
def forward(self, idx, targets=None):
B, T = idx.shape
# idx和targets都是整数的(B,T)张量
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
这里用 x 保存 token 嵌入和位置嵌入的加和，但由于目前只是一个简单的 bigram 模型，不同位置是具有平移不变性的，但当我们使用注意力机制时就不一样了。
v4. self-attention#
我们当前所做的只是简单的平均，但现实中，每个 token 的意义都不是不同的，也就是数据相关的。比如对于一个元音，它就想知道在自己前面传递信息的辅音们是什么。这就是自注意力机制解决的问题。
每个 token 都将发出两个向量，一个 query（查询：我要找什么，对啥感兴趣）和一个key（键：我包含什么信息，和谁相似）。
我们得到序列中这些 token 之间亲和力（权重）的方法，基本上就是在键和查询之间做点积。如果一个键与查询非常匹配或对齐，这个键对应的权重会很高。因此，模型的注意力会集中在这个键对应的值上 —— 即模型会更加关注这个特定的信息（或 token ），而不是序列中的其他任何信息。
也就是说，注意力机制通过计算查询和键之间的匹配度并据此分配权重，使模型能够聚焦于最相关的信息，从而提高了处理和理解序列数据的能力。
我们还需要一个 Value（值：你对我感兴趣的话，我会贡献给你的信息），最终聚合的是经过一个线性层传播的 x 而不是直接聚合 x。
现在我们来实现这个单头注意力机制:
# 版本 4: self-attention!
torch.manual_seed(1337)
B,T,C = 4,8,32 # batch, time, channels
x = torch.randn(B,T,C)
# 让我们看看单头自注意力是如何运作的
head_size = 16
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)
# 现在key和query对x进行前向传播
k = key(x) # (B, T, 16)
q = query(x) # (B, T, 16)
# 点积得到亲和力（权重）
wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)
tril = torch.tril(torch.ones(T, T))
# wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
v = value(x)
out = wei @ v
# out = wei @ x
可以看到权重矩阵 wei 中，并非像以前一样完全平均了，而是数据依赖的：亲和力高的 token 会在加权聚合中为当前 token 提供更多信息
注意力总结#
-
注意力是一种通信机制。可以看作有向图中的节点相互观察，并通过指向它们的所有节点的加权和来聚合信息，权重是数据相关的。
-
本身并没有空间概念。注意力只是作用在一组向量上。这就是为什么我们需要对标记进行位置编码。
-
每个批次维度上的示例完全独立处理，永远不会互相通信。
-
注意力机制并不在乎你是否只关心过去的信息。我们这里的实现中，当前 token 对于未来的信息是屏蔽的，但只要在 “编码器” 注意力模块中删除使用
masked_fill进行掩码的代码就能允许所有 token 互相通信。这里称为 “解码器” 模块，因为它具有三角矩阵掩码。注意力机制支持任意节点之间的连接。 -
"自注意力" 意味着键和值是从与查询相同的源（x）产生的。在 " 交叉注意力（cross-attention）" 中，查询仍然由 x 生成，但键和值来自其他某些外部来源（例如一个编码器模块）
-
"Scaled" Attention：额外将
wei除以 $\sqrt {head_size}$ 。
原因是，我们当前只有标准高斯分布 （均值为 0，方差为 1）输入时，会发现简单的加权得到的 wei 的方差其实是 head size 的数量级（我们的实现中正是 16），
加上这个 normalization 后，权重的方差就是 1 了：
为啥这一步很重要呢？
# 我们的 wei 会经过 softmax
wei = F.softmax(wei, dim=-1)
对于 softmax 来说，一个性质是经过 softmax 后分布中绝对值较大的元素会向 1 趋近：
可以看到，同样的分布在乘上 8 后，变得趋于 one-hot 分布了（one-hot：一个元素是 1 其它全 0），即在初始阶段导致分布过于尖峭，基本上只是从单个节点获取信息。
因此，“缩放” 注意力机制说白了就是控制初始化时的方差，通过将 wei 除以 1/√(头大小) 来进行额外调整。这样，当输入 Q（查询）和 K（键）的方差为 1 时，wei也将保持单位方差，这意味着 Softmax 将保持分散，而不会过度饱和。使得在应用 Softmax 函数之前，权重分布不会因为过大的值而导致梯度消失或梯度爆炸，进而保持了模型的稳定性和效果。
代码实现#
class Head(nn.Module):
""" one head of self-attention """
def __init__(self, head_size):
super().__init__() # 人们一般不在这里用 bias
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)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B,T,C = x.shape
k = self.key(x) # (B,T,C)
q = self.query(x) # (B,T,C)
# 计算注意力分数 ("affinities")，这里应用了上面提到的 scaled
wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
# 执行加权聚合
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
构造函数中， tril 并非 nn.Module 的参数，这在 PyTorch 命名约定中被称为缓冲区（buffer），要调用的话必须使用 register_buffer来将它赋值给 nn.Module，
Multi-Head Attention#
论文中还有这部分没有复现，多头注意力就是并行地应用多个注意力，并将它们的结果连接起来：
代码实现#
在 PyTorch 中，我们可以简单地通过创建多个 head 来做到这一点。
class MultiHeadAttention(nn.Module):
""" 自注意力中的多头并行。 """
def __init__(self, num_heads, head_size):
super().__init__()
# 在一个列表中并行运行，然后将输出连接起来
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
def forward(self, x):
return torch.cat([h(x) for h in self.heads], dim=-1) # 在 chanmel 维度上连接
现在我们有了 4 个并行的通信 channel 而不是一个，单个通道会相应地变小。嵌入维度是 32，对应的就有 8 维度的 self-attention，连接起来又得到 32 个，这就是原始嵌入。这里有点类似于 群卷积 ( group convolution )，不进行一个大卷积而是分组卷积。
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每个 token 直接从查找表中读取下一个 token 的 logit
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.sa_head = nn.MultiHeadAttention(4, n_embd//4)
self.lm_head = nn.Linear(n_embd, vocab_size)
Blocks#
对于下图，也就是论文所展示的网络结构，我们不会实现对编码器的交叉注意力。但图中还有一个前馈部分，它被分组成一个 block，一个不断重复（$N\times$）的块。
前馈网络#
这个前馈部分只是一个简单的 MLP：
注意论文这里说输入和输出的维数是 512，前馈的内层维数是 2048，所以前馈的内层通道大小应该乘以 4
class FeedFoward(nn.Module):
""" 一个简单的线性层，后面跟着一个非线性函数 """
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
""" Transformer block: 分散了通信和计算
通信：多头注意力
计算：前馈网络在所有 token 上独立完成
"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我们想要的 head 个数，这里是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
#
self.ffwd = FeedFoward(n_embd) # 计算
def forward(self, x):
x = self.sa(x)
x = self.ffwd(x)
return x
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每个 token 直接从查找表中读取下一个 token 的 logit
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=4),
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
)
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx 和 targets 都是 (B,T) 大小的整数 tensor
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
...
尝试解码（decode），发现结果并没有改善多少。原因是我们现在得到了一个相当深的神经网络，会受到优化问题影响，我们还需要从 Transformer 的论文中借鉴一个解决这个问题的方法。
现在有两种能极大地提高网络深度的同时，确保网络保持可优化状态的方法：
残差连接#
block 中红色圈起的部分（箭头和 Add）就是残差连接（residual connection），这个概念由 Deep Residual Learning for Image Recognition论文提出。
Andrej 的原话是 “你转换数据，然后与先前特征进行跳连接并相加”，我们这里详细解释一下（可以配合图片食用）：
-
你转换数据（You transform data）：深度神经网络中的每一层，输入数据会通过权重矩阵的乘法和激活函数的非线性变换等操作进行 “转换”，从而学习到数据的抽象表示。
-
但然后你有一个跳跃连接（But then you have a skip connection）：在传统的深度网络中，数据的这种转换是连续且线性的。残余连接通过引入 “跳过连接” 打破了这种模式。跳过连接直接将某一层的输入连接到后面的某一层（通常是相隔一层或几层的输出），这样做的目的是为了将前面层的特征直接传递给后面的层。
-
与之前的特征连接并相加（With addition from the previous features）：跳跃连接的实现通常是通过将跳过的输入和目标层的输出进行元素级的加法操作来完成的。这种加法操作保证了原始特征信息能够在网络中直接传递，而不会被后续层的变换所 “稀释”。
残余连接的引入使得网络能够更容易地学习到恒等映射（identity mapping），这对于训练非常深的网络是非常有用的。实际上，残余连接允许梯度直接流过网络，有助于缓解梯度消失或梯度爆炸的问题，从而使得训练深层网络变得更加可行和高效。
我们在 micrograd 中提到过，神经网络中加法节点的作用就是将梯度均匀分配给所有输入（因为加法操作对于每个输入来说是线性的，并且每个输入对输出的贡献是独立的） 看到这里的时候不禁感叹：鲜活的知识点居然又串起来了。
class MultiHeadAttention(nn.Module):
""" 自注意力中的多头并行。 """
def __init__(self, num_heads, head_size):
super().__init__()
# 在一个列表中并行运行，然后将输出连接起来
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(n_embd, n_embd) # 引入投射
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1) # 在 chanmel 维度上连接
out = self.proj(out) # 残差路径的投影就是 out 的线性变换
return out
class FeedFoward(nn.Module):
""" 一个简单的线性层，后面跟着一个非线性函数 """
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(n_embd, 4 * n_embd), # 投射回残差通路
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
""" Transformer block: 分散了通信和计算"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我们想要的 head 个数，这里是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
self.ffwd = FeedFoward(n_embd) # 计算
def forward(self, x):
x = x + self.sa(x)
x = x + self.ffwd(x)
return x
Layer Norm#
第二个优化方法是这里的 Norm，指的是一种叫做层归一化（Layer Norm）的东西：
Layer Norm 和我们之前实现过的 Batch Norm （确保在 batch 维度上，任意神经元都是标准正态分布）非常相似，唯一与 Batch Norm 不同的是，Layer Norm 不是在批次维度上进行归一化，而是在特征维度上进行归一化。这意味着，对于网络中的每个样本，Layer Norm 会计算该样本所有特征的均值和标准差，并用这些统计量来归一化该样本的所有特征。
这里需要注意的是，Transformer 论文发布到现在内部的细节并没有太多变动，但我们这里的实现和原论文略有出入，可以看到论文中 Add&Nrom 是在 Transform 后添加的，但现在更普遍的做法是在 Transform 之前应用 Layer Norm，这被称为 Pre-norm 公式。
现在我们已经拥有了一个相当完整的 Transformer 了（只有解码器 decoder）
class Block(nn.Module):
""" Transformer block: 分散了通信和计算
通信：多头注意力
计算：前馈网络在所有 token 上独立完成
"""
def __init__(self, n_embd, n_head):
# n_embd: embedding dimension, n_head: 我们想要的 head 个数，这里是 8
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size) # 通信
#
self.ffwd = FeedFoward(n_embd) # 计算
self.ln1 = nn.LayerNorm(embd)
self.ln2 = nn.LayerNorm(embd)
def forward(self, x):
x = self.sa(self.ln1(x))
x = self.ffwd(self.ln2(x))
return x
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# 每个 token 直接从查找表中读取下一个 token 的 logit
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=4),
Block(n_embd, n_head=4),
Block(n_embd, n_head=4),
nn.LayerNorm(n_embd),
)
'''
self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)]) # 扩大模型，指定 block 层数
self.ln_f = nn.LayerNorm(n_embd) # 最终的 Layer Norm
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx 和 targets 都是 (B,T) 大小的整数 tensor
tok_emb = self.token_embedding_table(idx) # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb # (B,T,C)
x = self.blocks(x) # (B,T,C)
x = self.ln_f(x) # (B,T,C)
logits = self.lm_head(x) # (B,T,vocab_size)
Dropout#
Dropout 是可以在残差连接回到残差路径之前添加的东西，由Dropout: A Simple Way to Prevent Neural Networks from Overfitting论文提出，它基本上让你的神经网络在每次 forward-backward pass 随机关闭一些神经元（也就是降为 0，不参加后面的训练）
这里只需要知道它是一种正则化技术就可以了
# 超参数（可以用 Colab V100 跑，CPU的话太慢了，或者调低超参数）
batch_size = 64
block_size = 256 # 增加上下文长度，预测第 257 个 token
max_iters = 5000
eval_interval = 100
learning_rate = 1e-3 # 网络变大，降低学习率
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6 # 384 / 6
n_layer = 4 # 4 层 Block
dropout = 0.2 # 每次 forward-backward pass 都有 20% 概率丢弃神经元
class Head(nn.Module):
""" one head of self-attention """
def __init__(self, head_size):
super().__init__() # 人们一般不在这里用 bias
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)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x):
B,T,C = x.shape
k = self.key(x) # (B,T,C)
q = self.query(x) # (B,T,C)
# 计算注意力分数 ("affinities")，这里应用了上面提到的 scaled
wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
wei = F.softmax(wei, dim=-1) # (B, T, T)
wei = self.dropout(wei) # Dropout
# 执行加权聚合
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
class MultiHeadAttention(nn.Module):
def __init__(self, num_heads, head_size):
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) # Dropout
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
class FeedFoward(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout), # Dropout
)
def forward(self, x):
return self.net(x)
现在的结果就像是莎士比亚风格的胡言乱语，不过显然很可观了。
图中左半边所示的编码器（encoder）和右边红圈中的交叉注意力并没有在这里实现。
我们只用了解码器的原因是，我们只是生成不受任何条件限制的文本，就像最后的结果一样，只是根据给定的莎士比亚数据集胡言乱语，我们通过三角掩码来实现注意力机制，使其具有自回归的性质便于抽样进行语言模型建模。
而原始论文采用 encoder-decoder 的结构是因为它是在机器翻译领域的，模型期待的输入是外文编码 token（如法语），然后解码翻译为英语，如下图所示：
这里编码器就是取感兴趣的法语句子来创建 token ， 在上面使用 Transformer 结构但不用三角掩码，让所有的 tokens 都尽可能多地相互通信。负责语言建模的解码器会连接编码完成后的输出（上面架构图中的左半部分的最上面），这是通过交叉注意力完成的。实际上所做的就是对解码进行制约，不仅仅是在当前解码过去信息，更是在完整的、完全编码的法语 tokens 上进行。而我们实现的是 decoder-only 的版本。
本节知识对应的项目：karpathy/nanoGPT，同样是只关注预训练部分的实现。
重回 ChatGPT#
要训练 ChatGPT 大致有两个阶段：预训练和微调阶段。
预训练#
在大量互联网语料上训练，试图得到一个 encoder-only 的 Transformer。我们现在已经完成了一个很小的预训练步骤。还有一点不同是， OpenAI 的训练使用的是 tokenizer，也就是词汇表不是单个字符而是字符块，我们使用的莎士比亚数据集对应的大概会有 30 万 token，我们在上面训练了约 1000 万个参数，而 GPT3 的 Transformer 最多拥有 1750 亿个参数，在 3000 亿个 token 上进行了训练
在完成这一步后，你并不能向模型提问，因为它目前只会创建互联网上的新闻等信息，也就是只有补全序列的作用。
微调#
这个阶段就是把它调教为一个语言模型助手。
第一步是收集成千上万的文档，其格式为 “问题：答案”，对模型进行微调对齐（align），在对大模型中微调过程中的样本效率是很高的。
第二步，评分者根据模型的回应进行排名，用此来训练一个奖励模型（reward model）。
第三步，运行 PPO（ Proximal Policy Optimization，一种策略梯度强化学习优化器）来微调抽样策略，把模型从一个文档补全器变成了一个问题回答器。
当然，这些部分个人是基本无法复刻的，大公司才行。
关于 GPT 的详细讲述，Andrej 在 2023 年 3 月份的 Microsoft Build 演讲中进行了全面的讲述，详见GPT 的现状。