快速 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. (作為典範轉移式重新定義的範例被引用)。