Vol. III · Learning Machinery Plate 1 of 10embedding, embedding table, vector, cosine similarity, nearest neighbours, representationAL

Embeddings.

A token id is just a row number. The row it selects — a vector of learned numbers — is where the model’s knowledge of the token begins.

A model cannot do arithmetic on token ids: id 4 512 is not “bigger” than id 300 in any useful sense. So the first thing it does is look each id up in an embedding table, a matrix with one row per vocabulary entry, and replace the token by that row, a vector of d numbers. In the course model micro, d is 128; in large models it is several thousand. The rows start random and are trained along with everything else.

Nothing tells the model what the vectors should mean. But tokens that appear in similar contexts receive similar gradients, so their vectors drift together. After training, the angle between two rows — measured by cosine similarity, one for the same direction and zero for unrelated ones — says how alike the model treats the two tokens. Colours end up near colours and organs near organs, without anyone having said what a colour is.

Below: left, 29 words from micro’s embedding table projected onto the two directions along which they vary most. Right, the eight tokens in the whole 8 192-token vocabulary whose rows are closest to the word typed.

What you will learn

Describe the embedding table and what an embedding lookup does.
Measure the similarity of two tokens by the cosine of their vectors.
Find a token’s nearest neighbours in embedding space.
Explain why training places related tokens close together.

Workshop

Left: micro’s embeddings of 29 single-token words, projected onto their first two principal components and coloured by kind. Right: the eight nearest tokens to the typed word, by cosine similarity over the whole vocabulary.

Meaning as position

In micro’s table, red and blue have a cosine of 0.79, three and four 0.92, he and she 0.55, but red and heart only 0.46. Search the whole vocabulary for the rows nearest heart and the answers are other organs and parts of the body; nearest three, other numbers; nearest 细胞 (cell), other words of biology in Chinese. The projection on the left squeezes 128 dimensions into two, so it distorts, yet the kinds of word still separate. Nobody labelled anything: the clusters are what next-token prediction on the course corpus made useful.

Two cautions. First, an embedding row is only the model’s starting point for a token: the layers above rewrite it according to context, so the vector that finally represents bank in a sentence about rivers differs from the one in a sentence about money (Vol. V). Second, tied models such as micro use the same table to score outputs (plate V.7), so a row is shaped both by how a token is read and by when it is predicted. In large models the table has tens of thousands of rows of thousands of numbers — Llama 3 8B’s is 128 256 × 4 096 — and works the same way.

Word vectors became famous with word2vec (Mikolov and colleagues, 2013), whose vectors supported analogies such as king − man + woman ≈ queen; in a language model the embedding table is simply the first layer, trained by the same loss as everything else.

Terms

Embedding table: A V × d matrix with one learned row per token.
Embedding lookup: Replacing a token id by its row of the table; equivalent to multiplying a one-hot vector by the matrix.
Cosine similarity: The cosine of the angle between two vectors: 1 for the same direction, 0 for perpendicular.
Nearest neighbours: The tokens whose vectors are most similar to a given token’s.

An embedding is a learned row per token. Training places tokens used alike near one another, so similarity of vectors comes to stand for similarity of use.