Token Embeddings

Starting from where we left off in the tokenisation deep dive, after tokenisation, we have turned our prompt into a list of token IDs.

The dog fetched the ball

The791 ␣dog5679 ␣fetched42542 ␣the279 ␣ball5041

[791, 5679, 42542, 279, 5041]

Each token has become a token id. From this point on, the model works with numbers rather than the prompt text.

Neural networks, and in particular LLMs, operate on vectors, which can be represented as points in multidimensional space. Therefore we need to transform our token IDs into vectors.

There are many reasons why vectors are the natural objects for LLMs to work with, including:

vectors are lists of numbers, which makes them convenient for the types of mathematical operations that neural networks need to perform
vector spaces with a large number of dimensions have lots of “room” in which to encode complex features of language, such as semantic meaning and grammatical structure
vectors have geometry. The distance between two learned vectors can express relationships, such as whether their tokens often appear in similar contexts

Real LLMs use vectors with many thousands of dimensions, but we’ll use 6-dimensional vectors to make them easier to show on the page.

The embedding matrix

In order to map each token ID to a vector, we define a matrix of size n by 6, where n is the vocabulary size:

token ID	dimension 1	dimension 2	dimension 3	dimension 4	dimension 5	dimension 6
0	0.12	-0.31	0.08	0.44	-0.17	0.25
1	0.34	0.06	-0.22	0.11	0.29	-0.40
2	-0.09	0.27	0.38	-0.15	0.04	0.19
⋮	⋮	⋮	⋮	⋮	⋮	⋮
n − 3	0.21	-0.07	0.13	-0.36	0.18	0.02
n − 2	-0.28	0.16	0.05	0.31	-0.12	0.24
n − 1	0.07	0.33	-0.19	0.14	0.22	-0.26

Each horizontal row stores one 6-dimensional vector. The vertical dots stand in for all the other rows in the vocabulary.

This acts as a lookup table; the first row is the vector that the token with ID 0 will be transformed to, the second is the vector that token ID 1 will be transformed to, and so on:

token ID1

0.340.06-0.220.110.29-0.40

Looking up token ID 1 selects the whole second row, turning one number into a 6-dimensional vector.

How do we decide the actual numbers that go into the embedding matrix? In other words, how do we decide which vector each token gets transformed to? The numbers in the matrix are called weights (there will be many more weights in the full LLM!), and these are optimised to produce meaningful output from the LLM during model training. Before training, we usually initialise the embedding matrix with random weights.

Now we have a way to turn token IDs into vectors.

import torch

vocabulary_size = 100_000
embedding_size = 6

token_embedding_matrix = torch.nn.Embedding(
    vocabulary_size,
    embedding_size,
)

token_ids = torch.tensor([791, 5679, 42542, 279, 5041])
token_vectors = token_embedding_matrix(token_ids)

# one 6-dimensional vector for each of the 5 token IDs
# token_vectors.shape == torch.Size([5, 6])

Positional embeddings

The same token can appear in different positions within text.

position	0	1	2	3	4
token	the	dog	fetched	the	ball
token ID	42	8	19	42	31

The repeated token has the same token ID at positions 0 and 3, so the token embedding matrix gives it the same vector both times.

Sometimes these different positions can be meaningful within the context of a sentence, but the above embedding scheme can’t capture this meaning; it maps the same token to the same vector, no matter the position.

For this reason, LLMs employ a second embedding matrix: the positional embedding matrix.
This maps each possible token position to a vector. This matrix has size m by 6, where m is the maximum context length (i.e., the maximum number of tokens that will be allowed as input) of the LLM.

position	dimension 1	dimension 2	dimension 3	dimension 4	dimension 5	dimension 6
0	0.17	-0.25	-0.21	-0.02	0.20	-0.19
1	-0.17	-0.43	-0.59	-0.06	-0.33	0.09
2	0.18	-0.63	-0.48	0.02	-0.31	0.06
⋮	⋮	⋮	⋮	⋮	⋮	⋮
m − 2	0.09	0.21	-0.14	0.32	-0.07	0.12
m − 1	-0.11	0.05	0.28	-0.16	0.24	-0.03

Each possible position has its own 6-dimensional vector. The number of rows is the maximum context length, m.

We now compute the vector embedding of a token ID as the sum of these two embeddings.

maximum_context_length = 2_048
position_embedding_matrix = torch.nn.Embedding(
    maximum_context_length,
    embedding_size,
)

positions = torch.arange(len(token_ids))
position_vectors = position_embedding_matrix(positions)

vector_embeddings = token_vectors + position_vectors
# vector_embeddings.shape == torch.Size([5, 6])

This allows an LLM to encode information about both the token itself and the token’s position within the given input text.