2 — Tokens become vectors

Chapter 1 ended with The Time Machine as 45,259 token IDs, and with the problem that an ID is only a label. ID 3862 is " Time" and ID 3863 is " maybe". Doing arithmetic on those numbers directly would treat " maybe" as almost the same as " Time", and " time" (ID 640) as very different.

The fix is to give every token a list of numbers of its own: a vector. The model starts with random vectors and adjusts them during training, so tokens used in similar ways drift toward similar vectors. This chapter builds that lookup, adds a second lookup for position, and then cuts the book into the training examples the model will learn from.

All runs here were on an Apple M3 Pro with 18 GB of memory, with Python 3.12, PyTorch 2.15.0 and tiktoken 0.14.0, on the CPU. Seeds are given with each run.

An embedding table is a lookup

An embedding table is a matrix with one row per token in the vocabulary. The row for token i is that token’s vector. Turning a list of IDs into vectors means picking out those rows.

A toy version has a vocabulary of 6 tokens and vectors of 4 numbers each:

import torch

torch.manual_seed(123)
tok_emb = torch.nn.Embedding(num_embeddings=6, embedding_dim=4)
print(tok_emb.weight)

Parameter containing:
tensor([[ 0.34, -0.18, -0.30, -0.59],
        [ 0.35,  0.66, -0.22, -0.38],
        [-0.16, -0.40,  0.70, -1.81],
        [ 1.90, -0.18,  1.37, -1.60],
        [-0.78, -1.41, -0.41,  0.80],
        [ 1.00,  0.22,  1.83, -0.34]], requires_grad=True)

A few PyTorch basics appear here. A tensor is PyTorch’s array: a grid of numbers with a shape, like a NumPy array, that can also live on a GPU and keep track of how it was computed. torch.manual_seed(123) fixes the random number generator so the same “random” values come out every run. torch.nn.Embedding(6, 4) creates a 6-by-4 table and fills it with random numbers drawn from a normal distribution. Its weight is a parameter, a tensor the training process is allowed to change, which is what requires_grad=True means. (Values are printed to two decimal places throughout.)

Looking up three tokens, IDs 2, 0 and 4:

ids = torch.tensor([2, 0, 4])
print(tok_emb(ids))

tensor([[-0.16, -0.40,  0.70, -1.81],
        [ 0.34, -0.18, -0.30, -0.59],
        [-0.78, -1.41, -0.41,  0.80]], grad_fn=<EmbeddingBackward0>)

torch.tensor builds a tensor from a Python list. Calling the embedding on it returns rows 2, 0 and 4 of the table, in that order. Three IDs went in; a 3-by-4 tensor came out. The grad_fn note means PyTorch recorded how this result was made, so it can later work out how to adjust the table.

The same thing as a matrix multiply

A lookup sounds like a different kind of operation from the matrix arithmetic in the rest of a network, but it isn’t. Write each ID as a one-hot vector: as long as the vocabulary, all zeros except a single 1 at the ID’s position. Multiplying that by the table picks out exactly one row, because every other row is multiplied by zero.

one_hot = torch.nn.functional.one_hot(ids, num_classes=6).float()
print(one_hot)
print(one_hot @ tok_emb.weight)
print(torch.equal(one_hot @ tok_emb.weight, tok_emb(ids)))

tensor([[0., 0., 1., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0.]])
tensor([[-0.16, -0.40,  0.70, -1.81],
        [ 0.34, -0.18, -0.30, -0.59],
        [-0.78, -1.41, -0.41,  0.80]], grad_fn=<MmBackward0>)
True

one_hot builds the 0/1 rows, .float() converts them from integers to floating point so they can be multiplied with the table, and @ is matrix multiplication. A (3 × 6) matrix times a (6 × 4) matrix gives a (3 × 4) matrix, and torch.equal confirms it matches the lookup exactly.

By hand, for the first row: the one-hot row is [0, 0, 1, 0, 0, 0], so the first column of the product is

which is the first number of row 2. So an embedding layer is a linear layer with no bias, fed one-hot inputs. nn.Embedding exists because building a 50,257-wide one-hot vector per token, and multiplying mostly by zero, would be pure waste.

The equivalence also explains how the table learns. When a result is used in a computation and PyTorch works backward to find how each parameter affected it (.backward()), only the rows that were looked up receive any adjustment:

tok_emb(ids).sum().backward()
print(tok_emb.weight.grad)

