Embeddings

Key idea

Embedding matrix E (V × d); token i → row i of E, which equals one-hot(i) × E; rows are learned, similar tokens end up with similar rows; often the same E scores the output (tied)

Real world

A library catalogue card. The book’s number on its own tells you nothing, but its card lists subject, period, language and length. Books on similar topics have similar cards, even if their numbers are far apart.

In an LLM

In GPT-2, token 3797 (“ cat”) is looked up as row 3797 of a 50,257 × 768 table. Its 768 numbers are closer to the rows for “ dog” and “ kitten” than to those for “ France” or “ seven”, because training put them there.

From ID to vector

After tokenization, the text is a list of ID numbers. The numbers themselves carry no meaning: “ cat” is 3797 and “ dog” is 3290 in GPT-2, but that doesn’t make them similar, and 3797 isn’t “more” than 3290. The model needs, for each token, a list of numbers that it can do arithmetic on: a vector, called the token’s Embedding.

The simplest way to get one is a table with one row per token in the vocabulary: the embedding matrix $E$, with $V$ rows (the vocabulary size) and $d$ columns (the model width). Token $i$ gets row $i$. That’s all the embedding layer does.

It is often written as a matrix multiplication. Write ID $i$ as a One-hot vector vector: all zeros except a 1 in position $i$. Multiplying it by $E$ adds up the rows of $E$ weighted by the one-hot vector, and every weight is 0 except one, so the result is row $i$:

Step through it with the toy model used on the attention pages:

In code nobody actually multiplies by a one-hot vector, which would waste $V$ multiplications per number. nn.Embedding just indexes the row. But thinking of it as a multiplication explains how the table is trained: the gradient that reaches $E$ only touches the rows of the tokens that were in the batch, which is why rare tokens get so little training.

The real table

GPT-2 small has $V = 50{,}257$ and $d = 768$: 38.6 million numbers, almost a third of the whole 124M-parameter model. Here is a corner of it:

At the start of training the table is random noise. Gradient descent then moves each row so that the rest of the network can predict the next token better. Nobody says what the 768 columns should mean, and individual columns rarely mean anything you could name. Meaning shows up in the directions between rows: which rows are close to which.

A map of meaning

To compare two embeddings we use Cosine similarity: the cosine of the angle between them, 1 for the same direction, 0 for unrelated, −1 for opposite. Here are 48 words from six groups, with their real GPT-2 vectors:

The groups mostly hold together. “ cat” is closest to the other animals, “ three” to the other numbers (cosine 0.89 with “ four”), “ France” to the other countries. GPT-2 was never told these categories. Words that appear in similar contexts lead to similar predictions, so gradient descent gives them similar rows: “the ___ sat on the mat” is about the same whether the blank is cat or dog.

Two cautions about reading the picture:

• The 2D map is a squashed picture of 768 dimensions. PCA shows the two directions in which these 48 vectors vary most and loses the other 766; t-SNE keeps each word’s near neighbours together but distorts everything else, so tidy clusters partly come from the method. The cosine list, computed on all 768 numbers, is the real measure.
• All GPT-2 embeddings point roughly the same way. Two random tokens have a cosine of about 0.27 on average, not 0. Compare similarities with each other.

King − man + woman?

A famous demonstration from word2vec (2013) is vector arithmetic: take the vector for “king”, subtract “man”, add “woman”, and the nearest word is “queen”. Does it work in GPT-2’s input embeddings? Partly, and the details matter:

What the real numbers show:

• With the usual trick it often works. Excluding the three input words, the nearest token to king − man + woman is “ queen” (cosine 0.71), and Paris − France + Germany gives “ Berlin”.
• Without the trick, an input word usually wins. The nearest token to king − man + woman is “ king” itself (0.78). Many published demos quietly exclude the inputs.
• The answer is often already a neighbour. “ queen” has cosine 0.66 with “ king” before any arithmetic: it is king’s third-nearest token. Subtracting man and adding woman nudges the result in the right direction, but the answer was close already.
• Some fail. doctor − man + woman gives “ physician”, not “ nurse”: here the arithmetic stays in the doctor neighbourhood.

So the embedding table does hold consistent directions for some relations (country to capital, singular to plural), but it is not a clean parallelogram world. Most of a model’s knowledge is built up later, in the layers above. Embeddings are only the starting point.

The same table at the output

At the top of the network, the model must turn its final vector into a score for every token: the unembedding. That is another $V \times d$ matrix, and many models simply reuse $E$ for it. The score for token $j$ is then the dot product of the final vector with row $j$ of $E$: “how much does this vector point towards token $j$’s embedding?”. This is called tying the embeddings. GPT-2 does it, as do Gemma and small Llama models such as Llama 3.2 1B; it saves $V \times d$ parameters, which matters most in small models with big vocabularies. Larger models such as Llama 3 8B and DeepSeek-V3 keep separate input and output tables.

Position comes next

The embedding of “ cat” is the same wherever the word appears: first, last, or twice in the same sentence. Something must tell the model where each token is. That is the subject of the next page.

import torch, torch.nn as nn
E = nn.Embedding(50257, 768)        # weight: (50257, 768), random at first
ids = torch.tensor([[464, 3797, 3332]])  # "The cat sat"
x = E(ids)                          # (1, 3, 768): rows 464, 3797, 3332
# identical, but wasteful:
onehot = nn.functional.one_hot(ids, 50257).float()
x2 = onehot @ E.weight              # (1, 3, 768)