Fast WordPiece 分词：一种适用于 BERT 及其后续模型的线性时间算法

目录

1. 引言

分词是几乎所有自然语言处理任务的基础预处理步骤，从情感分析到机器翻译皆然。BERT、GPT-3、XLNet 等现代模型依赖于子词分词算法，特别是 WordPiece 这类算法，将文本分解为有意义的单元，同时处理词汇表外词。本文解决了一个关键瓶颈：WordPiece 分词算法本身的效率问题。

作者指出，WordPiece 中使用的标准贪婪最长匹配优先策略具有次优的时间复杂度——$O(n^2)$ 或 $O(nk)$，其中 $n$ 是输入长度，$k$ 是最大词元长度。他们提出了 LinMaxMatch，一种新颖的算法，为单词语义分词实现了严格的 $O(n)$ 线性时间复杂度，并提出了一种用于通用文本的集成式 端到端 WordPiece 算法。

2. 背景与相关工作

BERT 中使用的 WordPiece 分词涉及两个步骤：(1) 根据空白和标点符号将文本预切分为单词，(2) 对每个单词应用 MaxMatch 算法，对照固定词汇表进行匹配。MaxMatch 算法也用于中文分词，它迭代地选择剩余字符串中与词汇表词元匹配的最长前缀。

先前的高效算法仅限于 $O(n^2)$ 或 $O(nk)$。$k$ 因子是有问题的，因为词汇表可能包含很长的词元（例如 "supercalifragilisticexpialidocious"）。本文的贡献在于完全消除了这个依赖于词汇表的因子。

3. LinMaxMatch 算法

3.1 核心思想与灵感来源

LinMaxMatch 的灵感来源于 Aho-Corasick 字符串匹配算法。关键见解是将词汇表视为一组模式，将输入单词视为文本。通过在词汇表字典树上预计算辅助数据结构，该算法可以在匹配过程中避免回溯，而回溯正是朴素 MaxMatch 算法产生二次复杂度的根源。

3.2 数据结构：字典树、失效链接与失效弹出

该算法在离线预处理阶段使用从词汇表构建的三个核心数据结构：

• 字典树：一种标准的前缀树，其中每个节点代表一个字符，从根节点出发的路径代表词汇表词元。
• 失效链接（类似于 Aho-Corasick）：对于代表字符串 $s$ 的节点，其失效链接指向代表 $s$ 的最长真后缀的节点，该后缀同时也是某个词汇表词元的前缀。这使得算法能够高效地进行“故障转移”。
• 失效弹出：一项新颖的补充。当通过失效链接从一个节点遍历到另一个节点时，算法可能会“弹出”一个或多个在该节点结束的已匹配词元。这个列表是预先计算并存储的。

3.3 算法流程详解

对于一个长度为 $n$ 的输入单词，算法从左到右单次遍历处理字符。它维护字典树中的一个当前节点。对于每个新字符：

1. 尝试从当前节点沿匹配边前进。
2. 如果成功，移动到子节点。如果该节点是词汇表词元的终点，则将其记录为潜在匹配。
3. 如果不存在匹配边，则跟随当前节点的失效链接。收集关联的“失效弹出”列表中存储的任何词元。然后，从新节点重新尝试匹配相同的输入字符。

这个过程确保每个字符只被处理一次，从而保证了 $O(n)$ 的时间复杂度。

4. 端到端 WordPiece 分词

为了对完整句子或文档进行分词，作者提出将预分词步骤（基于空白/标点符号切分）与 LinMaxMatch 集成到一个统一的算法中。这种端到端 WordPiece 算法相对于总文本长度也以线性时间运行，避免了分阶段处理的开销。它巧妙地在同一个字典树遍历框架内，将边界检测与子词匹配交织在一起。

5. 实验结果与性能

本文提供了令人信服的基准测试，将其实现与两个广泛使用的库进行了比较：HuggingFace Tokenizers 和 TensorFlow Text。

结果表明，理论上的 $O(n)$ 优势转化为了显著的现实世界性能提升，特别是对于较长文本或在移动端推理等延迟敏感的环境中。

6. 核心见解与分析

核心见解

本文的根本突破不仅仅是一个更快的分词器；它正式地将分词识别为一个可通过经典自动机理论解决的字符串匹配问题。通过将 WordPiece 的 MaxMatch 从朴素的贪婪搜索重新定义为 Aho-Corasick 的变体，作者弥合了理论计算机科学与应用 NLP 工程之间长达数十年的鸿沟。这让人联想到 CycleGAN 论文如何将图像翻译重新定义为循环一致性对抗学习问题，从而创造了一个新范式。

逻辑脉络

论证在逻辑上是严密的：1) 识别瓶颈（$O(n^2)$/$O(nk)$ 复杂度）。2) 从经过验证的最优算法（Aho-Corasick）中汲取灵感。3) 使其适应 MaxMatch 的特定约束（需要输出所有匹配的词元，而不仅仅是判断是否存在）。4) 引入新颖的数据结构（失效弹出）来处理这些约束。5) 将解决方案扩展到完整流程（端到端）。从问题到灵感，再到适配和泛化的脉络，堪称应用研究的典范。

优势与不足

