Similarity and Embeddings

Let's try to tackle the generalization problems outlined in the previous chapter.

What should our model do when encountering a context it wasn't trained on? Given that we have no way of suddenly obtaining additional training data for such a context, and that backing off and only considering part of the context would effectively mean ignoring part of the prompt, it seems that the best thing our chatbot could possibly do is derive inspiration from the contexts it does know and make an educated guess. Of course, most of the contexts known to our model won't be very useful in forming that guess, as most of the training data is bound to be completely unrelated to the prompt.

But if our training dataset is large enough, there should be at least some contexts in it which are similar to the new, unknown context in question. Ideally, we'd like our model to use a next-token probability distribution inspired by the next-token distributions of these similar contexts. For example, if our model is prompted with  Cats   make   for   great   pets   because   and this phrase never occurred in the training data while the context  Dogs   make   for   good   pets   because   did, then we'd like our model to somehow have learned that these are very similar contexts and choose the next token accordingly.

What about the output problem? The most straightforward way to enable our n-gram model to create new sequences it didn't encounter during training would be to simply scale down the probabilities of the next-token continuations it did encounter during training and to distribute the freed-up probability equally (or according to lower-order n-gram counts) over the other tokens in the vocabulary. The result would be a model that follows the n-gram distribution from the training data most of the time but sometimes goes rogue in the name of diversity and draws the next token from some fixed distribution that does not depend on what token(s) followed the context in the training data.

But that's a bad solution, of course. A smarter model would take into account the nature of the training-data continuation(s) and would preferentially pick tokens similar to the token(s) which followed the context during training. For example, if a given context is only followed by the token  dog in the training data, we'd still like our model to assign some probability to tokens like  labradoodle,  dachshund, or even  cat, given how similar the meaning of these tokens is to that of  dog, but we wouldn't want it to assign significant probability to the token  spaceship because that would almost certainly be non-sensical.

So for both the input and the output problem, the notion of Similarity seems central to any solution we might come up with.

Features, scores, embeddings

So to build a better model, let's derive some inspiration from how we think about similarity in everyday life.

Let's say you are making small talk with someone and they ask you how similar the city you are currently living in is to the one that you grew up in. What do you do?

Assuming you are taking the task seriously, you'll likely come up with a set of Features to compare the two cities by. For example, you might start commenting on their size, typical weather, cost of living, proximity to the ocean, and so on. As long as you are only asked about two cities and the conversation is casual, these comparisons would likely be qualitative, meaning you would say things like "city A is bigger than city B", or "living in city A is much more expensive than living in city B".

But what if that weren't the case? What if you were instead working for the government of your home country and had to compare all its major cities in order to determine how certain resources will be allocated to them? You'd still have to come up with a bunch of features to compare them by, just by virtue of trying to be specific, but given the stakes, you'd probably start to be much more quantitative. For example, you might say that city A has a population of 3.1 million people, city B has a population of 1.4 million people, city C has a population of 2.2 million people, and so on. Generally speaking, you will assign some type of Score to each city for each individual feature, and then say that cities whose scores are more alike should be treated similarly in terms of resource allocation, simply because they are more alike.

Of course, there is nothing fundamentally special about comparing cities in particular, and comparing things using feature scores like this seems very intuitive in general. In fact, feature scores like this pop up in many contexts across daily life, ranging from sports video games to college applications.

This seems promising, so let's think about it a bit more rigorously. From a mathematical point of view, we can think of performing this scoring trick as choosing to represent a set of discrete entities with no obvious relation (cities, athletes, applicants, tokens, contexts) by tuples of real numbers. The intuitive way to visualize tuples of real numbers is to think of them either as points or as vectors in a space that has as many dimensions as there are numbers in the tuples.

So if we score a bunch of discrete entities on a number of features like this, we are effectively finding point/vector representations for them and embedding them in a continuous space. For this reason, such tuple representations are usually referred to as Embeddings in machine learning.

Of course, if we have many feature scores, these embeddings will live in a high-dimensional abstract space that is hard to visualize, but if we choose a simpler situation we can plot them directly. For example, we could choose to represent students who took the SAT exam by the two-number tuples representing their SAT scores for the verbal and math sections, respectively. We could then think of these tuples as embeddings in 2D space:

Because the embeddings represent feature scores, more similar entities are bound to have more similar embeddings, meaning they will be represented by more similar points/vectors in embedding space. For example, in the above scenario, students having similar academic strengths (e.g. students that are better at math than at verbal reasoning) are bound to be represented by more similar points/vectors in 2D space, and if we could somehow figure out how to map our contexts and tokens to proper embeddings, more similar contexts/tokens would be represented by more similar points/vectors as well.

In and of itself, this does not immediately change the game – instead of trying to consider the similarity of two entities (e.g. our tokens), we now have to consider the similarity of their embeddings. But given that we ultimately need to be quantitative about similarity in order to have our probabilistic model leverage it, this mapping to points/vectors is a genuine game-changer since it fundamentally enables us to use spatial intuition and powerful mathematical machinery instead of relying on vague qualitative statements.⁠[note] Even though the answer is far from obvious, it simply seems much easier to quantitatively tackle a question like "How similar are the tuples $(0.3, -0.5, 1.2)$ and $(0.6, 0.2, 1.4)$?" than it does to quantitatively tackle a question like "How similar are  running and  spoon?"

So far, so good. But this still leaves us with three major sub-problems:

1. We need to create embeddings for our contexts and tokens.
2. We need to actually figure out how to quantify the similarity of these embeddings.
3. We need to figure out how to adjust our language model to leverage this similarity.

Let's tackle these issues one by one.

Creating embeddings

Contexts are themselves comprised of tokens and seem like somewhat of a higher-level problem, so let's start with the token embeddings.

Token embeddings

The token embeddings represent the meaning of the various tokens (think: words). In the spirit of machine learning, we would ideally like our model to learn this meaning from the training data, much like human infants learn the meaning of words based on context clues. If we could somehow pull that off, it would not only ensure that the embeddings are created automatically but also that the training data remains the only "ground truth" our model's capabilities are based on and its performance doesn't depend on any additional subjective interventions attempting to specify the meaning of invidiual words.

But jumping from not having any embeddings at all to trying to automatically learn them from the training data seems like a huge challenge in its own right and would dramatically derail us from our similarity-focussed journey right now. So let's just add this to our list of things to revisit down the line for now:

For the time being, let's just keep the conceptual ball rolling by using a much simpler solution: Let's do the same thing we did in the small-talk example above, which is to come up with a bunch of (hopefully reasonable) features and manually perform the scoring.

Since we'd like to use the token embeddings to generalize the next-token prediction of our model beyond the training data, the features we come up with to score the tokens should be features that we believe are helpful in achieving that generalization.

For example, if a given context is followed in the training data only by tokens with positive meaning, we'd probably like to generalize preferentially to other tokens that also have positive meaning, so "positivity of meaning" might be a useful feature. Similarly, it might be useful to have a feature that scores how commonly a token appears in the training data, or in the English language in general, if we believe that more common words tend to make for better substitutes for other common words. Lots of binary classification features such as "Can the token represent a noun?" or "Is the token a punctuation mark?", for which the scores would be one of two real numbers (e.g. 0 and 1), would likely also be helpful, as would features scoring where a token lies on the continuum from informal to formal English, and so on.

Of course, none of these individual features will hold much predictive power, and manual scoring like this comes with a fair bit of subjectivity and arbitrariness for some of the continuous features (such as positivity of meaning) in particular, which is part of the reason that we would like to come up with a more sophisticated approach in the long run. But it's nonetheless true that if we come up with enough features, more similar tokens will be represented by more similar scores on average, and so even these suboptimal embeddings should help make our model better than one that generalizes ignorantly.

Since this average effect is all we're after for the time being, the exact choice of features and the scores themselves actually aren't very important to our conceptual invention journey at this point, so let's just assume we have performed this incredibly tedious manual scoring and have thus created embeddings for all the tokens in our vocabulary.

Context embeddings

To address the input generalization problem, we will additionally need a way to create embeddings for entire contexts in order to be able to generalize from known contexts unknown contexts of a similar nature.

In principle, we could do this the same way as for the tokens, i.e. by coming up with a bunch of features we suspect might be useful in identifying whether two contexts should lead to similar next-token predictions and then manually scoring any context we encounter.

