最笨的模型：先让训练循环转起来

这一篇不求聪明，只求循环转起来

第一篇说过，训练是个循环：喂真实文字、看模型猜错在哪、按错误方向拧旋钮，反复上万遍。前两篇我们把数据准备好了。这一篇要做的，是把这个循环第一次真正搭起来、让它转动。

为了把注意力集中在"循环"本身，这一篇的模型会故意做得非常笨——它只看前一个字来猜下一个字。这种模型叫 Bigram（二元）模型。它学不到什么真本事，但没关系。这一篇的目标不是聪明，是让"喂数据、算损失、拧旋钮"这套机器先转起来。聪明，是第五篇之后的事。

一次喂多少：批次和上下文长度

训练时，我们不会把整份语料一次性塞给模型，那既装不下也没必要。每一步只喂一小块。这就牵出两个数字。

一个是上下文长度 block_size：模型一次最多看多长的文字。我们先设 8，表示模型每次根据最多 8 个字来预测。

另一个是批次大小 batch_size：为了效率，每一步同时喂好几小块，让模型并行处理。设 32，表示一次喂 32 小块。

关键在于，每一小块要配一个"正确答案"。我们从语料里随机抠出一段长 block_size 的文字作为输入 x，再把这段文字整体往后挪一个字作为答案 y。这样一来，x 的每一个位置，它对应的"下一个字"正好就是 y 同一位置上的字。

import torch

torch.manual_seed(1337)        # 固定随机种子，让结果可复现

block_size = 8
batch_size = 32


def get_batch(split: str):
    """随机抠出 batch_size 个小块，返回输入 x 和答案 y"""
    data = train_data if split == "train" else val_data
    # 随机选 batch_size 个起始位置
    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")
print(f"输入 x 的形状：{xb.shape}")      # (32, 8)
print(f"答案 y 的形状：{yb.shape}")      # (32, 8)
print(f"第一小块 x：{xb[0].tolist()}")
print(f"第一小块 y：{yb[0].tolist()}")

对照打印出来的第一小块，你会发现 y 就是 x 错开一位的结果。x 的第 0 个号码后面应该接什么？答案是 y 的第 0 个号码。x 的第 1 个号码后面接什么？看 y 的第 1 个。一小块长度 8，里面其实藏了 8 道"预测下一个字"的题。

最笨的模型：Bigram

现在搭模型。在 PyTorch 里，模型是一个继承 nn.Module 的类。Bigram 模型只需要一样东西：一张表。

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


class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size: int):
        super().__init__()
        # 一张 vocab_size × vocab_size 的表，这就是模型的全部"旋钮"
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # idx 形状 (B, T)，每个元素是一个字的号码
        logits = self.token_embedding_table(idx)   # 查表，得到 (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

nn.Embedding(vocab_size, vocab_size) 是一张 vocab_size 行、vocab_size 列的表。每个字的号码对应其中一行，这一行有 vocab_size 个数，代表"看到这个字之后，下一个字是词表里各个字的可能性打分"。

这就是它笨的地方：预测只查了当前这一个字对应的那一行，完全没看再往前的内容。"今天天气真"这五个字，它只拿最后一个"真"去查表，前面四个字直接忽略。这台机器只有一个字的记忆。

代码里 B, T, C 是三个维度：B 是批次大小（batch，32），T 是时间步也就是序列长度（time，8），C 是通道数也就是每个位置的打分个数（channel，等于 vocab_size）。这套 (B, T, C) 记法后面每一篇都会用，先混个眼熟。

模型怎么"打分":logits、softmax 与损失

forward 里出现了三个新词，得讲清楚，它们是训练的核心。

第一个是 logits。模型对"下一个字"的预测，不是直接报一个字，而是给词表里每个字打一个分。这一串原始分数就叫 logits。分数有高有低，可正可负。

第二个是 softmax。logits 是原始分数，不直观。softmax 是一个函数，它把这一串分数换算成一组概率——全都在 0 到 1 之间、加起来正好等于 1。这样我们就能说"下一个字是'好'的概率是 0.42"。

第三个是损失（loss）。我们需要一个数字来衡量"模型这次猜得有多差"。这个数字就是损失。这里用的损失函数叫交叉熵（cross entropy），它的逻辑很朴素：看模型给"正确答案那个字"打出的概率有多高。概率给得高，损失就低;概率给得低，损失就高。模型把正确答案的概率压到接近 0，损失会冲到很大。

整个训练，就是想方设法让这个损失数字变小。损失小，等于模型给正确答案的概率高，等于它猜得准。

