Character Prefix Conditioning using Back Tokenization

I recently came across a problem published in cursor's blog named Character Prefix Conditioning. It opens with the following context:

When using a language model for code completion, we typically want the model to produce a completion that begins with what the user has typed. However, modern language models operate on sequences of tokens, not characters, so naively tokenizing the user's input and sending it to the model produces wrong results if the user's cursor doesn't happen to lie on a token boundary.

To understand this, let's consider the following completion example: Here, the user has typed "The company hired an intermed" and the model is supposed to suggest a completion "iary to negotiate." There's a problem here. Let's look at the token representation of the user's current state and the actual sentence to get a better understanding of the problem:

User's current state: Actual sentence: If we were to naively tokenize the user's current state and send it to the model, we would have to sample the following tokens to reach the actual sentence: This would hurt the performance of the model as the sequence of these tokens does not match the actual sentence's token sequence.

You can already tell that a simple approach won't solve it. Even token healing won't fully help here as if we backtrack a token and generate from there, we would start sampling from The company hired an inter which won't get us to the token sequence of the actual sentence which need us sampling from The company hired an.

I am going to use the Byte-Pair Encoding (BPE) tokenization algorithm to solve this problem. Specifically the cl100k_base used for GPT-4 tokenizer.

The Byte-Pair Encoding tokenizer works by starting with a vocabulary of individual characters and then iteratively merging the most frequent pair of adjacent tokens to form a new token. This process continues until a desired vocabulary size is reached or no frequent pairs remain. The important thing to understand here is that when tokenizing a new string, we apply our merge rules in the exact order they were learned during training. Unlike the training phase where we repeatedly count frequencies, during inference we simply apply the rules in sequence.

This is a simplified example - in practice, the cl100k_base tokenizer used by GPT-4 has a much larger vocabulary (100,000+ tokens) and more sophisticated merge rules. The GPT-2 paper highlights one of things done additionally to improve the performance of the tokenizer:

We observed BPE including many versions of common words like dog since they occur in many variations such as dog. dog! dog? . This results in a sub-optimal allocation of limited vocabulary slots and model capacity. To avoid this, we prevent BPE from merging across character categories for any byte sequence.

They enforce this prevention of merging by using a regex pattern that identifies boundaries where merges shouldn't occur.



I am going to use this regex to only search and validate over subset of the overall user input.

When implementing CPC with self-hosted language models, there are two main approaches to handling multiple completion candidates, with significant performance implications:

• Sequential autoregressive generation for each candidate path
• Batched processing of multiple candidates together

Most of the time spent on this problem went into preparing the dataset. I tried to think of and find as many different cases as possible.

Here are some of the cases I considered, ordered by increasing complexity:

• 1. We have to implement search over the token space in some form to find possible generation paths
• 2. We cannot make the assumption of taking the longest matching token. The actual one would be very much dependent on the context of the user's input and may not be the longest one
• 3. The search might have to be initiated at multiple points in the user's input but we need to be efficient about it
• 4. We cannot base this search on the existing token boundaries constructed on the user's input as the token boundaries might be different for the final sentence
• 5. We need to prune the search space to make it manageable while still covering all the cases

If you've been able to follow along so far, the solution would seem obvious.

The strategy involves four main steps:

1. Create a trie data structure from the tokenizer vocabulary with the ability to take a string and get all tokens that has the the string as the prefix
2. For a given user input, search from the right to left of the last work determined by the regex pattern. Move character by character from left to right accumulating the prefix and for each prefix, find possible token completions that have the prefix
3. Filter token candidates to ensure they maintain valid token boundaries
4. Return a list of possible completion paths for the model to consider

At the core of our solution is a trie (prefix tree) data structure optimized for byte-level searching.
Each node in the trie represents a byte and the children of the node represent the next byte in the token.
The trie also stores the original token that ends at that node.

So given a prefix, we can traverse the trie and collect all the tokens that have the prefix. We first take the user's input and split it into words using the regex pattern from the tokenizer library.

For the dataset above, the last words would be: Now, we search the trie for each of these prefixes. For example, if we consider the prefix " indivi", we would initiate search for the following prefixes: The above search yields many potential candidates, forming a superset of the combinations we want to explore. We need an effective pruning strategy to filter these combinations. Let's walk through an example using the prefix "divi" from our search candidates above: Since the re-encoded tokens don't match our expected sequence, we discard "division" as a candidate. By applying this validation process across all candidate tokens, we can effectively filter out invalid completions while preserving the ones that maintain proper token boundaries.

Let's see the algorithm in action with some of our test cases:

The implementation includes several optimizations:

• Trie-based search: O(k) lookup time instead of O(nk) for linear vocabulary search
• Early pruning: Using the regex pattern to focus search on relevant portion of text
• Iterative DFS: Using queue-based iteration to avoid stack overflow on deep trie traversals
• Token validation: Filtering out invalid candidates early to reduce downstream processing

I am sure there are many other ways to solve this problem. If you have alternative approaches, optimizations, or spot any issues in my implementation, please feel free to open an issue or submit a pull request on GitHub. I welcome contributions and feedback from the community to make this solution even better.

The code is available here.