But actually implementing this is completely unrealistic due to the same combinatorics that brought about the curse of dimensionality in the first place — it's already ambitious to score all the tokens in our vocabulary, but doing it for all contexts contained in the training data is completely absurd for the amounts of training data we are striving to use. Making matters even worse, even if we had somehow magically pulled that off, we'd have to halt the model whenever it encounters an unprecedented context in a prompt and score the new context from scratch before continuing. In other words, manual context scoring is completely infeasible and we have to create our context embeddings automatically.

Since contexts are just sequences of tokens, an intuitive approach is to use the token embeddings we already have to do so.

The simplest way of doing this is to simply string the individual token embeddings together. For example, if we embed each token in 3D space and want to find an embedding for a 20-token context, we can simply concatenate the embeddings for the individual tokens in that context and represent it with a tuple of $20 \cdot 3 = 60$ numbers. Of course, in practice, these numbers would be much larger since scoring our tokens on more features will yield more useful embeddings.

Simple concatenation like this is not an optimal solution because it will only result in similar embeddings for contexts which are similar on a token-by-token basis. For example, this type of embedding would yield very different representations for different contexts created from the same sequence, such as  that   the   average   November   day   in   Ireland   was   and   average   November   day   in   Ireland   was   cold   and  , even though they have similar meanings and should likely be followed by similar tokens. This is another area for further improvement we should return to down the line.

That being said, this simple token-wise embedding is a solution we can readily implement, and if we have enough training data, taking inspiration from other contexts that are similar to the new one in this specific fashion should already go a long way towards addressing the generalization problem, so let's adopt this approach for now.

Measuring similarity

So far, so good. But to actually address the input and output generalization problems, we will need to define next-token distributions for unprecedented contexts based on those for similar contexts, and spread some probability from the next-token continuations seen during training to other tokens. Simply knowing that embeddings of more similar contexts or tokens are more similar isn't enough to pull this off. We need a Measure of Similarity†We are referring to the colloquial meaning of the term "measure" here, not the proper mathematical one.[note] that quantifies this notion.

Inspired by the observation that we can think of the embeddings either as points or as vectors, two types of options come to mind:

1. Thinking of the embeddings as points, it's intuitive to define the similarity of two tokens or contexts based on the distance $d$ between the corresponding points $E_1$ and $E_2$ in embedding space. The simplest possible option would be to use

   $$\text{sim}({\small \text{thing}_1}, {\small \text{thing}_2}) \;=\; - d$$

   which would ensure that larger distance corresponds to smaller (more negative) similarity. But this yields the value 0 if the embeddings are identical, which doesn't feel very intuitive, and it often turns out to be easier to work with square distances in math to begin with, so a slightly better option might be

   $$\text{sim}({\small \text{thing}_1}, {\small \text{thing}_2}) \;=\; 1 - d^2$$

   which yields a similarity of 1 if the two embeddings are the same and goes towards $-\infty\,$ if the two embedding points are very far away from each other. If we'd like the scale to go from 0 to 1 instead, we could alternatively use something like

   $$\text{sim}({\small \text{thing}_1}, {\small \text{thing}_2}) \;=\; e^{-d^2}$$

   and there are many other options we could come up with. Here is a quick visualization of what this looks like in 1D, 2D and 3D, even though we ultimately want to apply our similarity dimensions to higher-dimensional spaces where visualization is more challenging:

2. Thinking of the embeddings as vectors, we might instead define the similarity of two tokens or contexts based on how aligned of the corresponding embedding vectors†Following standard ML practice, we will denote vectors in bold. So $\vct{e}$ means $\vec{e}$.[note] $\vct{e}_1$ and $\vct{e}_2$ are. An intuitive way to do this is to use the cosine of the angle $\theta$ between them, i.e.

   $$\text{sim}({\small \text{thing}_1}, {\small \text{thing}_2}) \;=\; \cos(\theta)$$

   since this gives a similarity of 1 for two perfectly aligned vectors pointing in the same direction, a similarity of 0 for two vectors that are orthogonal and a similarity of $-1$ for two anti-aligned vectors pointing in opposite directions. This is aptly called Cosine Similarity and is closely related to the dot product of the two vectors, which is given by

   $$\vct{e}_1 \cdot \vct{e}_2 \;=\; \lVert\vct{e}_1\rVert \lVert\vct{e}_2\rVert \cos(\theta)$$

   and seems like yet another way we could define our similarity measure:

   $$\text{sim}({\small \text{thing}_1}, {\small \text{thing}_2}) \;=\; \vct{e}_1 \cdot \vct{e}_2$$

   This is called Dot-Product Similarity and can yield any similarity value ranging from large positive numbers (gigantic vectors pointing in the same direction) to large negative numbers (gigantic vectors pointing in opposite directions). Once again, here is a quick visualization for the 1D, 2D and 3D cases:

Since we are much more used to thinking about distances than about angles, using a distance-based similarity measure might seem much more intuitive, and indeed, the distance-based definitions yield a property that we would instinctively deem desirable: The measure of similarity is maximized only if the two embeddings are identical. But we might well argue that the distance-based approach has its shortcomings as well — any distance-based measure would conclude, for example, that the numbers 100 and 102 are just as similar as the numbers $+1$ and $-1$, whereas we might intuitively consider 100 and 102 roughly the same and $+1$ and $-1$ literal opposites.

What type of measure seems more intuitive boils down to whether we instinctively think of our feature scores in absolute terms (in which case thinking of them as points and using distance-based measures seems more appropriate) or if we think of them as relative trends with respect to some neutral middle (in which case thinking of them as vectors and using angle-based measures seems more appropriate). Of course, that same thinking is already going to impact the creation of our embeddings itself, so the choice of similarity measure might be less important than ensuring consistency between how we score our tokens and how we judge their similarity.

Another argument why the choice might not be critical is that the difference between the various definitions becomes much less significant if all the embedding vectors are roughly of the same size. In this case, the dot-product similarity is just a scaled version of the cos-similarity⁠[show me] and $1\!-\!d^2$ equals a scaled version of the dot product plus some constant.⁠[show me] If we come up with enough features, it seems reasonable to expect that most embedding vectors will end up being of similar length (in relative terms), as the varying scores over all the many features will tend to average out, thereby diminishing the importance of our choice of metric.

Still, we have to choose one. Given that we do not have an obvious winner, let's simply go with the one that is simplest to calculate. This seems like a good idea not just because it feels right philosophically but because we ultimately want to train our model at massive scale, meaning compute will be incredibly precious and when in doubt, we should always go for the compute-saving strategy.†Indeed, one reason that most of ML can be grasped without knowing much advanced math is that simpler approaches are often preferrable to more sophisticated ones simply because they require less compute.[note]

Of the measures we have, the dot product is the simplest to compute both conceptually and in terms of required operations,⁠[details] so let's use that.

This is also the go-to standard in ML and will turn out to have more advantages down the road.

Similarity-based generalization

With our embeddings and similarity measure in place, we can finally attempt to tackle the generalization problems.

Let's start with the output problem, which we said we'd like to address by taking some of the probability assigned to any given token (for a given context) and spreading it to similar tokens across the vocabulary.

Spreading probability to other tokens

Assuming the given context appears only once in the training data, which we said will be the case for large enough context windows, our n-gram model will learn to pick the token $t_{\text{data}}$ that the context is succeeded by in the training data with probability 1.

We can then use our token embeddings to calculate how similar any other token $t_i$ is to that training-data token. Using dot-product similarity yields

$$\text{sim}(t_i, t_{\text{data}}) \;=\; \vct{t}_i \cdot \vct{t}_{\text{data}}$$

where $\vct{t}_i$ and $\vct{t}_\text{data}$ represent the embeddings of the tokens $t_i$ and $t_\text{data}$.

Doing this for all tokens in the vocabulary yields a large number of token-similarity scores. The remaining challenge is to figure out how to translate these scores, which could theoretically range from $-\infty$ to $\infty$, into next-token probabilities that range from 0 to 1 and sum to 1.

The intuitive way to ensure the summing-to-1 part is to simply calculate some non-negative pseudo-probability values that are free to range from 0 to $\infty$ first, and to then normalize them by dividing by their sum:

$$p_\text{gen}(t_i \,|\, c) \;=\; \frac{p_{\text{pseudo}}(t_i \,|\, c)}{\sum_k^{V} p_{\text{pseudo}}(t_k \,|\, c)}$$

where $p_\text{gen}(t_i \,|\, c)$ is the post-spread, generalized next-token probability, $c$ is the given context and $V$ is the size of the vocabulary.

Using this approach, all we have left to do is to pick a function that maps the token-similarity scores, which can range from $-\infty$ to $\infty$, to pseudo-probabilities ranging from 0 to $\infty$.

