Fast W ord P iece Tokenization (2021)

Abstract: 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.

## introduction tokenization is the process of splitting text into smaller units called tokens (e.g., words). it is a fundamental preprocessing step for almost all nlp applications: sentiment analysis, question answering, machine translation, information retrieval, etc. modern nlp models like bert @xcite , @xcite , and xl-net @xcite ) tokenize text into subword units @xcite @xcite . as a midpoint between words and characters, subword units retain linguistic meaning (like morphemes), while alleviating out-of-vocabulary situations even with a relatively small-size vocabulary. in this paper, we propose efficient algorithms for wordpiece, the subword tokenization used in bert @xcite . given unicode text that has already been cleaned up and normalized, wordpiece has two steps: (1) pre-tokenize the text into words (by splitting on punctuation and whitespaces), and (2) tokenize each word into wordpieces. for single-word tokenization, wordpiece uses a greedy longest-match-first strategy: iteratively pick the longest prefix of the remaining text that matches a vocabulary token. this is well-known as maximum matching or maxmatch @xcite , which has also been used for chinese word segmentation since 1980s @xcite . despite its wide use in nlp for decades, to the best of our knowledge, the most efficient max-match algorithms so far are ( 2 ) (where is the input word length) or ( ) (where is the maximum vocabulary token length) (see section 2). it's worth noting that the latter has a vocabularyspecific multiplicative factor , which can be large when the vocabulary contains long words. we propose linmaxmatch, a novel maxmatch algorithm for wordpiece tokenization, whose tokenization time is strictly ( ) without any vocabulary-specific multiplicative factors. inspired by the aho-corasick algorithm @xcite , we organize vocabulary tokens in a trie @xcite and introduce precomputed failure links and failure pops. during tokenization, if an input character does not match any trie edge, we perform smart transitions to avoid backtracking to earlier input characters. this involves collecting the recognized tokens (i.e., failure pops) and moving to a trie node (via the failure link), from where we continue to match the same character (section 3). for general text tokenization, referred to as end-to-end tokenization in this paper, we propose e2e wordpiece, an end-to-end algorithm that combines pre-tokenization and wordpiece tokenization into a single, linear-time pass (section 4). experimental results show that our method is 8.2x faster than huggingface tokenizers (hug-gingface, 2020) and 5.1x faster than tensor-flow text @xcite on average for general text tokenization (section 5). although tokenization is relatively faster than other steps, it's still worth improving the performance: tokenization is a prerequisite step for almost all nlp tasks, and any improvement on its efficiency helps reduce the latency of the entire inference. one potential impact of the work, for example, is on mobile nlp applications. ondevice models are generally highly optimized for reducing latency, e.g., by distilling or compressing larger models. thus, the impact of tokenization can be significant here. another impact is on aggregate computational savings for web services like google, facebook, twitter, etc. for example, google uses bert to power its web search nowadays. 1 google serves billions of search queries per day, and it processes hundreds of trillions of web pages in index building. by employing a faster tokenization system, the aggregate computational savings would be material, which also benefits the environment (for less power consumption). this paper also makes a theoretical contribution. the proposed linmaxmatch algorithm solves the decades-old maxmatch problem in the optimal ( ) time, and the idea is applicable to other string matching or rewriting problems (section 3.6). the code will be available at https://www. tensorflow.org/text . ## related work maximum matching (or maxmatch) has been used for chinese word segmentation (cws) since the 1980s @xcite . recent cws work focuses on machine learning-based segmentation approaches, but maxmatch remains a commonly referenced baseline @xcite . more recently, subword tokenization techniques have become a near-universal feature of modern nlp models, including bert @xcite , @xcite , xlnet @xcite , etc. common subword tokenization techniques include byte-pair encoding (bpe) @xcite , sen-tencepiece @xcite (based on unigram lan-1 https://blog.google/products/search/ search-language-understanding-bert/ guage modeling), and wordpiece @xcite . the widely-adopted maxmatch algorithm, which is used in the original wordpiece algorithm @xcite , starts from the longest possible prefix and decrements the length in search of the longest-matching token @xcite . a variant starts from the shortest substring and increases the length @xcite @xcite . the worst-case time complexity of the previous algorithms are ( foot_0 ) or ( ) or even higher than that. 23 for example, the complexity of @xcite is ( ) (in our notations), since lookup(t,c,i,n) (figure 1 in their paper) may take ( ) time (which is similar to the analysis in section 3.2 of this paper). @xcite recognizes maximum matching tokens using regular expressions in the context of compilers; their complexity is (| | ), where | | is the number of states in the automaton built from the grammar/vocabulary. if applied to wordpiece tokenization, since vocabulary tokens are finite strings, their complexity can be refined as ( ). our algorithm is inspired by the aho-corasick algorithm @xcite , but the two algorithms are designed to address different problems. aho-corasick is not optimal for the max-match problem. in the worst-case scenario where every substring in the input matches a vocabulary token, aho-corasick finds a quadratic number of matches, resulting in an overall quadratic complexity for maxmatch. by comparison, our algorithm achieves the worst-case linear complexity for max-match due to a novel definition of failure links, the newly-introduced failure pops, as well as a different way of emitting tokens. it's worth clarifying the difference between our failure links and the tabular solution of @xcite . in their work, a table called failed_previously is used to store whether a @xmath0,@xmath1. @xcite uses that table to avoid wasteful revisits of the same state. the table @xmath2,@xmath3. in contrast, our failure links capture which state to transit to when trie matching cannot continue (definition 1), and they are precomputed based on the vocabulary only, independent of the input. finally, we discuss the complexity of algorithms for byte-pair encoding (bpe) @xcite and sentence-piece @xcite . note that they are different problems from maxmatch (the topic of this paper). sentencepiece is based on unigram language modeling, and the optimal segmentation can be found in ( ) time with the viterbi algorithm @xcite . bpe algorithms can be implemented in two ways. one is to enumerate the symbol pairs in the order that they were added to the vocabulary in the building phase. for each symbol pair, we scan the current sequence and replace all their occurrences with the merged symbol. the complexity is (| | ), where | | is the size of the vocabulary. the other approach is to repeatedly select the pair of symbols from the current sequence that has the highest priority (e.g., the maximum frequency). using a heap, this approach can be done in ( ). ## linear-time single-word tokenization in this section, we present linmaxmatch, an ( ) algorithm for single-word wordpiece tokenization. ## linear-time end-to-end tokenization the existing bert tokenization implementations @xcite pre-tokenize the input text (splitting it into words by punctuation and whitespace characters) and then call wordpiece tokenization on each resulting word. for example, the text john johanson's may be split into @xcite . we propose an end-to-end wordpiece tokenizer that combines pre-tokenization and wordpiece into a single, linear-time pass. it uses the linmaxmatch trie matching and failure transition loop as much as possible and only checks for punctuation and whitespace characters among the relatively few input characters that are not handled by the loop. it is more efficient as it traverses the input only once, performs fewer punctuation / whitespace checks, and skips the creation of intermediate words. precomputation we use the same process as in section 3.4, with several differences: after the trie is constructed, we remove all trie links labeled with a punctuation character. 7 then, for every possible punctuation character , we add a trie data node with no descendants, and a trie link from the root to with label . if is part of the vocabulary, we @xmath4, @xmath5. the resulting trie matches all punctuation characters, as either themselves or @xmath6, depending on the vocabulary. punctuation characters are not part of longer tokens, and there is no suffix token for a punctuation character. this reflects the fact that each punctuation character is a word by itself. we then run the rest of algorithm 2 to compute the failure pops and failure links. finally, for punctuation nodes, we set their failure links to a special node ; their failure pops are not changed. the special node has no parent and no descendants, @xmath7, ( ) = ∅. node indicates that a punctuation character was matched. tokenization algorithm 3 tokenizes general text into wordpieces. it starts by appending a whitespace at the end of the input (line 1). in each iteration, it recognizes wordpieces for the current word by employing (almost) the same routine as in single-word tokenization (lines 3-7 in algorithm 3 versus lines 1-5 in algorithm 1). 8 when returning from matchloop(), algorithm 3 must have met a character that cannot be consumed after attempting failure transitions, such as a whitespace, a punctuation, or some unseen character. lines 4-5 examine whether the current word can be tokenized (by checking if the current position is at a word boundary and where the node lands at) and reset the tokens as appropriate (see related discussions in section 3.3). lines 6-7 further handle the corner case that the word happens to be the suffix indicator itself (in the same way as algorithm 1, see section 3.3). note that normally the suffix indicator contains only punctuation characters (e.g., ## in bert); in that case lines 6-7 can be saved, because the suffix indicator itself is not be tokenized as a single word. the tokens of the current word are then appened to the result (line 8). finally, the algorithm moves the cursor past the boundary of the current word (lines 9-10) and skips any following whitespaces (lines 11-12) to process the next word. it can be shown that algorithm 3 is consistent with @xcite for general text tokenization, and the time complexity is ( ). ## experiments experimental setup we benchmark our method against two widely-adopted wordpiece tokenization implementations: algorithm 3: end-to-end tokenization function e2ewordpiece(text): function iswdbndry( , ): in both cases, we use pre-tokenization and word-piece tokenization, and skip other steps provided by those libraries (text cleanup, normalization, etc) for fair comparison. both libraries use the original wordpiece tokenization algorithm @xcite . they both generate not only the numeric ids of the tokens, but also the token strings and start/end offsets of the input word. we modify both libraries to generate only the token ids, foot_6 for two reasons: (1) most downstream models (e.g., bert) consume only the token ids, and (2) we want to focus on the core tokenization work, not on, e.g., string copying. we implement linmaxmatch and e2e word-piece and made them return the numeric ids of the tokens, leveraging a double array-based trie library @xcite . and tensorflow text on a large corpus (several million words) and found that the tokenization results are identical for both single-word and end-to-end tokenization. in the rest of this section, we focus on the tokenization speed. all experiments are conducted on a linux desktop with a six-core intel xeon @ 3.60ghz cpu and 64gb memory. we iterate each benchmark (after warming up) until it ran for a long-enough period of time, repeat each experiment 10 times, and report the average results. our method is implemented and benchmarked in c++; so is tensorflow text. huggingface uses (and is benchmarked in) rust. we use the wordpiece vocabulary released with the bert-base, multilingual cased model, a model that supports 104 languages @xcite . to generate the test data, we sample 1,000 sentences from the multilingual wikipedia dataset, covering 82 languages including english, chinese, french, russian, etc. on average, each word has 4 characters, and each sentence has 82 characters or 17 words. we found this dataset large enough: a much larger dataset (consisting of hundreds of thousands of sentences) generated similar results. we run bert's basictokenizer @xcite to clean up and normalize each sentence, including unicode clean-up and normalization. following the guidance for the bert-base multilingual cased model @xcite , we do not instruct basictokenizer to do lower casing or accent stripping. in addition, preprocessing adds spaces around every cjk character, and thus chinese is effectively character-tokenized. for simplicity, we keep chinese in the test set, but keep in mind that each chinese word is just one chinese character, and any wordpiece implementation is efficient on such short words. using a dataset with long words would emphasize the speed advantage of our algorithm even more than indicated below. for single-word tokenization, we further used basictokenizer to pre-tokenize each sentence on punctuation and whitespace characters. this results in 17,223 words, 8,508 of them unique. ## conclusion we proposed linmaxmatch for single-word word-piece tokenization, which is asymptoticallyoptimal linear-time with respect to the input length, without a vocabulary-specific multiplicative factor. we also proposed e2e wordpiece that combines pre-tokenization and wordpiece tokenziation into a single, linear-time pass for even higher efficiency. experimental results show that our approach is 8.2x faster than huggingface and 5.1x faster than ten-sorflow text on average for general text tokenization. for future work, we will adapt the proposed methods to more text processing techniques.