A Problem in Continuation Tasks for Language Models with Large Vocabularies, and a Countermeasure

大词表语言模型在续写任务上的一个问题及对策 苏剑林 · 2023-09-13 · Machine Learning probability problem language model

For LLMs, increasing the size of the tokenizer's vocabulary in order to improve the compression ratio, and thereby shorten sequence length and reduce decoding cost, is something everyone is happy to see. After all, enlarging the vocabulary only requires enlarging the embedding layer and the final dense layer, and the extra computation this introduces is barely noticeable, while the speedup in decoding brought by shorter sequences is very real. Of course, increasing the vocabulary size can also bring some negative effects on model performance, so it's not something that can be scaled up without limit. This post analyzes a problem that arises in continuation tasks when the vocabulary is enlarged, and proposes a reference solution.

Pros and Cons

The benefits of enlarging the vocabulary are obvious. On one hand, since LLMs are autoregressive, decoding inevitably slows down as sequence length grows; "larger vocabulary → higher compression ratio → shorter sequence length" means that the same piece of text now corresponds to fewer tokens, i.e. fewer decoding steps, which speeds up decoding. On the other hand, language models are trained via Teacher Forcing, and shortening the sequence length can alleviate the Exposure Bias problem caused by Teacher Forcing, potentially improving model performance.

That said, the drawbacks of enlarging the vocabulary are also clear. The most direct one is that it breaks the character-level connections between tokens, which may hurt generalization or even degrade the model's ability to perform certain tasks. For example, if both "太阳能" (solar energy) and "太阳" (sun) are single tokens in the vocabulary, the model has no way of knowing that "太阳能" is composed of "太阳" and "能," nor that "太阳" is composed of "太" and "阳." This makes subword-related tasks quite difficult — the classic example being the question "What do you get if you reverse '太阳能'?" The expected answer is "能阳太," but since the model doesn't know it's made up of the three characters "太," "阳," "能," it's very hard for it to answer correctly.

The Continuation Problem

Recently @Armen Aghajanyan shared another problem. When training a code model with an extremely large vocabulary, they found that common commands such as "import numpy as np" ended up as a single token — and as a result, when a user typed "import numpy," the model could not continue with " as np." The reason is simple: since "import numpy as np" is treated as a single token, whenever "import numpy" appears on its own, the model finds that it is never followed by " as np" (because every instance followed by " as np" has been merged into the single token "import numpy as np"), so naturally it cannot produce that continuation.

This phenomenon is indeed quite classic, and it's not limited to code models — it also shows up in common natural language models. For instance, if both "太阳能" and "太阳" become independent tokens, then after a user types "太阳" (sun), the next character the model continues with will basically never be "能" (energy), which may not match what the user expects. Similarly, if "白云" (white cloud / Baiyun), "白云山" (Baiyun Mountain), and "白云机场" (Baiyun Airport) are all independent tokens, then after a user types "广州的白云" (Guangzhou's Baiyun), the model will almost never continue with "广州的白云机场" or "广州的白云山," and so on.

A Reference Countermeasure

However, I would argue that the phenomenon Armen Aghajanyan pointed out isn't really a drawback of enlarging the vocabulary — in fact, with a little bit of handling, it can actually become an advantage of a larger vocabulary. The problem is actually quite simple: back before LLMs existed, we could already perform a certain amount of completion using "vocabulary + prefix search." Now that we have LLMs, why should we be confined to LLMs alone, instead of combining LLM-based continuation with vocabulary-based continuation?

Take the same example: suppose the user inputs "广州的白云" (Guangzhou's Baiyun), and the tokenizer splits it into "广州/的/白云." If we directly convert these three words into ids and feed them into the model, we won't be able to generate continuations like "广州的白云机场." This is fundamentally because the tokenizer cannot foresee future text, leading to an incorrect tokenization (of course, one could also consider using a tokenization algorithm with some randomness during training, in which case "白云机场" might sometimes appear as a single word and sometimes as "白云/机场" — in that case the tokenization result wouldn't severely affect downstream performance, and might even improve generalization; see Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates).

So, can we anticipate the future text in some way? Suppose the tokenization gives us "广州/的/白云"; we can then back up one step and use "白云" to perform a prefix search over the vocabulary. Let's say the search returns four words: "白云," "白云机场," "白云山," and "白云路." This search is done purely based on the vocabulary, and its computational cost is negligible compared to running the LLM. Once we have the search results, we use the LLM to compute:

\begin{equation}\begin{aligned} p(\text{Baiyun}|\text{Guangzhou},\text{of}) \\p(\text{Baiyun Airport}|\text{Guangzhou},\text{of}) \\ p(\text{Baiyun Mountain}|\text{Guangzhou},\text{of}) \\ p(\text{Baiyun Road}|\text{Guangzhou},\text{of}) \\ \end{aligned}\end{equation}

Since the input is the same in all four cases, computing these four conditional probabilities only requires a single forward pass of the LLM. Once we have these four conditional probabilities, we renormalize them and sample. If the sampled result is "白云," we continue generation as if the tokenization were "广州/的/白云"; if we sample "白云机场," we can output "机场" and continue as if the tokenization were "广州/的/白云机场"; and so on. This neatly resolves the problem raised by Armen Aghajanyan, and even turns the drawback into an advantage (when the compression ratio is high, even though we back up one step, the word found by the prefix search may be quite long, allowing more characters to be generated at once). In particular, the backtracking operation only needs to happen at the very first sampling step — its sole purpose is to avoid tokenization errors caused by incomplete input — and from the second step onward no backtracking is needed, so the extra computation introduced is minimal.

It's worth mentioning that Microsoft has a library called "guidance" that proposes exactly the same trick (see here). Furthermore, considering more general scenarios, sometimes backing up just one step isn't enough. Take the "import numpy as np" example: when only "import numpy" is entered, it might be tokenized as "import/ numpy," in which case we'd need to back up at least two steps to get a complete, reasonable sequence. But this is not fundamentally different — just a bit more complicated in the details — so I won't elaborate further here; readers can work out the specifics themselves when deploying inference models.

Summary

This post described a problem that can arise in text continuation tasks for LLMs with extremely large vocabularies, and shared a reference solution.