Token Alignment via Character Matching for Subword Completion (2024)

Abstract: Generative models, widely utilized in various applications, can often struggle with prompts corresponding to partial tokens. This struggle stems from tokenization, where partial tokens fall out of distribution during inference, leading to incorrect or nonsensical outputs. This paper examines a technique to alleviate the tokenization artifact on text completion in generative models, maintaining performance even in regular non-subword cases. The method, termed token alignment, involves backtracking to the last complete tokens and ensuring the model’s generation aligns with the prompt. This approach showcases marked improvement across many partial token scenarios, including nuanced cases like space-prefix and partial indentation, with only a minor time increase. The technique and analysis detailed in this paper contribute to the continuous advancement of generative models in handling partial inputs, bearing relevance for applications like code completion and text.

## introduction generative models have shown remarkable efficacy in a range of applications. however, they have been observed to falter when dealing with partially provided inputs or subwords during text completion. for instance, a generative model might struggle to predict the remaining part of the word where a prompt ending in a subword often leads to incorrect or nonsensical outputs. this issue arises due to the artifact of tokenization where a partial token can be out-of-distribution during inference. in this paper, we introduce a method to address and rectify these shortcomings in generative models, ensuring they maintain their performance on complete contexts. our approach involves backtracking to the last complete tokens and aligning the model's generation to match with the given prefix, as shown in @xmath0, after which the completion is shown. without token alignment, the model fails to predict "return" correctly, as the sequence of tokens "re", "turn" is out-ofdistribution. token alignment alleviates this constraint by backtracking to full tokens before "re", then align subsequent generations with the prompt. figure 2 . for example, if the end of the prefix is "sys" we guide the model's generation to ensure compatibility with this prefix where all possible tokens match with "sys". this method not only improves the accuracy of the model's predictions on subword data metrics but also retains the performance on generic evaluation sets. to illustrate our approach, consider the following programming code snippet where the end of the prompt is right after the incomplete token "re" in figure 1 . a conventional model might generate an incorrect or suboptimal output such as "@xmath1, while the correct prediction should be "return". our approach solves this issue by matching the characters starting from the last complete pretoken and aligning the subsequent tokens with the prefix. the results of our approach, as shown in various examples, demonstrate model's ability to reliably generate the correct output regardless of partial context. figure 2 : an illustration of token alignment process. we first tokenize the prompt where the last tokens may correspond to partial tokens. we use the backtracked prompt as model input and use the alignment prefix to filter possible tokens at byte or character level. we then mask out non-matched tokens by zeroing their probabilities, which is later used for sampling to select the next token. our evaluation of the token alignment approach indicates noticeable improvement in handling partial token contexts, outperforming the traditional decoding approach without any subword handling. we present comprehensive evaluation results on many partial token scenarios such as subword of natural words, as well as much less obvious cases of partial tokens related to punctuations, or artifacts of contiguous white spaces and space prefix. our approach also does not incur much extra latency, with an average increase of only 3-7 ms for using token alignment due to an efficient trie-based lookup table. we believe our findings provide a meaningful and practical contribution to the ongoing work in improving the performance and robustness of generative models, especially related to handling partial inputs. we release our code and datasets at https://github.com/ amazon-science/tokenalign . ## related work recent works have identified severe robustness issues for large language models. for instance, simple perturbations like synonym substitution on one word can significantly change outputs and fool the models @xcite . most studies focus on general perturbations especially at random locations of prompts for both text @xcite @xcite and code models @xcite @xcite . however, we see a large gap of partial versus full tokens at the end of prompts with only a few characters missed or added (e.g., with or without a white space) due to the artifact of tokenization. this robustness problem is barely in-vestigated. subword regularization tried to partially mitigate such problems by randomly introducing sub-tokens in training and improve robustness behavior @xcite . however, this requires to retrain the models and also leads to inference latency increase. our approach, on the other hand, can work for any models directly with negligible latency increase. we acknowledge concurrent work in the form of a blog post namely token healing @xcite , which attempts to similarly solve the partial token problems on inference time. however, we independently develop our own token alignment algorithm in a more general design. we also thoroughly validate and analyze the effectiveness of token alignment with extensive empirical results on various tasks and many partial token scenarios. other related previous work includes approximate marginalization @xcite for beam search. our method however is developed for language models which primarily use nucleus sampling @xcite for next token prediction. ## methodology generative language models have been remarkably successful in a variety of applications, including code completion in integrated development environments (ides) and text autocompletion in email or productivity tools. however, these models often encounter difficulties when dealing with incomplete inputs that end with partial tokens, which we refer to as the partial token issue. in such cases, the models tend to generate outputs that are not compatible with the given context due to the constraint imposed by tokenization. this limitation significantly hampers the utility and user experience of algorithm 1 tokenalign algorithm in our methodology, we first define the last complete token by examining the tokenization of the input sequence. given the tokenization of the last line as t 1 , t 2 , ..., t n , we start generating from the -b th token, i.e., t n-b+1 , where b is the number of tokens we need to backtrack. the rest of the context is the sequence t n-b+1 , . . . , t n , which we refer to as the alignment prefix where the subsequent generation needs to match at a byte level (or character level if the tokenizer is not a byte-level bpe). 1 the token alignment step ensures that the next token starts with the alignment prefix or is a prefix of the alignment prefix itself by masking out probabilities of tokens that do not match. after each step where the next token is selected, we deduct the generated bytes from the alignment prefix until the alignment prefix is empty, after which we finish the token alignment process. the efficiency of our approach is supported by two key techniques: a character-trie for fast prefix matching, and a mask cache. the character-trie is a tree data structure, where each node represents a character and the paths from the root node to the leaf nodes represent complete tokens in the vocabulary. we pre-build this trie using the model's vocabulary before starting the generation process, thereby enabling a fast and space-efficient representation of the vocabulary. the mask cache is used to further accelerate the lookup process by avoiding unnecessary trie lookup. for common cases, such as a single space as context, we pre-build and cache a boolean mask. this cached mask is then used for quickly filtering tokens that share a common prefix with the given context. these techniques significantly reduce the latency, enabling our method to be highly efficient. the use of a trie-based lookup table and the pre-built mask cache results in a time saving compared to a naive implementation using a dictionary lookup. in sum, our methodology presents an effective and efficient solution to the partial token issue in generative models, enhancing their performance in auto-completion applications in code and text domains. ## partial token scenarios we give an overview of common partial token scenarios where the model is susceptible to degenerate behavior when the prompt ends with partial tokens. such partial token scenarios range from obvious cases such as natural subwords to less obvious cases such as space prefix. examples of these scenarios are illustrated in figure 3 in the appendix. subword. this case refers to natural language subword. for instance, a possible subword of "banana" is 'banan' or 'bana'. note that while it is possible for a tokenizer to use a subword as an actual token, such as "beautifully" can be represented as 'beautiful' + 'ly', not all possible subwords will correspond to a token, e.g., 'beautifu' is unlikely to be a full token. punctuation. it is possible that multiple punctuations are grouped together as a token during the tokenizer training process, which can lead to some punctuations being a partial token in a non-obvious way. for example, in python, it is possible that the code snippet if (@xmath2:', which can lead to a suboptimal completion behavior if presented to a model. in the case of punctuation, it is generally quite hard to judge which case would correspond to a partial punctuation. we separate out this case to emphasize the importance of token alignment as a universal method that can handle both prompt ending in either partial token or non partial token in a prompt-agnostic way. space prefix. most tokenizers use the space-prefix schema where a whitespace is often grouped as part of a word as a prefix @xcite @xcite . this schema is primarily aimed to reduce the number of tokens required to represent text. for example, with spaceprefix, 'i like' can be tokenized into two tokens, 'i' and ' like', instead of three ('i', ' ', and 'like'). if a prompt 'i ' is presented to the model, such model can have difficulty predicting coherent text since the prompt is a sequence of tokens ('i', ' '), which is out-of-distribution compared to the training stage which observes ('i', ' like'). the space prefix's constraint is also quite pronounced in the presence surrounding indentation block. for instance, @xmath3, ' x', '='). however, if the prompt ends after the full block of indentation , this becomes a constraint since (' ', 'x') is outof-distribution compared to (' ', ' x') seen during training. we refer to this as the space prefix with indentation case. contiguous spaces. modern tokenizers often intentionally group whitespaces together for improved compression rates. for instance, codex @xcite extends the gpt-3 tokenizer @xcite to include contiguous spaces of different lengths (up to 65 contiguous spaces in some version of tiktoken (openai, 2023), for instance). other models such as llama @xcite , starcoder @xcite , codegen @xcite , to name a few, also use such contiguous white spaces. while such grouping of whitespaces is compression efficient, it can lead to out-of-distribution behavior if the prompt ends in partial token of such contiguous whitespaces. for instance, consider the following prompt ' if true:\n ' where represents a space character. the correct syntax for python is such that what comes after the if clause requires another level of indentation, meaning that the model should predict two extra characters. however, the model will have a hard time completing the next two . this is because during training the model always observes four contiguous whitespaces if it were to complete the next level of indentation. ## evaluation in this section, we empirically demonstrate that each of the partial token scenarios described in section 4 can unnecessarily constrain the model due to the artifact of tokenization and lead to significant drops in evaluation scores. we perform evaluations on token completion tasks in both code and text generation. we show that token alignment method graciously handles the constraint due to all cases listed in section 4, making a universal partial token handling approach for language model inference. ## ablation study and generality of token alignment in this section, we discuss an alternative approach to subword completion by processing training data with subword regularization. we show that our token alignment method is complementary and can offer a robust approach without further training to handle subwords compatible with subword regularization. we also discuss the latency as well as ablation study on the number of backtrack tokens. ## conclusion this paper explores the complex challenge of handling partial tokens in generative models, introducing specific scenarios such as subword, punctuation, space prefix, and contiguous spaces. through the application of token alignment, we demonstrated noticeable improvements across these scenarios. in the future, as tokens become longer due to possible advancement in tokenizer compression, the token alignment may become increasingly important as more cases may subtly fall into the categories of partial tokens. the findings of this study emphasize the importance of understanding tokenization artifacts and present a meaningful advancement in enhancing the robustness and practical applicability of generative models in various real-world tasks.