It's clear that this function should map large negative scores to vanishing pseudo-probabilities, since we don't want our model to spread any probability to extremely dissimilar tokens. But it is less obvious how the function should behave for large positive inputs — a function whose output approaches some constant value for large enough inputs, for example, would result in roughly equal probabilities being assigned to sufficiently similar tokens. A function whose output keeps increasing for larger and larger inputs, on the other hand, would always translate larger token-similarity scores into larger probabilities, with the nature of the increase determining how strongly the model prioritizes more similar tokens.

So which type of behavior do we want? We can't know for sure, but it seems reasonable to err on the side of caution. The token embeddings themselves are ignorant to the specific context at hand and the similarity scores represent an average over all features, so even if two tokens are fairly similar by our dot-product measure, substituting one with the other could result in non-sensical output rather quickly depending on the context at hand and the features that caused the token $t_i$ to be a reasonable choice. Put simply, we want to spread our probability, but we don't want to spread it too far. And even if we restrict ourselves to spreading it to the few most synonymous tokens only, that will already enable the model to generate sequences not contained in the training data.

So, all things considered, we need a function that

1. maps $(-\infty, \infty)$ to $(0, \infty)$
2. approaches $0$ for large negative inputs
3. increases rather dramatically for larger and larger inputs

Of all the simple "standard" functions we learn in school and constantly have to deal with when doing math or physics, $f(x) = e^x$ is the one that matches these criteria:

So let's go with that one and define our pseudo-probabilities as the exponentials of the token-similarity scores:

$$p_{\text{pseudo}}(t_i \,|\, c) \;=\; e^{\text{sim}(t_i ,\, t_{\text{data}})} \;=\; e^{\vct{t}_i\, \cdot \, \vct{t}_{\text{data}}}$$

If we plug this into the normalization formula, we get

$$p_\text{gen}(t_i \,|\, c) \;=\; \frac{e^{\text{sim}(t_i,\, t_{\text{data}})}}{\sum_k^{V} e^{\text{sim}(t_k,\, t_{\text{data}})}} \;=\; \frac{e^{\vct{t}_i\, \cdot \, \vct{t}_{\text{data}}}}{\sum_k^{V} e^{\vct{t}_k\, \cdot \, \vct{t}_{\text{data}}}}$$

In machine learning, things are usually framed in terms of vectors and matrices whenever possible in order to compress notation, so this sort of calculation would be thought of as calculating an output vector, whose elements in this case are the probabilities, from an input vector, in this case the similarity-score vector $\vct{s}$, whose elements are the token-similarity scores $s_i = \vct{t}_{\text{i}} \cdot \vct{t}_\text{data}$.

The specific vector-to-vector mapping represented by the prescription above is referred to as the Softmax function.⁠[misnomer] Formally, the $i$-th element of the softmax output for some input vector $\vct{z} = (z_1,\,\dots, z_N)$ is defined as

$$\text{softmax}(\vct{z})_i \;\coloneqq\; \frac{e^{z_i}}{\sum_k^{N} e^{z_k}}$$

The softmax pops up a lot in the field, usually when some set of scores has to be converted into a set of probabilities.

So following ML lingo, we can alternatively frame our calculation of the next-token probabilities as simply applying the softmax function to the vector of similarity scores:

$$p_\text{gen}(t_i \,|\, c) \;=\; \frac{e^{\vct{t}_i\, \cdot \, \vct{t}_{\text{data}}}}{\sum_k^{V} e^{\vct{t}_k\, \cdot \, \vct{t}_{\text{data}}}} \;=\; \frac{e^{s_i}}{\sum_k^{V} e^{s_k}} \;=\; \text{softmax}(\vct{s})_i$$

If we test-run this softmax approach using a handful of three-dimensional dummy embeddings, it looks like this:

Interestingly, this probability distribution is quite flat. Even though we chose an exponential specifically to ensure that the model ends up very picky and strongly prioritizes high similarity, the actual next token from the training data only ends up with a slightly higher probability than many others in this case.

One reason this is happening is that we only used three dimensions in this dummy example, meaning the embedding vectors are comprised of only three elements and it is relatively likely that at least some of our random dummy embeddings are well aligned by sheer coincidence, yielding closely spaced dot-product similarity values. If we were to use radomly generated dummy vectors in a much higher dimension, this problem should become less severe. In fact, one of the many cool facts about high-dimensional spaces is that two high-dimensional random vectors tend to be orthogonal, meaning if their length is fixed, their dot product tends to be zero.⁠[intuition]