tensor([[1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [1., 1., 1., 1.],
        [0., 0., 0., 0.]])

.grad holds the gradient: for each number in the table, how much the result would change if that number were nudged up. Here the “result” is just the sum of the three looked-up vectors, so each used entry has gradient 1 and every unused row has 0. In real training the result is the model’s error, and each step changes only the rows of the tokens that appeared in the batch. Rare tokens get updated rarely.

What trained embeddings look like

Random rows mean nothing. After training, they do. GPT-2’s published weights (the 124-million-parameter model, from the openai-community/gpt2 repository on Hugging Face, which chapter 6 loads in full) include its trained 50,257 × 768 token table. Comparing rows by cosine similarity, which measures whether two vectors point the same way (1 for the same direction, 0 for unrelated, −1 for opposite):

the nearest rows to a few tokens are:

Nobody told GPT-2 that " cat", " Cat" and “cat” are the same word with different spacing and case, which chapter 1 left as four unrelated IDs. Training put them close together because they appear in similar places. That is the whole point of the table: the IDs stay labels, but the vectors behind them carry meaning.

Order gets lost

Embedding each token on its own has a blind spot. Take a three-token vocabulary, dog = 0, bites = 1, man = 2, and embed two sentences with the same table:

dog_bites_man = tok_emb(torch.tensor([0, 1, 2]))
man_bites_dog = tok_emb(torch.tensor([2, 1, 0]))

dog bites man:                         man bites dog:
[[ 0.34, -0.18, -0.30, -0.59],         [[-0.16, -0.40,  0.70, -1.81],
 [ 0.35,  0.66, -0.22, -0.38],          [ 0.35,  0.66, -0.22, -0.38],
 [-0.16, -0.40,  0.70, -1.81]]          [ 0.34, -0.18, -0.30, -0.59]]

The two results are the same three vectors with the first and last rows swapped. The vector for “dog” is identical whether the dog did the biting or got bitten. Nothing about any single vector records where it sat.

That would be fine if the next layer read the rows strictly in order, the way a recurrent network does. The layer the next chapter builds, attention, doesn’t. It compares every vector with every other vector and combines them with weighted sums, and a sum doesn’t care about order. Chapter 3 shows this directly: feed attention the rows in reverse, with no position information, and it returns exactly the same outputs in reverse. To that layer, “dog bites man” and “man bites dog” are the same bag of three vectors.

Adding position

The fix GPT-2 uses is a second embedding table, indexed by position instead of token ID. Row 0 is a vector meaning “first position”, row 1 “second position”, and so on. Each token’s position vector is added to its token vector, element by element:

pos_emb = torch.nn.Embedding(num_embeddings=3, embedding_dim=4)
positions = torch.arange(3)
print(positions)
print(pos_emb(positions))

tensor([0, 1, 2])
tensor([[-0.03,  1.42, -0.39, -0.89],
        [-0.96, -0.41,  1.08, -0.40],
        [-0.73, -0.12, -0.48, -0.52]], grad_fn=<EmbeddingBackward0>)

torch.arange(3) is the tensor [0, 1, 2], the positions of a three-token input. (This table was created right after the token table, so the seed-123 generator simply continued.) Adding:

a = dog_bites_man + pos_emb(positions)
b = man_bites_dog + pos_emb(positions)

dog bites man + positions:             man bites dog + positions:
[[ 0.31,  1.24, -0.69, -1.48],         [[-0.19,  1.02,  0.31, -2.70],
 [-0.61,  0.25,  0.86, -0.78],          [-0.61,  0.25,  0.86, -0.78],
 [-0.89, -0.52,  0.22, -2.32]]          [-0.39, -0.30, -0.78, -1.10]]

By hand, for “dog” at position 0, first number: 0.34 + (−0.03) = 0.31. For “dog” at position 2: 0.34 + (−0.73) = −0.39. The same token now gives a different vector depending on where it is. “Bites” sits at position 1 in both sentences, so its row is the same in both, as it should be. The two sentences are no longer the same set of rows, so a layer that ignores order can still tell them apart.

Adding, rather than gluing the position vector onto the end, keeps the width the same (4 here, 768 in GPT-2). It sounds as if adding would scramble the two kinds of information together. In practice the model has hundreds of dimensions to work with, and training arranges the two tables so that both kinds of information can still be read back out.

Learned positions, as in GPT-2

There are two common ways to fill the position table. The original transformer paper, “Attention Is All You Need”, used fixed sine and cosine waves of different frequencies, computed by formula. It also tried learned position vectors, reported that the two “produced nearly identical results”, and chose the fixed version because it might extend to sequences longer than those seen in training.

GPT-2 learns them. In OpenAI’s model code there are two tables, wte (token embeddings, 50,257 × 768) and wpe (position embeddings, 1,024 × 768), both initialized with small random numbers, and the input to the first layer is the token rows plus the position rows, exactly as above. Training shapes the position table like any other weight.

A learned table has a hard limit: it has rows for positions 0 to 1,023 and nothing beyond. Asking for position 1,024 is an error:

pos_emb = torch.nn.Embedding(1024, 768)
pos_emb(torch.tensor([1024]))
# IndexError: index out of range in self

That is why GPT-2’s context is exactly 1,024 tokens. The trained table also has visible structure. In GPT-2’s wpe, neighboring positions are nearly identical while distant ones are not:

Training discovered on its own that position 100 and position 101 should look alike, a smooth pattern that the fixed sine-wave version builds in by design.

Real sizes

At GPT-2’s size the two tables are:

tok_emb = torch.nn.Embedding(50257, 768)
pos_emb = torch.nn.Embedding(1024, 768)
print(tok_emb.weight.numel(), pos_emb.weight.numel())

38597376 786432

numel() counts the numbers in a tensor. The token table alone has 38,597,376 parameters, nearly a third of GPT-2’s roughly 124 million, and takes 154 MB as 32-bit floats. The position table is small next to it: 786,432 parameters.

Training pairs from a long text

Before any of these vectors can be trained, the model needs examples to learn from. A GPT is trained to predict the next token, and the nice thing about that task is that any text supplies its own answers. No labels are needed.

Take the first sentence of the story:

text = "The Time Traveller (for so it will be convenient to speak of him)"
ids = encode(text)
print(ids, len(ids))
for t in range(1, 6):
    print(f"{decode(ids[:t])!r:35} -> {decode([ids[t]])!r}")

[464, 3862, 43662, 6051, 357, 1640, 523, 340, 481, 307, 11282, 284, 2740, 286, 683, 8] 16
'The'                               -> ' Time'
'The Time'                          -> ' Trave'
'The Time Trave'                    -> 'ller'
'The Time Traveller'                -> ' ('
'The Time Traveller ('              -> 'for'

(encode and decode are the two tiktoken helpers from chapter 1.) Every prefix of the sentence is a question, and the token after it is the answer. Sixteen tokens give fifteen such questions.

A model doesn’t need fifteen separate examples for that. Shift the sequence by one: the inputs are tokens 0 to 14, and the targets are tokens 1 to 15.

x: [464, 3862, 43662, 6051,  357, 1640, ...,  286,  683]
y: [3862, 43662, 6051,  357, 1640,  523, ...,  683,    8]

Position t of y is the token that follows position t of x. Because the model (as the next chapter arranges) only lets position t see positions 0 to t, one pass over x produces fifteen predictions at once, one per position, each checked against the matching entry of y. One sequence of length T gives T training examples.

A sliding window

A real text is far longer than the model’s context, so it is cut into windows of a fixed length. A window of context_length tokens slides along the text, moving stride tokens each time. Each window is an input, and the same window shifted one token right is its target.

PyTorch’s standard way to package this is a Dataset, a class that says how many examples there are (__len__) and how to fetch example i (__getitem__):

from torch.utils.data import Dataset

class TokenWindows(Dataset):
    def __init__(self, ids, context_length, stride):
        self.ids = torch.tensor(ids, dtype=torch.long)
        self.context_length = context_length
        # Every start that leaves room for context_length inputs plus one more target token.
        self.starts = list(range(0, len(ids) - context_length, stride))

    def __len__(self):
        return len(self.starts)

    def __getitem__(self, i):
        s = self.starts[i]
        x = self.ids[s : s + self.context_length]
        y = self.ids[s + 1 : s + self.context_length + 1]
        return x, y

dtype=torch.long stores the IDs as 64-bit integers, the type an embedding lookup expects. The whole text is turned into one tensor once, and each example is just two slices of it, so nothing is copied up front.

On the 16-token sentence, with windows of 4 tokens and a stride of 4:

ds = TokenWindows(ids, context_length=4, stride=4)
print(len(ds))
for i in range(len(ds)):
    x, y = ds[i]
    print(x.tolist(), y.tolist())

3
[464, 3862, 43662, 6051] [3862, 43662, 6051, 357]
[357, 1640, 523, 340] [1640, 523, 340, 481]
[481, 307, 11282, 284] [307, 11282, 284, 2740]

The windows start at tokens 0, 4 and 8. A fourth window would start at 12 and need a target running to token 16, which doesn’t exist, so it’s dropped. The last few tokens of a text can be lost this way; across a whole book that’s negligible.

The stride controls overlap:

Overlapping windows give more examples from the same text, but they are mostly repeats, which makes it easier for a model to memorize a small text. Setting the stride equal to the context length uses each token once per pass, and that’s the default used here.

Batches with a DataLoader

Training works on many windows at a time, a batch, because the hardware does a stack of identical computations far more efficiently than one at a time. PyTorch’s DataLoader takes a Dataset and hands out batches, stacking the individual x and y tensors along a new first dimension:

from torch.utils.data import DataLoader

def make_loader(text, batch_size=8, context_length=256, stride=256,
                shuffle=True, drop_last=True, num_workers=0, seed=None):
    ds = TokenWindows(encode(text), context_length, stride)
    gen = torch.Generator().manual_seed(seed) if seed is not None else None
    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
                      drop_last=drop_last, num_workers=num_workers, generator=gen)

