[论文解读] Linear-Time WordPiece Tokenization.

Xinying Song, Alex Salcianu+3 more|arXiv (Cornell University)|Dec 31, 2020

Natural Language Processing Techniques参考文献 14被引用 5

一句话总结

本文提出 LinMaxMatch，一种基于优化的 tries 数据结构并受 Aho-Corasick 算法启发的额外链接机制的线性时间 WordPiece 分词算法。该算法在推理速度上比 HuggingFace 和 TensorFlow Text 系统快 3 倍，且在长尾输入上实现 4.5 倍的加速，为 MaxMatch 问题提供了实际效率与理论最优性兼具的解决方案。

ABSTRACT

WordPiece tokenization is a subword-based tokenization schema adopted by BERT: it segments the input text via a longest-match-first tokenization strategy, known as Maximum Matching or MaxMatch. To the best of our knowledge, all published MaxMatch algorithms are quadratic (or higher). In this paper, we propose LinMaxMatch, a novel linear-time algorithm for MaxMatch and WordPiece tokenization. Inspired by the Aho-Corasick algorithm, we introduce additional linkages on top of the trie built from the vocabulary, allowing smart transitions when the trie matching cannot continue. Experimental results show that our algorithm is 3x faster on average than two production systems by HuggingFace and TensorFlow Text. Regarding long-tail inputs, our algorithm is 4.5x faster at the 95 percentile. This work has immediate practical value (reducing inference latency, saving compute resources, etc.) and is of theoretical interest by providing an optimal complexity solution to the decades-old MaxMatch problem.

研究动机与目标

• 解决现有 MaxMatch 算法长期存在的低效问题，尽管其在 BERT 及相关模型中被广泛使用，但其时间复杂度仍为 O(n²) 或更差。
• 为子词分词中的 MaxMatch 问题提出一个可证明最优的线性时间解决方案，这是自然语言处理领域数十年来的挑战。
• 通过实现更快、可扩展的分词机制，降低实际 NLP 系统中的推理延迟和计算成本。
• 通过实现 MaxMatch 问题的最优时间复杂度，提供理论洞见，从而解决自然语言处理工程领域长期存在的开放问题。

提出的方法

• 从词汇表构建一个 tries 数据结构，以表示 WordPiece 分词中使用的所有可能子词单元。
• 通过引入额外的失败链接（受 Aho-Corasick 算法启发），增强 tries 结构，实现在当前路径匹配失败时的高效转移。
• 采用智能的状态转移机制，利用这些链接避免回溯，从而保持线性时间性能。
• 将增强后的 tries 集成到最长匹配优先（MaxMatch）分词流水线中，确保每一步都实现最大子词分割。
• 设计算法以单次遍历处理输入标记，保持 O(n) 的时间复杂度，其中 n 为输入长度。
• 通过智能链接传播机制最小化冗余的 tries 遍历，从而在平均情况和长尾输入上均实现性能优化。

实验结果

研究问题

• RQ1 能否设计一种针对基于 MaxMatch 的 WordPiece 分词的线性时间算法，打破先前方法的 O(n²) 时间复杂度瓶颈？
• RQ2 与 HuggingFace 和 TensorFlow Text 等生产系统相比，所提出的 LinMaxMatch 算法在速度和可扩展性方面表现如何？
• RQ3 该算法在长尾输入（如罕见或复杂序列）上的性能提升程度如何？
• RQ4 是否可以通过基于 tries 的结构并辅以 Aho-Corasick 风格的失败链接，实现 MaxMatch 问题的最优时间复杂度？
• RQ5 该算法在实际 NLP 部署中对推理延迟和计算资源使用量的实际影响是什么？

主要发现

• LinMaxMatch 实现了 MaxMatch 分词的线性时间复杂度 O(n)，为这一数十年来的问题提供了首个最优解。
• 在各种输入分布下，LinMaxMatch 的平均速度比 HuggingFace 和 TensorFlow Text 的生产级分词器快 3 倍。
• 在长尾输入的第 95 百分位数下，LinMaxMatch 的速度比基线系统快 4.5 倍，表明其在复杂序列上的卓越可扩展性。
• 该算法在短输入和长输入上均保持高效率，在实际推理工作负载中表现出一致的性能提升。
• 通过在 tries 上使用失败链接，实现了智能且无需回溯的转移，这是实现线性时间性能的关键。
• 所提出的方法在理论上最优且可实际部署，为低延迟 NLP 系统提供了立竿见影的收益。