优势： $O(n)$ 的保证在理论上优雅，在实践中也得到了验证。预处理成本是一次性开销。正如作者所指出的，并得到 Google AI 等机构设备端 AI 研究趋势的支持，该算法对移动/边缘 NLP 的影响可能是巨大的。

潜在不足/遗漏： 本文简要提及但未深入分析增强字典树（失效链接和弹出）的内存开销。对于大规模词汇表（例如多语言模型），这可能不容忽视。此外，基准测试虽然令人印象深刻，但仅针对特定实现。与其他优化的 MaxMatch 变体（例如使用后缀数组）进行比较，将能强化其主张。

可操作的见解

对于 NLP 从业者：立即分析你的分词延迟。 如果它占你总推理时间的 1-2% 以上，特别是在流式或移动应用中，采用像 LinMaxMatch 这样的算法是一项高投资回报率的优化。对于研究人员：这项工作打开了大门。类似的基于自动机的方法能否优化字节对编码或 SentencePiece？其核心思想——用预计算的状态转移替换回溯——广泛适用于 NLP 中其他贪婪分割算法。

7. 技术细节与数学形式化

LinMaxMatch 的核心可以形式化。令：

• $V$ = 词汇表，一组字符串。
• $T$ = 由 $V$ 构建的字典树。
• 对于字典树节点 $u$，令 $L(u)$ 为从根节点到 $u$ 的字符形成的字符串。
• 失效函数 $f(u)$： 节点 $v$，使得 $L(v)$ 是 $L(u)$ 的最长真后缀，且 $v$ 也是 $T$ 中的一个节点。
• 失效弹出 $P(u)$： 词汇表词元 $t \in V$ 的集合，其中 $t$ 是 $L(u)$ 的后缀，并且存在从根节点到 $f(u)$（或 $u$ 本身）路径上的一个节点 $w$，使得 $L(w) = t$。这对所有节点都是预先计算的。

输入字符串 $S[1..n]$ 的分词过程通过维护一个当前节点 $u$（初始为根节点）和一个索引 $i$ 来进行。当 $i \leq n$ 时：

1. 如果存在 $u$ 的一个子节点 $v$，其边标签为 $S[i]$，则设置 $u = v$，$i = i+1$。如果 $u$ 标记了一个词汇表词元的结尾，则记录它。
2. 否则，输出 $P(u)$ 中的所有词元。设置 $u = f(u)$。不递增 $i$。 如果 $u$ 是根节点且没有子节点匹配 $S[i]$，则 $i = i+1$（处理未知字符）。

这个循环的 $O(n)$ 复杂度是显而易见的，因为 $i$ 要么增加，要么 $u$ 跟随失效链接，而失效链接的跟随可以分摊到每一步为常数时间。

8. 分析框架：案例研究

场景： 为智能手机上的实时语音助手优化分词。当前流程使用来自流行库的标准 WordPiece 分词器。

框架应用：

1. 基线测量： 性能分析显示，每个查询的分词消耗约 15 毫秒，占总流程延迟（125 毫秒）的 12%。目标是 <100 毫秒。
2. 瓶颈分析： 查询通常较长（口语句子）。$O(nk)$ 因子影响显著，因为词汇表中包含技术性复合词。
3. 解决方案设计： 用 LinMaxMatch 替换分词器。在模型部署期间，构建带有失效链接/弹出功能的增强字典树的一次性成本是可以接受的。
4. 预期结果： 基于本文的 5-8 倍加速，分词时间应降至约 2-3 毫秒，将总延迟减少到约 112 毫秒，这是朝着目标迈出的坚实一步。常数时间的保证也提高了延迟的可预测性。

这个案例研究说明了该算法的理论特性如何直接应对现代 AI 系统中的实际约束。

9. 应用前景与未来方向

高效、线性时间分词的影响超越了 BERT：

• 设备端与边缘 AI： 随着类似 GPT 架构的模型被蒸馏以适应手机，预处理中节省的每一毫秒都会累积。这项工作是下一代响应迅速、具备离线能力的 NLP 应用的基础。
• 流式分词： 对于实时翻译或实时字幕，$O(n)$ 复杂度确保了可预测的性能。未来的工作可以探索针对无限流的增量分词。
• 多语言与大词汇表： 扩展到涵盖 100 多种语言、包含 100 万以上词元的词汇表。需要研究如何优化预计算结构的内存占用，或许可以使用压缩字典树或近似方法。
• 与模型架构的集成： 分词是否可以部分学习或与 Transformer 的第一层合并？基于上下文的“动态分词”仍然是一个开放的挑战，但像这样的高效基础算法是此类探索的先决条件。

10. 参考文献

1. Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL-HLT.
2. Aho, A. V., & Corasick, M. J. (1975). Efficient string matching: an aid to bibliographic search. Communications of the ACM.
3. Kudo, T. (2018). Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. ACL.
4. Wolf, T., et al. (2020). HuggingFace's Transformers: State-of-the-art Natural Language Processing. arXiv preprint arXiv:1910.03771.
5. Google. (2020). TensorFlow Text. https://www.tensorflow.org/text
6. Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. NeurIPS.
7. Zhu, J.-Y., Park, T., Isola, P., & Efros, A. A. (2017). Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks. ICCV. (作为范式转换式重构的例证被引用)。