N-grams and the Curse of Dimensionality

Having defined the vocabulary our chatbot will work with, we now face the task of making it learn how to chain all these tokens together in a fashion reminiscent of the training data. This is an extremely daunting challenge, which we will spend much of our journey tackling. We have no idea how to overcome it all at once, so let’s start with a dumb approach and see if we can make it smarter step by step.

Picking tokens at random

Arguably the simplest possible way to generate text would be to randomly choose a token from the vocabulary at each step, with all tokens being equally likely to be chosen. A primitive chatbot using this approach produces output like this:

The only thing learned from the training data here is the vocabulary itself, and the output is complete gibberish. Notably, it does not even contain spaces, which makes sense given our tokenization — since the whitespace is just one of our 20,000 tokens and each token is equally likely to be chosen at every step, the chatbot will rarely output a space.

A unigram model

Choosing every token with equal probability is, of course, rather senseless. In English text, some tokens appear much more often than others. If the goal is for our chatbot to mimic the training data, our next-token probability distribution should reflect this fact. For example, if about 2% of the tokens in the text data we feed our learning program are instances of the article “the”, then we would like our chatbot to output the token  the about 2% of the time. We can ensure this by going through the training data, counting the number of times each token appears and adjusting the probability that our chatbot picks that token accordingly:

The resulting next-token probability distribution can be thought of as a (very naive) probabilistic representation of the language in the training data. In machine learning, such probabilistic representations of language which can be used to generate text output by repeatedly drawing from the probability distribution in question are generally referred to as Language Models.⁠[note]

The simple count-based next-token probability distribution we calculated here is called a Unigram Model, since it models the training data on a single-token level. For our simple War-and-Peace training-data example, the output of such a unigram model looks something like this:

As expected, this reflects the rates at which various tokens appear in the training data, yielding lots of spaces and plenty of common words such as pronouns and prepositions. But it’s still gibberish.

A bigram model

This lack of coherence is unsurprising, since our model still draws each token independently without taking into account any of the preceding tokens. In machine learning, the sequence of proceding tokens a given model considers when generating output is referred to as the Context.

Let’s ensure that our model considers some context by making another incremental improvement: Whenever we generate a new token, let’s take into account the token immediately preceding it. For example, if the previous token is  said, and the word “said” is followed by a period 5% of the time in the training data, let’s make sure there is a 5% chance our chatbot generates the token  . next.

To ensure this, we need to figure out the set of probabilities with which any one token follows any other in the training data. This sounds complicated but is actually straightforward — we can simply count the number of times any given two-token combination occurs and divide it by the number of times the first token occurs: 

Formally, the resulting probability is the Conditional Probability that the next token will be token $t_2$ given that the previous token was token $t_1$. The proper mathematical notation for this probability is $p(T_{\text{next}} = t_2 \,|\, T_{\text{previous}} = t_1)$, even though we will often just write things like $p(t_2 \,|\, t_1)$ for short. For example, in the scenario above, our model would learn that $p(T_{\text{next}} =$  .$\,|\, T_{\text{previous}} =$  said $) = 0.05$ simply because 5% of the occurrences of the word “said” in the training data are followed by a period.

The probabilistic model given by all of these conditional probability distributions is called a Bigram Model, as it reflects learning on a two-token level.

During output generation, our chatbot can now simply use the most recently generated token as its new input, draw a next token from the corresponding conditional probability distribution, use that next token as its new input, and so on. The concept of repeatedly re-attaching the output to the input like this is called Autoregression†From the Greek “autos” meaning “self” and the Latin “regressus” meaning “to go back”[note], and models that keep eating their previous output as input in this manner are referred to as autoregressive models.

Because the bigram model consists of conditional probability distributions, we necessarily need to provide our bigram chatbot with a first token to kick things off. If we want it to generate language from nothing, the simplest solution is to draw a first token from the unigram distribution described above and to use the bigram method from there:

In chatbot practice, we don’t want the model to generate language from nothing though. We want it to respond to a given prompt. To get a good response, we will necessarily have to figure out how to feed that entire prompt to our model eventually, but let's start small. Our simple bigram model can only consider a single token of context, so we should at least pick that one token from the tokens contained in the prompt.

An intuitive choice is to simply take the end of the prompt and use that to kick off our bigram chatbot's output generation, so let's just do that. Here is what this looks like if we use an arbitrary sentence from War and Peace as our prompt and simply let our bigram bot carry on:

Regardless of whether we jumpstart our bigram bot with a prompt token or with one drawn from the unigram distribution, the results are still far from proper English. But, crucially, they do look better than the output of the unigram model. In particular, words are actually properly spaced out now, which makes sense given that the bigram model considers the previous token and that double spaces are exceedingly rare in the training text — whenever the most recently generated token is a space, the conditional probability for another space to follow is virtually zero and the bigram model will output a word. If the previous token is a word, on the other hand, the conditional probability for another word to follow is zero, and the bigram model will output a space or punctuation mark.

