#分词器(Tokenizer)与位置编码
系列导读:在前面的教程里,我们从 第一篇(从 Bigram 到 GPT-2) 出发,在 nanoGPT 架构上尝试了 第二篇(钢琴曲生成) 与 第三篇(算子加速与多模态),验证了不同模态都可以离散化为 Token 算条件概率,但在生成质量上并不理想。从本篇开始,我们回到语言模型主线,跟着 LLM 发展的脚步,用相同的数据集、几乎一样大小的模型参数量、一样的训练步数,一步步看每一次改进到底能给模型带来什么。
#1. 基础模型(Baseline)
我们先搭一个约 12M(1,199 万)参数的 GPT-2 基线模型(6 层,6 头,隐藏层维度 384,上下文窗口 block_size = 512)。训练数据采用 chinese-law-corpus 里的 93 份法律相关文件(5,761 条法条)和 278 个最高法指导案例,拼成纯文本一共约 160 万字符。
这里我们没有使用 GPT-2 官方训练好的 BPE Tokenizer,因为原版词表基本是在英文上训练的,对中文支持很差,一个汉字往往会被拆成 2 到 3 个底层的字节 Token。所以基线版本里,我们先退回到最简单的按字切分(Character-level):直接把语料里出现的 3,062 个唯一字符做成词表,一个汉字对应一个编号。
在 T4 GPU 上训练 2,000 步(约 4 分钟,验证集 Loss 停在 2.572)后,我们喂几个提示词看看模型的效果。
输入 犯故意伤害罪,判处,模型生成:
犯故意伤害罪,判处有期徒刑二年,并处罚金人民币七万元。 被告人王某元犯侵犯公民个人信息罪,判处有期徒刑三年零五年,并处罚金人民币三万元;被告人王某某某在判决生效后十日内赔偿人民币25100元。宣判后,双方当事人均未提出上诉。北京市第级人民法院于20年3月1日作出(202)京01民终131号民事判决,驳回上诉,维持原判。 法院生效裁判认为:本案的争议焦点为,原审被告王某明是否属
再试一条法条续写,输入 《中华人民共和国民法典》第五百七十七条:当事人一方不履行合同义务,模型生成:
《中华人民共和国民法典》第五百七十七条:当事人一方不履行合同义务的,依照法律、行政法规的规定代表人享有的权利。 《中华人民共和国民法典》第五百七十八条:合同的当事人约定代表示或者履行义务的真实意思表示,代表示不得以拒绝接受代表示的权利。 《中华人民共和国民法典》第五百七十九条:代表人的行为不得代表示的意思表示。
可以看到,模型已经摸到了判决书和法条的格式套路,甚至知道从“第五百七十七条”往下数到“第五百七十八条”。但“有期徒刑三年零五年”、“北京市第级人民法院”、“于20年3月1日作出(202)”这些地方还是不通顺,还会把“代表人”和“意思表示”生硬缝合成“代表示”。仅凭这些样例,还不能把问题归因于字符级分词。字符级表示有一个很直接的限制:block_size = 512 就只能装下 512 个字符,连标点和换行也算在里面。接下来尝试 BPE,看看合并常见字符片段后会有什么变化。
#2. 训练自己的 BPE Tokenizer
这次我们采用 BPE 的高频相邻符号合并方法,在法律语料上训练一个以字符为基础的分词器。
具体做法很简单:我们保留语料里出现的 3,062 个唯一字符作为基础词表,这样就不用先从 UTF-8 字节合并出汉字。接着先用正则把连续汉字、数字、字母和标点符号切开,防止 BPE 把标点和词语粘在一起(比如把“合同,”和“合同。”当成两个不同的 Token);然后统计语料中相邻符号对的出现频率,循环合并最高频的一对,直到词表扩充到 vocab_size = 4,096(即新增了 1,034 个常见法律词组与片段,例如“人民”、“或者”、“处罚”、“有期徒刑”、“人民法院”、“当事人”等)。BPE 只管出现频率,合出来的不一定是完整的词。
换上这套 BPE 分词器后,词表从 3,062 变成 4,096,模型总参数量从 11.99M 增到 12.39M。训练集的 1,441,680 个字符变成了 917,987 个 BPE Token,平均一个 Token 对应约 1.57 个字符。同样的 block_size = 512 窗口,现在平均能看约 800 个字符的原文。
注意:当前 notebook 用整份文本训练 BPE,验证集也参与了合并频率的统计。更严格的做法是先划分文本,只用训练集学习合并规则 BPE,再把同一套分词器用在训练集和验证集上。以避免 BPE 泄露信息。
保持其余训练设置不变,训练 2,000 步:
step 1 ( 1.2s) | batch_loss=8.393 | train=8.053 | val=8.041
step 200 ( 25.7s) | batch_loss=4.950 | train=4.955 | val=5.207
step 600 ( 80.5s) | batch_loss=3.679 | train=3.608 | val=4.415
step 1000 (144.9s) | batch_loss=2.811 | train=2.806 | val=4.107
step 1400 (208.7s) | batch_loss=2.424 | train=2.233 | val=4.036
step 1800 (273.3s) | batch_loss=1.965 | train=1.776 | val=4.076
step 2000 (305.5s) | batch_loss=1.956 | train=1.573 | val=4.149
再用同样的提示词测试生成效果。下面的样例都来自第 2,000 步的模型,当时的 val loss 是 4.149。
输入 犯故意伤害罪,判处:
犯故意伤害罪,判处有期徒刑二年六个月,并处罚金人民币二万元。二、被告人陈某元不服,向最高人民法院申请再审。最高人民法院于2019年12月26日作出(2018)最高法刑再9号刑事裁定:驳回陈某元的再审申请。 本案的争议焦点有二: 一、关于本案主要证据 一、关于陈某2是否构成危险驾驶罪 一、危险驾驶犯罪。二是从一、醉酒后如何具体情况 一、醉酒的问题 一、关于“醉酒危险驾驶人醉酒后是否应当是否应当负有醉酒的危险驾驶行为是否是否具有下列情形之一;二是醉酒驾驶行为是否是否具有下列情形之一的行驶的作用作用,其是否构成危险驾驶行为是否应当予以判断。
输入 《中华人民共和国刑法》第二百六十四条:盗窃公私财物,:
《中华人民共和国刑法》第二百六十四条:盗窃公私财物,数额较大或者有其他严重情节的,处三年以上十年以下有期徒刑,并处罚金;数额特别巨大或者有其他特别严重情节的,处十年以上有期徒刑或者无期徒刑,并处罚金或者没收财产。 伪造、买卖武装部队车辆运输的,依照本法第二百六十四条的规定定罪处罚。 《中华人民共和国刑法》第二百六十五条:非法出售、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、运输、物品的,处三年以下有期徒刑或者拘役,并处或者单处罚金;情节特别严重的,处三年以上十年以下有期徒刑,并处罚金。
对比单字基线,有两个地方需要分开看。
#为什么验证集 Loss 反而从 2.572 涨到了 4.036?
换上 BPE 后,最佳 val loss 从 2.572 变成了 4.036。但两个数的单位已经不一样了。
交叉熵(Cross-Entropy)算的是平均每个 Token 的预测损失。在单字模型里,1 个 Token 就是 1 个字符。学会常见表达后,看到“有期徒”接着预测“刑”,通常会比决定下一个词容易。这些词内位置也参与了平均。
而在 BPE 模型里,“中华人民共和国”、“有期徒刑”、“人民法院”都可以作为一个 Token,模型一次就要预测出这一整段。预测覆盖的文本长度变了,词表也变大了,原始 Loss 自然不能直接和单字模型比较。
可以把损失换算成每字符的值:用每 Token 的 Loss 乘以参与评估的 Token 数,再除以这些 Token 对应的字符数。这次实际抽到的 5 个验证 batch 里,共有 81,920 个目标 Token,对应 128,743 个字符,所以最低的 4.036 换算后是 $4.036 \times 81920 / 128743 \approx 2.568$,和单字基线的 2.572 很接近。到第 2,000 步,4.149 换算后约为 2.640,反而稍高一些。
两版抽到的原文窗口不同,2.568 和 2.572 这点差别还不足以说明 BPE 更好。至少可以看出,原始 Loss 从 2 点多变成 4 点多,并没有看起来那么夸张。
#为什么 1,400 步后开始过拟合,且生成结果喜忧参半?
再看生成文本,变化也没有那么整齐。
第一段判词里的“有期徒刑二年六个月”和法院名称、日期、案号,比前面单字版的例子顺了不少。但后半段又出现了“运输、运输、运输”、“是否是否”和“作用作用”。常见片段合成一个 Token 后,片段内部不用再逐字拼接,片段之间怎么接仍然要靠模型学。
日志里能看清的是过拟合:第 1,400 步后,训练 Loss 继续下降,验证 Loss 却从 4.036 回升到了 4.149。我们的训练集只有约 144 万字符,BPE 编码后剩下约 91.8 万个 Token。同样训练 2,000 步,模型总共预测 $2000 \times 32 \times 512 = 32,768,000$ 个 Token,按训练集长度折算,单字版相当于看了约 22.7 遍,BPE 版约 35.7 遍。代码是随机抽取文本窗口,并非真的按顺序遍历,但这个数能说明:步数相同,BPE 版会更频繁地看到同一批文本。不过,仅凭几段输出,还不能确定复读具体是哪个环节造成的。
BPE 让同样的窗口能装下更多原文。接下来保留这套 vocab_size = 4,096 的分词器,只换位置编码,看看这一步会带来什么变化。
#3. 位置编码:从绝对位置编码(wpe)到旋转位置编码(RoPE)
在原版 GPT-2 中,为了让注意力机制感知语序,模型在最底层的输入端维护了一张形状为 (block_size, n_embd) 的可学习绝对位置编码表 wpe,在进入 Transformer 层之前直接把词向量与位置向量相加:
$$x_t = \text{wte}(\text{idx}_t) + \text{wpe}(t)$$
这种做法给每个绝对位置分配了一个向量。一段话从第 10 个位置挪到第 300 个位置,虽然内部各个词之间的距离没变,加上去的位置向量却换了一套。模型虽然能靠海量数据硬学,但 wpe 本身并不知道这两个位置上的词距其实是同一回事。
这次换成苏剑林等人在 RoFormer 论文中提出的 RoPE(Rotary Position Embedding,旋转位置编码),直接在每层 Attention 里引入相对位置信息。
具体做法是:去掉输入端的 wpe,在每一层算自注意力之前,对每个 Token 的 Q 向量和 K 向量做旋转(V 保持不变)。
在我们的代码里,单个注意力头的向量长度是 $d = 384 / 6 = 64$。要在二维平面上旋转,每次需要两个数当横纵坐标,所以我们先把这 64 个数两两配对,分成 32 组:第 $0$ 组是 $(q_0, q_1)$,第 $1$ 组是 $(q_2, q_3)$,一直到第 $31$ 组 $(q_{62}, q_{63})$。
如果当前 Token 排在输入窗口内的第 $m$ 个位置($m = 0, 1, \dots, 511$),就把它的第 $i$ 组坐标($i = 0, 1, \dots, 31$)旋转 $m\theta_i$ 的角度:
$$\begin{pmatrix} \hat{q_{2i}} \\ \hat{q_{2i+1}} \end{pmatrix} = \begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix} \begin{pmatrix} q_{2i} \\ q_{2i+1} \end{pmatrix}$$
即
$$\hat{q_{2i}} = q_{2i} \cos m\theta_i - q_{2i+1} \sin m\theta_i$$
$$\hat{q_{2i+1}} = q_{2i} \sin m\theta_i + q_{2i+1} \cos m\theta_i$$
同样地,如果另一个 Token 排在输入窗口内的第 $n$ 个位置,它的 K 向量也按完全相同的规则,把第 $i$ 组坐标 $(k_{2i}, k_{2i+1})$ 旋转 $n\theta_i$ 的角度。
注意公式里的 $m$ 和 $i$ 是两个完全不同方向的下标,互不相干:
- $m$(或 $n$)是输入窗口内的 Token 位置(横向上是第几个 Token,范围 $0 \sim 511$),只有它才代表文字的先后顺序;
- $i$ 是单个注意力头内 Q、K 向量的组号(纵向上 64 个特征里的第几对,范围 $0 \sim 31$)。每一组都有自己固定的基础步进角 $\theta_i = 10000^{-2i/d}$:$i$ 越小,每前进一个 Token 转过的角度越大,像秒针;$i$ 越大,角度变化越慢,像时针。不同组用不同的速度表示位置变化。
用 $R_m$ 表示位置 $m$ 对应的旋转矩阵,位置 $m$ 的 Q 与位置 $n$ 的 K 旋转后再做内积,结果只取决于两个位置的夹角差:
$$(R_m q)^T(R_n k) = q^T R_{n-m} k$$
这样,位置就通过相对距离 $n-m$ 进入内积了。保持旋转前的 Q、K 不变,把位置从 8 和 5 同时挪到 408 和 405,内积不会变。这里的前提不能省:如果一句话换了上下文,深层算出来的 Q、K 也可能改变,不能说只要距离相同,注意力分数就一定相同。
删掉 wpe 表还省下了 $512 \times 384 = 196,608$ 个参数,模型参数量从 12.39M 降到 12.19M。
- BPE + RoPE 完整源码:llm_tut4.2_nano_gpt_bpe_rope.ipynb (Colab)
下面是 RoPE 版先前训练 2,000 步的日志,vocab_size = 4,096 的 BPE 分词器和超参数与上一节相同。链接里的 notebook 后来把 max_iters 改成了 1200,用来做后面的补充对比。
step 1 ( 0.9s) | batch_loss=8.394 | train=8.023 | val=8.021
step 200 ( 33.4s) | batch_loss=4.596 | train=4.583 | val=4.902
step 600 (109.7s) | batch_loss=3.264 | train=3.164 | val=4.056
step 1000 (184.8s) | batch_loss=2.419 | train=2.385 | val=3.842
step 1200 (222.2s) | batch_loss=2.471 | train=2.086 | val=3.820
step 1400 (260.0s) | batch_loss=2.057 | train=1.823 | val=3.825
step 1800 (334.8s) | batch_loss=1.621 | train=1.376 | val=3.930
step 2000 (372.4s) | batch_loss=1.615 | train=1.184 | val=4.002
因为第 2 节和第 3 节用的是同一套 BPE 分词器,验证时抽到的 Token 窗口也相同,这次两个模型的 Loss 可以直接横向对比。
这次运行里,RoPE 在第 600 步的验证 Loss 为 4.056,已经接近上一节 wpe 在第 1,400 步达到的最佳值 4.036。到了第 1,200 步,RoPE 的验证 Loss 降到 3.820,比 wpe 的最佳值低 0.216。按同一批验证 Token 对应的字符数换算,$3.820 \times 81920 / 128743 \approx 2.431$,也低于单字基线的 2.572。不过,字符版和 BPE 版的窗口没有对齐,和单字基线的比较只能作为参考。
RoPE 后面也开始过拟合了:第 2,000 步时,验证 Loss 回升到 4.002,仍低于 wpe 同一步的 4.149。这次两版都只跑了一个随机种子,可以先记下这个结果,还不能当成每次都会有的提升。另外,RoPE 版跑完约 372 秒,wpe 版约 306 秒,训练步数上的优势也不等于每一步算得更快。
再来看看同样提示词下的生成结果。我们先看训练结束时(第 2,000 步)的模型输出,后面再对比一下第 1,200 步(验证 Loss 最低点)的权重。
先看前面用过的民法典提示词,截取开头一段:
《中华人民共和国民法典》第五百七十七条:当事人一方不履行合同义务或者不履行对方义务义务的义务义务义务的,对方可以要求其履行该义务。对方不履行预付义务致使不能实现合同目的的,对方可以解除合同并可以随时解除合同,但是可以获得相应的违约责任等民事法律行为的效力。
这段一开头就卡在“义务”上,复读显然还在。
输入 《中华人民共和国刑法》第二百六十四条:盗窃公私财物,:
《中华人民共和国刑法》第二百六十四条:盗窃公私财物,数额较大的,处二年以下有期徒刑或者拘役,并处或者单处罚金;数额巨大或者有其他严重情节的,处二年以上十年以下有期徒刑,并处罚金。 前款规定的公司、企业或者其他单位的工作人员滥用公务员,情节严重的,依照前款的规定处罚。 《中华人民共和国刑法》第二百二十五条:有下列情形之一,进行信用卡诈骗活动的,处三年以下有期徒刑或者拘役,并处或者单处二万元以上十万元以下罚金;数额巨大或者有其他严重情节的,处三年以上十年以下有期徒刑,并处二万元以上二十万元以下罚金;数额特别巨大或者有其他特别严重情节的,处十年以上有期徒刑或者无期徒刑,并处五万元以上五十万元以下罚金或者没收财产: (一)使用伪造、变造的单据、文件、印章的; (二)伪造、变造、买卖或者盗窃、买卖或者使用伪造的货币的; (三)出售伪造、变造、买卖国家机关、人民团体印章的;
输入 《中华人民共和国劳动合同法》第三十九条:劳动者有下列情形之一的,用人单位可以解除劳动合同::
《中华人民共和国劳动合同法》第三十九条:劳动者有下列情形之一的,用人单位可以解除劳动合同: (一)劳动者与用人单位协商一致或者协商内容,协商一致的; (二)用人单位依照本法第三十条规定解除劳动合同的; (三)未依法解除或者终止劳动合同的; (四)用人单位的规章制度违反法律、法规的规定,给劳动者造成损害的; (五)劳动者严重损害劳动者权益的; (六)用人单位以欺诈、胁迫的手段或者乘人之危,给劳动者造成损害的。 《中华人民共和国劳动合同法》第三十条:用人单位违反本法规定解除或者终止劳动合同的,用人单位应当承担赔偿责任;劳动者依法解除或者终止劳动合同的,劳动者应当继续履行。 《中华人民共和国劳动合同法》第三十一条:用人单位违反本法规定解除或者终止劳动合同的,应当承担赔偿责任。 《中华人民共和国劳动合同法》第三十二条:劳动者有下列情形之一的,用人单位可以解除劳动合同:
输入 犯故意伤害罪,判处:
犯故意伤害罪,判处有期徒刑二年六个月,并处罚金人民币三万元。宣判后,被告人金某均未上诉,判决已发生法律效力。 法院生效裁判认为,被告人金某均未上诉,称金某均系自首,且不构成危险驾驶罪。《中华人民共和国消费者权益保护法》第二十条第三款规定:“数罪应当判处死刑,但是应当判处死刑,但是应当判处死刑即达到死刑缓期执行的,一般应当判处死刑,死刑缓期执行的,应当依照死刑缓期执行的刑期,剥夺政治权利。”《最高人民法院关于审理死刑缓期执行期间应准死刑缓期执行期间,缓期执行期满后依法减为无期徒刑的,应当在死刑缓期执行的期间以后依法减为无期徒刑后,终身监禁刑期自死刑缓期执行期满之日起计算。 本案中,被告人金某在上诉期内曾收到上诉的,应当在
这几个样例里,《刑法》没有再卡在“运输”上,《劳动合同法》也列完了 (一) 到 (六),接着写出 第三十条、第三十一条、第三十二条,这几处比上一版整齐。不过,它还是从提示词里的第三十九条退回了第三十条;民法典样例反复念“义务”,最后一段判词也反复念“应当判处死刑”。只能说部分样例的列表和编号更完整,复读的问题还没解决。
至于最后一段把死刑规定安到《消费者权益保护法》头上,还是一本正经地胡说八道。这次 Loss 降低了,不代表法律内容就写对了。也不能从这几段文字断定,错误全是参数太少或者语料不够造成的。
上面用的是第 2,000 步已经开始过拟合的模型。我们又按相同设置重新训练到第 1,200 步,验证 Loss 同样是 3.820,再用相同的提示词和采样设置试了一次。
比如同样输入 《中华人民共和国劳动合同法》第三十九条:劳动者有下列情形之一的,用人单位可以解除劳动合同:,第 1,200 步生成:
《中华人民共和国劳动合同法》第三十九条:劳动者有下列情形之一的,用人单位可以解除劳动合同: (一)用人单位与用人单位协商一致或者协商解决,劳动者协商解决或者解除劳动合同的; (二)用人单位在劳动合同的; …… (八)劳动合同的劳动合同的; (九)用人单位依照本法第三十八条规定解除劳动合同的; (十)用人单位在劳动合同的其他、劳动合同制度或者本单位解除劳动合同的; (十三)用人单位与劳动者订立或者终止劳动合同的劳动关系终止劳动合同的。 用人单位的会与或者用人单位订立劳动合同的劳动合同,用人单位应当向劳动者支付劳动合同的劳动合同的劳动合同解除
输入 犯故意伤害罪,判处,第 1,200 步生成:
犯故意伤害罪,判处有期徒刑二年六个月,并处罚金人民币二十万元;被告人严格报告执行有期徒刑三年,剥夺政治权利一年,并处罚金人民币一万元;被告人刘海涛犯开设赌场罪,判处有期徒刑一年,并处罚金人民币五万元;被告人刘海涛犯开设赌场罪,判处有期徒刑一年,并处罚金人民币一万元。宣判后,陈淑娟和赵延海犯开设赌场罪,判处有期徒刑二年,并处罚金人民币六万元;被告人刘海涛犯开设赌场罪,判处有期徒刑二年三个月,并处罚金人民币一万元。
第 1,200 步的输出里,《民法典》开头不再卡在“义务义务”上,《消费者权益保护法》判死刑那段也没出现。但“(二)用人单位在劳动合同的”、“(八)劳动合同的劳动合同的”还是不通顺,末尾依然卡在“劳动合同的劳动合同的”上,判词里也念叨了三次“被告人刘海涛犯开设赌场罪”。这次第 1,200 步的模型在验证集上预测得更好,但不保证每个提示词生成的结果都更好,复读也没有消失。
#4. 小结
把这三次运行的结果放在一起。表里的最佳值来自每 200 步的评估记录。
| 实验版本 | 分词词表 | 位置编码 | 参数量 | 最佳验证集 Loss(Token 级) | 折算每字符 Loss(约) | 这次观察到的变化 |
|---|---|---|---|---|---|---|
| 4.0 Baseline | 单字 (3,062) | 绝对位置 wpe | 11.99M | 2.572 | 2.572 | 能模仿句式,仍有“代表示”、“第级”这样的错误;窗口覆盖 512 个字符 |
| 4.1 + BPE | BPE (4,096) | 绝对位置 wpe | 12.39M | 4.036 | 2.568 | 窗口平均覆盖约 800 个字符;第 1,400 步后验证 Loss 回升,生成仍有复读 |
| 4.2 + BPE + RoPE | BPE (4,096) | 旋转位置 RoPE | 12.19M | 3.820 | 2.431 | 少了约 20 万参数,最佳验证 Loss 降低 0.216;部分列表更整齐,复读和内容错误仍在 |
下一篇继续改 Transformer Block:试试用 RMSNorm 替代 LayerNorm、用 SwiGLU 替代 GELU MLP,每次只换一处,看结果怎么变。推理部分再加入 KV Cache,看看能省下多少重复计算。