从零实现一个极简 GPT 系列开到第二篇说实话百感交集。上一篇我们还在纠结怎么把文本变成数字这一篇终于要接触真正的模型了。但别急着把 Transformer、多头注意力这些庞然大物搬进来——在动手写 attention 之前我强烈建议先实现一个结构简单到极点、但五脏俱全的 Bigram 语言模型也就是所谓的二元语法模型。它只用一个 Embedding 层没有任何上下文窗口却能把数据管线、损失函数、训练循环、采样生成这条主链路完整跑通一遍。等这条链路在你手里变得像呼吸一样自然后面再加注意力机制时你才会知道每一步改动到底在解决什么问题。我写这篇内容时的心态是假设你已经能用 PyTorch 写出一个基础的全连接网络对 tensor 的 shape 操作有一定直觉但没碰过任何语言模型。接下来所有代码你都可以直接复制到 Jupyter 或脚本里跑数据用任何纯文本都行。1. 为什么第一个模型选 Bigram一个不需要注意力的语言模型1.1 语言模型的最小定义语言模型的任务一句话就能说清楚给定一段已经出现的字符序列预测下一个字符是什么。形式化地讲我们希望建模一个条件概率分布P(next_char | previous_chars)。比如说你给模型看 I like cheese它应该给 I 后面的所有候选字符分配一个概率其中像空格、逗号这类在训练语料中高频出现在 I 后的字符概率就会偏大。那么这个任务的最小实现是什么如果我们彻底放弃previous_chars里面所有历史信息只保留最后一个字符模型就退化成只看当前字符来预测下一个字符。这就是 Bigram 名称的来源它等价于统计语料中每两个相邻字符的共现频率。比如语料里 t 后面跟着 h 的次数很多模型就会学到 t 之后输出 h 概率较高。在代码层面Bigram 模型甚至可以不用神经网络——手写一个二维计数矩阵就能完成。但为了后续平滑过渡到 GPT我们会用 PyTorch 的nn.Embedding来把它写成一个标准模型。这样它和最终 Transformer 模型的输入输出接口完全一致输入一个形状为(Batch, Time)的 token 序列输出一个形状为(Batch, Time, VocabSize)的 logits 序列。1.2 Bigram 的局限与它的不可替代价值你肯定会问这种只能看一个字符的模型能有什么用它连最基本的语法结构都学不会更不用说理解语义了。确实Bigram 生成的文本几乎就是随机字符排列最多在字符共现频率上显得像人话。但它的价值不在生成质量而在于它把整个训练链路压缩到了最小可复现的范围。我见过很多直接上手实现 Transformer 的同学一上来就被多头注意力的 shape 变换劝退。问题往往不是注意力机制本身难理解而是他们还没有一个能快速反馈的完整工具链数据加载有问题时会误以为是模型写错了损失下降不下去时也分不清是优化器问题还是模型设计问题。先用 Bigram 跑通一遍等于先给整个流程做了一次冒烟测试。后面无论增加什么模块数据部分和训练框架都不需要再动我们只需要替换模型内部的计算逻辑。这个价值在调试复杂模型时会体现得淋漓尽致。此外Bigram 也是我们之后验证梯度是否正确的好参照物。当 Transformer 版模型损失怎么也降不到 Bigram 以下时至少能证明底层数据管线没问题。2. 数据管线把一段文本变成模型能吃的 (输入, 目标) 对2.1 字符级 tokenizer 的实现我们需要一个能反复调用的文本处理器。我的做法和大多数入门教程一致使用字符级 tokenizer也就是把每个字符当成一个 token。这样做的好处是词汇表极小方便可视化而且不会引入 BPE 等分词算法的额外复杂度。对于 1MB 左右的英文语料字符集合通常只有 50~100 个足够我们研究模型行为。先把原始文本读进来然后构建两个映射表字符到整数的stoi整数到字符的itos。Python 代码一般长这样with open(input.txt, r, encodingutf-8) as f: text f.read() chars sorted(list(set(text))) vocab_size len(chars) print(f语料总字符数: {len(text)}去重后字符数: {vocab_size}) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} def encode(s): 字符串 - token 列表 return [stoi[c] for c in s] def decode(l): token 列表 - 字符串 return .join([itos[i] for i in l])这里有一个值得注意的细节用set(text)去重会丢失字符顺序信息但因为我们之后用sorted重新排序所以结果可复现。另一个细节是encode函数需要处理数据中遇到的每个字符一旦语料里出现未登录字符就会报 KeyError因此如果你换数据集务必确保chars是基于新文本重新生成的而不是复用旧词汇表。2.2 构造输入序列与目标序列核心是 shift 一位把整个语料编码成一个长整数张量后就该构造模型训练用的 (x, y) 对。语言模型的监督信号并不是显式标注好的而是从文本自身生成的对于任意位置的字符它的正确答案就是紧接着它的下一个字符。所以我们需要把原始序列和它的移位版本配对。假设原始 token 序列是[1, 2, 3, 4, 5]如果我们想让模型通过[1, 2, 3]预测[2, 3, 4]那么输入就是data[0:3]目标就是data[1:4]。在 PyTorch 里我们通常一次处理一批形状为(batch_size, block_size)的序列。data torch.tensor(encode(text), dtypetorch.long) # 分训练集和验证集 n int(0.9 * len(data)) train_data data[:n] val_data data[n:] batch_size 32 block_size 8 # 上下文长度Bigram 模型其实用不到但留着方便后续扩展 def get_batch(split): 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:iblock_size] for i in ix]) y torch.stack([data[i1:iblock_size1] for i in ix]) return x, y这段代码中的torch.randint是异步随机采样不用像有些教程那样手动 shuffle 整个语料。需要说明的是我们并没有把语料切成互不重叠的 batch而是每次随机选一段连续子序列所以同一个字符会以不同上下文出现在不同训练步里这极大地增加了样本多样性也符合现代语言模型的做法。还要提醒一下block_size在这里虽然是 8但 Bigram 模型并不真正利用它——因为nn.Embedding对每个位置的输出只由该位置的 token 决定不涉及跨位置的信息混合。但保留这个维度能让代码接口上和后续的 Transformer 保持一致到时候只需要改模型内部不需要改数据加载。3. Bigram 模型的前向实现一张表和一堆 logits3.1 用 nn.Embedding 实现查表现在到了最核心的部分定义一个前向传播能直接输出 logits 的模型。对于 Bigram我们要做的只是把每个输入 token 映射到一个长度等于vocab_size的向量这个向量里的数值是缩放前的分数也就是 logits。用nn.Embedding(vocab_size, vocab_size)一行就能搞定Embedding 本质上就是一张可学习的vocab_size × vocab_size的权重矩阵取第i行就是输入 tokeni对应的输出向量。import torch import torch.nn as nn import torch.nn.functional as F class BigramLanguageModel(nn.Module): def __init__(self, vocab_size): super().__init__() self.token_embedding_table nn.Embedding(vocab_size, vocab_size) def forward(self, idx, targetsNone): logits self.token_embedding_table(idx) # logits shape: (B, T, C) if targets is not None: # 交叉熵需要把 (B, T, C) 变成 (B*T, C)targets 变成 (B*T) B, T, C logits.shape logits logits.view(B * T, C) targets targets.view(B * T) loss F.cross_entropy(logits, targets) else: loss None return logits, loss def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): logits, loss self(idx) # 取最后一个时间步的输出 logits logits[:, -1, :] probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx这里最反直觉的地方在于明明说是 Bigram为什么forward接受一个长度为block_size8的序列而不是单个字符原因很简单——参数共享。上面那张 Embedding 表对所有位置都是同一张所以无论我们输入长度为 1 还是长度 8本质上模型在做的事情没变对于输入序列的每个位置查表得到对应输出。换句话说block_size8只是并行计算了 8 个单独的 Bigram 预测而已。3.2 为什么没有 softmax 也能训练交叉熵内部处理很多刚接触代码的同学会有一个疑问语言模型的输出不应该是概率吗怎么代码里只有 logits没有看到softmax呢这个疑问其实源于对损失函数实现细节的不熟悉。F.cross_entropy这个函数在 PyTorch 里的标准用法就是直接接收 logits而不是经过 softmax 之后的概率向量。原因很简单交叉熵损失定义为-log(softmax(logits)[target])其中涉及一个 log 和一个 softmax而数学上这两步可以合并成一个数值上更稳定的计算过程能避免因为中间概率值过小导致的对数溢出。所以训练阶段我们不需要手动加 softmax模型输出的 logits 直接喂给损失函数就行。至于generate函数里的F.softmax和torch.multinomial那是为了在推理时把 logits 变成概率分布再从分布中随机采样一个 token。采样的随机性很重要如果直接每次都取概率最大的 token模型会陷入重复输出反之如果按概率采样就能让模型偶尔说出一些意料之外但概率合理的内容这正是生成式模型的魅力所在。这个模型总参数量就是vocab_size * vocab_size。如果词汇表是 65 个字符那么参数只有 4225 个小到甚至可以在 CPU 上瞬间完成推理。4. 训练循环loss 从 4.8 降到 2.5 的过程中发生了什么4.1 数据按 batch 切分我们已经在get_batch里实现了随机 batch 采样。接下来就是标准的训练循环。这里有一个容易踩的坑我见过有人为了追求简单一次性把整个训练集喂给模型做梯度下降。对于语言模型这种做法几乎不可行因为全语料的损失是一个超级大的 batch计算图会占用大量内存。所以我们一定用 mini-batch 随机采样每个 step 只喂32 × 8 256个字符让模型在一个 batch 内同时看到 32 个不同位置的上下文片段。model BigramLanguageModel(vocab_size) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) for step in range(10000): xb, yb get_batch(train) logits, loss model(xb, yb) optimizer.zero_grad(set_to_noneTrue) loss.backward() optimizer.step() if step % 1000 0: print(fstep {step}: train_loss {loss.item():.4f})设置zero_grad(set_to_noneTrue)是比zero_grad()更高效的做法因为直接把梯度置为 None 能减少不必要的内存清零操作这在模型规模变大后收益明显。4.2 优化器与超参数我选择的优化器是AdamW它比原版 Adam 做了权重衰减解耦是目前训练语言模型最常用的默认选择。学习率1e-3对于这么小的模型是一个相对安全的初始值。你可以自己尝试把学习率改成1e-2大概率会出现 loss 震荡甚至直接发散改成1e-4则会发现 loss 下降变得异常缓慢。对于这个 Bigram 模型1e-3是我实测比较稳的位置。另一个超参数是 batch size。我这里用了 32你可以按需调整。但要注意batch size 变大后每个 step 的计算量会增加而 loss 的噪声会减小反之batch size 变小虽然噪声大但 step 数相同的情况下更不容易过拟合。对于这个玩具级模型32 完全够用。4.3 训练日志与 loss 变化的意义当 vocab_size 为 65 时模型初始化的 Embedding 权重是从标准正态分布随机采样的。此时每个 token 对应的输出 logits 大致随机经过 softmax 后每个类别的概率近似相等所以初始 loss 约等于ln(65) ≈ 4.17。如果考虑字符频率不均真正的初始 loss 可能略高或略低但大致在这个范围。训练几分钟后CPU 上 1 万步大概需要 3~5 分钟loss 通常会稳定下降到 2.5 左右。这一步下降说明模型学到了什么它学到的不是语法而是字符共现频率比如字母q后面大概率是u句号后面大概率是空格空格后面大概率是某个常见字母。这类统计规律恰恰是 Bigram 模型能够捕捉的最有价值信息。如果你拿验证集去算 loss会发现验证 loss 也在 2.5 附近说明这个模型还没有严重过拟合——毕竟容量太小想记下整个语料也记不住。在训练过程中我建议你刻意观察每个 step 的 loss 变化曲线。你会发现 loss 不是单调下降的而是伴随着大量小波动这是 mini-batch 随机采样带来的正常现象。不要因为某个 step 的 loss 突然升高就怀疑代码写错了连续观察几百步的整体趋势才有意义。5. 文本采样让模型输出第一个看起来像英语的句子5.1 采样函数与 temperature训练结束后我们最想做的就是让模型开口说话。给定一个初始字符或者干脆用换行符作为起始 token不断调用generate把上一次输出的新 token 拼回输入序列再预测下一个。这个过程会持续max_new_tokens次。为了让生成结果更有趣可以在采样前加入一个 temperature 参数把 logits 除以 temperature再送进 softmax。temperature 大于 1 会让概率分布变得更平坦生成结果更随机小于 1 会让分布更尖锐更倾向于选择高频 token。对于 Bigram 模型我建议 temperature 设在 0.8 到 1.2 之间因为原始概率分布本来就已经很平坦了温度太低反而会生成大量重复的常见字符序列。def generate_with_temp(idx, max_new_tokens, temperature1.0): for _ in range(max_new_tokens): logits, _ model(idx) logits logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx5.2 第一个生成结果解析我拿一段 1MB 左右的英文语料做实验训练 5000 步后以换行符作为起始 token设置 temperature1.0生成的文本大概像这样the fout far the terther I the the foung the and wor, and ton thy t. Wher ber tat the the thar ton说实话这段文本在人类眼里几乎就是乱码。但请你看几个现象它已经能产生许多合法的英文单词片段比如 the、and、ton、thy而且空格的分布位置也基本符合英文文本习惯。你还能注意到大量 th、he、an 这样的高频双字符组合。这说明模型确实学到了字符级的转移概率只是它没有能力组织成一句话因为它根本看不到超过一个字符的历史信息。如果你把 temperature 调低到 0.5生成结果可能会变成一堆 the the the the 这样的重复因为高频字符组合被过度放大调高到 1.5则会看到大量低频字符甚至是标点符号的混杂。这些现象都是理解语言模型采样机制的最好素材。6. 从 Bigram 走向 Transformer下一步的进化点和容易踩的坑6.1 为什么 Bigram 本质上是一张静态表现在我们已经跑通了第一个语言模型是时候跳出代码站在更高维度审视它。Bigram 模型的参数就是一个vocab_size × vocab_size的矩阵其中每行i表示如果当前字符是i下一字符各个候选的分数是多少。这个矩阵一旦训练完成就固定下来永远不会根据上下文动态变化。换句话说无论输入是 I love cheese 还是 I hate cheese只要当前要预测的下一个位置前面的字符是同一个比如空格模型给出的概率分布就是完全相同的。它完全忽略了空格之前的love和hate这两个词在语义上的巨大差异。这也意味着Bigram 模型没有能力处理任何需要长期依赖的任务比如代词指代、语法一致性、主题连贯性等等。要解决这个问题必须让模型在预测时看到更长的历史。最直接的想法就是把最近若干个字符拼接成一个向量喂给网络但如果只是简单拼接那也不过是一个更长上下文的 n-gram仍然无法处理任意长度的依赖。真正能解决这个问题的是注意力机制它允许每个位置动态地从所有历史位置中聚合信息而且这种聚合是由当前内容决定的不是静态查表。6.2 引入上下文能力的路线图很多从零实现 GPT 的教程会在这一步直接跳到多头注意力我觉得这个跨度还是有点大。比较平滑的过渡方式是先实现一个 Block 类包含一个单头注意力 一个 MLP验证它在小型语料上的损失能明显低于 Bigram 后再扩展为多头加上残差连接和 LayerNorm。下一篇内容我们会沿着这条路线走但在动手之前我强烈建议你先自己完成一个简单的实验手工把 Bigram 的 loss 记下来当作 baseline后续所有模型改进的 loss 都必须低于这个值否则说明新增的模块根本没有起到作用。在这个过程中有几个高频坑必须提前说。第一个坑是get_batch中把x和y的索引写反导致模型输入输出完全错位训练 loss 永远降不动。解决方法是在写完后用几个已知的短文本测试打印decode(x[0])和decode(y[0])确认逐字符 shift 一位。第二个坑是generate时将idx直接传入模型忘记block_size限制——一旦输入序列超过训练时的上下文长度后续真正引入注意力时就会出现 shape 不一致。虽然 Bigram 当前不依赖上下文但这个好习惯现在就要养成生成时一定只取最后一个block_size的 token。第三个坑是直接用model(idx)而不是model(idx, None)导致前向返回的 loss 为 None 时被误用从而报错。我个人在实际操作中的体会是Bigram 模型虽然简单但它是检验你对语言模型基本概念是否透彻的最佳尺子。如果你能不看任何参考代码完全靠自己的理解写出训练和采样流程并且能把每一行 shape 变化都解释清楚那后面面对 Transformer 时你的心态会完全不一样。下一步我们就要在这个模型里加入真正的上下文了到时候你再回头看这张 Embedding 表会明白它只是完整模型的第一层。
从零实现一个极简 GPT 系列开到第二篇说实话百感交集。上一篇我们还在纠结怎么把文本变成数字这一篇终于要接触真正的模型了。但别急着把 Transformer、多头注意力这些庞然大物搬进来——在动手写 attention 之前我强烈建议先实现一个结构简单到极点、但五脏俱全的 Bigram 语言模型也就是所谓的二元语法模型。它只用一个 Embedding 层没有任何上下文窗口却能把数据管线、损失函数、训练循环、采样生成这条主链路完整跑通一遍。等这条链路在你手里变得像呼吸一样自然后面再加注意力机制时你才会知道每一步改动到底在解决什么问题。我写这篇内容时的心态是假设你已经能用 PyTorch 写出一个基础的全连接网络对 tensor 的 shape 操作有一定直觉但没碰过任何语言模型。接下来所有代码你都可以直接复制到 Jupyter 或脚本里跑数据用任何纯文本都行。1. 为什么第一个模型选 Bigram一个不需要注意力的语言模型1.1 语言模型的最小定义语言模型的任务一句话就能说清楚给定一段已经出现的字符序列预测下一个字符是什么。形式化地讲我们希望建模一个条件概率分布P(next_char | previous_chars)。比如说你给模型看 I like cheese它应该给 I 后面的所有候选字符分配一个概率其中像空格、逗号这类在训练语料中高频出现在 I 后的字符概率就会偏大。那么这个任务的最小实现是什么如果我们彻底放弃previous_chars里面所有历史信息只保留最后一个字符模型就退化成只看当前字符来预测下一个字符。这就是 Bigram 名称的来源它等价于统计语料中每两个相邻字符的共现频率。比如语料里 t 后面跟着 h 的次数很多模型就会学到 t 之后输出 h 概率较高。在代码层面Bigram 模型甚至可以不用神经网络——手写一个二维计数矩阵就能完成。但为了后续平滑过渡到 GPT我们会用 PyTorch 的nn.Embedding来把它写成一个标准模型。这样它和最终 Transformer 模型的输入输出接口完全一致输入一个形状为(Batch, Time)的 token 序列输出一个形状为(Batch, Time, VocabSize)的 logits 序列。1.2 Bigram 的局限与它的不可替代价值你肯定会问这种只能看一个字符的模型能有什么用它连最基本的语法结构都学不会更不用说理解语义了。确实Bigram 生成的文本几乎就是随机字符排列最多在字符共现频率上显得像人话。但它的价值不在生成质量而在于它把整个训练链路压缩到了最小可复现的范围。我见过很多直接上手实现 Transformer 的同学一上来就被多头注意力的 shape 变换劝退。问题往往不是注意力机制本身难理解而是他们还没有一个能快速反馈的完整工具链数据加载有问题时会误以为是模型写错了损失下降不下去时也分不清是优化器问题还是模型设计问题。先用 Bigram 跑通一遍等于先给整个流程做了一次冒烟测试。后面无论增加什么模块数据部分和训练框架都不需要再动我们只需要替换模型内部的计算逻辑。这个价值在调试复杂模型时会体现得淋漓尽致。此外Bigram 也是我们之后验证梯度是否正确的好参照物。当 Transformer 版模型损失怎么也降不到 Bigram 以下时至少能证明底层数据管线没问题。2. 数据管线把一段文本变成模型能吃的 (输入, 目标) 对2.1 字符级 tokenizer 的实现我们需要一个能反复调用的文本处理器。我的做法和大多数入门教程一致使用字符级 tokenizer也就是把每个字符当成一个 token。这样做的好处是词汇表极小方便可视化而且不会引入 BPE 等分词算法的额外复杂度。对于 1MB 左右的英文语料字符集合通常只有 50~100 个足够我们研究模型行为。先把原始文本读进来然后构建两个映射表字符到整数的stoi整数到字符的itos。Python 代码一般长这样with open(input.txt, r, encodingutf-8) as f: text f.read() chars sorted(list(set(text))) vocab_size len(chars) print(f语料总字符数: {len(text)}去重后字符数: {vocab_size}) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} def encode(s): 字符串 - token 列表 return [stoi[c] for c in s] def decode(l): token 列表 - 字符串 return .join([itos[i] for i in l])这里有一个值得注意的细节用set(text)去重会丢失字符顺序信息但因为我们之后用sorted重新排序所以结果可复现。另一个细节是encode函数需要处理数据中遇到的每个字符一旦语料里出现未登录字符就会报 KeyError因此如果你换数据集务必确保chars是基于新文本重新生成的而不是复用旧词汇表。2.2 构造输入序列与目标序列核心是 shift 一位把整个语料编码成一个长整数张量后就该构造模型训练用的 (x, y) 对。语言模型的监督信号并不是显式标注好的而是从文本自身生成的对于任意位置的字符它的正确答案就是紧接着它的下一个字符。所以我们需要把原始序列和它的移位版本配对。假设原始 token 序列是[1, 2, 3, 4, 5]如果我们想让模型通过[1, 2, 3]预测[2, 3, 4]那么输入就是data[0:3]目标就是data[1:4]。在 PyTorch 里我们通常一次处理一批形状为(batch_size, block_size)的序列。data torch.tensor(encode(text), dtypetorch.long) # 分训练集和验证集 n int(0.9 * len(data)) train_data data[:n] val_data data[n:] batch_size 32 block_size 8 # 上下文长度Bigram 模型其实用不到但留着方便后续扩展 def get_batch(split): 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:iblock_size] for i in ix]) y torch.stack([data[i1:iblock_size1] for i in ix]) return x, y这段代码中的torch.randint是异步随机采样不用像有些教程那样手动 shuffle 整个语料。需要说明的是我们并没有把语料切成互不重叠的 batch而是每次随机选一段连续子序列所以同一个字符会以不同上下文出现在不同训练步里这极大地增加了样本多样性也符合现代语言模型的做法。还要提醒一下block_size在这里虽然是 8但 Bigram 模型并不真正利用它——因为nn.Embedding对每个位置的输出只由该位置的 token 决定不涉及跨位置的信息混合。但保留这个维度能让代码接口上和后续的 Transformer 保持一致到时候只需要改模型内部不需要改数据加载。3. Bigram 模型的前向实现一张表和一堆 logits3.1 用 nn.Embedding 实现查表现在到了最核心的部分定义一个前向传播能直接输出 logits 的模型。对于 Bigram我们要做的只是把每个输入 token 映射到一个长度等于vocab_size的向量这个向量里的数值是缩放前的分数也就是 logits。用nn.Embedding(vocab_size, vocab_size)一行就能搞定Embedding 本质上就是一张可学习的vocab_size × vocab_size的权重矩阵取第i行就是输入 tokeni对应的输出向量。import torch import torch.nn as nn import torch.nn.functional as F class BigramLanguageModel(nn.Module): def __init__(self, vocab_size): super().__init__() self.token_embedding_table nn.Embedding(vocab_size, vocab_size) def forward(self, idx, targetsNone): logits self.token_embedding_table(idx) # logits shape: (B, T, C) if targets is not None: # 交叉熵需要把 (B, T, C) 变成 (B*T, C)targets 变成 (B*T) B, T, C logits.shape logits logits.view(B * T, C) targets targets.view(B * T) loss F.cross_entropy(logits, targets) else: loss None return logits, loss def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): logits, loss self(idx) # 取最后一个时间步的输出 logits logits[:, -1, :] probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx这里最反直觉的地方在于明明说是 Bigram为什么forward接受一个长度为block_size8的序列而不是单个字符原因很简单——参数共享。上面那张 Embedding 表对所有位置都是同一张所以无论我们输入长度为 1 还是长度 8本质上模型在做的事情没变对于输入序列的每个位置查表得到对应输出。换句话说block_size8只是并行计算了 8 个单独的 Bigram 预测而已。3.2 为什么没有 softmax 也能训练交叉熵内部处理很多刚接触代码的同学会有一个疑问语言模型的输出不应该是概率吗怎么代码里只有 logits没有看到softmax呢这个疑问其实源于对损失函数实现细节的不熟悉。F.cross_entropy这个函数在 PyTorch 里的标准用法就是直接接收 logits而不是经过 softmax 之后的概率向量。原因很简单交叉熵损失定义为-log(softmax(logits)[target])其中涉及一个 log 和一个 softmax而数学上这两步可以合并成一个数值上更稳定的计算过程能避免因为中间概率值过小导致的对数溢出。所以训练阶段我们不需要手动加 softmax模型输出的 logits 直接喂给损失函数就行。至于generate函数里的F.softmax和torch.multinomial那是为了在推理时把 logits 变成概率分布再从分布中随机采样一个 token。采样的随机性很重要如果直接每次都取概率最大的 token模型会陷入重复输出反之如果按概率采样就能让模型偶尔说出一些意料之外但概率合理的内容这正是生成式模型的魅力所在。这个模型总参数量就是vocab_size * vocab_size。如果词汇表是 65 个字符那么参数只有 4225 个小到甚至可以在 CPU 上瞬间完成推理。4. 训练循环loss 从 4.8 降到 2.5 的过程中发生了什么4.1 数据按 batch 切分我们已经在get_batch里实现了随机 batch 采样。接下来就是标准的训练循环。这里有一个容易踩的坑我见过有人为了追求简单一次性把整个训练集喂给模型做梯度下降。对于语言模型这种做法几乎不可行因为全语料的损失是一个超级大的 batch计算图会占用大量内存。所以我们一定用 mini-batch 随机采样每个 step 只喂32 × 8 256个字符让模型在一个 batch 内同时看到 32 个不同位置的上下文片段。model BigramLanguageModel(vocab_size) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) for step in range(10000): xb, yb get_batch(train) logits, loss model(xb, yb) optimizer.zero_grad(set_to_noneTrue) loss.backward() optimizer.step() if step % 1000 0: print(fstep {step}: train_loss {loss.item():.4f})设置zero_grad(set_to_noneTrue)是比zero_grad()更高效的做法因为直接把梯度置为 None 能减少不必要的内存清零操作这在模型规模变大后收益明显。4.2 优化器与超参数我选择的优化器是AdamW它比原版 Adam 做了权重衰减解耦是目前训练语言模型最常用的默认选择。学习率1e-3对于这么小的模型是一个相对安全的初始值。你可以自己尝试把学习率改成1e-2大概率会出现 loss 震荡甚至直接发散改成1e-4则会发现 loss 下降变得异常缓慢。对于这个 Bigram 模型1e-3是我实测比较稳的位置。另一个超参数是 batch size。我这里用了 32你可以按需调整。但要注意batch size 变大后每个 step 的计算量会增加而 loss 的噪声会减小反之batch size 变小虽然噪声大但 step 数相同的情况下更不容易过拟合。对于这个玩具级模型32 完全够用。4.3 训练日志与 loss 变化的意义当 vocab_size 为 65 时模型初始化的 Embedding 权重是从标准正态分布随机采样的。此时每个 token 对应的输出 logits 大致随机经过 softmax 后每个类别的概率近似相等所以初始 loss 约等于ln(65) ≈ 4.17。如果考虑字符频率不均真正的初始 loss 可能略高或略低但大致在这个范围。训练几分钟后CPU 上 1 万步大概需要 3~5 分钟loss 通常会稳定下降到 2.5 左右。这一步下降说明模型学到了什么它学到的不是语法而是字符共现频率比如字母q后面大概率是u句号后面大概率是空格空格后面大概率是某个常见字母。这类统计规律恰恰是 Bigram 模型能够捕捉的最有价值信息。如果你拿验证集去算 loss会发现验证 loss 也在 2.5 附近说明这个模型还没有严重过拟合——毕竟容量太小想记下整个语料也记不住。在训练过程中我建议你刻意观察每个 step 的 loss 变化曲线。你会发现 loss 不是单调下降的而是伴随着大量小波动这是 mini-batch 随机采样带来的正常现象。不要因为某个 step 的 loss 突然升高就怀疑代码写错了连续观察几百步的整体趋势才有意义。5. 文本采样让模型输出第一个看起来像英语的句子5.1 采样函数与 temperature训练结束后我们最想做的就是让模型开口说话。给定一个初始字符或者干脆用换行符作为起始 token不断调用generate把上一次输出的新 token 拼回输入序列再预测下一个。这个过程会持续max_new_tokens次。为了让生成结果更有趣可以在采样前加入一个 temperature 参数把 logits 除以 temperature再送进 softmax。temperature 大于 1 会让概率分布变得更平坦生成结果更随机小于 1 会让分布更尖锐更倾向于选择高频 token。对于 Bigram 模型我建议 temperature 设在 0.8 到 1.2 之间因为原始概率分布本来就已经很平坦了温度太低反而会生成大量重复的常见字符序列。def generate_with_temp(idx, max_new_tokens, temperature1.0): for _ in range(max_new_tokens): logits, _ model(idx) logits logits[:, -1, :] / temperature probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idx5.2 第一个生成结果解析我拿一段 1MB 左右的英文语料做实验训练 5000 步后以换行符作为起始 token设置 temperature1.0生成的文本大概像这样the fout far the terther I the the foung the and wor, and ton thy t. Wher ber tat the the thar ton说实话这段文本在人类眼里几乎就是乱码。但请你看几个现象它已经能产生许多合法的英文单词片段比如 the、and、ton、thy而且空格的分布位置也基本符合英文文本习惯。你还能注意到大量 th、he、an 这样的高频双字符组合。这说明模型确实学到了字符级的转移概率只是它没有能力组织成一句话因为它根本看不到超过一个字符的历史信息。如果你把 temperature 调低到 0.5生成结果可能会变成一堆 the the the the 这样的重复因为高频字符组合被过度放大调高到 1.5则会看到大量低频字符甚至是标点符号的混杂。这些现象都是理解语言模型采样机制的最好素材。6. 从 Bigram 走向 Transformer下一步的进化点和容易踩的坑6.1 为什么 Bigram 本质上是一张静态表现在我们已经跑通了第一个语言模型是时候跳出代码站在更高维度审视它。Bigram 模型的参数就是一个vocab_size × vocab_size的矩阵其中每行i表示如果当前字符是i下一字符各个候选的分数是多少。这个矩阵一旦训练完成就固定下来永远不会根据上下文动态变化。换句话说无论输入是 I love cheese 还是 I hate cheese只要当前要预测的下一个位置前面的字符是同一个比如空格模型给出的概率分布就是完全相同的。它完全忽略了空格之前的love和hate这两个词在语义上的巨大差异。这也意味着Bigram 模型没有能力处理任何需要长期依赖的任务比如代词指代、语法一致性、主题连贯性等等。要解决这个问题必须让模型在预测时看到更长的历史。最直接的想法就是把最近若干个字符拼接成一个向量喂给网络但如果只是简单拼接那也不过是一个更长上下文的 n-gram仍然无法处理任意长度的依赖。真正能解决这个问题的是注意力机制它允许每个位置动态地从所有历史位置中聚合信息而且这种聚合是由当前内容决定的不是静态查表。6.2 引入上下文能力的路线图很多从零实现 GPT 的教程会在这一步直接跳到多头注意力我觉得这个跨度还是有点大。比较平滑的过渡方式是先实现一个 Block 类包含一个单头注意力 一个 MLP验证它在小型语料上的损失能明显低于 Bigram 后再扩展为多头加上残差连接和 LayerNorm。下一篇内容我们会沿着这条路线走但在动手之前我强烈建议你先自己完成一个简单的实验手工把 Bigram 的 loss 记下来当作 baseline后续所有模型改进的 loss 都必须低于这个值否则说明新增的模块根本没有起到作用。在这个过程中有几个高频坑必须提前说。第一个坑是get_batch中把x和y的索引写反导致模型输入输出完全错位训练 loss 永远降不动。解决方法是在写完后用几个已知的短文本测试打印decode(x[0])和decode(y[0])确认逐字符 shift 一位。第二个坑是generate时将idx直接传入模型忘记block_size限制——一旦输入序列超过训练时的上下文长度后续真正引入注意力时就会出现 shape 不一致。虽然 Bigram 当前不依赖上下文但这个好习惯现在就要养成生成时一定只取最后一个block_size的 token。第三个坑是直接用model(idx)而不是model(idx, None)导致前向返回的 loss 为 None 时被误用从而报错。我个人在实际操作中的体会是Bigram 模型虽然简单但它是检验你对语言模型基本概念是否透彻的最佳尺子。如果你能不看任何参考代码完全靠自己的理解写出训练和采样流程并且能把每一行 shape 变化都解释清楚那后面面对 Transformer 时你的心态会完全不一样。下一步我们就要在这个模型里加入真正的上下文了到时候你再回头看这张 Embedding 表会明白它只是完整模型的第一层。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
