PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast WordPiece Tokenization

Xinying Song, Alex Salcianu, Yang Song, Dave Dopson, Denny Zhou

arXiv:2012.15524v3cs.CL

TL;DR

Prior WordPiece MaxMatch tokenization algorithms can require O(n^2) or O(nm) time. The paper introduces LinMaxMatch and E2E WordPiece for linear-time tokenization, reporting 8.2x and 5.1x average speedups over common systems.

  • Problem

    Existing MaxMatch algorithms for WordPiece require O(n^2) or O(nm) time, and the vocabulary-dependent factor m can be large.

  • Method

    LinMaxMatch adds precomputed failure links and failure pops to a vocabulary trie, while E2E WordPiece combines pre-tokenization and WordPiece tokenization in one pass.

  • Results

    8.2x faster than HuggingFace and 5.1x faster than TensorFlow Text on average for general text tokenization.

  • Takeaways & Limitations

    The algorithms provide linear-time WordPiece tokenization for single words and general text without a vocabulary-specific multiplicative factor.

  • Takeaways & Limitations

    Failure-pop construction requires O(Mm) total size and straightforward offline precomputation, with optimized implementations deferred to future work.

Abstract

from arXiv · show

Tokenization is a fundamental preprocessing step for almost all NLP tasks. In this paper, we propose efficient algorithms for the WordPiece tokenization used in BERT, from single-word tokenization to general text (e.g., sentence) tokenization. When tokenizing a single word, WordPiece uses a longest-match-first strategy, known as maximum matching. The best known algorithms so far are O(n^2) (where n is the input length) or O(nm) (where m is the maximum vocabulary token length). We propose a novel algorithm whose tokenization complexity is strictly O(n). Our method is 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. For general text, we further propose an algorithm that combines pre-tokenization (splitting the text into words) and our linear-time WordPiece method into a single pass. Experimental results show that our method is 8.2x faster than HuggingFace Tokenizers and 5.1x faster than TensorFlow Text on average for general text tokenization.

1 Introduction

The paper targets WordPiece tokenization, a fundamental NLP preprocessing step whose existing MaxMatch algorithms can scale quadratically or with vocabulary token length. It introduces linear-time single-word and end-to-end methods, reporting substantial speedups over established tokenizers.

  • WordPiece first pre-tokenizes normalized Unicode text into words, then tokenizes each word into wordpieces.
  • MaxMatch greedily selects the longest vocabulary-matching prefix of the remaining word.This longest-match-first strategy is used for single-word WordPiece tokenization.
  • Existing MaxMatch algorithms require O(n^2) or O(nm) time, with m potentially large for vocabularies containing long words.
  • LinMaxMatch achieves strictly O(n) tokenization without vocabulary-specific multiplicative factors by adding trie failure links and failure pops for smart transitions.The transitions avoid backtracking when trie matching cannot continue.
  • 8.2x faster than HuggingFace Tokenizers and 5.1x faster than TensorFlow Text on average for general text tokenization.

2 Related Work

Related work covers MaxMatch, subword tokenization, and alternative string-processing algorithms. The paper distinguishes LinMaxMatch from prior quadratic or vocabulary-dependent approaches and from methods addressing different segmentation objectives.

  • MaxMatch is a longstanding baseline for Chinese word segmentation and is used in the original WordPiece algorithm.
  • Prior MaxMatch algorithms have worst-case complexity O(n^2), O(nm), or higher, with m reflecting the maximum vocabulary token length.
  • Aho-Corasick is not optimal for MaxMatch because it can find quadratically many matches, whereas the proposed algorithm achieves worst-case linear complexity.
  • SentencePiece and BPE address different problems from MaxMatch, using unigram modeling or symbol-pair procedures rather than longest-match tokenization.

3 Linear-Time Single-Word Tokenization

