语言模型的工作原理一直是人们好奇的焦点。本文将深入探讨如何使用PyTorch从零开始构建一个简易版的GPT模型（Mini-GPT），剖析其核心机制。灵感来源于nanoGPT项目，并参考了“Attention is All You Need”论文中的注意力机制，我们将一步步地搭建起这个小型大模型，一窥其内部运作的奥秘。
数据集的准备：模型之基
构建任何大模型的第一步都是准备高质量的数据集。数据是模型学习的基石，如同植物生长的种子。在这个Mini-GPT项目中，我们使用莎士比亚作品集作为训练数据。这个未经处理的文本文件包含了莎士比亚的所有作品，可以从GitHub上获取。
# Load the text
with open('input.txt', 'r', encoding='utf-8') as f:
text = f.read()
模型需求分析：预测下一个Token
我们的目标是构建一个语言模型，能够根据前文的上下文预测下一个Token。为了实现这个目标，我们需要：
- 分词器 (Tokenizer)：将文本转换为模型可以理解的数字表示。
- 上下文窗口 (Context Window / block_size)：模型一次处理的Token序列长度。
- 批处理逻辑 (Batching Logic)：高效地加载和处理训练数据。
- 嵌入层 (Embedding Layers)：将Token转换为向量表示。
- Transformer架构 (Transformer Architecture)：模型的核心结构，使用注意力机制学习Token之间的关系。
- 输出层 (Output Layer)：预测下一个Token的概率分布。
- 损失函数 (Loss Function)：衡量模型预测的准确性。
- 优化器 (Optimizer)：调整模型参数以最小化损失。
- 采样逻辑 (Sampling Logic / Generation)：生成新的文本序列。
- 训练循环 (Training Loop)：迭代训练模型。
超参数和分词器：模型的微调
超参数是控制模型训练过程的关键参数。以下是我们Mini-GPT模型中使用的一些超参数：
batch_size = 16 # 每个训练步骤并行处理的序列数量
block_size = 32 # 模型预测的最大上下文长度
max_iters = 5000 # 模型训练的迭代次数
eval_interval = 100 # 每100步检查验证损失
learning_rate = 1e-3 # 学习率
device = 'cuda' if torch.cuda.is_available() else 'cpu' # 使用GPU或CPU
eval_iters = 200 # 验证期间平均的批次数
n_embd = 64 # Token向量表示的维度
n_head = 4 # 注意力头的数量
n_layer = 4 # Transformer层的数量
dropout = 0.0 # Dropout概率
batch_size: 决定了每次训练迭代中使用的样本数量。更大的batch_size可以提供更稳定的梯度，但也会增加VRAM的使用。block_size: 定义了模型在预测下一个Token时考虑的上下文长度。n_embd: 设置了每个Token的向量表示维度。更大的维度可以容纳更复杂的模式，但也需要更多的计算资源。dropout: 是一种正则化技术，可以防止过拟合。
分词器是将原始文本转换为模型可理解的数字形式的关键组件。在这个例子中，我们使用基于字符的分词器，将每个字符映射到一个唯一的整数ID。
# here are all the unique characters that occur in this text
chars = sorted(list(set(text)))
vocab_size = len(chars)
# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
# encoder: take a string, output a list of integers
encode = lambda s: [stoi[c] for c in s]
# decoder: take a list of integers, output a string
decode = lambda l: ''.join([itos[i] for i in l])
# Encode Full Dataset into a 1D tensor of integers
data = torch.tensor(encode(text), dtype=torch.long)
# first 90% will be train, rest val
n = int(0.9*len(data))
train_data = data[:n]
val_data = data[n:]
例如，如果我们的词汇表是['h', 'e', 'l', 'o']，那么"hello"将被编码为[0, 1, 2, 2, 3]。
批处理逻辑：数据的高效利用
批处理逻辑负责准备用于训练的Mini-batch数据，包括输入序列 x 和目标序列 y。模型通过学习 x 来预测 y，即“给定前n个Token，下一个Token是什么？”
def get_batch(split):
# generate a small batch of data of inputs x and targets 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])
x, y = x.to(device), y.to(device)
return x, y
get_batch 函数从训练或验证数据集中随机抽取一批数据。它首先随机选择 batch_size 个起始索引，然后从这些索引开始提取长度为 block_size 的序列作为输入 x。目标 y 是 x 的下一个Token序列。
Transformer核心：自注意力机制
自注意力机制是 Transformer 架构的核心，它允许模型学习Token之间的关系。在Mini-GPT模型中，我们使用单个注意力头来实现自注意力。
那么，为什么我们需要自注意力头？在神经网络和语言模型的背景下，什么是注意力？
注意力是一种机制，允许模型在进行预测时，决定输入中的哪些部分是重要的。举个例子，当模型生成句子“The cat sat on the mat”时，在预测单词 “mat”时，模型会学习关注：
- “sat” (因为动词和宾语相关联)
- “on” (符合语法规则)
- 甚至 “cat” (提供上下文信息)
注意力机制模拟了“在确定接下来应该是什么时，我应该关注这个句子的其他哪些部分？”的思考过程。
class Head(nn.Module):
""" one head of self-attention """
def __init__(self, head_size):
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)
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) # (B,T,C)
q = self.query(x) # (B,T,C)
# compute attention scores ("affinities")
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)
# perform the weighted aggregation of the values
v = self.value(x) # (B,T,C)
out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
return out
-
初始化：连接注意力神经元
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)
对于每个Token：
- query: 我想寻找什么？
- key: 我提供什么？
- value: 我携带什么？
register_buffer()函数创建一个下三角矩阵，以强制执行因果掩码。
-
生成Key和Query
k = self.key(x)q = self.query(x)
我们将输入投影到Key和Query空间。现在，每个Token都是一个Key和一个Query。
-
计算注意力得分矩阵
wei = q @ k.transpose(-2, -1) * C**-0.5
这行代码计算了注意力得分矩阵，表示每个Token应该关注其他Token的程度。它包括：
- Query和Key的点积。
- 除以
1/√C以提高训练稳定性（来自原始的“Attention Is All You Need”论文）。
-
Softmax 函数
wei = F.softmax(wei, dim=-1)
将注意力得分转换为概率（注意力权重）。
总而言之，这个注意力机制可以这样理解：将每个Token想象成一个好奇的聚会参与者。它四处张望，通过注意力得分来判断谁比较有趣，然后倾听他们的讲话（values），并将这些信息融入到自己的新状态中。
基于Transformer的因果语言模型
这是一个基于 Transformer 的因果语言模型，是 GPT 的一个精简版本。虽然名字是 BigramLanguageModel，但它不再只是一个 Bigram 模型。它是一个真正的 Transformer，能够从上下文中学习。
class BigramLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# each token directly reads off the logits for the next token from a lookup table
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=n_head) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd) # final layer norm
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# idx and targets are both (B,T) tensor of integers
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)
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
def generate(self, idx, max_new_tokens):
# idx is (B, T) array of indices in the current context
for _ in range(max_new_tokens):
# crop idx to the last block_size tokens
idx_cond = idx[:, -block_size:]
# get the predictions
logits, loss = self(idx_cond)
# focus only on the last time step
logits = logits[:, -1, :] # becomes (B, C)
# apply softmax to get probabilities
probs = F.softmax(logits, dim=-1) # (B, C)
# sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
# append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
return idx
model = BigramLanguageModel()
m = model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')
# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
-
剧院中的角色
self.token_embedding_table = ...self.position_embedding_table = ...
将文本中的每个 Token (字符) 想象成剧院中的一个演员。模型首先给每个演员一个角色 (通过词嵌入) 和一个舞台位置 (通过位置嵌入)，以便了解他们的身份以及他们在剧本中的位置。
-
每个演员关注其他人
self.blocks = nn.Sequential([...Block...])
在每个场景中，每个演员不仅仅是盲目地念他们的台词，他们还会环顾四周，看看他们之前的其他演员在做什么。这就是自注意力的作用：每个 Token 都关注之前的 Token，以决定如何行动。
-
对场景的理解是分层的
就像伟大的表演涉及细微差别和层次一样，该模型通过多个 Transformer 块（注意力机制 + MLP + 归一化），一遍又一遍地改进每个 Token 的角色，以提高他们对上下文的感知。
-
最后的麦克风检查 (LayerNorm)
self.ln_f = nn.LayerNorm(...)
在传递他们的最后一行台词之前，每个演员都会进行层归一化处理，以确保音调、表达和传递的一致性。
-
观众预测接下来会发生什么
self.lm_head = nn.Linear(...)
在所有这些戏剧和互动之后，最终的输出会传递给一个 logits 投影仪。这就像问观众，“根据你刚才所看到的，你认为剧本中的下一行台词是什么？”
-
通过猜测进行学习
该模型通过反复玩猜测游戏来接受训练：
“这是演员们所说的内容，现在预测下一个单词。”如果它猜错了，它会从损失函数 (交叉熵) 中得到一点“纠正”，并使用优化器 (AdamW) 来更新自己。
-
改进演出：生成文本
generate(...)
经过训练后，该模型可以即时地即兴创作新的场景：
它会倾听一个简短的提示，选择下一个单词，将其添加到场景中，然后重复这个过程，从而一次生成一个新的莎士比亚式独白，一次一个 Token。
如前所述，此模型仅预测下一个 Token。为了使其具有对话能力，需要进行称为监督微调的过程，我们将在下一部分中实现。
模型评估：困惑度的衡量
在训练大模型的过程中，我们需要定期评估模型的性能。常用的评估指标是困惑度 (Perplexity)。困惑度越低，模型预测的准确性越高。困惑度可以理解为模型预测下一个Token时的平均选择数量。例如，如果模型的困惑度是20，那么模型在预测下一个Token时，平均需要在20个Token中做出选择。
模型部署和应用
训练好的Mini-GPT模型可以用于生成各种文本内容，例如：
- 莎士比亚风格的文本
- 代码
- 故事
- 对话
此外，还可以将Mini-GPT模型集成到各种应用中，例如：
- 聊天机器人
- 文本摘要工具
- 代码生成器
结论：大模型之路
通过从零开始构建Mini-GPT，我们深入了解了大模型的核心原理，包括注意力机制、Transformer 架构以及训练流程。虽然Mini-GPT只是一个简易版的 GPT 模型，但它包含了构建更复杂的大模型所需的所有关键组件。希望本文能够帮助读者更好地理解大模型的工作原理，并为进一步探索大模型领域奠定基础。
虽然这个Mini-GPT已经能够生成莎士比亚风格的文本，但是为了让它能够进行更有意义的对话，还需要进行监督微调等后续步骤。 大模型的未来充满无限可能，让我们一起探索吧！