But there is another factor playing into the closely spaced similarity values and the resulting flat-ish distribution as well, namely that the magnitudes of the dummy embeddings in the above example was chosen to be roughly 1. If we use larger-magnitude vectors, the resulting probability distribution remains much more sharply peaked even if the embedding vectors point in the same exact directions:

Investigating the softmax formula we use to calculate the probabilities, it's easy to see why: The ratio of the probability assigned to two tokens simply turns out to be $e^{\Delta s}$ where $\Delta s$ is the difference between their similarity scores with respect to $t_\text{data}$:

$$\frac{p_\text{gen}(t_i \,|\, c)}{p_\text{gen}(t_j \,|\, c)} \;=\; \frac{e^{s_i}}{e^{s_j}} \;=\; e^{s_i - s_j}$$

In other words, how strongly we are spreading probability across the vocabulary depends on the absolute spacing of the similarity scores, which will grow or shrink with their overall scale, which is in turn tied to the embedding vector magnitudes via the dot product formula.

In principle, we could try to figure out a target scale for our similarity scores that works well and attempt to create our embeddings with that in mind, but that feels like a clumsy and inflexible approach. If we are already thinking about using scale to affect the spacing of the similarity values, a much more natural framing is to simply introduce a scaling parameter to our similarity formula itself. This allows us to scale the similarity scores as we see fit and leaves us with a simple lever we can play with to encourage the model to be more or less aggressive in spreading probability across the vocabulary.

In machine learning, this scaling parameter is usually placed in the denominator and is called the Temperature due to analogies with physics we won't get into here:

$$\vct{t}_i \cdot \vct{t}_\text{data} \quad \rightarrow \quad \frac{\vct{t}_i \cdot \vct{t}_\text{data}}{T}$$

If the temperature value is very small, the spacing between the similarity scores gets larger and the model hardly spreads probability to other tokens at all. If the temperature is very high, the spacing between the similarity scores is compressed and probability is spread across the vocabulary more liberally.† Another way of thinking about this is that we are now using an exponential with basis $b=e^{1/T}$ and tuning the temperature simply allows us to change our basis.[note]

Let's play around with this a little bit to get a better feel for it:

This seems to work pretty well, so if the context $c$ is only followed by a single token $t_{\text{data}}$ in the training data, we can calculate the probability with which our model should choose some token $t_i$ using this recipe.

But what if the original probability distribution contains multiple tokens with non-zero probability, either because a context appears multiple times and just so happens to be followed by several different tokens in a huge dataset, or perhaps because whatever way we come up with to construct distributions for unknown contexts might result in such a situation?

In that case, we can simply apply this same approach many times: Instead of only calculating how much of the probability assigned to one specific token $t_{\text{data}}$ should be reallocated to $t_i$, we can calculate how much of the probability assigned to any token in the vocabulary should be spread to $t_i$ (using the same similarity+softmax strategy as before) and weigh the results by the amount of probability assigned to that token to begin with. To get the total probability the model should assign to $t_i$, we then simply sum over all these contributions.

Of course, this is likely not a perfect solution, but it fundamentally enables our chatbot to generate original output even for long context windows, and if the embeddings are good enough, it should go a long way towards solving the output generalization problem.

Unfortunately, if we test-run this approach, we find that creating original output immediately results in unknown contexts, meaning we run into the input generalization problem straight away:

Processing unknown contexts

So let's move on to tackling that input problem, which we said we'd like to do by assigning to each unknown context the model encounters a next-token probability distribution reminiscent of those associated with similar contexts seen during training.

Since our similarity measure relies on embeddings, the first thing this will require us to do when encountering an unknown context $c_{\text{new}}$ is to look up the embeddings for all the individual tokens and concatenate them (that is, string them together) into an embedding of $c_{\text{new}}$, as described earlier. Once we have that embedding, we can interpret it as a vector $\vct{c}_{\text{new}}$ and calculate the dot products of this vector with the individual embedding vectors $\vct{c}_j$ of all the training-data contexts, for each of which our n-gram model has learned a next-token distribution:

$$\text{sim}(c_{\text{new}}, c_j) = \vct{c}_{\text{new}} \cdot \vct{c}_j$$

