Creating Token Embeddings

Training a large language model on raw text starts with a simple question: given some words, which word comes next? Before any model can learn to answer that question, the text has to be turned into training examples. Each example pairs a stretch of tokens the model can see with the token that follows it. This article builds those pairs from a single token sequence, and then shows how the pairs are turned into vectors that a model can work with.

Keep that answer in mind. The rest of this article explains how the code produces exactly that pair, and why the same mechanism produces every other pair hidden in the sequence.

The input-target pairs for next-word prediction are built from two variables. The variable x holds the input tokens, which is the text the model sees. The variable y holds the targets, which are the same tokens shifted by one position. Each target token is the word that appears immediately after its matching input token in the original sequence.

The diagram uses a short sequence of five token IDs. The input row holds the first four, and the target row holds the same four positions moved one slot later. Read down any column and you find an input token together with the token that follows it. The last input position is paired with the fifth token of the sequence, which is why the target row ends with 50 even though 50 does not appear in the input row.

The two lines that matter are two slices of the token sequence. The first slice takes the opening block of the sequence. The second slice starts one position later and ends one position past the end of the first slice. The variable context_size controls how long each block is.

A single pair of slices is only one training example. To see every next-word prediction task hidden in one sequence, the sequence is read one step at a time. At each step the input grows by one token, and the target is the token that comes right after the end of that growing input. Looping over the growing prefixes of a sequence reveals every such task, one context and one target at a time.

The four printed pairs are exactly the pairs that the single slice expression produces when it takes the whole window at once. The loop shows them one at a time, while the slices show the whole window together. Once the pairs are visible as token IDs, they can be decoded back into words with the tokenizer to check that they make sense. Each context and each target can be passed through decode in the same loop.

Read the diagram as a count. Four steps produce four context-to-target pairs from one five-token sequence. In general, a sequence of n tokens contains n minus 1 such pairs when every prefix is used, because the final token can only ever be a target and never has a token after it.

Generating pairs one by one works for a demonstration, but training needs a data loader. Its job is to iterate over the input dataset and return inputs and targets as PyTorch tensors. A tensor can be thought of as a multidimensional array. The loader returns two tensors. The input tensor holds the text the LLM sees, and the target tensor holds the targets the LLM must predict.

The code operates on token IDs, not on words. The sample text is shown as words only to make it readable. The encode step of the BPE tokenizer performs tokenization and conversion to IDs in a single step. So the rows of the input tensor are rows of token IDs, and the rows of the target tensor are the same windows moved one position forward.

A tokenizer turns text into integers. Each integer is a token ID that names one entry in a fixed vocabulary. An integer by itself carries no information about how one token relates to another, so a model cannot do much arithmetic with it directly. The usual remedy is to replace each ID with a vector, a list of numbers that the model can learn to adjust during training. That replacement is called a token embedding.

In this lesson the token IDs come from the BPE tokenizer built earlier in the course, and the vocabulary has 50,257 entries. The embedding size is set to 256 dimensions. The source describes 256 as much smaller than the 12,288 dimensions used by the original GPT-3 model, yet still a reasonable size for experimentation on ordinary hardware. It is a practical choice for learning rather than a claim about what production models use.

The embedding layer is created with a constructor of the form torch.nn.Embedding(vocab_size, output_dim). In this lesson the first argument is 50,257 and the second is 256. These two numbers fully determine the layer's shape. The vocabulary size sets how many rows the layer holds, one for each possible token ID, and the output dimension sets how long each row is. A row is the vector that a given token ID maps to.

Because the table is indexed by token ID, looking up an ID means retrieving one row. A token with ID 7 is represented by row 7 of the table, and a token with ID 50256 is represented by the last row. The values in those rows are the numbers the model adjusts during training.

The batch passed to the layer is a two-dimensional grid of token IDs, with one row per text sample and one column per token position. The lesson uses a data loader to fetch a batch of 8 samples with 4 tokens each, so the input grid has shape 8 × 4. Every entry in that grid is one integer ID between 0 and 50,256.

Calling the embedding layer on that grid replaces every integer with its 256-value vector. The grid positions do not move. The batch and sequence dimensions are preserved, and a third dimension of length 256 is appended. The result has shape 8 × 4 × 256, which means every one of the 32 token IDs is now a 256-dimensional vector.

The token embedding layer treats each token by its identity alone. It does not record where in the sequence a token sits. A GPT-style absolute embedding approach needs a second embedding layer with the same embedding dimension, 256 in this lesson, so that the two outputs can be combined. That second layer is the subject of the next step, and both layers must share the 256-dimensional size.