Incorporating Context into Subword Vocabularies (2023)

Abstract: Most current popular subword tokenizers are trained based on word frequency statistics over a corpus, without considering information about co-occurrence or context. Nevertheless, the resulting vocabularies are used in language models’ highly contextualized settings. We present SaGe, a tokenizer that tailors subwords for their downstream use by baking in the contextualized signal at the vocabulary creation phase. We show that SaGe does a better job than current widespread tokenizers in keeping token contexts cohesive, while not incurring a large price in terms of encoding efficiency or domain robustness. SaGe improves performance on English GLUE classification tasks as well as on NER, and on Inference and NER in Turkish, demonstrating its robustness to language properties such as morphological exponence and agglutination.

## introduction much of the research space in current nlp is focused on advancing models: modifying pretraining objectives, improving network architectures, adding tasks and schemes for downstream evaluation. limited work is dedicated to a crucial step underlying all modern large language models (llms), namely the tokenization phase. in order to process a given string of text, an llm must first obtain a vector representation of the input by segmenting it into tokens. since out-of-vocabulary (oov) items inhibit the performance of models, current tokenizers produce tokens which are possibly proper subsegments of input words, known as subwords. this method, popularized by systems such as wordpiece @xcite , byte-pair encoding (bpe; @xcite and unigramlm @xcite , allows any word to be represented by one or more tokens, removing the oov problem while allowing more flexibility in determining the token vocabulary size, which ultimately affects model speed (mostly through bpe his son raj ash ri sud h ak ar has p enn ed dial og ues and songs for some films that were dubbed into telugu . sage his son raj ash r i sud h a k a r has penn e d dial ogues and songs for some films that were dubbed into telugu . bpe this gene is a pseud og ene in humans and most other prim ates . sage this gene is a pseud ogene in humans and most other prim ates . bpe the st o og es work for mir acle det ective agency , sage the st o o g e s work for mir acle det ective agency , one potential pitfall of both bpe and uni-gramlm, as well as their proposed variants @xcite , is that they are trained on word frequency statistics alone, without considering information about word co-occurrence or contexts. at the same time, the resulting vocabularies are used in highly contextualized settings, the llms, where a single subword such as og might appear in very different contexts derived from words like dial og ues and pseud og ene. we propose a system which prepares subwords for their downstream use by baking in the contextualized signal at the vocabulary creation step. our model, sage, uses the skipgram objective @xcite over a corpus as the basis for iteratively eliminating candidate subwords from an initial large vocabulary until the desired vocabulary size has been reached. as we present our algorithm, sage, which is predicated on iterative pruning of contextually noisy tokens from the vocabulary, and compare its effects on token properties and context cohesion with bpe both in-and out-of-domain, in english and turkish. we then evaluate its performance on downstream tasks by training a bert-based llm @xcite on a vocabulary produced by both tokenizers in both languages, demonstrating substantial improvements on most english glue tasks and on ner, as well as turkish nli and ner. we emphasize that as opposed to most current tokenizer variants, our model is a "plug and play" substitution for any subword token vocabulary, requiring no modification in the inference protocol (or code) when pre-training or applying an applicable llm from a popular shared library. foot_0 ## subword vocabulary creation the methods used to tokenize corpus in order to later assign tokens with continuous vectors, or embeddings, have evolved over the years. initially, each word in the corpus was assigned its own embedding @xcite . oovs, i.e. words not appearing in the original training corpus or below a certain frequency threshold, would receive a special (but identical) "unk" vector. subword tokenizers @xcite were introduced to alleviate this issue, allowing segmentation of all text into embeddable units (assuming no unseen characters, a much more relaxed constraint for languages using alphabetical scripts). the training process used to create a subword vocabulary from which the model then decodes text input involves optimizing an encoding objective over a large corpus. to date, all tokenizers used in practice in large models focus on efficiency and informationtheoretic objectives, and reduce the corpus to a unigram frequency count of space-delimited words, reducing calculation time but losing all contextual signal. sage reintroduces the contextual dependencies between words into vocabulary creation via a two-stage process, namely over-application of bpe followed by iterative pruning using ideas inspired by unigramlm and skipgram. we briefly present these algorithms before tying them together into sage. algorithm 1 byte-pair encoding vocabulary creation @xcite input: corpus c, vocabulary final size v . output: vocabulary v of size v (ordered). 1: procedure bpe(c, v ) 2: v ← all unique characters in c 3: v ← v ⊕ [tnew ] 7: c.replaceall(⟨tl, tr⟩, tnew ) 8: unigram language model. unigramlm offers a top-down vocabulary creation process, starting with an initial vocabulary of all substrings in the input corpus and pruning tokens iteratively until reaching the desired vocabulary size. the pruning procedure involves calculating the overall unigram likelihood of the corpus with the current vocabulary versus a vocabulary lacking the candidate pruning token (see algorithm 2 for details), which we refer to as the ablation objective. under this system, decoding is ideally performed by considering probabilities of all possible segmentations using, e.g., the viterbi algorithm; again, common practice is to use left-to-right greedy decoding. skipgram objective. the skipgram objective @xcite formalizes the relation between a target token t and its context, asking whether context tokens c within a window w t of pre-defined size can be predicted from t. these predictions are done via sigmoid activation over the inner product of embeddings trained for targets (e (t ) ) and contexts (e (c) ). when aggregated over algorithm 2 unigramlm vocabulary creation @xcite . n arg min x denotes the n bottomranked elements in x. input: corpus c, vocabulary final size v , pruning batch size k. output: vocabulary v of size v . 1: procedure unigramlm(c, v ) 2: v ← all substrings occurring more than once in c 3: as token vocabularies or their inference methods change, so do the target sequences and their contexts, resulting in differences in aggregated likelihood which can then act as scores comparing one tokenization to another. we use this behavior as the ablation objective for sage. ## sage vocabulary creation sage foot_1 is a top-down tokenizer, following un-igramlm's general procedure, incorporating a skipgram objective as its vocabulary trimming rule. given an initial vocabulary v and a corpus c, sage computes a skipgram embedding space over v which provides it with an overall likelihood over c as in (1). it then proceeds to calculate the loss of each token in the vocabulary were it to be removed, eliminating the tokens incurring minimal loss and re-tokenizing the corpus according to the updated vocabulary, repeating this procedure until reaching the desired vocabulary size v . having learned this vocabulary, downstream inference proceeds exactly as in the other segmentation-based methods, in a greedy left-to-right manner. sage can also be adapted to anticipate other decoding algorithm 3 sage vocabulary creation. n arg min x denotes the n bottom-ranked elements in x. input: corpus c, vocabulary final size v , basic tokenizer t , overshoot factor n, pruning batch size k, likelihood recalculation frequency m, size of pruning candidate set m , embedding recalculation frequency l. v bot ← m arg min @xmath0: else ▷ update losses for bottom set 15: in practice, applying the full process described above introduces multiple sources of considerable computational complexity: for example, calculating the ablation objective for each token in each iteration produces a quadratic amount of calculations over the entire corpus; recalculating embeddings for an updated vocabulary is similarly unreasonable to perform at each iteration. we ameliorate these and other sources of complexity using a series of heuristics found in preliminary experiments to be minimally disruptive to precision of likelihood calculations. we will now describe these heuristics, all depicted in algorithm 3. first, instead of initializing the vocabulary as the full set of possible character sequences in the corpus, as in unigramlm, we use any existing noncontextual tokenizer such as bpe to learn a vocabulary larger than v by a factor of n, and begin the pruning process from there. next, instead of removing a single token from the bottom of the loss-ranked vocabulary, we remove a batch of the k bottom tokens contextual loss. in order to calculate the pertoken skipgram likelihood loss, all sentences where a token t occurs need to be re-segmented according to v \ {t}, and their new likelihoods recorded. to support performing this calculation on a large scale, we maintain a mapping of tokens to sentences containing them, as well as these sentences' current likelihoods. this must be done at the sentence level rather than the window level, since a remaining suffix from an out-of-window re-tokenization may combine with in-window characters and form different token sequence replacements at a given stage. consider the example in negative sampling. the original skipgram objective uses negative samples to estimate context probabilities. since our application of skip-gram within the vocabulary creation algorithm (independent of the embeddings training procedure) includes only likelihood estimation with no parameter updates, we do not sample negative tokens, a process which would introduce substantial noise and complexity. ## sage vocabulary properties for an analysis of our modified algorithm's advantages, we trained vocabularies of a pre-determined size using both bpe and sage. we selected |v| = 16, 000, and obtained corpora for english (750,000 lines from the august 2022 wikipedia dump) and turkish (the entire text of the september 2022 wikipedia dump), opting for languages that share the latin alphabet but differ in family (indo-european vs. turkic) and, crucially, in morphological properties: english is a low-exponence, lowsynthesis language, while turkish features multiple inflectional exponence and high verbal synthesis, as well as highly agglutinative morphology @xcite . we used the following hyperparameter settings to compute the vocabularies: initial vocab size 20, 000 (or @xmath1.25), @xmath2, @xmath3, @xmath4, @xmath5. we used the gensim package to train the skipgram models @xcite , and sentencepiece @xcite to obtain the initial bpe vocabularies. 4 more hyperparameters are detailed in appendix a. we present an analysis of the resulting vocabularies, highlighting the advantages and trade-offs exhibited by context-based subword tokenization. generally speaking, most of the tokens discarded from sage's initial vocabulary appear in the baseline bpe's final vocabulary. among the differences between the vocabularies are many short tokens that appear in bpe's but not sage's, proper substrings of longer tokens also appearing in the bpe vocabulary. this is due to bpe's bottom-up merge table construction, which forces retention of the entire chain of tokens created: if the is part of the vocabulary, either th or he must also be there. while essential for the original intended decoding process, actual implementations of greedy decoding have no need for this property. sage's initial vocabulary shares this characteristic, but the trimming process allows any token to be ablated, more frequent in sage bpe e s es ic ing ist ings ff ation ate ations ates token length. figure 1 shows a histogram of token lengths (in characters) for the 16,000-token sage and bpe vocabularies in english (results on turkish are similar). sage clearly selects longer tokens for its vocabulary, again a sensible outcome given their higher chance of being contextually coherent. the difference is most stark with tokens of length 2 and 3; when considering only tokens appearing in exactly one of the final vocabularies, we find that 56% of bpe-only tokens are of length 2 and 3, while 55% of sage-only tokens are of length 5 and above. token frequency. we compute the frequency of tokens in the encoding form of the english training corpus, once using sage vocabulary and once using bpe's. in subword fertility. fertility, as defined in the statistical machine translation literature, refers to the average number of subwords produced per tokenized word. figure 2 exhibits a histogram of all english corpus words by their subword length, using the bpe vocabulary and the sage vocabulary. although sage retains more words as single tokens, it trades them off with more words having five subwords or more, compared with bpe's abundance of words with 2 and 3 subwords. this follows the trend described so far, of sage's preference for dismantling unknown words into meaningless single-character tokens rather than confusing, ambiguous length-2 and length-3 tokens. we believe that bpe's behavior harms text understanding in suggesting that these ambiguous fragments (consider "og") have some meaning that an llm can try and learn, whereas sage's single-character breakdown indicates a word that's truly unknown and cannot be inferred by composing constituent in-vocab subwords. fertility translates to a trade-off in encoding efficiency to sage's contextual advantage: a sample of 150k lines from english wikipedia is encoded by 4 million bpe tokens, optimizing only an information-theoretic objective, whereas sage produces 4.5 million. having said that, this inefficiency might be further offset during llm pretraining: we propose that contextually coherent tokens will require fewer update steps in order to achieve useful embedding parameters, helping the model converge faster compared to bpe tokens. we leave testing this hypothesis to future work. ## downstream evaluation in order to evaluate the utility of our tokenization algorithm for major nlp tasks, we compare sage to a bpe vocabulary of the same size by means of pre-training a bert-parameterized model @xcite using an expedited training scheme @xcite . we then evaluate the llm's performance both on sequence classification via the english glue benchmark @xcite and the turkish partition of xnli @xcite , and on named entity recognition in english @xcite and turkish @xcite . we use the default settings from huggingface's library implementations of the finetuning processes @xcite and do not perform hyperparameter tuning for either model. we present our results on sequence-level tasks in ## related work in recent years, a growing body of research has demonstrated the shortcomings of existing tokenization algorithms in the context of representing linguistic phenomena in different languages across different tasks @xcite @xcite @xcite @xcite , as well as the statistical properties affecting their downstream performance @xcite . our work addresses the concerns raised in this line of work by introducing an improved subword vocabulary creation method which leverages the contextual aspects of the main intended use case, namely llms. previous work towards this goal includes algorithms which offer robustness within an existing subword vocabulary @xcite @xcite , necessitating modification of either training, inference, or both procedures in the context of llms. others have considered tuning the size of a subword vocabulary @xcite , or selecting from an enlarged set of possible segmentations @xcite , for optimizing performance on downstream tasks. some alternative tokenization methods focus on the application of a model which considers the expected downstream tasks together with the pretraining corpus @xcite , to the degree of jointly optimizing the tokenizer with the downstream model @xcite . in addition to the massive changes in training and inference procedures this approach incurs, we note that it is difficult to apply to large contextualized models due to the long path from tokenization to prediction; sage overcomes this problem by "nudging" only the llm vocabulary itself towards a contextualization-friendly segmentation. the concept of subword tokenization made its rise alongside that of contextualized representations, meaning that little work exists where skip-gram or other static models are trained over proper subword segmentations. recently, @xcite did so for a proof-of-concept of a spelling prediction model, in lieu of training full llms. to our knowledge, no work to date has used a static embedding-based objective to score token sequence likelihood for a separate task (as we do for vocabulary trimming). finally, we acknowledge the recent efforts to do away with tokenization altogether, be it through character-only @xcite or byteonly @xcite models, or through encoding characters visually and passing them through a vision model @xcite . these represent an even more radical departure from the established application of llms, and we look forward to testing their abilities against our improved contextual subword tokenization methods. we note that while these models have been facing issues regarding scaling, mostly on the decoding side, sage vocabularies are ready to be used immediately within existing popular llm implementations. furthermore, recent work has shown the limited utility of character-level transformers in semantic tasks, even for morphologically rich languages with nontrivial orthography-morphology relations @xcite . ## conclusion in this work, we introduced sage, a context-aware tokenizer built using insights from bpe, uni-gramlm, and skipgram, and showed that it achieves better results when used in an llm-pretrain-then-fine-tune schema on two typologically distant languages on both the sequence and token levels. we believe that further investigation into incorporating context in tokenization models can improve results even further, and intend to also extend our efforts toward other languages and writing systems, as well as to multilingual tokenizers. for example, we plan to apply sage in the context of abjads like hebrew and arabic, as well as languages written in alphasyllabaries such as devanagari. within sage itself, there is room for improvement. the algorithm is still relatively slow, taking roughly a day to run on a strong cpu, making it difficult to apply to a truly large corpus, to start from a larger initial vocabulary, or to conduct exhaustive search over the hyperparameters. we intend to keep optimizing it, and continue evaluation against other subword and character-only schemas.