One subtlety to note here is that for this dot product to be defined, both context embeddings need to be of equal length, which is not necessarily a given since the length of the context our model takes in will vary depending on prompt and output lengths, and given that we create the context embedding $\vct{c}_{\text{new}}$ by concatenating the individual token embeddings, its length will vary as well.

But at least in theory, guaranteeing equal lengths of $\vct{c}_{\text{new}}$ and $\vct{c}_j$ is straightforward. If we create the training-data context embeddings $\vct{c}_j$ at runtime by combing through the training data, we can simply choose to use however many tokens $\vct{c}_{\text{new}}$ contains and directly build $\vct{c}_j$'s of the appropriate length. If we instead wish to store all the full-length training-data context embeddings once, we can simply truncate them to match the length of $\vct{c}_{\text{new}}$ or zero-pad $\vct{c}_{\text{new}}$ at earlier-token positions to match the full length of the $\vct{c}_j$'s, which will effectively yield the same result since the corresponding vector elements won't contribute to the dot product.

So let's just assume we have chosen one of these options and keep it moving for now.

Carrying out all these dot products yields a large number of context-similarity scores, just like our earlier calculation yielded a large number of token-similarity scores. The remaining challenge is to figure out how exactly we should use these context-similarity scores to define the next-token distribution for an unprecedented context.

The simplest approach would be to pick the single most similar context from the training data and to simply copy its next-token probability distribution, but this would leave a lot of information left unused. For example, if there were many other contexts that were almost as similar to $c_{\text{new}}$ as the top match, this approach would not take their next-token distributions into account at all. So let's avoid throwing away information like this and instead define the new next-token distribution as a weighted average of all the known ones:

$$p_\text{gen}(t_i \,|\, c_{\text{new}}) \;=\; \sum_{j}^{N_c} \; w_j \; p_{\text{data}}(t_i \,|\, c_{j})$$

Here, $N_c$ is the number of training-data contexts $c_j$, and the $w_i$'s are the weights. To ensure that the values $p_\text{gen}(t_i \,|\, c_{\text{new}})$ form a proper probability distribution, the weights $w_j$ should fulfill $0 \leq w_j \leq 1$ and sum to 1, i.e. $\sum_j w_j = 1$. And to ensure that this distribution represents a reasonable generalization from the training data, the known distributions corresponding to contexts that are more similar to the new context (i.e. those with higher context-similarity scores) should receive larger weights.

In other words, we have to translate all the similarity scores, which could in principle range from $-\infty$ to $\infty$ into a bunch of non-negative numbers that sum to 1.

But wait a minute, we already know how to do this! We simply use a softmax again:

Here $T_c$ is the temperature for this softmax calculation, which works the same way as the temperature on the output side: For $T\!\approx\!0$, our distribution for the new context will simply match that of the most similar context in the training data. For $T\!\to\!\infty$, the new distribution will be the equal-weight average of all the known distributions. For intermediate temperatures, the new distribution will be a context-similarity-weighted mix of the known ones.

If the embeddings are good enough and the new context is sufficiently similar to at least one context from the training-data, this new distribution should represent a reasonable-ish generalization.

If we visualize this process using some 3D dummy embeddings, as we did earlier for the output problem, it looks like this:

If we combine this context-similiarity-based method of creating distributions for unknown contexts with our existing method of spreading probability across tokens, we get a chatbot that can process unknown contexts and doesn't run the risk of regurgitating the training data word for word.

There are some minor design decisions to be made on how to combine the two approaches here, but we will keep changing our architecture anyway, so let's not overthink it. If we include this context-similarity-based generalization on the input side and simply always apply a bit of token-similarity-based generalization on the output side (regardless of whether we are dealing with a newly constructed distribution or with a known context from the training data), the output of our simple War-and-Peace 20-gram bot looks something like this:

Unsurprisingly, this is still a mess. After all, we are still dealing with a primitive 20-gram chatbot trained on only a single book, so there is no way we are going to get great responses here.

But we shouldn't let that take away from the huge wins we have achieved – our model is now able to process unknown contexts and can generate original output which deviates from the training data even for known contexts. This means we now have a model architecture which, at least in principle, allows us to use large context windows and is capable of at least some generalization beyond the training data.

Given that large context windows and proper generalization are strict necessities if we hope to build a useful model, this is a big leap forward.

Next Chapter