Module 1 · Foundations
Chapter 5 of 17 Run it in your browser

Chapter 5: Embeddings¶

Name: Teach Your Computer to Write: Build and Train an LLM from Zero
Author: Truong (Jack) Luu

Where we are in the big picture Figure 5.1: We have numbers, now we give them meaning.

In the previous chapter, we turned text into tokens (Figure 5.1), but tokens are arbitrary numbers. To predict what comes next, the model needs to understand how characters relate to each other. In this chapter you will:

Map simple token IDs to large vectors of numbers.
Learn how these vectors capture meaning as coordinates.
Add position information so the model knows the order of characters.

Words to Know

Vector: A list of numbers that works like coordinates on a map, giving a token its place among the others.
Embedding: A lookup table that gives each token ID its own vector.

Theory¶

The Problem with IDs¶

IDs have no math meaning Figure 5.2: Token IDs are arbitrary labels, not math quantities.

After tokenization, each character is an integer: A=13, B=14, a=39. But feeding these raw integers to a neural network creates a problem (Figure 5.2). The network learns by multiplying numbers together, so it would see Z (token 38) as almost three times as big as A (token 13). It might learn that Z is more important than A only because of this accident. But token IDs are arbitrary labels.

Think of it like an employee ID number, where Employee 105 is not "five times more employee" than Employee 21. The number alone carries no meaning about their role. We need a way to turn these arbitrary labels into a format the "Numbers Machine" can use to find patterns.

Vectors as Coordinates of Meaning¶

Looking up a token ID to find its embedding vector Figure 5.3: An embedding table maps each ID to a vector of numbers.

The solution is a lookup table called an embedding (Figure 5.3). It gives each token ID a vector, a list of decimal numbers that programmers call floating-point numbers. In our model, we use 128 numbers for each character.

Instead of seeing 39 for a, the model sees [0.55, 0.03, -0.89, ...]. These 128 numbers act like coordinates, and characters that appear in similar contexts will gradually be moved closer together in this 128-dimensional space during training. The network learns that capital and lowercase versions of a letter behave similarly, so their vectors become similar. This table of profiles is learned from scratch.

Position Embeddings¶

Adding token and position embeddings together Figure 5.4: The final representation combines what the token is with where it sits.

One more problem is that the model reads all characters at the same time. Without help, it cannot tell the difference between "cat" and "act" because they use the same letters. Order matters.

We fix this with positional embeddings (Figure 5.4), a second lookup table. Here you look up a row by where the character sits in the sequence (0, 1, 2, ...), not by its token ID. Position 0 gets its own 128-number vector, position 1 gets a different vector, and so on.

We then add the token embedding and the positional embedding together. The network learns to use these combined 128 numbers to represent both "what character is this?" and "where does it appear?". With the character and its position in one vector, the "Embeddings" stage in our map is complete (Figure 5.1). The text is now converted into math vectors, ready for the next stage.

In Business

When you build an assistant that writes in your company's house style, embeddings turn your archive into a map of meaning. Token embeddings capture the vocabulary, while the position embeddings give the model the order of the characters, which is what lets it pick up the shape of an invoice or a polite greeting.

Code¶

We use src/ch04_embeddings.py to create and combine these two tables.

src/ch04_embeddings.py (excerpt)
    token_emb = nn.Embedding(config.vocab_size, config.n_embd)
    token_embeddings = token_emb(token_ids)
    pos_emb = nn.Embedding(config.block_size, config.n_embd)
    positions = torch.arange(T)
    position_embeddings = pos_emb(positions)
    x = token_embeddings + position_embeddings

When we run it, the script prints this:

Terminal

$ python src/ch04_embeddings.py
--- 2. Looking up embeddings ---
Input token_ids shape : torch.Size([2, 5])
Token embeddings shape: torch.Size([2, 5, 128])
  (B=2, T=5, C=128)
...
--- 3. Positional embeddings ---
Position indices: [0, 1, 2, 3, 4]
Position embeddings shape: torch.Size([5, 128])

--- 4. Combining token + position embeddings ---
Final x shape: torch.Size([2, 5, 128])
  (B=2, T=5, C=128)

What just happened:

Line 1 created a lookup table for the 65 characters in our vocabulary.
Line 2 translated a batch of token IDs into their embedding vectors.
Line 3 created a second lookup table for the sequence positions.
Line 6 added the token and position vectors to create the final input x.

Shape Check: Table 5.1 shows the dimensions of our data structures.

Table 5.1: Tensor shapes before and after the embedding layer.

Variable	Shape	Meaning
`token_ids`	`[B, T]`	Batch size by Time (sequence length).
`token_embeddings`	`[B, T, C]`	`C` is the embedding dimension (128).
`position_embeddings`	`[T, C]`	One vector per position.
`x`	`[B, T, C]`	The combined input for the model.

Try It¶

Try It

In src/ch04_embeddings.py, add one more token ID to each row of token_ids (any number from 0 to 64) and change T from 5 to 6. Run the script again. The output shape grows from [2, 5, 128] to [2, 6, 128], but the parameter counts stay exactly the same. The lookup tables do not care how much text you process at once.

Watch Out

A common mistake is trying to look up a token ID that is not in the vocabulary. If your vocab_size is 65, the valid IDs are 0 to 64. Passing an ID of 65 will crash the program with an "index out of bounds" error.

Key Takeaways¶

Token IDs are arbitrary labels and cannot be used directly for math.
Embeddings are learned lookup tables that map token IDs to vectors.
These vectors act as coordinates that capture relationships and meaning.
Positional embeddings are added to tell the model the order of the characters.

Check Your Understanding¶

Why do neural networks struggle with raw token IDs?
How many numbers make up a single character's embedding vector in our model?
Why do we need positional embeddings in addition to token embeddings?
Does the size of the embedding table depend on the batch size?

Assignment¶

Assignment 5: Bigger ID, bigger numbers?

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

Run it. Run src/ch04_embeddings.py as it is. Copy the line that starts with First token's embedding. It shows the first 8 of the 128 numbers for the first token.
Change one number. Find the line token_ids = torch.tensor([[3, 14, 7, 2, 50], [10, 22, 45, 1, 8]]). These are ten token IDs. Change the first one, 3, to 38, so the line starts with token_ids = torch.tensor([[38, 14,. Run the script again.
Compare. Did the 8 numbers change? Are they bigger, now that the ID is bigger? Did any shape, or any count under --- Summary ---, change?
Explain it. In two or three sentences a manager could follow, say what the token ID is used for here, and whether the size of the ID matters.

Hand in: both First token's embedding lines, your answers to step 3, and your sentences.