In-Place Tokenizer Expansion for Pre-trained LLMs

LLM PTSAD MQIA NLP MLRCC CLKTA
在预训练初始阶段即固定下来的分词器(tokenizer),其词表规模按预训练语料库的分布比例进行分配,从而反映当时部署场景下的优先级。当这些优先级发生改变时,后续新增的语言往往被切分为更多子词(token),导致使用这些语言的用户面临更高的推理延迟、计算开销和能耗。云端大模型可以承载更庞大的词表,因为其嵌入层(embedding)与语言建模头(LM-head)矩阵仅占全部参数量的极小部分;而在紧凑型模型中,这两部分矩阵却占据了单次 token 解码所需带宽的显著比重,因此端侧部署的模型通常采用较小的词表,并默认接受超出预设语言集合之外的语言出现更严重的子词碎片化现象。本文提出“分词器扩展”(tokenizer expansion)方法——一种原位升级已预训练模型分词器的实用方案,适用于模型设计者对其架构拥有完全控制权的情形。该方法在多语种语料上延续原有分词器的字节对编码(BPE)合并过程,使得绝大多数原始词元(source tokens)得以完整保留、不作拆分,而每个新增词元均可被精确分解为若干原始子词。我们直接复用原始嵌入矩阵中对应保留词元的行向量,而对新增词元的嵌入向量,则初始化为其组成原始子词嵌入向量的算术平均值。随后通过两阶段适配:先仅微调嵌入层,再对整个模型开展持续预训练,即可恢复至原始检查点(source checkpoint)的性能水平。我们将该方法应用于持续预训练后的 LFM2-8B-A1B 检查点(一个含 80 亿参数的混合专家模型),成功构建出词表规模达 128K 的升级版本 LFM2.5-8B-A1B。扩展后的分词器将印地语和越南语分别压缩为原始分词结果的约 2.4 倍和 2.6 倍更少的 token(泰语甚至可达最高 4.0 倍的压缩比)。结合实测所得的大词表下每个 token 的解码开销,我们估算出:在各类基准设备上,上述语言的每字符解码速度可提升 2.2 至 3.7 倍。我们已开源该模型权重及扩展后的分词器,并如实报告了在方法探索过程中所发现的若干负面实验结果,这些结果亦深刻塑造了本方案的设计路径。
A tokenizer fixed at the start of pre-training allocates vocabulary in proportion to the pre-training corpus, reflecting the deployment priorities at that time. When those priorities shift, languages added later are split into many more tokens per word, which can raise latency, compute, and energy consumption for users of those languages. Cloud models can afford a broad vocabulary because the embedding and LM-head matrices are a small fraction of their parameters. On a compact model those matrices are a material share of per-token decode bandwidth, so on-device models ship small vocabularies and accept fragmentation outside a fixed language set. We present tokenizer expansion, an in-place recipe for upgrading a pre-trained model's tokenizer when the model producer controls its design. We continue the existing tokenizer's BPE merges on a multilingual corpus, so most source tokens carry over unchanged as single tokens and every new token has an exact decomposition into source tokens. We copy the carried-over embedding rows unchanged and initialize new rows as the mean of their source sub-token embeddings. A two-stage adaptation, embedding-only training then full-model continued pre-training, recovers source-checkpoint quality. We apply the recipe to a continued pre-trained checkpoint of LFM2-8B-A1B, an 8B-parameter Mixture-of-Experts model, to help produce LFM2.5-8B-A1B with a 128K tokenizer. The expanded tokenizer encodes Hindi and Vietnamese in roughly $2.4\times$ and $2.6\times$ fewer tokens than the source (up to $4.0\times$ on Thai). Combining these reductions with the measured per-token cost of the larger vocabulary, we estimate a $2.2$-$3.7\times$ per-character decode speedup for these languages across our reference devices. We release the model weights and the expanded tokenizer, and report the negative findings that shaped the recipe.
许愿