More generally, we can examine the sample output and verify that any pair of successive tokens, taken on its own, could very well stem from proper English text (as indeed it does, since the bigram model can only produce bigrams contained in the training data)! It’s only if we consider more than two successive tokens at a time, such as two words separated by a space, that the incoherence of the output becomes apparent.

Of course, in actual chatbot practice, the prompt will be an arbitrary prompt made up by the user, not a sentence drawn from the training data. So let’s play around with this:

This didn’t go so well. The problem is that our model was only trained on the training data, in this case the raw text of War and Peace. So if the prompt provided by the user ends in a token that never occurred in this data, the bigram model doesn’t recognize it and throws an error. This should become less of a problem as we make our training dataset larger and larger, as it will become increasingly unlikely that the prompt contains a token the model has never seen, but the problem will never go away completely — assuming we don’t want our bigram bot to throw an error if a user happens to end a prompt with a new word creation or an exotic typo, we need to make sure it has a way out if it doesn’t recognize the last token of the prompt.

A simple way of ensuring this is to revert to the unigram model for just one token. Since any token the unigram model can generate occurred in the training data and is therefore known to the bigram model, the bigram model can always carry on from there. This strategy is referred to as Backing Off — we attempt to use our more ambitious bigram model first, but if it fails, we “back off” and use the simpler unigram model which is guaranteed to work.

There are other many approaches we could use, such as replacing all tokens which only appear a very small number of times (or never) in the training data with a special “rare token” token and using that general token’s bigram distribution whenever we encounter something unknown, but obsessing over this problem doesn’t seem super instructive for now, so let’s just assume we use the back-off strategy to handle it and return our focus to the main challenge of producing more coherent language output.

A trigram model

The bigram model clearly fared better than the unigram model, so let’s keep going. If considering the previous token created coherence across successive tokens, let’s simply consider the previous two. To find the corresponding conditional probabilities, we simply count the number of times any given three-token combination occurs and divide it by the number of times the first two tokens occur together:

We have slightly adjusted our previous notation here to make it more generalizable — $p(T_{i+1} = t_3 \,|\, T_{i-1} = t_1,\, T_i = t_2)$ is the conditional probability that the next token will be token $t_3$ given that the previous two tokens were tokens $t_1$ and $t_2$. For example, if there is an 11% chance that the two-token sequence  will   is followed by the token  not in the training data, then $p(T_{i+1} =$  not$\,|\, T_{i-1} =$  will$,\, T_i =$   $) = 0.11$.

A collection of conditional probability distributions of this type is referred to as a Trigram Model, since it models the training-data language on a three-token level. Since the probability of any given token to be generated depends on the two tokens before it, a trigram chatbot necessarily needs to be provided with two initial tokens to kick things off.

Let’s stick with our previous strategy of linking the response to the prompt and use the last two tokens of the prompt as our initial tokens. Here is what this looks like if we use the same War and Peace sentence as before as our prompt:

This is slowly starting to feel like English and, as expected, looks indistinguishable from an English text if we only consider any one triplet of three successive tokens! Amazing.

Let’s play around with this and see how our trigram bot handles an arbitrary prompt that isn’t drawn from the training data:

Once again, this didn’t go so well. The problem here isn’t that one of the tokens in the prompt is unknown to the model, as was the case earlier. Rather, it’s that even though they are part of our vocabulary, the last two tokens of the prompt  Please   describe   the   duties   of   a   Russian   nobleman   in   the   nineteenth   century ., namely the word “century” and the period “.”, never occur together in the training data, meaning that our trigram bot never learned a probability distribution for this two-token context and therefore produces an error.

Of course, using vastly more training data would likely eradicate the problem with this specific prompt eventually, but there is no guarantee that all two-token combinations that might appear at the end of a prompt would show up even in a huge dataset, just as there was no guarantee that no prompt could ever end in an unrecognized token to begin with.

So for these problematic cases, let's just use the same solution we used when facing that problem with our bigram model and extend our back-off strategy: Let's try to use a trigram model first. If this fails, no matter if it’s because of an unknown individual token or because of an unknown combination of known tokens, we back off and try to use the bigram model. If this fails, too, we back off all the way to the unigram model. Since the unigram model is always guaranteed to work (it does not depend on the prompt at all), such a back-off trigram bot will never throw an error.

An n-gram model

Models like the ones we have discussed so far are generally referred to as N-gram Models, where $n$ is the total number of tokens considered including both the context of preceding tokens and the token to be generated, so $n\!=\!1$ for the unigram, $n\!=\!2$ for the bigram and $n\!=\!3$ for the trigram model.

Well then, enough with the baby steps already. If making n larger and larger results in coherence across longer and longer pieces of text and we will eventually need to make the context larger to take into account the entire prompt anyway, let’s just run with it and set $n$ equal to something much larger!

For sake of specificity, let's look at $n\!=\!20$. We can make a 20-gram model by counting the occurrences of any given 20-token sequence and dividing by the number of times the first 19 tokens occur together. Our chatbot can then use the resulting conditional probabilities $p(T_{i+1} = t_{20} \,|\, T_{i-18} = t_1,\, \dots,\, T_i = t_{19})$ to draw the next token at any given time step, just as before. Here is what this produces if we prompt it with a sentence from the War and Peace training data:

Wow! This actually sounds like Tolstoy, which seems... too good to be true. And indeed, if we check the original text we find that this doesn’t just sound like Tolstoy, it is Tolstoy! It’s exactly how the prompt is continued in the book.

On a superficial level, this might sound like we achieved our dream — we wanted to imitate language, and here we have a chatbot that perfectly imitates the great work of literature it was trained on. But that is not what we meant by imitate language, of course. We want our model to use the training data to learn to generate original output on its own, the way human infants do. We don’t just want it to regurgitate the exact texts it was exposed to.

In machine learning, this "going beyond the training data" is referred to as Generalization. It's pretty much what distinguishes real learning from mere memorization and is the implicit goal of virtually all AI endeavours.

So this output isn’t good. But how does it come about in the first place, given that we didn’t explicitly store the training text and that our chatbot only ever considers the previous 19 tokens?

The answer is simple: The 19-token sequence    The   spindles   hummed   steadily   and   ceaselessly   on   all   sides . appears in the War and Peace training text exactly once, meaning the corresponding 20-gram count is 1 and our model assigns a next-token probability of 100% to the whitespace    , which follows this 19-token sequence in the novel. The new 19-token context  The   spindles   hummed   steadily   and   ceaselessly   on   all   sides .   also appears in the novel only once, so the chatbot again generates the one token that follows it in the training text (the token  With), and so on and so forth.

Basically, even though we are not explicitly storing the training text as a whole, we are doing so implicitly, because every context appears only once in the novel and the conditional probability distribution for any given context collapses entirely onto the one token Tolstoy chose next.

This wasn’t a problem for the trigram model, because the context window only consisted of two tokens and many contexts therefore appeared multiple times in the training text, with multiple different continuations. As a result, the model spread its next-token probability for those contexts across several tokens, and by sheer randomness, our trigram chatbot was able to produce sequences that weren’t contained in the training data.

But for large enough context windows, any context that appears in the training data at all will appear in it only once, with one specific next-token continuation. This is true regardless of how much training data we use simply because the number of possible contexts explodes dramatically with increased context length – even if we fed the entire internet to our model, any context of a few hundred or thousands of tokens would still only occur once, unless there happen to be multiple copies of the corresponding text in the data, in which case the continuations for thos copies will all be the same anyway. As a result, an n-gram chatbot would effectively memorize the next-token continuation from the data like our 20-gram bot did here.

This explosion of possible contexts is based on sheer combinatorics⁠[note] and very much lies at the heart of the difficulty of language modeling. It means that our training data is bound to be but a drop in the ocean of what we ultimately want our chatbot to be able to process and generate, and strictly implies that successful language modeling requires extreme generalization. Because this combinatorial tragedy is so fundamental, people have given it an epic name: The Curse of Dimensionality.

There is no way around the immediate result of this curse — if we make our context window large enough, the training-data contexts will all be unique. The part of the problem we might be able to address is that our chatbot only learns the single continuation contained in the training data for each of those contexts, and cannot generalize to other reasonable next-token continuations.

For example, if a given context is followed by the token  April in the training-data, it would be nice if our chatbot could also generate reasonable substitute tokens such as  March or  May in addition to the original  April, as this would enable it to generate reasonable output that differs from the training data and thus address the memorization problem. But sadly, our simple n-gram bots have no way of doing this. Let's refer to this as the Output Generalization Problem in the following.

Next, let's once again try an arbitrary prompt:

Somewhat unsurprisingly, this causes the same problem we encountered earlier — our model only has next-token probabilities for the contexts it saw during training, so it doesn't know what to do with a prompt whose ending cannot be found in the training data. But the problem is way more severe now. For the trigram model, a significant fraction of the reasonable (two-token) contexts a prompt could end with was already contained in War and Peace, and most reasonable contexts would eventually pop up in a large enough training dataset. For the 20-gram model, this is not the case — an arbitrary prompt will almost never end in a 19-token sequence that happens to appear in War and Peace, so our novel-trained 20-gram model is essentially useless for arbitrary prompts, and even if we use a much, much larger dataset, we will always run into this same problem for sizeable context windows due to the curse of dimensionality.

For the trigram model, we decided to circumvent this issue using our back-off strategy, which ensures that our chatbot can always keep going by temporarily shortening the context window it considers. But this simply isn’t good enough if the problem occurs virtually all the time. After all, a 20-gram model that constantly has to back off to a trigram model is basically just a trigram model, and it will never produce the coherent language output we were chasing when defining the 20-gram model in the first place.

So we will have to actually address this problem of running into unprecedented contexts instead of avoiding it. Because the heart of the issue is that our simple n-gram models cannot generalize to contexts which weren’t contained in the training data and the context represents the input to our n-gram model, I will refer to this as the Input Generalization Problem.

Let's see if we can solve these problems.

Next Chapter