The arguments: shuffle=True visits the windows in a random order on each pass, so the model doesn’t see the book front to back every time. drop_last=True discards a final batch that would be smaller than the rest, so every batch has the same shape. num_workers=0 loads data in the main process, which is plenty for a pre-tokenized book. generator gives the shuffle its own seeded random source, so the order is repeatable.

On the whole book, with windows of 256 tokens:

loader = make_loader(book, batch_size=8, context_length=256, stride=256, shuffle=True, seed=123)
print(len(loader.dataset), len(loader))
x, y = next(iter(loader))
print(x.shape, y.shape, x.dtype)
print(x[0, :8]); print(y[0, :8])
print(repr(decode(x[0, :8].tolist())))

176 22
torch.Size([8, 256]) torch.Size([8, 256]) torch.int64
tensor([6071,   13,  679, 3332,  198, 1891,  287,  465])
tensor([  13,  679, 3332,  198, 1891,  287,  465, 5118])
' forth. He sat\nback in his'

The book’s 45,259 tokens make 176 windows (45,259 ÷ 256 is 176.8, and the partial window is dropped), which fill 22 batches of 8. next(iter(loader)) takes the first batch. Both x and y have shape [8, 256]: 8 windows, 256 token IDs each. Shuffling picked a window from the middle of the story, and y is visibly x shifted by one.