LinMaxMatch tokenizes single words with WordPiece’s longest-match-first behavior while avoiding the repeated reprocessing that makes earlier trie-based approaches slower. Precomputed failure pops and links enable strictly O(n) tokenization, which is asymptotically optimal.

  • 3.1 Background and Notations: WordPiece segments a word by repeatedly selecting the longest matching prefix, with middle wordpieces marked by a suffix indicator such as ##.If the word cannot be tokenized, WordPiece maps the entire word to <unk>.
  • 3.2 Intuition: A vocabulary trie supports prefix matching, but trie failure can reveal the longest token several characters earlier and force repeated reprocessing, yielding O(nm) time.The example reprocesses previously scanned characters after matching abcd but recognizing only a.
  • 3.2 Intuition: LinMaxMatch precomputes failure pops and failure links so a failed trie match emits recognized tokens, moves to the remaining suffix state, and continues without backtracking.Failure pops are the longest-match tokens removed from the matched string; the failure link identifies the trie node for the remaining suffix.
  • 3.3 LinMaxMatch Tokenization: The algorithm is consistent with the original WordPiece algorithm and processes input through MATCHLOOP using trie edges, failure transitions, and token emission.MATCHLOOP processes one input character per step and resets the result to [<unk>] when the word cannot be tokenized.
  • 3.5 Complexity Analysis: O(n) total tokenization time follows because failure transitions are no more numerous than normal transitions and the total number of output tokens is at most n.The implementation makes one state transition on each input character, despite requiring precomputed results.
  • 3.5 Complexity Analysis: LinMaxMatch is asymptotically optimal because reading an input of length n already requires at least n operations, and the method has no vocabulary-specific multiplicative factor.The paper presents this as the first strictly O(n) optimal-complexity result for MaxMatch.

4 Linear-Time End-to-End Tokenization

E2E WordPiece combines pre-tokenization and WordPiece tokenization in a single linear-time pass. It preserves the existing tokenization behavior while avoiding intermediate words and repeated input traversal.

  • 4 Linear-Time End-to-End Tokenization: E2E WordPiece combines pre-tokenization and WordPiece tokenization into a single, linear-time pass.It reuses LinMaxMatch trie matching and failure transitions while checking punctuation and whitespace only when needed.
  • 4 Linear-Time End-to-End Tokenization: Punctuation is represented as individual trie matches, either as itself or as <unk> when absent from the vocabulary.Punctuation is not part of longer tokens and has no suffix token because each punctuation character is treated as a word by itself.
  • 4 Linear-Time End-to-End Tokenization: The algorithm appends whitespace, processes the current word with the single-word matching routine, and handles word boundaries after matching stops.Boundary handling checks whether the current word can be tokenized and resets tokens as appropriate.
  • 4 Linear-Time End-to-End Tokenization: After completing a word, the algorithm appends its tokens, advances past the word boundary, and skips following whitespace.A special case handles inputs consisting exactly of the suffix indicator.
  • 4 Linear-Time End-to-End Tokenization: The resulting algorithm is consistent with Google’s general-text tokenization and has O(n) time complexity.Its single-pass design also avoids creating intermediate words and reduces punctuation and whitespace checks.

5 Experiments

The experiments compare the proposed implementations with HuggingFace Tokenizers and TensorFlow Text under matched token-id-only settings. The proposed method matches their tokenization results and is faster for both single words and general text.

  • 5 Experiments: The evaluation compares the proposed implementations with HuggingFace Tokenizers and TensorFlow Text.The comparison skips cleanup and normalization and modifies systems to return only numeric token ids.
  • 5 Experiments: Tokenization results are identical to HuggingFace and TensorFlow Text for both single-word and end-to-end tokenization.The remaining analysis focuses on tokenization speed.
  • 5 Experiments: The benchmark uses the BERT-Base Multilingual Cased vocabulary and 1,000 multilingual Wikipedia sentences covering 82 languages.The sampled sentences average 82 characters or 17 words, and a much larger dataset produced similar results.
  • 5 Experiments: The experiments report mean and 95th-percentile running times for single-word and end-to-end tokenization.Each benchmark is repeated 10 times after warm-up, with average results reported; implementations use C++ except HuggingFace’s Rust implementation.
  • 5 Experiments: 3x faster average single-word tokenization is reported for the proposed system, with greater speedup on long-tail inputs.Figure 2 plots average running time against input length for single-word tokenization.
  • 5 Experiments: 8.2x faster than HuggingFace and 5.1x faster than TensorFlow Text is reported on average for general text tokenization.Table 3 reports running times in nanoseconds for the evaluated systems.

6 Conclusion

