基于印度民间故事训练小型GPT：分词器、掩码与损失函数
从零构建的900百万参数GPT：为何GPT-2的分词器无法处理民间故事中的词汇，BPE与因果掩码的运作原理，以及为何一个批次会产生4,096个下一个词元的损失值。
来源: Building LLM, Part 1 — Training: Predicting the Next Token
基于印度民间故事训练的小型GPT
任何构建过视觉变换器的人大致都已了解语言模型70%的构成。嵌入层、注意力机制、残差连接结构以及训练循环在两者中都是相同的组件。剩下的差异部分将通过从零开始构建小型GPT并在一台笔记本电脑的GPU上对其进行训练来加以说明。该模型的训练语料为印度民间故事，包括本生故事、潘查坦特拉风格的寓言以及旁遮普和孟加拉地区的传说。
早在初期就出现了真正的难题。GPT-2原有的分词器因从未见过“Brahman”和“Rajah”这样的词汇，会将它们拆分成碎片，因此人们转而使用从语料库中训练得到的分词器。最终形成的模型拥有900万个参数，已从随机噪声转变为经过训练的权重，任何人只要运行大约十分钟的训练任务就能得到相同的参数值。
为何视觉技术先出现
既然聊天机器人处理的是文字，那么先看图片再看文字似乎顺序有误。但实际上，变换器并不关心输入的是图像块还是句子片段——注意力机制、LayerNorm、跳过连接、多层感知机以及优化循环都保持不变。更恰当的问题是，在哪种场景下这种通用架构最容易学习，而图像技术因四个原因占据了优势。
- 令牌其实就是从图像中截取的固定网格区域，因此无需词汇表、分词器，也不存在需要先解决的生词问题。
- 没有因果掩码。ViT可以一次性看到整幅图像，所以普通的注意力机制就能直接使用，无需遵循语言中关于不能预看未来的额外规则。
- 其行为是可见的。这些图像块、注意力图以及模型关注的区域都可以被绘制出来，与文本相比，像素具有更直观的可感知性。
- 它可以在几分钟内就在单块GPU上从零开始训练（针对CIFAR-10数据集），因此每个数值都是你自己计算得出的，而非直接采用的预设值。
归根结底，一旦构建了ViT，完整的变换器就已经存在了，只是为视觉任务做了适配。语言模型所做的只是将文本转换为标记（通过查找而非计算）、添加因果掩码以及进行下一个标记的预测。而最复杂、最抽象的部分——语言本身——在机器熟悉之后就变得容易理解多了。
训练其实很简单：预测下一个标记
在训练过程中，语言模型所做的一切归结为读取已有的标记并预测接下来的那个标记。
该图展示了如何根据前面的词元来预测下一个词元。输入开头的“在某个村庄里，住着一位贫穷的婆罗门”，系统就会输出每个可能接续词元的概率。句号、逗号，或是谁、名叫这样的词都是符合语法的。无需任何人进行标注：正确答案就是文本中紧接着出现的那个词元。正因如此，仅凭原始文本就能训练语言模型，而且它就是全部的学习信号。其他所有组件都是为了辅助这一预测而存在的，ViT已经提供了其中的大部分。
在代码中，这一思路只需两行：运行模型，然后将它的得分标准化为下一个词元的概率分布。
logits, _ = model(ids) # ids = the tokens so far, shape (1, T)
probs = F.softmax(logits[0, -1], dim=-1) # one probability per vocabulary token
其余部分介绍了如何构建和训练model，从而使这些概率值达到理想水平。
为什么分词器很重要
ViT会将图像分割成块，而语言模型则将文本分割成词元。块是由像素值计算得出的，而文本词元则来自固定的词汇表。必须有人事先构建这个词汇表，而且构建者决定了模型认为哪些内容是容易处理的。
最初的尝试重用了GPT-2现成的分词器r50k_base，其词汇表源自大约2019年的西方网络文本且无法修改。下图展示了该分词器对故事中词语的处理效果。
Brahman会拆分为Brah和man，Rajah会拆分为Raj和ah，而Panchatantra则会分裂成五部分。这些词都并非生僻词汇。Rajah在语料库中出现了709次，但分词器在其训练数据中并未频繁见到这类单词，因此没有为每个词单独分配标记。这种碎片化现象会拉长序列长度，迫使模型从Brah和man重新组合出Brahman，而man这个词还会与woman和human共享。鉴于训练数据仅有约一百万个标记，这无疑是个巨大障碍。
解决方案是一种基于这些故事、采用字节对编码（BPE）训练的分词器。GPT也使用相同的算法，只是将其应用于更匹配的语料库上。这些原始材料来自古腾堡计划的十二部公共领域的印度民间故事集，包括两本本生经、两册印度童话、《旧德干时代》、《孟加拉民间故事》、《会说话的画眉鸟》、《旁遮普故事集》、《二十二个妖精的故事》、两部桑塔尔族故事集以及《南印度太阳故事集》，总字数约为809,000字。
该图展示了用于训练分词器的字节对合并过程。BPE从256个原始字节值开始，因此不存在未知元素，随后会不断将最常见的相邻字节对合并为一个新标记。经过对这些故事进行7,935次合并后，词汇表中共有8,192个条目，其中还包括诸如Rajah、Brahman和jackal这样的完整词汇。
有一个设计选择需要明确说明：GPT-2的分词器是被完全替换，而非简单扩展。
图表显示了由8,192个标记构成的自定义词汇表。当需要保留大型预训练模型原有的英语处理能力时，如OpenHathi项目所示，将r50k现有的50,257个标记再加以扩展是一种合理的做法。而基于小型专业语料库从零开始训练则是另一种情况——此时r50k的西方网络词汇表只会占用空间，而这些故事仅涵盖了其中37%的内容。全新的8,000个标记构成的词汇表规模更小，使用率为96.6%，且是专为这些故事设计的，因此进行替换才是最佳选择。
使用HuggingFace的tokenizers库只需几行代码即可完成分词器的训练。该脚本会读取每一本原始文本书籍，去除Project Gutenberg的页眉和页脚信息，进而训练一个BPE模型。
import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from llm_transformer.prep_tales import strip_gutenberg, RAW_DIR
# read every raw .txt book and strip its Project Gutenberg header/footer,
# leaving just the story text - one big string per book
tales_texts = [strip_gutenberg(open(p, encoding="utf-8", errors="ignore").read())
for p in sorted(glob.glob(f"{RAW_DIR}/*.txt"))] # 12 books, ~809k words
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) # start from bytes
trainer = trainers.BpeTrainer(vocab_size=8192, special_tokens=["<|endoftext|>"],
initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(tales_texts, trainer) # learn the 7,935 merges
tok.save("tales_bpe.json")
生成的 .json 文件即为最终的分词器，是一份固定的合并列表。它不会再被训练，而是被加载后用于将文本转换为编号，再将编号转换回文本。只需一次 .encode() 调用，即可将原始文本与神经网络关联起来。
tok = Tokenizer.from_file("tales_bpe.json") # reload the trained tokenizer
ids = tok.encode("Once upon a time").ids # -> [412, 987, 15, 733] (integers)
idx = torch.tensor(ids).unsqueeze(0) # (1, 4) — a batch of one sequence
vocab_size = tok.get_vocab_size() # 8192 - sets the NUMBER OF ROWS below
只有这些整数会传送到网络中，原始文本已不复存在。每个ID仅仅是表格中的一个位置，而vocab_size则决定了嵌入层会预留多少条记录。
从ID到向量
当文本变为一系列ID后，每个ID都会转换成一个向量。这也是大语言模型比视觉变换器更简单的部分。
该图将分词器的索引与嵌入向量联系起来。在视觉变换器中，通过小型矩阵乘法根据像素计算出每个块的token；而在大语言模型中，token是一个离散的ID，因此向量仅仅是学习得到的表格中的行号ID，只需简单查找即可。例如，Token 789会对应第789行，这是一個256维的向量，再像视觉变换器那样加上位置嵌入后，输入阶段就完成了。
在代码中，nn.Embedding层相当于一个表格，调用该层时是通过索引进行查找而非乘法运算。它的.weight是一个可学习的8192×256维张量，包含210万个权重，这些权重会通过反向传播算法进行训练；只有当对应标记出现时，该层的某一行才会发生变化。
self.token_embed = nn.Embedding(vocab_size, n_embed) # the 8,192 x 256 table
# .weight IS a learnable (8192, 256) parameter tensor — 2.1M weights, trained by backprop;
# each id's row is nudged only on steps where that token appears in the batch
...
tok_emb = self.token_embed(idx) # idx (B, T) -> vectors (B, T, 256), one row per id
从这里开始的所有内容，包括注意力机制、残差流以及多层感知器块，都与视觉变换器相同，只有一个例外。
因果掩码
视觉变换器允许每个图像块关注所有其他块。而语言模型无法做到这一点，因为它的任务是预测接下来的内容。如果位置 t 能看到位置 t+1，它就会直接复制答案。解决方法是使用掩码限制注意力范围，使得每个标记只能看到自身以及更早的标记。
图中的掩码为下三角形结构。在应用软阈值函数之前，所有后续位置的权重都会被设为负无穷大，软阈值函数会将这些值转换为恰好为零的权重，因此被屏蔽的标记不会对结果产生任何影响。这种三角形结构是视觉变换器注意力机制与大型语言模型注意力机制之间唯一的架构差异。点积运算、按头部尺寸平方根进行的缩放、软阈值函数以及权重求和等步骤则完全一致。
在代码中，它只是在普通注意力机制的基础上多出了一行。首先构建一个全为1的下三角矩阵，然后像往常一样计算得分，在应用softmax函数之前将值为0的位置替换为负无穷大。
# tril = torch.tril(torch.ones(T, T)) — a lower-triangular matrix of 1s, made once
attn = q @ k.transpose(-2, -1) / math.sqrt(head_size) # scores, (B, T, T)
attn = attn.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # blank out the future
attn = F.softmax(attn, dim=-1) # -inf -> weight 0
out = attn @ v
一次遍历即可处理所有上下文长度
那条单行掩码其实隐藏了值得深入探讨的细节。人们往往认为语言模型的工作方式就是像生成文本那样：在不断扩展的上下文中添加标记，并在每一步都重新计算注意力。这在推理阶段确实如此，但在训练阶段并非如此。在训练时，句子内容是完全已知的，因此所有T个标记的Q、K和V值会一次性共同计算出来，从而在单一步骤内形成完整的T×T注意力矩阵。
正是由于因果掩码的存在，这种单次遍历才可行。以掩码矩阵的第行为例：对角线右侧的所有元素均为负无穷大，经过softmax处理后会变为零。此时计算out = attn @ V，第行就会仅通过前t个值向量得到加权求和结果。这样一来，所有上下文长度都能同时被处理。
该图展示了训练过程中不同的迭代次数以及掩码的形状。请逐行阅读：
- 第1行仅关注标记1，因此其输出是在仅一个标记的上下文下预测标记2。
- 第3行关注标记1至3，从而在三个标记的上下文下预测标记4。
- 第6行关注全部六个标记，进而基于完整上下文预测标记7。
这是另一种视角。每一行的可见上下文构成一个下三角矩阵，其规模每次增加1——第一个标记对应1×1的矩阵，接着是2×2、3×3……这些就是大矩阵中用橙色标出的嵌套左上角块，因此一次6×6的运算就能处理所有较小的序列片段，而底部的条带则负责将它们展开。所有这些都在一个attn @ V操作中生成，Python代码中并没有对各个位置进行循环处理。这就解释了Transformer的训练效率：一个包含T个标记的序列相当于T个训练样本，上下文长度从1到T不等，但可以并行计算得分。用于生成这些样本的同一三角掩码还能防止在生成过程中任何位置向前查看，生成过程是按行进行的，通过缓存之前的键值（即KV缓存）来保持每一步的计算成本较低。
通过注意力机制处理实数
目前每个单元格仅表示“参加”或“不参加”。展开多行显示后，可以看到针对相同三次运行的实际数值，这些数值来自经过训练的模型在句子“Once there lived a clever jackal”上的第0块、第0头。
该图表展示了各次运行中Q、K、V的计算过程。在同一行中，注意力权重为正数且总和为1，因此输出结果是各可见值向量的凸组合，即加权平均值。
- 第一次运行中只有一个标记，其权重为1.00，因此
out1 = V1。由于还没有其他内容可混合，第一个位置直接复制自身的值。 - 在第三次运行中，输出结果通过约0.20的权重混合了第一个标记的值，0.06的权重混合了第二个标记的值，0.74的权重混合了第三个标记的值。它主要依赖自身的标记“lived”，但也会吸收“Once”五分之一的内容，因此属于真正的混合结果。
这涵盖了注意力机制的所有功能：查询项和键项产生的分数决定了权重，而这些权重又用于混合各个数值。掩码机制确保混合过程不会超出当前标记的范围，因此每一行都代表了迄今为止所读取内容的公平概括。实际的头部结构具有32个数值维度，模型在六层中每层都包含八个这样的头部。这些条状图展示了单个头部的前8个维度，不过整个机制在所有头部中都是相同的。
T预测值，T损失值
下一个标记预测带来了ViT所缺乏的优势。ViT对每张图像仅进行一次预测，即一个类别标签。而因果掩码让语言模型能够同时猜测所有位置上的下一个标记，因此一个序列所能提供的学习信号数量与其标记数相同。
该图显示了每次运行中T个标记输入的T次预测结果与损失值的平均值。每个位置都会产生一个交叉熵损失值，用以衡量模型对实际下一个标记的惊讶程度。有些情况的损失较低，比如在“lived”之后出现“a”的损失仅为0.58；而有些情况的损失则很高，比如形容词“clever”的损失为7.17，因为有很多其他单词也可能出现在该位置。将这六次的损失值取平均后为3.59，随后通过与ViT相同的链式法则机制，反向传播会不断降低这个数值。
在代码中，所有的 T 损失都源自一次 cross_entropy 调用。形状为 (B, T) 的位置矩阵会被展平为一个长列表，列表中的每个元素都会与其真实的下一个标记进行对比评分。
logits = self.lm_head(x) # (B, T, vocab)
B, T, C = logits.shape
loss = F.cross_entropy(logits.reshape(B*T, C), # every position ...
targets.reshape(B*T)) # ... vs the token that actually followed
每次损失都是一次词汇分类
明确cross_entropy所比较的内容非常重要，因为预测值和目标值都是向量，而非单个单词。在每个位置上，lm_head会生成一个包含8,192个元素的logits向量，每个词汇标记对应一个分数，随后通过softmax函数将其转换为所有标记的概率分布。目标值则是相同大小的独热向量，实际出现的单词对应位置为1，其余位置均为0。因此，预测下一个标记实际上就是一个分类问题，其中词汇表中的单词即为各个类别。
该图对比了大型语言模型与视觉变换器中的对数概率值。它阐明了一个细微之处：为了便于阅读，条形图是按概率大小排序的，而实际输出则是具有固定结构的向量，其中第个位置始终对应令牌ID。无论是在训练还是推理过程中，“son”这个词在每次前向传播时的索引始终为498，而softmax函数仅改变各位置的概率值大小，不会改变其顺序。这也就解释了为什么独热编码目标向量中唯一的1位于索引498而非最开头的位置：分词器在初次为每个令牌分配身份时，就确定了其对应的位置。
这正是ViT的分类器头部，它会对类别得分应用softmax函数，并通过交叉熵与独热编码的标签进行对比评分。两者存在两处差异：类别数量（8,192个词汇标记 versus 大约10个图像标签）以及出现频率（在所有T位置上出现 versus 每张图像出现一次）。面对独热目标时，交叉熵可简化为负的对正确单词概率的对数，因此每个位置的损失值看起来就像一个惊讶度分数。出现概率为9.8%的“son”对应的损失值为2.32，而模型原本预期这个单词的损失值会低得多；真正的意外情况则会导致更高的损失值。
批量的构建方式
reshape(B*T)这一操作表明一个批量包含B个序列，每个序列都会贡献T个位置的损失值。那么这些B个序列是从哪里来的？它们的长度是否都相同？
人们常常忽略一个细节：预训练序列并未进行填充，且所有序列的上下文长度都完全相同。这些语料被合并成一个连续的标记流，不同文档之间用 <|endoftext|> 标记分隔，然后从随机位置截取固定长度的窗口。由于所有窗口都来自同一个连续流，因此不存在较短的窗口，也就无需填充标记。
该流程展示了从训练语料到数据批次，再到损失值的路径。数据流中有一个特殊的标记 <|endoftext|>（EOT），它用作分隔符而非填充符，用来标识一个民间故事结束的位置以及下一个故事的起始位置。窗口可能会跨越这样的分隔线并在中间包含EOT标记，模型正是通过这种方式得知一个故事已经结束，新的故事即将开始。在微调阶段，当样本长度不同时，填充符会出现，此时被填充的位置会在损失计算中被屏蔽，这与紧凑预训练的设置有所不同。
损失是大型语言模型与视觉变换器不同的第二个方面，而这种差异其实是值得欢迎的。视觉变换器每处理一张图像只会产生一个预测结果和一个损失值。而128个标记的序列则会生成128个损失值，这些值会被平均起来作为该样本的损失值（图中样本1的损失值为3.78）；同一批次中的32个样本损失值则会被平均，得到3.92的批次损失值，随后通过反向传播来最小化这一数值。经过这两轮平均处理后，每一步就能降低4,096个下一个标记损失值的平均值，相比每张图像仅一个损失值而言，每步能获得的信号量要大得多。
观察损失值下降的过程
将这个损失值连接到优化器上并重复该过程，这就是在故事语料库上进行的实际训练。
该图表展示了不同训练轮次下的训练损失与保留的验证损失。训练损失起始值为9.2，对应ln(8192)，即在使用包含8,192个标记的词汇表时随机猜测所产生的损失；而最佳的保留验证损失则为4.73。两条曲线始终非常接近，相差约1.0，这说明在该有限的领域内，模型具备泛化能力而不仅仅是死记硬背。该模型的结构本身也相当简洁。
该图还将训练过程与可学习权重的数量联系起来。该模型拥有900万个参数，在配备6GB显存的笔记本电脑GPU上大约需要十分钟即可完成训练。其中略多于一半的参数位于六个Transformer块中，其余则属于嵌入表和输出层。训练过程实际上就是不断调整这900万个数值，从而让下一个标记的概率更倾向于真正出现过的那个标记。
在代码中，循环结构与预期一致，也和ViT的格式相同：随机抽取上下文及下一个词元的窗口，执行前向传播以得到一个均值损失值，将梯度置零，再进行反向传播并更新参数。
for step in range(train_steps):
xb, yb = get_batch(train_data) # random (context, next-token) windows
_, loss = model(xb, yb) # forward: T predictions -> one mean loss
opt.zero_grad(); loss.backward() # backprop the loss into all 9M parameters
opt.step() # nudge them downhill
完整脚本增加了余弦学习率调度机制以及评估功能。
新变化与未变内容
与视觉Transformer相比，大语言模型仅有少量新变化：用作查找表的令牌嵌入、因果掩码，以及每个位置上对整个词汇表进行分类的下一令牌损失函数。还有一个重要的实践经验值得注意，那就是分词器属于模型的一部分，因此使用领域特定的分词器会带来显著收益。注意力机制、位置嵌入、残差流、LayerNorm、MLP模块以及完整的反向传播和优化器循环均保持不变。
实数揭示了其内部运作机制：因果掩码如何让单次前向传播同时处理所有上下文长度，attn @ V会产生什么结果（每行输出都是该标记可见的所有值向量的混合），训练批次是如何通过打包而非填充来生成的，以及某个位置的损失——即对词汇表的分类结果——如何与其他位置的损失合并为单一平均值。
权重已确定，但训练好的模型在运行之前什么都不会做。接下来是推理阶段：生成循环、用于提升速度的KV缓存，以及当模型获得故事开头信息时会输出什么内容。