Learn Your Tokens: Word-Pooled Tokenization for Language Modeling (2023)

Abstract: Language models typically tokenize text into subwords, using a deterministic, hand-engineered heuristic of combining characters into longer surface-level strings such as ‘ing’ or whole words. Recent literature has repeatedly shown the limitations of such a tokenization strategy, particularly for documents not written in English and for representing numbers. On the other extreme, byte/character-level language models are much less restricted but suffer from increased sequence description lengths and a subsequent quadratic expansion in self-attention computation. Recent attempts to compress and limit these context lengths with fixed size convolutions is helpful but completely ignores the word boundary. This paper considers an alternative ‘learn your tokens’ scheme which utilizes the word boundary to pool bytes/characters into word representations, which are fed to the primary language model, before again decoding individual characters/bytes per word in parallel. We find that our moderately expressive and moderately fast end-to-end tokenizer outperform by over ‘300%‘ both subwords and byte/character models over the intrinsic language modeling metric of next-word prediction across datasets. It particularly outshines on rare words, outperforming by a factor of 30! We extensively study the language modeling setup for all three categories of tokenizers and theoretically analyze how our end-to-end models can also be a strong trade-off in efficiency and robustness.

## introduction almost all natural language processing (nlp) begins with tokenization @xcite . sequences of characters are (mostly deterministically) segmented into discrete tokens, each of which has a lookup embedding in an enormous vocabulary matrix. statistical nlp methods, similar to other forms of machine learning at the time, relied on feature extraction from these tokens, in the form of n-gram occurrences or part-of-speech tags or other representations of syntax. all of these pipelines have over time been replaced with end-to-end learning using recurrent neural networks (rnns) or transformers, however the tokenization schemes remain static, deterministic, and manually engineered. state-of-the-art approaches include subword tokenization schemes such as wordpiece @xcite , byte pair encoding or bpe @xcite , and unigram @xcite , all of which are statistical methods for preprocessing a large unlabeled corpus of text to yield a fixed vocabulary, midway between characters or bytes at one end and whole words at the other. this results in a convenient trade-off in sequence description length while avoiding the unk token, that is, a fallback mechanism for handling rare words. however, it is not obvious why these hand-engineered algorithms would be the optimal forms of tokenization and whether there exists a possibility for end-to-end models to also include this crucial stage of the nlp pipeline. recent work has shown countless limitations with subword embeddings. several languages contain diverse morphological features whereas subword segmentation is mostly apt at only identifying suffixes and prefixes @xcite . technical domains such as biomedical documents often need to pre-train their own tokenizer for improved vocabulary @xcite . finally, numbers are often inconsistently segmented into subwords, leading to decreased arithmetic @xcite and estimation @xcite skills. the extent of these numeric limitations is so dire that @xcite has an explicit workaround of adding all numbers from 0 to 999 as individual tokens to the model's vocabulary. recently, several language models have been proposed which remove the tokenizer vocabulary entirely, beginning with a character @xcite or byte-level @xcite vocabulary and often compressing them into fixed units of around four tokens each @xcite @xcite . while these zeroassumption methods are useful in compressing text and consequently expand context windows, they completely ignore the word boundary. besides, the so-called 'tokenizer-free' byte-based models are not entirely bias-free since the unicode-8 encoding they use is itself biased towards representing latin scripts with a single byte each, whereas some scripts 1 like bammum (africa), meetei (india), and cherokee (north america) may require four bytes to represent a single character. the concept of words is a fundamental feature of nearly all human languages, including those written in chinese or japanese scripts that do not explicitly delineate words by whitespaces. this paper empirically studies the case where tokenizers lose their subword segmentation algorithms but utilize the word boundary for a multi-level model with added efficiency. more concretely, we use the word boundary to compress the base tokens of bytes or characters into word representations, which are then fed into the underlying language model (here, a small version of gpt @xcite ). our end-to-end learned tokenization undoubtedly has several limitations. it is not faster than subwords. it does not allow characters/bytes within one word to directly attend to those in another word. it relies on the word boundary, which is not straightforward to find for most internet-scale datasets. nevertheless, we believe this empirical deep-dive into tokenizers for language modeling offers the following contributions: we will publicly release all code (see supplementary material) and checkpoints upon acceptance. next, (step 2) the pooled per-word embeddings are passed onto the main language model, in our case, a vanilla transformer decoder like gpt @xcite . finally, the contextualized word embeddings are fed through another transformer decoder to autoregressively decode the next word, one base unit (character or byte) at a time. note that we call this method an end-to-end 'tokenizer' since it compresses the many units into a few embeddings per word, just like subwords, except the compression is learned from scratch. finally, at decoding stage (step 3), the contextualized word representations are unrolled with another transformer decoder to autoregressively predict one base token (character/byte) at a time. 3 note how we achieve our purported trade-off between subwords and byte/character models. the cls representations learnt are unconstrained by a deterministic mapping as in subwords. they are also efficient to compute and decode from, since the first and last steps only allow intra-word attention. for a tokenizer-free model, roughly 80% of the memory bottleneck 4 is spent on tokens from one word attending to tokens on another word, which we contest is of questionable importance relative to the overhead incurred. formally, we begin with a sequence of words w 0 , w 1 , . . . , w n each of which is comprised of an ordered set of base units (character/bytes) w @xmath0, c 1 i , . . . , c m i i where m i + 1 is the length of the ith word. the task is autoregressive language modeling, i.e. given the previously seen words w 0 , w 1 , . . . , w i-1 as well as the previously seen units in w i (the current word): character/byte level models ignore the wordboundary and directly model the task as: this paper uses 4 cls tokens per word, except in section 4.2 where we ablate with 1 cls per word. subword segmentation maps the base units deterministically into fewer subwords per word, i.e. , our end-to-end models instead follow a threestep process to (1) pool base units into a fixed set of embeddings per word, (2) autoregressively predicting the next word embedding, and (3) autoregressively predicting individual unit embeddings per word: here, encoder refers to a transformer bert-like encoder and decoder refers to a transformer gptlike decoder. from an implementation standpoint, we prefix a fixed number (@xmath1. the word-level contextualized representations obtained on the other end are collectively depicted here as w i . figure 2 is a visualization of how our end-to-end model saves on self-attention computation bottleneck by only allowing intra-word attention at the first step, before allowing contextualization of information across the word boundary in step 2 using the base decoder model. finally step 3 again restricts the individual characters/bytes to be predicted using only the single word-level predicted embeddings. 5 3 experiments there are numerous nlp tasks that can benefit from improved tokenization, such as machine translation, question answering, and text classification. however, the scope of our preliminary analysis is not to cast a wide net over every downstream application. instead, we choose to analyze in depth the most commonly used pre-training task in nlp i.e. language modeling. we pretrain autoregressive language models from scratch using different tokenizers described in the previous section, on different datasets described in section 3.2. ## related work some recent work has challenged the subword tokenization schemes. character/byte-level byt5 @xcite , canine @xcite , and subchar @xcite propose using very small fixed-length units such as characters, bytes, or glyph strokes instead of dynamic-length subwords or words. this often comes at the expense of larger sequence lengths and more compute requirements, especially for a transformer architecture which typically has a complexity of o(n 2 ) in number of input tokens. beyond word level codebpe @xcite and multi word expressions (kumar @xcite @xcite show promise in yet larger tokens that cross word boundaries, e.g., a vocabulary with single tokens for the strings "for i in range" or "new york city" respectively. visual segmentation yet another line of work @xcite ) renders text as images before feeding them to cnns, doing away with tokenization altogether and showing gains in robustness to spelling or printing errors. learnt subword segmentation finally, some methods @xcite @xcite @xcite parameterize the process of tokenization by pooling character n-grams or randomly choosing one of the many ways to segment a given word. a recent preprint on machine translation by @xcite proposes a method called wsf, perhaps closest to ours, except that they only use the word boundary fusion at encoder stage. our independent analysis focuses on language modeling instead and also generates text in parallel using end-to-end attention based tokenization. ## conclusion subword tokenization is efficient but too rigid and deterministic. character/byte-level models on the other hand are too expressive, which leads to inefficient training and inference. we propose a wordboundary-informed tokenizer that efficiently and robustly performs language modeling in a hierarchical, end-to-end, learned model. we show that it outperforms by over 300% both extremes: subwords and character/byte models. we also analyze its trade-offs in training and inference efficiency. despite its many flaws including reliance on a word boundary signal and moderate efficiency as well as moderate expressiveness, we expect this preliminary study to pose an interesting trade-off tokenization for truly end-to-end language modeling. our code is released on github.