寻找最优分词器

# 寻找最优分词器 来源：https://blog.aqnichol.com/2026/06/10/optimal-tokenizers/ 在这篇文章中，我将介绍一种算法，它能够在某些设置下计算出最优分词器。这项成果很酷，因为最优分词在理论上被认为是难解的（https://arxiv.org/abs/2411.08671），但在实践中似乎可以求解。我的发现与旅行商问题（TSP）的多种结果非常相似，即即使是很困难的实例，也可以通过割平面法（https://en.wikipedia.org/wiki/Cutting-plane_method）技术来最优求解。 需要指出的是，虽然这个结果很酷，但它并不一定*有用*，原因有几点。首先，现有的最优方法已经相当接近最优（通常差距在1%以内）。其次，即使一个分词器在*训练数据*上是最优的，它在留出测试数据上的泛化能力可能不如其他分词器。最后，低效的分词器基本上也没问题：你可以通过稍微增加词汇表大小来弥补低效分词器带来的成本。 尽管存在上述注意事项，我在这个项目上的工作过程非常愉快，也希望其他人能对这个问题的前沿探索产生兴趣。 ## 背景：分词器 前沿的大语言模型通常使用称为*tokens*（词元）的整数序列进行训练。每个词元对应某个字节序列，这些字节序列通常对应常见的单词。例如，在GPT-5分词器中，词元290对应字节“ the”，6602对应“ token”，因此文本“ the token”可以被编码为序列[290, 6602]。 从词元到字节的映射，即“词汇表”，是在大语言模型训练之前就固定的。通常，我们会尝试找到一个能压缩一部分训练数据的词汇表。具体来说，我们希望选择一个固定大小的词汇表，使编码数据所需的词元数量最小化。寻找此类词汇表的主流技术是字节对编码（BPE）（http://www.pennelynn.com/Documents/CUJ/HTML/94HTML/19940045.HTM），这是一种已有数十年历史的贪心压缩算法。 ## 分词作为整数线性规划 在最近的一篇论文中，Tempus 等人（https://arxiv.org/abs/2605.22821）将分词与整数线性规划联系起来。他们方法的基本思想是将整个数据集的分词表示为一组整数变量。 在这种表述中，每个可能的词汇表条目都有一个“颜色”变量。具体来说，我们为数据集中每一个唯一的子串创建一个颜色变量。如果对应的字节序列在词汇表中，则颜色变量为1，否则为0。我们添加一个约束，强制所有颜色变量的总和等于目标词汇表大小。 一种颜色对应某个字节序列，但给定的字节序列可能在数据集中出现多次。对于每次出现的颜色，都有一个单独的“边”变量。这些边共同编码了数据集的实际分词。如果一条边为1，则在该特定位置使用该边对应的词元。我们线性规划的目标是最小化所有边变量的总和，即编码数据集所使用的词元数量。 例如，在下图中，我们将单词“Queue”分词为词元[“Q”, “ue”, “ue”]。我们也可以将其分词为[“Qu”, “e”, “ue”]，但这并不是当前ILP解所指示的分词，因为初始的“Qu”和“e”边的变量为0。 一幅图展示了单词Queue的整数分词，包含边变量和颜色变量。 我们以两种方式约束这个线性规划。首先，如果某个词元不在词汇表中，我们就不能使用它。为此，我们约束每个边变量小于或等于其对应的颜色变量。其次，我们要确保以唯一有效的方式对数据集进行分词。为此，我们添加流量约束：对于数据集中的每个字节位置，我们希望流入该位置的边之和等于流出该位置的边之和，边界位置除外。对于第一个和最后一个位置，我们要求流出或流入为1。在整数解中，你可以将流量约束理解为以下断言：*任何边进入的位置都必须有一条边离开，除了第一个和最后一个位置*。 如果所有变量都是整数且约束在[0, 1]范围内，那么这个线性规划就足以编码最优分词。然而，由于我们无法高效求解任意的整数线性规划，Tempus 等人将ILP松弛为连续LP，并使用优化良好的求解器进行求解。 连续LP的解通常不是整数解。下面的例子展示了这一点：单词“Queue”有两种重叠的分词方式：要么编码为[“Q”, “ue”, “ue”]，要么编码为[“Qu”, “e”, “ue”]。这个解的问题在于，我们的颜色变量之和为2.5，但实际上我们使用了总共四种颜色，因此并没有找到大小为3的最优词汇表。一般来说，我们最终可能会得到比目标词汇表大小多得多的非零颜色变量。 一幅图展示了单词Queue的分数分词，包含分数边变量和分数颜色变量。 Tempus 等人提出用几种不同的方式对颜色变量进行“舍入”，从而得到一个整数但次优的ILP解。连续LP的解给出了最优解词元数量的*下界*，而舍入后的分词器给出了上界。 关于这项工作，我还想提一点：为了使其可行，我们对数据集进行预分词（拆分成单词），并合并重复单词（根据单词出现次数在目标函数中赋予相应权重）。这大大减少了LP中的变量数量，但这也意味着我们的解只是在预分词器下的“近似最优”。今天，我不会试图去除这个限制，但这将是未来工作的一个有趣方向。 ## 割平面 去年我花了一些时间研究旅行商问题（TSP），它同样可以表述为ILP。我们通常可以使用割平面法（https://en.wikipedia.org/wiki/Cutting-plane_method）来求解这个ILP：首先，将ILP转化为连续LP，然后添加额外的约束，直到最优解为整数。这些约束必须是可证明“有效”的——即对于实际的整数解来说，它们永远不会被违反。理论上，任何ILP都可以通过添加额外约束转化为连续LP，但神奇的额外约束可能难以找到。TSP求解器使用多种启发式方法，在大多数实际情况下高效地找到这样的约束。Concorde（一个TSP求解器）的作者写了一整本书来介绍寻找有用割的技术。 在阅读Tempus等人的论文后，我想知道是否可以将割平面法应用于分词ILP。这种方法的工作流程如下：首先，求解初始LP，得到最优分词的下界和上界；然后，不断向LP添加有效割并重新求解，使这些界限越来越接近，直到它们相遇于最优解。 提出对ILP可能有用的“割族”需要大量工作和创造力，所以我没有自己绞尽脑汁，而是让Codex来完成这个任务。起初，它几乎没找到什么——有些割稍微改善了LP下界，但大多数尝试都停留在表面层次的单词启发式。 然后我尝试了另一种方法：暴力搜索。一个“割”是某个被所有整数解满足，但被当前分数LP解违反的约束。我们可以通过构造一个辅助线性规划（每行对应一个可能的整数解）并优化它以最大化对分数解的违反程度来寻找割。我们不能对整个LP这样做，因为行数会指数级增长，但我们可以对LP中小的、有趣的“投影”这样做。Codex建议查看具有共同分数颜色的单词对或三元组中的所有变量。 上述技术找到了非常好的割，它们改进了舍入分词器并提升了下界。然而，这种方法的效率非常低，因为它需要为大量单词对求解（相当大的）辅助LP。接下来的技巧是让Codex查看我们实际找到的割。 通过观察暴力搜索得到的割，Codex发现了几个可以更高效找到的割模板。最有效的族似乎是Codex命名的“循环约束”。这种技术会找出当前LP解中重叠的分数边对。例如，我们可能会找到颜色A和B的重叠（即冲突）边对。然后我们找出几对共享共同颜色的边，例如颜色B和C的另一对，以及颜色C和A的另一对。接着，我们可以根据相应的边和颜色变量创建一个约束，这个约束常常被连续LP解违反，但不会被任何合法的整数解违反。 找到冲突边对AB、BC、CA构成的循环可以通过一个巧妙的技巧完成：构造一个图，其中顶点是颜色，并将当前解中作为分数边重叠的任何一对颜色连接起来。得到这个图后，运行DFS寻找其中的循环。Codex自主实现了这一切，尽管我确信这并非原创技巧。 ## 实验设置 这个项目的硬件资源相当有限，我只有Mac Studio和Mac mini。这些硬件上没有很好的GPU加速LP求解器，所以我主要依赖HiGHS（https://highs.dev/）单核单纯形求解器。遗憾的是，我发现这个求解器有时会卡住，尤其是在后期迭代中施加了大量（可能退化的）割之后。 为了在这些硬件上在合理时间内运行实验，我研究了单本电子书。我需要保持LP足够小以便在CPU上求解，因此保留了Tempus等人的预分词方法。 最后，我采用了Tempus等人的一些启发式方法来缩小LP，例如删除出现次数少于5次的子串的颜色变量。我还对颜色的字节长度施加了限制——本例中为16字节。我发现这相对于8字节限制有改进，后者得到的最优分词效果稍差。 ## 结果 我至少在一些小玩具问题上找到了可证明的最优分词器。我最引以为豪的是为书籍《傲慢与偏见》找到的词汇表大小为512的最优分词器。算法经过大约十几轮迭代收敛，耗时一天多一点。 我在相同问题上尝试将词汇表大小从512增加到1024，发现仅靠循环约束不足以找到最优解。在我添加其他割族后，下界继续显著移动，尽管我最新的运行尚未完成。毫无疑问，这里还有待发现的其他割族，其中一些甚至可能是解决1024词汇表问题所必需的。 ## 未来工作 目前，我实验中的主要瓶颈是LP求解时间。在许多实验中，每次LP求解可能需要数小时到数天。我尝试了几个求解器（HiGHS、SCIP中的求解器以及OR-Tools PDLP），它们都在我的高度约束LP上开始变得吃力。我怀疑我的割平面方法产生了退化的LP，这可能是潜在的改进方向。 总的来说，我非常希望看到有人继续将这项工作扩展到更大的语料库。我怀疑我所探索的割族对于更困难的问题并不足够，而且肯定存在丰富的探索空间。 我也希望看到有人去除预分词器。这目前使得LP相当大，因为我们无法合并重复单词。去除预分词器也会消除使用基于单词的割策略的能力。例如，我的一些割策略枚举了每个单词的所有有效整数解，然后将这些组合投影到一个变量子集上。这些策略需要为“单个超大单词”数据集完全重新设计。 ## 结论 这是一个很棒的项目，看到Codex只需我少量指导就能完成整个研究循环，非常有趣。我真的很希望继续玩下去，但这取决于能否找到解决*慢LP*问题的办法。 这个项目极其粗糙的Codex实现可在Github（https://github.com/unixpickle/tokenizer_lp）上找到。供参考，我找到的《傲慢与偏见》（https://www.gutenberg.org/cache/epub/37431/pg37431.txt）的最优词汇表在此（https://gist.github.com/unixpickle/55079366aca7818f65697da36aca0652）（注意，词汇表实际上是510，因为代码库保留了两种特殊词元）。