From a batch of IDs to a batch of vectors

The two tables turn that batch into the tensor the rest of the model works on:

torch.manual_seed(123)
tok_emb = torch.nn.Embedding(50257, 768)
pos_emb = torch.nn.Embedding(1024, 768)
h = tok_emb(x) + pos_emb(torch.arange(x.shape[1]))
print(tok_emb.weight.shape, pos_emb.weight.shape, h.shape)

torch.Size([50257, 768]) torch.Size([1024, 768]) torch.Size([8, 256, 768])

tok_emb(x) looks up every ID in the [8, 256] batch and returns [8, 256, 768]: one 768-number vector per token. pos_emb(torch.arange(256)) is [256, 768], one vector per position. The two shapes differ, and PyTorch’s broadcasting handles it: when one tensor is missing a leading dimension, it is reused along that dimension, so the same 256 position vectors are added to each of the 8 windows.

That three-dimensional shape, batch × tokens × width, written (B, T, C), is the shape every layer from here on takes in and gives back.

Where this leaves things

Text is now numbers in the form a network can use. Each token is a 768-number vector that training can move, so related tokens can end up near each other, as GPT-2’s " cat", " Cat" and " cats" did. A position vector added to each one means the same word at a different place is a different input. And the book is now 22 batches of 8 windows, with targets that are the inputs shifted by one.

But look at what each vector actually contains: the token’s own row and its position’s row. Nothing else. The vector for " bank" is the same whether the sentence is about rivers or money. The vector for " it" in “The Time Traveller (for so it will be convenient to speak of him)” has no idea what “it” refers to. Every vector was computed alone, and each one knows nothing about its neighbors.

To predict the next token, a position needs to gather information from the tokens before it, and to decide which of them matter. That mechanism is attention, and chapter 3 builds it slowly, starting from three words and a 3 × 3 table.