The conclusion presents LinMaxMatch as an optimal linear-time solution for single-word WordPiece tokenization and E2E WordPiece as a single-pass extension for general text. Experiments report substantial average speedups over two established implementations.

  • 6 Conclusion: LinMaxMatch achieves asymptotically optimal O(n) single-word tokenization without a vocabulary-specific multiplicative factor.The conclusion identifies this as the paper’s single-word WordPiece contribution.
  • 6 Conclusion: E2E WordPiece combines pre-tokenization and WordPiece tokenization into a single, linear-time pass for higher efficiency.The conclusion frames this as the general-text extension of the proposed approach.
  • 6 Conclusion: 8.2x faster than HuggingFace and 5.1x faster than TensorFlow Text is reported on average for general text tokenization.The paper identifies adaptation to other text-processing techniques as future work.

A Mathematical Formulations and Proofs of LinMaxMatch

This appendix formalizes LinMaxMatch’s notation and MaxMatch recursion, then develops the lemmas used in its correctness proof. It defines vocabulary prefixes, suffix transformations, and the special suffix-indicator behavior.

  • A Mathematical Formulations and Proofs of LinMaxMatch: The appendix presents mathematical formulations of LinMaxMatch and proves its correctness.The formal development introduces additional notation before stating the MaxMatch definitions and lemmas.
  • A Mathematical Formulations and Proofs of LinMaxMatch: The length of a string is its character count, except that a leading suffix indicator is excluded from the length.Examples assign lengths 3 to abc, 1 to ##d, and 0 to ε or ##.
  • A Mathematical Formulations and Proofs of LinMaxMatch: p_w is the longest non-empty vocabulary prefix of w, or ε when no such prefix exists.If w starts with the suffix indicator, the selected prefix also starts with it unless empty.
  • A Mathematical Formulations and Proofs of LinMaxMatch: q_w replaces p_w in w with the suffix indicator and retains the remaining suffix.For V={a, ab, ##c}, the appendix gives p_abcd=ab and q_abcd=##cd.
  • A Mathematical Formulations and Proofs of LinMaxMatch: If appending the final character c produces a string outside the vocabulary, p_wc=p_w and q_wc=q_w c.The proof derives this by showing that p_wc cannot include c as its final character.
  • A Mathematical Formulations and Proofs of LinMaxMatch: The suffix-indicator example shows that p_##a can be ε rather than the single-character # when ##a is not in the vocabulary.This illustrates the special prefix constraint for strings beginning with the suffix indicator.
  • A.1 MaxMatch / Definition 5. MaxMatch: MaxMatch recursively emits p_w, handles empty-prefix and terminal cases with <unk> or an empty list, and continues on q_w.When the input is exactly the suffix indicator, the formal definition may differ from the original algorithm, so the original algorithm is used instead.

A.2 MinPop Matching

MinPop Matching computes MaxMatch by minimally removing longest-matching prefixes until the remaining string reaches a trie node. This equivalence supports tokenization outcomes, including unknown-token handling.

  • A.2 MinPop Matching: MinPop Matching repeatedly pops the fewest longest-matching prefixes from a string until the remainder matches a trie node.This provides an alternative formulation of MaxMatch.
  • A.2 MinPop Matching: The function g(w) returns a trie node for w or for a suffix obtained after minimally popping consecutive longest-matching prefixes.If no such node exists, it returns ∅.
  • A.2 MinPop Matching: The function G(w) records the consecutive longest-matching prefix tokens popped while computing g(w).When w is already on the trie, no popping is needed.
  • A.2 MinPop Matching: MaxMatch M(w) can be computed from g(w) and G(w) on an augmented trie containing whitespace and the suffix indicator ♯.The added nodes are not vocabulary tokens but make the equivalence hold for tokenizable strings.
  • A.2 MinPop Matching: If g(w)=∅, M(w) returns [<unk>]; otherwise, M(w) returns G(w).For example, M(abcdx)=[abcdx], whereas M(z)=[<unk>] in the example vocabulary.

A.3 One-Step MinPop Matching

One-Step MinPop Matching updates MinPop state one character at a time. Recursive transitions compute the state efficiently, while failure-link implementations reduce storage compared with full transition tables.

  • A.3 One-Step MinPop Matching: One-Step MinPop Matching computes g(wc) and G(wc) from the state for w and the next character c.The transition functions h(u,c) and H(u,c) perform minimal popping of longest-matching prefixes.
  • A.3 One-Step MinPop Matching: Lemma 3 recursively computes g(w) and G(w) by splitting a string into prefix w and final character c.The base cases are ε and ♯, whose popped-token lists are empty.
  • A.3 One-Step MinPop Matching: The recursive transition has three cases covering direct trie matches, empty prefix remainders, and nonempty remainders requiring recursive failure handling.The proof proceeds by induction on the prefix length.
  • A.3 One-Step MinPop Matching: Precomputing h(u,c) and H(u,c) enables efficient computation but requires O(|T|·|Σ|) space for the h table.Here |T| is trie size and |Σ| is alphabet size.
  • A.3 One-Step MinPop Matching: Failure links f(v) and failure pops F(v) compute the same transitions with O(|T|) link-table space while maintaining overall linear time complexity.The paper presents this as a more practical approach.

A.4 Failure links and Failure Pops

Failure links and failure pops are defined over trie nodes to encode fallback transitions and the tokens removed during those transitions. They support recursive computation of one-step matching states.

  • A.4 Failure links and Failure Pops: For a trie node v, failure links f(v) and failure pops F(v) are defined from the node’s represented string and its parent transition.The definitions continue the earlier failure-link construction.
  • A.4 Failure links and Failure Pops: Lemma 4 computes h(u,c) and H(u,c) recursively using f(·) and F(·).This connects node-level failure information to one-character MinPop transitions.
  • A.4 Failure links and Failure Pops: When a transition cannot continue and the current prefix is empty, the failure link and failure-pop outputs are ∅ and [ ].The resulting one-step transition also returns ∅ and [ ].
  • A.4 Failure links and Failure Pops: The paper illustrates failure-transition computation with a trie-node example using node 6 and character z.The example is tied to the augmented trie in Figure 4.

A.5 Tokenization and its Correctness

Algorithm 1 incrementally computes the MinPop state and returns MaxMatch tokens, including [<unk>] for untokenizable inputs. The paper establishes correctness from the preceding lemmas under the stated input assumption.

  • A.5 Tokenization and its Correctness: The correctness discussion assumes that the original input string is not the suffix indicator ♯ itself.This assumption is stated in the accompanying note.
  • A.5 Tokenization and its Correctness: Algorithm 1 incrementally computes h(u,s[i]), H(u,s[i]), g(s), and G(s) while scanning the input string.When g(s)≠∅, the resulting tokens are G(s).
  • A.5 Tokenization and its Correctness: For tokenizable inputs, MATCHLOOP returns tokens equal to G(w), which equals MaxMatch M(w).This behavior is shown for both the augmented and original tries.
  • A.5 Tokenization and its Correctness: For untokenizable inputs, MATCHLOOP follows the same failure behavior and Algorithm 1 returns [<unk>].The augmented-trie formulation has g(w)=∅ in this case.
  • A.5 Tokenization and its Correctness: Algorithm 2 precomputes failure links and failure pops for trie nodes using Lemma 5.The root and augmented-root nodes have empty failure links and pops.

A.7 LinMaxMatch as a Finite-State Transducer (FST)

LinMaxMatch can be implemented as a finite-state transducer by precomputing transition and output functions, eliminating failure transitions and retaining linear time complexity.

  • A.7 LinMaxMatch as a Finite-State Transducer (FST): LinMaxMatch becomes an FST by precomputing the transition function δ′(u, c) and output function σ′(u, c).These functions eliminate failure transitions during matching.
  • A.7 LinMaxMatch as a Finite-State Transducer (FST): The precomputed functions are stored as h(u, c) and H(u, c) to rewrite LinMaxMatch as Algorithm 4.The rewritten algorithm differs from Algorithm 1 at lines 8–11.
  • A.7 LinMaxMatch as a Finite-State Transducer (FST): Algorithm 4 presents LinMaxMatch as an FST.
  • A.7 LinMaxMatch as a Finite-State Transducer (FST): The FST implementation eliminates failure transitions, making LinMaxMatch’s time complexity trivially linear.
Loading 2012.15524v3…