Cosine similarity and large language models

A companion to the language identification lab

Computing the angle between feature vectors is a ubiquitous technique in machine learning. One important example shows up in large language models.

Background

First a little of the relevant background on how current large language models work. The basic unit of language is a token. Most of the time, tokens are words, but more complex words are usually broken up. For instance, in GPT-2,

cosine = cos + ine

The basic task of a model is to take a token, or a sequence of tokens, and generate the next one. To do this, each token must first be turned into something we can compute with, a vector. Then some computations must happen, at the end of which another vector is produced, and this vector has to be turned into a token.

token"cat" embedding layertokens becomevectors many layers ofcomputation(magic, for now) unembedding layervectors turn backinto tokens token"kitten" Two vectors are compared with the dot product at the unembedding step, exactly as in the digraph lab.

Embedding vectors

The embedding layer is crucial. It associates a vector with each token. We can graph this association, drawing just the point at the tip of each vector. Each point in the image below represents one of the 50,257 tokens in GPT-2, the OG large language model. A few points are labeled.

catCatkittentheThecheeseburgercosineSchwarzeneggerSolidGoldMagikarpFriday2019Japanpurplecalculusing
Figure 1. The embedding layer of GPT-2, one point per token, with eighteen points labeled.

Note: In the actual GPT-2 model, each vector is 768-dimensional. Our image is two-dimensional. How this simplification is accomplished, that is, which two dimensions are picked to display, is an extremely important question. A method is usually seen in a second course in linear algebra. Even better methods use fancy geometry.

Directions

Since embedding vectors are vectors, their direction should mean something. And it does. It turns out that vectors in similar directions have similar meanings. See Figure 2 below. And, relating this to our lab, similarity is measured by computing the angle using the dot product!

catdogkittenpuppyhorsecowliontigerMondayTuesdayWednesdayThursdayFridaySaturdaySundayFranceGermanyItalySpainJapanChinaBrazilredbluegreenyellowpurpleorangepinktheofandtoina
Figure 2. The same picture with the vectors drawn in for 37 tokens in five categories: animals, colors, days of the week, countries, and common function words. Vectors in the same category point in similar directions.

Magnitudes

The meaning of the magnitudes of the embedding vectors tends to vary from model to model. In many, long vectors correspond to infrequent tokens and short ones are more common. So if you wanted to use an embedding as a fancy thesaurus, you could try to replace a pedestrian word with something more fancy by looking for tokens that point in roughly the same direction as your word, but of a much longer magnitude.

Other ways to visualize embeddings

The above images were made by using principal component analysis to turn a high-dimensional vector into one that we can visualize in two dimensions. There are other, more modern methods. The image below is made using UMAP, which distorts angles. To compare embeddings of tokens, we have to use both magnitude and direction, so cosine similarity is no longer applicable.

animals days countries colors numbers function
Figure 3. A UMAP embedding of 4,000 GPT-2 tokens. Similar words form tight “clusters.”

The UMAP embedding uses position instead of direction. The next two figures color points by part of speech. The UMAP embedding separates these beautifully; the original embedding from the first part of the demo, shown in Figure 5, does not.

noun proper noun verb adjective adverb function word not in tagged corpus
Figure 4. The map of Figure 3 with each token colored by its most common part of speech in a tagged corpus of English. Tokens that do not appear in that corpus are shown in light gray.

For comparison, here is the PCA picture of Figure 1 with the same coloring. Since what matters is the angle, the meaning of relative positions is more muddled.

noun proper noun verb adjective adverb function word not in tagged corpus
Figure 5. The tokens of Figure 1, colored by part of speech as in Figure 4.