Token Embeddings and Vocabulary Mapping

An LLM never works directly with words. Before any computation happens, each token in the input text is turned into a numeric identifier. That identifier is just a whole number, and it does not carry any information about what the token means. The model still needs something it can compute with, so the identifier must be converted into a vector of numbers. The component responsible for this conversion is the embedding layer.

Think of the embedding layer as a lookup structure. You give it a token, and it returns the vector assigned to that token. Everything in this article is about how that lookup structure is sized, what it contains when it is first created, and how to tell when its shape is wrong.

A lookup structure has to be sized before it can be used, and two numbers decide that size. The first is how many distinct tokens the layer must cover, which is the vocabulary size. The second is how many numbers make up each token's vector, which is the output dimension. Once these two values are fixed, the layer has a definite shape and a definite amount of storage.

An embedding layer is created by calling a PyTorch module class and passing it the two sizes. The signature has exactly two positional parameters, and their order matters. The first argument is the vocabulary size and the second is the output dimension. Swapping them does not raise an error in the way a typo might, but it gives the layer a different shape than you intended, so the order must be treated as part of the design.

The two constructor arguments map onto the two axes of the weight matrix. The vocabulary size sets the number of rows, one row per token. The output dimension sets the number of columns, one column per embedding feature. The weight matrix therefore has the shape vocab_size by output_dim. Written as math, a layer built with a vocabulary of $V$ tokens and an output dimension of $d$ has a weight matrix of shape $V \times d$.

Once the layer exists, a token ID arrives and selects a row. Suppose the token ID is 5. Because the vocabulary has 6 tokens, IDs run from 0 to 5, so ID 5 is the last row of the weight matrix. The layer does not search for anything. It uses the ID as a position and returns the row stored at that position. That returned row is the vector for that token.

Notice what this means for the whole table. The lookup never mixes rows. Token 5 always reads row 5, and no other token's row is touched. This is why the vocabulary mapping is a one-to-one association between each token ID and one row.

The matrix is a block of stored numbers with a fixed layout. Its rows are indexed by tokens and its columns by embedding dimensions. Because the two sizes are fixed when the layer is created, the layer has a definite amount of storage from the start. Every token row has the same width, and the total number of entries is always the vocabulary size multiplied by the output dimension.

When the layer is created, every entry of the weight matrix is filled with a small random number. These numbers are not chosen to represent any particular word. At creation time they carry no linguistic meaning. The row for token 5 is simply some small random values, and so is every other row. They are a starting point from which training can begin.

The values are not fixed for good. They are adjusted during LLM training, as part of the same optimization process that trains the rest of the model. Over many updates, the rows that correspond to tokens are nudged toward values that help the model predict text well. This is why the initial values are a beginning rather than a result. A trained embedding table and a freshly created one have the same shape, but very different contents.

The weights are marked as trainable parameters, which is why the printed representation of the layer reports that gradients are required. This flag tells PyTorch to track operations on the matrix so that it can be updated during optimization.

The embedding layer is one stage in a longer chain. Text is first split into tokens and converted to numeric IDs. Those IDs are passed to the embedding layer, which returns vectors. Those vectors are then passed on to the rest of the model, which expects them to have a particular shape. A problem at any stage shows up as a shape that the next stage does not accept.

Random initialization creates a practical problem. If every run produced different starting numbers, two people running the same code could see different results and find it hard to compare notes. A random seed solves this. Setting a seed makes the random generator produce the same sequence every time, so the starting weights are reproducible.

A fixed seed makes the starting values reproducible without changing the shape. The seed controls which numbers fill the matrix, not how many there are. A layer built with a vocabulary of 6 and an output dimension of 3 still has 6 rows and 3 columns whether or not a seed is set.

When an embedding layer behaves unexpectedly, the fastest route to the cause is to compare the observed weight shape against the shape you predicted from the arguments. A mismatch almost always traces back to one of three places: the values passed in, the order in which they are passed, or the moment the seed is set relative to layer creation. Checking the shape first isolates the first two causes quickly. Checking the seed ordering isolates the third.