One-Hot Encoding for Tokens

Before a model learns anything, every token in a text has already been reduced to an integer. That integer is the token ID. It tells the model which entry of its vocabulary a token corresponds to, but it carries no information about meaning. The ID 3 is just a label, in the same way a locker number tells you which locker to open without telling you what is inside it.

An embedding layer is the part of the model that turns that bare label into a vector of numbers the model can actually compute with. The question this article answers is how it performs that conversion. The short answer is that it does not compute anything complicated. It looks up a row.

Think of the embedding layer as holding a weight matrix with one row per vocabulary entry. The token ID is used as a position. The layer moves to that position in the weight matrix and reads the row it finds there. Rows at other positions are not touched for this lookup. Only one row is read, and the vector it contains is what comes back.

Notice the numbering in the diagram. Token ID 3 retrieves the row in position 3, which is the fourth row, because positions start counting at zero. The other rows are shown only to make clear that they are present in the matrix but are not read for this lookup.

Token IDs are zero-based row indices. The first row is numbered 0, the second is numbered 1, and so on, up to the last row of the vocabulary. Token ID 3 therefore retrieves the fourth row, not the third. This is the most common source of confusion when reading embedding results, because people naturally count from one.

The returned embedding vector is identical to the stored row it came from. This gives you a simple verification procedure. Take the vector returned for a given token ID, find the row at that zero-based position in the embedding matrix, and compare the two. If they match exactly, the lookup worked as intended. If they do not match, the problem lies in the index that was used for the lookup, not in the stored values.

Readers who already know one-hot encoding may recognize the lookup as a shortcut. One-hot encoding represents a token ID as a vector whose length equals the vocabulary size. That vector contains a single 1 at the position of the token and 0 everywhere else. The one-hot vector is then multiplied by a weight matrix in a fully connected layer.

In this expression, $\mathbf{x}$ is the one-hot vector for token ID 3, $W$ is the weight matrix, and $W_{3,:}$ is the row at zero-based position 3. Because every entry of $\mathbf{x}$ except the one at position 3 is zero, every other term in each sum is multiplied by zero and drops out. Only row 3 survives.

The lookup is equivalent to one-hot encoding followed by matrix multiplication in a fully connected layer, and it is the more efficient implementation. The one-hot route performs a multiplication against a vector as long as the whole vocabulary, even though all but one of those products are zero. The lookup skips that work and reads the single row that the one-hot route would have kept. Because the two are equivalent, the embedding layer can be treated as a neural network layer. The weights of its matrix can be optimized through backpropagation during training, just like the weights of other layers.

Suppose the vector printed for token ID 3 does not match the fourth row of the embedding matrix. The useful question is where in the path from token ID to returned vector the actual result stopped matching the row the lookup should return. The path has a small number of stages, and checking them in order narrows the problem quickly.

The returned vector is identical to the stored row, so a mismatch points to the index used for the lookup rather than to the stored values themselves. In practice, the first place to look is the index: was the token ID 3 read as a zero-based position, and was the fourth row the one actually read? If the index is correct and the vector still differs from the stored row, the comparison itself needs to be checked against the correct row.