Lesson 2.3: Embeddings, vectors and token positions
Give every token a learned vector, see similar words land close together, and add position information so a model knows word order, with code and numbers.
By the end of this lesson you will be able to explain what an embedding is and why a model needs one, show that an embedding lookup is a matrix product with a one-hot vector, factor a bigram model through a small embedding and train it, measure how close two token vectors are with cosine similarity, explain why a model that only adds up token vectors cannot tell word order, and build both a learned and a sinusoidal table of positions. You will show that a bigram model with 3 numbers per token matches the full model's loss (1.162), that after training the vectors for "cat" and "dog" have a cosine similarity of 0.979 while "cat" and "chased" have 0.068, that the sentences "the dog chased the cat" and "the cat chased the dog" have identical average vectors, and that the dot product of two position vectors depends only on the distance between them. You should have read Lessons 1.4 and 2.1 to 2.2. The mathematics is dot products and a sine. This course is education about how language models work, not a guide to any commercial product.
After Lesson 2.2, a text is a list of integer ids. But an integer is a poor thing to compute with. The id of "cat" is 4 and the id of "dog" is 11 because of alphabetical order, and the number 7 is not halfway between 6 and 8 in any useful sense. The model needs a representation in which similar tokens look similar and the arithmetic of the network can use that. That representation is a vector of numbers for each token, called an embedding.
The embedding table
An embedding table is a matrix E with one row for each token in the vocabulary and d columns, where d is the width of the model. To embed a token with id i, take row i:
python
E = rng.standardnormal((V, d)) 0.1 # (34, 8) in our toy model
vec = E[ids] # ids has shape (N,) -> vec has shape (N, d)
This looks like a lookup, and it is one. It is also a matrix multiplication, which matters because it shows that the lookup is something the gradient can pass through. Write the id as a one-hot vector: all zeros with a single one at position i. Then
onehot(i) @ E = E[i]
and we can check it:
python
onehot = np.eye(V)[x[:5]] # (5, 34)
print(np.allclose(onehot @ E, E[x[:5]])) # True
An embedding is therefore a linear layer applied to a one-hot input, which is why the gradient formulas of Lesson 1.3 apply unchanged: the gradient on E is onehotᵀ @ dvec, which in code is a sum of the gradients of each use of a row, np.add.at(dE, ids, dvec). The numbers in E are parameters like any other. They start random and are changed by gradient descent so that the loss falls. Nobody decides what the columns mean.
The table has V × d parameters. For GPT-2 small that is 50,257 × 768 = 38,597,376, which is 31 per cent of the whole model (Lesson 1.1).
A bigram model through a bottleneck
To see embeddings learned, we change the model of Lesson 1.4 in one way. There, the score row for each previous token was a free row of 34 numbers. Now we insist that the model first turns the token into a vector of only d numbers, and then produces the 34 scores from that vector by a second matrix:
logits = E[x] @ U, where E has shape (V, d) and U has shape (d, V)
The number of parameters falls from V × V = 1,156 to 2 × V × d, which is 544 for d = 8. The vector E[x] is a bottleneck: everything the model knows about the previous token must pass through those d numbers. We train it exactly as before, with the gradients dU = E[x]ᵀ @ dlogits and dE[x] = dlogits @ Uᵀ:
python
def trainembedbigram(d, steps=3000, lr=10.0, seed=0):
rng = np.random.defaultrng(seed)
E = rng.standardnormal((V, d)) 0.1
U = rng.standardnormal((d, V)) 0.1
N = len(ytrain)
for step in range(steps + 1):
H = E[xtrain] # (N, d): the embeddings of the inputs
P = softmax(H @ U) # (N, V)
dlogits = P.copy()
dlogits[np.arange(N), ytrain] -= 1
dlogits /= N
dU = H.T @ dlogits # (d, V)
dH = dlogits @ U.T # (N, d)
dE = np.zeroslike(E)
np.add.at(dE, xtrain, dH) # add into the rows that were used
E -= lr dE
U -= lr dU
return E, U
Here are the results for several widths, each trained for 3,000 steps from the same seed and scored on the validation set:
Width d · Parameters · Training loss · Validation loss1 · 68 · 2.0316 · 2.0230
2 · 136 · 1.1664 · 1.1652
3 · 204 · 1.1627 · 1.1619
4 · 272 · 1.1618 · 1.1620
6 · 408 · 1.1614 · 1.1621
8 · 544 · 1.1615 · 1.1626
16 · 1,088 · 1.1612 · 1.1613
34 · 2,312 · 1.1612 · 1.1615
Full table (Lesson 1.4) · 1,156 · 1.161 · 1.162
The rows for widths 1 and 2 needed a different setting. They would not train at the learning rate of 10: the loss overflowed to nan, which is the divergence of Lesson 1.3, because with so few numbers per token the gradient on each is large. With a learning rate of 3 and 4,000 steps they trained, and a single number per token is clearly too few (validation loss 2.023), while two numbers are almost enough (1.165).
The result is surprising at first. A bottleneck of only three numbers per token loses nothing: the validation loss is 1.162 whether the model has three numbers per token or thirty-four. Why? Because the language we are modelling has only a handful of kinds of word. All eight subjects ("the cat", "the dog", "a fox" ...) are followed by the same verbs with the same probabilities. All six verbs are followed by the same objects. The model has to remember "this is a subject-like word", not which subject it was. A few numbers per token are enough to say which group a token belongs to and how it behaves inside the group.
That is the central idea of embeddings, and it is why they generalise. The table of counts had a separate row for each token and could share nothing between them. The embedding model has to compress, and compression forces it to find what the tokens have in common. The width d sets how much it can keep. A real language is far richer than our toy, so real models need hundreds or thousands of numbers per token, but the principle is the same.
What the vectors learned
We trained the width-8 model and look at the vectors. A natural measure of how alike two vectors are in direction is the cosine similarity: the dot product divided by the product of their lengths. It is 1 for vectors pointing the same way, 0 for vectors at right angles and −1 for opposite vectors.
python
def cosine(a, b):
return float(a @ b / (np.linalg.norm(a) np.linalg.norm(b)))
vec = lambda w: E[stoi[w]]
print(cosine(vec("cat"), vec("dog"))) # 0.979
print(cosine(vec("cat"), vec("chased"))) # 0.068
Pair · Cosine similarity · Relationshipcat – dog · 0.979 · Both subjects
cat – fox · 0.985 · Both subjects
cat – baker · 0.977 · Both subjects
chased – watched · 0.986 · Both verbs
chased – found · 0.995 · Both verbs
quickly – slowly · 0.999 · Both adverbs
ball – fish · 0.961 · Both objects
a – the · −0.164 · Different roles
red – small · 0.139 · Different adjectives
cat – chased · 0.068 · Subject and verb
quickly – cat · −0.071 · Adverb and subject
cat – ball · −0.277 · Subject and object
The nearest neighbours of a word confirm the pattern: the four closest words to "cat" are farmer, teacher, bird and fox (0.98 to 0.99). The nearest to "chased" are liked, followed, saw and found (all 1.00). The nearest to "fish" are bus, tree, ball and drum (0.96 to 0.98), and to "quickly" they are slowly, again, today and night (0.99 to 1.00).
Two details are instructive. First, "a" and "the" are not similar (−0.164), although they are both articles. In this language they are used differently: "a" precedes only bird, fox, fish, small and loud, while "the" precedes many more. A model learns what the data says, and the data says that they behave differently. Second, "red" and "small" are barely related (0.139) because "red" is always followed by "bus", and "small" by "boat".
This is the distributional hypothesis, expressed by the linguist John Firth as "you shall know a word by the company it keeps". The embedding for a token ends up describing the company: which words tend to appear before and after it. A model never needs a dictionary. It needs enough text for the company to be visible.
What the sum of vectors cannot do
Embeddings give us a vector per token, but a sentence is a sequence, and the order of the words matters. Consider the simplest possible way to combine a sentence's vectors: add them, or equivalently average them. Take two sentences made from the same words:
python
a = [stoi[w] for w in "the dog chased the cat".split()]
b = [stoi[w] for w in "the cat chased the dog".split()]
print(np.allclose(E[a].mean(axis=0), E[b].mean(axis=0))) # True
The averages are identical, and always will be, whatever E contains: addition does not care about order. The same is true of attention, the mechanism we will build in Module 3. As we will see, attention treats its input as a set of tokens: it computes a weighted average of their values, and if you shuffle the inputs, the outputs are shuffled in the same way. Without something extra, the model would be unable to distinguish who chased whom.
The extra something is information about position. We give each position in the window its own vector and add it to the token vector:
x = E[ids] + P[0 : T]
where P is a table of position vectors of the same width as the token vectors, and T is the length of the window. A token at position 0 now looks different from the same token at position 5, and the model can learn to use the difference. There are two common ways to build P.
Learned positions
The simplest choice is to treat P like E: a table with one row per position, up to the longest window the model will ever see, learned by gradient descent. GPT-2 does this: its table has 1,024 positions of width 768, which is 786,432 parameters, less than 1 per cent of the model.
The disadvantages are plain. The table has no row for position 1,025, so the model cannot be run on longer windows than it was trained on. And each row is learned separately, so position 500 and 501 have no built-in link: the model has to discover that they are neighbours.
Sinusoidal positions
The original transformer paper proposed a fixed table built from sines and cosines, with no parameters to learn. For position pos and an even index 2i, with model width d:
P[pos, 2i] = sin(pos / 10000^(2i/d))
P[pos, 2i+1] = cos(pos / 10000^(2i/d))
Each pair of columns is a sine and a cosine of the same angle, and the angle grows with the position at a rate that depends on the pair. The first pair has the fastest rate (a full cycle every 6.3 positions); later pairs are slower. In code:
python
def sinusoidal(T, d):
pos = np.arange(T)[:, None] # (T, 1)
i = np.arange(0, d, 2)[None, :] # (1, d/2)
angle = pos / (10000 (i / d)) # (T, d/2)
P = np.zeros((T, d))
P[:, 0::2] = np.sin(angle)
P[:, 1::2] = np.cos(angle)
return P
P = sinusoidal(64, 16)
print(np.round(P[0, :6], 3)) # [0. 1. 0. 1. 0. 1.]
print(np.round(P[1, :6], 3)) # [0.841 0.54 0.311 0.95 0.1 0.995]
print(np.round(P[5, :6], 3)) # [-0.959 0.284 1. -0.01 0.479 0.878]
Position 0 is 0, 1, 0, 1, ..., since sin 0 = 0 and cos 0 = 1. Position 1 starts with sin 1 = 0.841 and cos 1 = 0.540, then moves more slowly in the following pairs, 0.311 and 0.950 (an angle of 0.316), and 0.100 and 0.995 (an angle of 0.1). The columns are a set of clocks that tick at different speeds: the first pair is like the second hand, and the last like the hour hand. Together they identify the position, much as the hands of a clock identify the time.
At the widths used in practice, the clocks are spread over a huge range. For d = 768, the wavelength of the fastest pair is 6.3 positions, the hundred-and-first pair has a wavelength of 69, the two-hundred-and-first one of 761, and the slowest pair has a wavelength of 61,343 positions.
The property that makes them useful
The reason for choosing sines and cosines is a property that the model can exploit. For any fixed offset k, the vector at position pos + k is obtained from the vector at pos by a fixed rotation in each (sine, cosine) pair, whatever pos is. We can verify it numerically: for offset k = 5 from position 7 and width 16, rotating each pair at position 7 by its own angle reproduces the vector at position 12.
python
d, k, pos = 16, 5, 7
ok = True
for j in range(d // 2):
w = 1 / (10000 (2 j / d))
R = np.array([[np.cos(w k), np.sin(w k)],
[-np.sin(w k), np.cos(w k)]])
a = P[pos, 2 j : 2 j + 2]
b = P[pos + k, 2 j : 2 j + 2]
ok = ok and np.allclose(R @ a, b)
print(ok) # True
A consequence is that the dot product of two position vectors depends only on how far apart they are, never on where they are. We compute it for width 16:
Offset between positions · Dot product, starting at position 10 · Dot product, starting at position 300 · 8.000 · 8.000
1 · 7.485 · 7.485
2 · 6.368 · 6.368
4 · 5.560 · 5.560
8 · 4.697 · 4.697
16 · 4.214 · 4.214
32 · 3.541 · 3.541
The columns for the two starting positions are identical. Nearby positions are alike (8.0 at distance 0 against 7.5 at distance 1) and distant ones are less alike (3.5 at distance 32). The distance between the vectors of positions 0 and 1 is 1.015, and between positions 0 and 32 it is 2.986. A mechanism that compares two tokens by taking a dot product, as attention will, can therefore tell how far apart they are with no extra machinery. And since a sinusoid exists for every position, the table can be extended to windows longer than any seen in training, although how well a trained model uses the longer positions is a separate question.
Neither choice has won outright. Learned positions are simple and are used by GPT-2. Sinusoidal tables were used by the original transformer. Most recent models use rotary position embeddings, which apply the rotation we verified above to the queries and keys inside attention, so that the relative distance appears in the attention scores directly. We use learned positions in this course because the code is shorter, and mention rotary embeddings again in Module 8.
Putting it together
The input to the transformer we build in Module 4 is, for a window of T token ids:
python
tok = E[ids] # (T, d) what each token is
pos = P[:T] # (T, d) where it is
x = tok + pos # (T, d) both, added
Two vectors with the same job are added, not joined side by side, so the width stays at d. It may look as if adding mixes the two sources of information beyond recovery. In practice the model learns to keep them in different directions, and a vector of hundreds of numbers has room for both.
Key takeaways
An embedding table E of shape (V, d) gives every token a vector of d numbers. Looking up row i equals multiplying the one-hot vector for i by E, so its gradient is a sum into the rows that were used.
The table is learned like any other parameter. Nobody assigns meaning to the columns.
Passing the token through a bottleneck of d numbers forces the model to find what tokens have in common. In the toy language, three numbers per token reached the full model's loss of 1.162, while one number per token left it at 2.02.
Words in the same role end up with nearly the same vector: cosine similarity 0.979 for cat and dog, 0.995 for chased and found, against 0.068 for cat and chased and −0.277 for cat and ball.
Adding or averaging token vectors ignores order: "the dog chased the cat" and "the cat chased the dog" have the same average.
Position information is added to the token vectors: x = E[ids] + P[:T]. The table can be learned (GPT-2: 786,432 parameters) or fixed (sinusoidal).
In the sinusoidal table, the dot product of two positions depends only on their distance: 8.0 at distance 0, 7.5 at 1, 3.5 at 32, for width 16, wherever the pair starts. Position pos + k is a fixed rotation of position pos.
Common mistakes
Treating ids as numbers. Never feed an id to a layer as if it were a quantity. Look up its vector.
Reading meaning into single columns. An embedding's directions are mixed. Compare whole vectors with cosine similarity.
Forgetting positions. A model that has only token vectors and attention cannot see word order, and will fail on any task where order matters.
Joining instead of adding. Concatenating token and position vectors doubles the width. Add them.
Sizing the position table by hope. A learned table is fixed at its training length. Choose it for the longest window you will ever use.
Comparing vectors from different models. Each model has its own coordinates. Vectors can only be compared within one model.
Check yourself
1. A vocabulary has 32,000 tokens and the width is 512. How many parameters are in the token table?
32,000 × 512 = 16,384,000.
2. Two vectors are (3, 4) and (4, 3). What is their cosine similarity?
Dot product = 12 + 12 = 24. Both lengths are 5. Cosine = 24 / 25 = 0.96.
3. Why does adding the average of the token vectors lose the word order?
Addition is commutative: the sum is the same whatever order the terms are in.
4. A learned position table covers 512 positions. What happens if you give the model 600 tokens?
There is no row for positions 512 to 599, so the lookup fails or must be handled in some other way. A learned table cannot extend beyond its training length.
5. Why are sines and cosines used in pairs, rather than a sine alone?
With a sine and cosine of the same angle, a shift in position is a rotation of the pair, which is a fixed linear operation. A sine alone repeats and is ambiguous.
Try this
Set aside two hours. You need Python and NumPy, and the code from Lessons 1.4 and 2.3.
1. Train the factored bigram model for d = 2, 3 and 8. Check that your validation losses are within 0.001 of the table.
2. Print the cosine similarity of every pair of the eight subjects. What are the smallest and largest values?
3. Find the nearest five neighbours of "bus", "red" and "." in your model. Explain each list from the generator.
4. Reduce the training steps to 100. How do the cosine similarities of cat and dog change? What does that say about when embeddings form?
5. Build the sinusoidal table for d = 16 and 64 positions. Reproduce the dot-product table for offsets 0, 1, 2, 4, 8, 16 and 32.
6. Plot the 16 columns of the sinusoidal table against position. Which column is the fastest? Which is the slowest?
A good result is a model with a 3-number bottleneck that matches the full table, a list of similarities that makes sense from the generator, and a position table whose dot products depend on distance only. The next module introduces attention, the mechanism that lets each token use the others.
Read more
Tomas Mikolov and colleagues' 2013 paper "Efficient Estimation of Word Representations in Vector Space" introduced word2vec and the idea that vector arithmetic captures meaning. Yoshua Bengio and colleagues' 2003 paper "A Neural Probabilistic Language Model" used learned embeddings in a language model. Vaswani and colleagues' 2017 paper "Attention Is All You Need" defines the sinusoidal encoding. Jianlin Su and colleagues' 2021 paper "RoFormer" introduces rotary position embeddings. John Firth's remark on "the company it keeps" is from his 1957 Papers in Linguistics.
Tags: Build an LLM Embeddings Positional encoding Cosine similarity Representation