Creating Token Embeddings for LLM Training

Imagine you want to build a model that predicts the next piece of text. Before you can decide how that model should be built, you need to answer a more basic question: what kind of data can the model actually learn from? This article approaches that question from the training side. We start by asking what happens during training, and then work backward to see why the data the model reads has to take a particular form.

The central claim is short. A continuous vector representation, also called an embedding, is necessary because GPT-like LLMs are deep neural networks trained with the backpropagation algorithm. Each part of that sentence does work. The architecture determines what the model is made of, the training algorithm determines how the model improves, and the combination of the two is what makes the vector form necessary.

Training a network means repeatedly comparing its output with a desired result, measuring the error, and then nudging the network's internal settings so the error gets a little smaller. Backpropagation is the algorithm that works out how much each internal setting should be nudged, by passing information about the error backward through the network. A nudge only makes sense when the quantities being adjusted are numbers that can change by small amounts.

The consequence for representation follows directly. If the model is trained by backpropagation, then the values it processes must also be numeric and continuous, so that small changes in the settings lead to small, measurable changes in the computation. A continuous vector gives the training process a space where gradual adjustment is possible. Without that space, there is nothing for backpropagation to measure and adjust in a graded way.

A deep neural network is built as a sequence of layers. Each layer takes the output of the one before it and passes a transformed result forward. The word deep refers to having many such layers stacked together. GPT-like LLMs belong to this family, which is why they are described as deep neural networks rather than as simple one-step models.

The layered structure is the reason depth matters for training. Each layer has internal settings that can be adjusted, and the whole stack has to be adjusted together so that the final output moves toward the desired result. Keeping track of how a change in an early layer affects the final output is exactly the job that backpropagation handles.

Follow a single training step in order. The step begins with a forward pass, where values enter at the input and move through each layer until the network produces an output. The output is compared with what it should have been, which gives an error. The error is then sent backward through the layers, and each internal setting receives an adjustment. The next forward pass uses the updated settings, and the cycle repeats many times.

The adjustments happen at the end of the pass, in the settings that backpropagation reaches. The forward pass itself does not change those settings. This is the point in the cycle where numeric values change, and it only works because those values are numbers that can move by small amounts.

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, which is a list of numbers that the model can learn to adjust during training. That replacement is called a token embedding.

In the lesson these 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 this 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. The 256 figure is a practical choice for learning, not a claim about what production models use.

The embedding layer holds a table with one row for every possible token ID. Looking up a token ID means selecting the single row with that index. The row is the vector that the token ID maps to, and its length is set by the output dimension. In the lesson, every row has 256 values, so a token ID that picks a row returns a 256-value vector.

The token embedding layer treats each token by its identity alone. It does not record where in the sequence a token sits. Position is a separate matter: the source notes that a GPT-style absolute embedding approach needs a second embedding layer with the same embedding dimension as the token embedding layer, and both layers must share the 256-dimensional size so their outputs can be combined.

The layer is defined by two numbers: the vocabulary size, which is the input dimension, and the output dimension, which is the embedding size. The constructor call takes the form torch.nn.Embedding(vocab_size, output_dim). In the lesson the first argument is 50,257 and the second is 256. These two numbers fully determine the layer's shape.

Multiplying those two numbers gives the count of trainable values in the table: 50,257 times 256 is 12,865,792 numbers. This is arithmetic from the lesson's dimensions, not a separate figure from the source. It shows why the vocabulary size and output dimension together control how large the layer is.

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 \times 4$. Every entry in that grid is one integer ID between 0 and 50,256.

Calling the embedding layer on the 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 is a three-dimensional tensor of shape $8 \times 4 \times 256$, which means every one of the 32 token IDs is now a 256-dimensional vector.

When checking embedding output, verify the input shape first, then confirm that the final axis equals the output dimension. If the input grid has an unexpected shape, the mismatch exists before the embedding layer is called, so the embedding layer is not the place to look first. If the input is correct but the final axis is not the output dimension, then the layer's output dimension is the place to check.