寻找最优分词器：从整数线性规划到切割平面法

背景

前沿大语言模型（Frontier LLMs）通常是在由整数序列组成的“标记（tokens）”上进行训练的。每个标记对应一段字节序列，这些字节序列通常对应常见的单词或子词。例如，在 GPT-5 的分词器中，标记 290 对应字节序列 “ the”，而 6602 对应 “ token”，因此文本 “ the token” 可以被编码为序列 [290, 6602]。

从标记到字节的映射被称为“词汇表（vocabulary）”，它在 LLM 训练之前就已经固定。通常，我们的目标是找到一个能最大程度压缩训练数据切片部分的词汇表。具体来说，我们希望选择一个固定大小的词汇表，以最小化编码数据所需的标记数量。目前，寻找此类词汇表的主导技术是字节对编码（Byte-Pair Encoding, BPE），这是一种拥有数十年历史的贪心压缩算法。

核心内容

将分词问题转化为整数线性规划

Tempus 等人最近发表的一篇论文将分词问题与整数线性规划（Integer Linear Programming, ILP）联系起来。其基本思路是将整个数据集的分词过程表示为一组整数变量。

在这种表述中，每个可能的词汇表条目都有一个“颜色（color）”变量。具体而言，我们为数据集中的每个唯一子串创建一个颜色变量。如果对应的字节序列在词汇表中，则该颜色变量为 1，否则为 0。我们添加一个约束条件，强制颜色变量的总和等于目标词汇表大小。

虽然一个颜色对应某些字节序列，但这些字节序列在数据集中可能出现多次。对于颜色的每次出现，都有一个单独的“边（edge）”变量。这些边共同编码数据集的实际分词。如果某个边为 1，则表示在该特定位置使用了该边对应的标记。线性规划的目标是最小化所有边变量的总和，即用于编码数据集的标记数量。

例如，在以下情境中，我们将单词 “Queue” 分词为标记 [“Q”, “ue”, “ue”]。我们本可以将其分词为 [“Qu”, “e”, “ue”]，但这并不是当前 ILP 解决方案指示的分词方式，因为初始 “Qu” 和 “e” 边的边变量为 0。

我们通过两种方式对线性规划（LP）进行约束：

1. 词汇表约束：如果某个标记不在词汇表中，就不能使用它。为此，我们将每个边变量约束为小于或等于其对应的颜色变量。
2. 流约束（Flow Constraints）：确保我们以唯一有效的方式对数据集进行分词。为此，我们添加流约束：对于数据集中的每个字节位置，流入该位置的边之和应等于流出该位置的边之和，边界除外。对于第一个和最后一个位置，我们希望流出或流入为 1。在整数解中，流约束断言：任何边进入的点必须有一个边从该点流出，第一个和最后一个位置除外。

如果所有变量都是整数且被约束在 [0, 1] 范围内，则该线性规划足以编码最优分词。然而，由于我们无法高效地解决任意整数线性规划，Tempus 等人将 ILP 松弛为连续 LP，并使用经过良好优化的求解器进行求解。

连续解的局限性与舍入

连续 LP 的解通常不是整数。我们可以看到一个例子，其中单词 “Queue” 有两个重叠的分词：要么编码为 [“Q”, “ue”, “ue”]，要么编码为 [“Qu”, “e”, “ue”]。该解的问题在于，我们的颜色变量总和为 2.5，但实际上我们使用了四个总颜色，因此我们并没有真正找到大小为 3 的最优词汇表。一般来说，我们最终得到的非零颜色变量数量可能远多于我们目标词汇表的大小。

Tempus 等人提出以几种不同方式“舍入（round）”颜色变量，从而获得 ILP 的整数但次优解。连续 LP 的解给出了最优解标记数量的下界，而舍入后的分词器给出了上界。

需要提到的一个注意事项是：为了使问题具有可处理性，我们需要对数据集进行预处理分词（将其拆分为单词）并合并重复的单词（在目标函数中根据单词出现的次数赋予相应的权重）。这极大地减少了 LP 中的变量数量，但这意味着我们的解仅在预处理分词器下是“近最优”的。

切割平面法的应用

作者去年花时间学习了旅行商问题（Traveling Salesman Problem, TSP），该问题也可以表述为 ILP。我们通常可以使用切割平面法（Cutting Planes）来解决这个 ILP：首先，我们将 ILP 转化为连续 LP，然后添加额外的约束，直到最优解变为整数。这些约束必须是证明为“有效（valid）”的——即对于实际的整数解永远不会被违反。理论上，任何 ILP 都可以通过额外约束“转化”为连续 LP，但那些神奇的额外约束可能难以找到。TSP 求解器使用多种启发式方法在大多数实际情况下高效地找到此类约束。

在阅读 Tempus 等人的论文后，作者思考是否可以将切割平面法应用于分词 ILP。该方法的工作原理如下：首先，求解初始 LP 以获得最优分词的一些上下界；然后，不断向 LP 添加有效的切割并重新求解，使这些界限越来越接近，直到它们在最优解处相遇。

利用 Codex 发现切割

提出可能对 ILP 有用的“切割族（cut families）”需要大量的工作和创造力。作者没有亲自尝试，而是让 Codex 承担这项任务。起初，Codex 几乎没找到什么——有些切割使 LP 界限略有改善，但大多数尝试只是表面上的单词启发式方法。

随后，作者尝试了另一种方法：暴力搜索。一个“切割”是满足所有整数解但违反当前分数 LP 解的约束。我们可以通过构建一个辅助线性规划来找到切割，该规划为每个可能的整数解设置一个约束，并优化它以最大化对分数解的违反。由于行数会指数级爆炸，我们不能对整个 LP 这样做，但可以对 LP 的小型有趣“投影”进行此操作。Codex 建议查看具有公共分数颜色的单词对或三元组中的所有变量。

上述技术找到了非常好的切割，改善了舍入分词器并提高了下界。然而，这种方法效率很低，因为它涉及为大量单词对求解（相当大的）辅助 LP。接下来的技巧是让 Codex 查看我们正在实际发现的切割。

关键要点

• 理论困境与实践突破：最优分词在理论上是不可处理的（intractable），但在实践中似乎可以通过特定算法求解。这一发现类似于旅行商问题（TSP），即使是非常困难的实例，也可以使用切割平面技术获得最优解。
• ILP 建模方法：
  • 将分词问题建模为整数线性规划（ILP）。
  • 引入“颜色变量”表示词汇表条目（0 或 1），引入“边变量”表示数据集中特定位置是否使用该标记。
  • 通过约束颜色变量总和等于目标词汇表大小，以及流约束确保分词的有效性，最小化边变量总和（即标记数量）。
• 松弛与舍入：
  • 直接求解 ILP 困难，因此将其松弛为连续 LP 求解。
  • 连续 LP 解通常非整数，需通过“舍入”获得整数解，但这会导致次优结果。
  • 连续解提供下界，舍入解提供上界。
• 预处理限制：为了降低计算复杂度，该方法需要对数据进行预处理分词和去重，这意味着得到的解仅在预处理后的数据分布下是“近最优”的。
• 切割平面法的引入：
  • 借鉴 TSP 求解经验，通过不断添加有效的“切割”约束来收紧 LP 的上下界，直至找到整数最优解。
  • 利用 AI 助手（Codex）辅助发现有效的切割族，从最初的启发式方法转向基于暴力搜索和变量投影的更深层分析。
• 实用性的局限性：
  • 现有最先进技术（如 BPE）已经非常接近最优（通常误差在 1% 以内）。
  • 在训练数据上最优的分词器，在保留测试数据上的泛化能力