这里有个能拿来对账的细节。训练刚开始，旋钮是随机的，模型对词表里每个字几乎一视同仁地瞎猜，给正确答案的概率大约是 1/vocab_size。此时交叉熵损失约等于 ln(vocab_size)。假设你的 vocab_size 是 3000，那一开始的损失应该在 8 上下（ln(3000) ≈ 8）。等会儿跑起来，如果初始损失和这个数对得上，说明一切正常。

让它生成文字

模型还要能写字。给 Bigram 加一个 generate 方法——就是第一篇说的"接龙":预测下一个字，接上，再拿去预测，循环。

    def generate(self, idx, max_new_tokens: int):
        # idx 是当前已有的文字（号码形式），形状 (B, T)
        for _ in range(max_new_tokens):
            logits, _ = self(idx)              # 前向，拿到打分
            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

注意 torch.multinomial 这一步：它不是死板地选概率最高的字，而是按概率随机抽。概率 0.42 的字有 42% 的机会被抽中。这点随机性让模型每次生成的内容不完全一样，也更自然。

训练循环：核心就五行

万事俱备。训练循环本身，核心其实只有五行：

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

for step in range(5000):
    xb, yb = get_batch("train")              # 1. 取一批数据
    logits, loss = model(xb, yb)             # 2. 前向：算预测、算损失
    optimizer.zero_grad(set_to_none=True)    # 3. 清空上一轮的梯度
    loss.backward()                          # 4. 反向：算出每个旋钮该往哪拧
    optimizer.step()                         # 5. 真正拧动旋钮

    if step % 500 == 0:
        print(f"step {step}: loss {loss.item():.4f}")

逐行说。第 1 行取一批数据。第 2 行把数据喂进模型，得到打分和损失——这一步叫前向传播。第 3 行先把上一轮残留的梯度清零（为什么要清，下一篇讲）。第 4 行 loss.backward() 是关键，它会算出"每一个旋钮，应该往哪个方向、拧多少",这个过程叫反向传播。第 5 行 optimizer.step() 按算出来的方向，真正把所有旋钮拧动一次。

optimizer 是优化器，负责"拧"这个动作。lr=1e-3 是学习率，控制每次拧多大幅度。loss.backward() 和优化器内部到底怎么工作的，正是下一篇——反向传播——的全部内容。这一篇你先把它当成"按错误方向拧旋钮"的黑盒，让循环转起来。

跑起来看看

把前两篇的数据准备代码、加上这一篇的全部代码，存成一个文件跑起来。你会看到损失打印出来，从一开始的 8 左右，一步步往下掉，降到 3 到 4 之间稳住。

损失在下降，就证明那套机器真的转起来了——模型在被一步步拧得"猜得更准"。

训练完，让它生成一段看看：

context = torch.zeros((1, 1), dtype=torch.long)   # 从号码 0 这个字起头
generated = model.generate(context, max_new_tokens=300)[0].tolist()
print(decode(generated))

别期待太高。Bigram 生成的东西，还是基本读不通的。但仔细看，它和训练前纯粹的乱码已经不同了——常见的字出现得多了，标点的位置不再完全离谱，单看相邻两个字偶尔还算搭。这正是它的能力上限：它只有一个字的记忆，能学到的也就是"哪个字后面常跟哪个字"这种最表层的规律。

但这不重要。重要的是：喂数据、算损失、反向传播、拧旋钮——这套循环，已经完整地、正确地转起来了。后面几篇要做的，全部是把模型从"只看一个字"升级成"看得见整段上下文",而这套训练循环的骨架，几乎一行都不用再改。

本篇要点

• 训练时数据分小块喂入，block_size 是一次看多长，batch_size 是一次喂几块。
• 输入 x 和答案 y 的关系是"整体错开一位",一小块里藏着多道"预测下一个字"的题。
• Bigram 模型只用一张 vocab_size × vocab_size 的表，只看前一个字，所以很笨。
• logits 是模型对每个字的原始打分，softmax 把它换算成概率，交叉熵损失衡量"猜得有多差"。
• 训练循环核心五行：取数据、前向算损失、清梯度、反向传播、拧旋钮。
• Bigram 生成的文字仍读不通，但训练循环已完整转起来，这就是这一篇的目标。

下一篇

这一篇里，loss.backward() 和 optimizer.step() 是当黑盒用的。但它们恰恰是"模型怎么知道该往哪改"的答案所在。下一篇专门拆开这个黑盒，讲反向传播——整个训练的心脏。

参考资料

• PyTorch nn.Module 教程
• 交叉熵损失（PyTorch 文档）
• Andrej Karpathy: Let's build GPT