大模型都在用 Token 切词,但 Meta 和华盛顿大学这次偏要反着来。一句话总结:让小模型直接从字节学老师,上限反而更高,End-Of-Token 方法预测平均准确率达 52.4%,比传统 Token 蒸馏高出 4 个百分点。
说真的我没想到,学得慢的那个,反而是跑得远的那个。你想啊,别人都是让学生背整块的词,老师说 isu 是一个整体,你就记 isu 是一个整体。这篇论文说,不如把词拆成字母,一个字节一个字节学。今天我们就把这事掰开讲清楚。
先说结论:字节蒸馏预测上限52.4%,超Token蒸馏4个百分点
学习单位更小,上限反而更高,这是这篇论文最反直觉的结论。
团队以 Llama 3-8B 做教师,学生 Transformer 层统一按 约 12.8 亿参数 对齐,再按 Scaling 曲线外推能力上限。三组蒸馏的预测平均准确率是:Token 蒸馏 48.4%,Marginalize-It 字节蒸馏 50.5%,End-Of-Token 字节蒸馏 52.4%。你没看错,最好的字节方法比 Token 蒸馏高出 4 个百分点,也是六组里最高的。
代价和便宜都写得很实在。便宜的是数据和硬盘:实际文本量约 Token 方案的六分之一,分布存储量约五分之一。贵的是计算:End-Of-Token 处理同样文本要多花约 30.94% 计算,推理也没省下来。说白了,这不是省钱方案,是花算力换上限的方案。
Token蒸馏的瓶颈:12.8万候选与top-k截断
蒸馏不是只告诉学生正确答案,是把老师的犹豫也教给学生。普通训练只学“下一个词是哪个”,蒸馏还要学“老师觉得每个候选各有几成把握”。
问题就卡在这里。Llama 3-8B 的词表有 128256 个 Token,每个位置都是十几万选一。要把完整概率存下来做离线蒸馏,硬盘根本扛不住,只能保留概率最高的 top-k,比如 top-600,剩下的直接扔掉。
打个比方,这就像老师改卷子写了满满一页评语,复印机太贵,你只复印了分数最高的那几行。学生看到的老师,其实是个删减版。
字节模型把这道题变简单了。一个字节只有 256 种取值,加几个特殊符号也就两百多个候选。每个位置的预测空间小了三个数量级,完整分布不用截断就能存下来。复印整页评语突然变得便宜了,学生终于能看到老师完整的犹豫了。
两种字节级转换方法:Marginalize-It与End-Of-Token
老师预测的是整块 Token,学生预测的是单个字节,怎么对上?团队拿 Tiramisu 举了个例子,它可能被切成 T、iram、isu 几块。
假设老师觉得下一个可能是 isu、isk 和 is,这三个都以字节 i 开头。那学首个字节时,就把所有以 i 开头的候选概率加起来。等真实前缀走到 i,再只看前缀匹配的候选,继续算下一个字节。顺着真实字节序列一步步拆,Token 分布就变成了逐字节目标。
两种方法的分歧,在 Token 长短不一的地方。走到前缀 is 时,isu 和 isk 后面还有料,is 这个 Token 已经结束了,后面接什么,得再问一次老师,计算量会炸。
Marginalize-It 选择不再往后追,把剩余候选重新归一化。 好处是只需要一次教师前向计算,坏处是扔掉了一部分已结束 Token 的概率,本质是近似。
End-Of-Token 给结束本身留了一个预测位,在每个 Token 末尾加一个特殊符号。 走到 is 时,学生可以预测 u,可以预测 k,也可以预测结束符,原来悬空的那部分概率有了明确去处。同样只需一次教师前向,还保留了 Token 边界,映射更完整。
你可以理解成查快递:一种是包裹签收了就不管了,另一种是要求收件人必须回个“已签收”,物流单才算走完。回执多了约三成计算,但链条完整了。
实验设置:Llama 3-8B为师,三类学生六组对照
实验设计得很干净,变量一次只动一个。
学生分三类:Token、普通字节、带结束符的字节,Transformer 层都是约 12.8 亿参数。算上词表后,总参数是 Token 学生约 18.1 亿,字节学生约 12.8 亿。每类再分监督训练和蒸馏训练,一共六组,既能看蒸馏有没有用,也能看换学习单位有没有用。
训练量给得很足,字节模型训练覆盖约万亿字节规模,再拿选择题问答、语言生成、机器翻译三类任务、八项评测打分。
训练过程出现了一个有意思的分叉。Token 模型前期学得更快,但很快摸到平台期;字节模型起步更慢,却随着计算量增加持续改善。 一个像短跑选手,开局就冲,一个像长跑选手,越到后面越稳。
这里还有个坑。只看 BPB,也就是统一到每字节的预测损失,几种字节模型很早就超过了 Token 模型,但真去做题,分数反而更低。因为 BPB 考的是文本预测准不准,下游考的是答案选得对不对,两回事。团队没拿 BPB 直接判胜负,而是拟合了训练计算量到 BPB 再到下游成绩的外推曲线。
实验结果与代价:文本六分之一、存储五分之一,计算多三成
外推结果就是开头那组数字:Token 蒸馏 48.4%,Marginalize-It 50.5%,End-Of-Token 52.4%。按曲线推,End-Of-Token 在约 6.33×10²² FLOPs 时追平 Token 蒸馏的上限,之后继续拉开。
有个细节很诚实。普通字节监督模型预测上限 51.2%,反而高于 Marginalize-It 蒸馏的 50.5%。 论文的解释是,Marginalize-It 为省计算扔掉的那部分概率,确实用成绩还了回来。而保留更完整的 End-Of-Token,同时超过了自己的监督版本和其他所有方案。
省下来的东西也很具体。在 Token 只存 top-600、字节存完整分布的设置下,字节处理的实际文本量约六分之一,存储量约五分之一。分词器变了也不怕,因为字节基础取值不随分词器改变,换词表不用重搞一套。
没省下来的也要说清楚。End-Of-Token 每个 Token 后多加一个符号,训练计算比普通字节高约 30.94%。 推理成本也没占便宜,论文承认还没做等推理成本下的公平比较。想兑现更高的上限,得先准备好更多的卡。
如何理解字节蒸馏这条小模型路径
过去我们默认词是最小学习单位,这篇工作换了个视角:词是人定的切分,字节才是机器层面更稳定的切分。
对做小模型的人来说,诱惑很直接。用更少的文本、更小的存储,换一个更高的天花板,部署时还没有十几万的大词表压着,12.8 亿级别的学生就有机会摸到原来够不着的分数。尤其分词器经常换的多语言、低资源场景,字节天然更抗折腾。
冷静的部分也有两句。上限是外推出来的,不是今天已经跑出来的,曲线要兑现,还得把那 6.33×10²² FLOPs 量级的计算 真投进去。而且推理变长、训练变贵,端侧实时场景未必划算。
留个好玩的收尾:大模型学了这么多年“背单词”,这次发现,把单词拆成字母表,一个字母一个字母抄老师的作业,反而抄出了更高的分数。下次你家小孩背单词背到哭,可以安慰他,Meta 证明了,拆开学,真有用。