Implementing Token Embeddings
The integers in input_ids are names for tokens. They are not the numeric features the model computes with. The embedding layer is what turns each of them into a vector.
Why a Model Needs Numbers
A language model does not read text directly. First, the text is split into tokens, and each token is represented by an integer token ID. Those integers are useful for identifying tokens, but they are not the numeric representations a model can compute with. Turning an integer token ID into an embedding vector is the job of an embedding layer, and this article follows that conversion one step at a time.
Throughout the article you will work with a small, concrete case: four input tokens whose IDs are 2, 3, 5, and 1. Each of these integers will be looked up by the embedding layer, and each lookup will return a vector of numbers that belongs to that token.
What do you think happens?
The embedding layer receives the token ID 3. Which row of its weight matrix do you think it selects: the third row or the fourth row?
Reveal answer
The fourth row, which is the row at index 3 when counting from zero.
Python starts counting at zero, so the index 3 names the fourth row. The source states that the embedding vector for token ID 3 is identical to the fourth row of the embedding matrix. The rest of this article explains why the count starts at zero and how to check it.
Storing the Token IDs in a Tensor
Before any lookup happens, the four token IDs are stored in a tensor. A tensor is the data container PyTorch uses for numbers, and torch.tensor() creates one from a Python list. In this example the tensor is held in a variable named input_ids. That variable holds the integers that the embedding layer will look up, so it is the starting point for every step that follows.
torch.tensor()builds a tensor from the Python list of four integersinput_idskeeps the token IDs that the embedding layer will look up- Each integer is a token ID, not an embedding vector
The integers in input_ids are names for tokens. They are not the numeric features the model computes with. The embedding layer is what turns each of them into a vector.
Tracing a Single Lookup
The next step applies the embedding layer to a single token ID. The layer is named embedding_layer in this example, and calling it like a function is how you use it. The call passes a tensor that contains only the ID 3. Before running it, notice what you are asking for: one token ID goes in, and the layer must return the vector that belongs to that token.
tensor([[ three numeric values ]], grad_fn=<EmbeddingBackward0>)single_idholds only the token ID 3- Calling
embedding_layerperforms the lookup and returns one embedding vector - The printed output includes the vector and metadata that PyTorch attaches to it
The returned result is displayed inside double square brackets. The outer pair of brackets means the value is a two-dimensional tensor with a single row. The three numbers inside that row are the components of the embedding vector for token ID 3. In this example, each embedding vector has three components.
Finding the row for token ID 3
The embedding layer receives the token ID 3. Which row of the weight matrix does it return, and how do you know?
- 1Read the token ID
The input is the single integer 3. It names a token, and it is not a vector.
- 2Convert to a zero-based row index
Python counts from zero, so the row index for token ID 3 is 3.
- 3Locate the row
Index 3 names the fourth row of the matrix, because the first row is index 0, the second is index 1, the third is index 2, and the fourth is index 3.
- 4Compare the returned vector
The vector returned for token ID 3 is identical to that fourth row. The layer read a stored row and did not compute a new vector.
Token ID 3 selects the row at index 3, the fourth row, and the layer returns the three numbers stored in that row.
From Token IDs to Embedding Vectors
The conversion has three stages. First, integer token IDs enter the embedding layer. Second, the layer uses each ID to select one row of its weight matrix. Third, the selected rows leave the layer as embedding vectors. The weight matrix holds one stored row per token ID, and the ID tells the layer which row to read.
Nothing in this path performs arithmetic on the token ID. The embedding layer's job is retrieval, so the same ID always returns the same stored row.
Token IDs and Embedding Vectors Are Different
Beginners often treat the token ID and the embedding vector as two versions of the same thing. They are different objects with different roles. The token ID is a single integer that names a token. The embedding vector is a list of numbers that the layer reads from a stored row of the weight matrix.
| Property | Token ID | Embedding vector |
|---|---|---|
| What it is | A single integer | A list of numbers |
| Role | Names a token | Represents a token numerically |
| Example in this article | 3 | Three numbers from the fourth row |
| Where it comes from | Produced before the layer runs | Read from the weight matrix by the layer |
Token IDs and embedding vectors differ in form and in role
Finding Where a Lookup Diverges
When a lookup result looks wrong, trace three things: the input token ID, the zero-based row index, and the vector returned by the layer. Compare the returned vector with the row you expected. The step where the two stop matching is where the expectation went wrong.
Common mistake Counting rows from one instead of zero
What it looks likeToken ID 3 is taken to mean the third row.
Why it failsPython starts with a zero index, so token ID 3 corresponds to the row at index 3, which is the fourth row.
Do this insteadConvert the token ID to a zero-based row index before you look for the row.
Common mistake Treating the token ID as the embedding vector
What it looks likeExpecting the lookup to return the integer 3 itself.
Why it failsThe token ID is one integer that names a token. The embedding vector is the list of numbers stored in the matching row.
Do this insteadKeep the two objects separate when you debug and compare the returned vector against the matrix row.
Common mistake Expecting the layer to compute a new vector
What it looks likeAssuming the output for an ID is calculated from the ID's value.
Why it failsThe embedding layer performs a lookup, not a calculation. It reads a stored row.
Do this insteadWhen a result differs from expectation, check which row was read rather than searching for a formula.
Common mistake Ignoring the output shape, so a single row inside double brackets is mistaken for several separate vectors
Practice Tracing a Lookup
Using the four token IDs 2, 3, 5, and 1, determine which zero-based row index the embedding layer selects for token ID 5. Then explain what you would check first if the returned vector matched the row at index 3 instead.
Hints
- Start from the token ID itself and convert it to a zero-based row index
- The source states that a token ID corresponds to the row at the same index, counted from zero
- If the returned vector matches a different row, compare the token ID, the row index, and the returned vector in that order
Key Takeaways
- A token ID is an integer, and the embedding layer uses it as a key to retrieve one stored row of its weight matrix.
- The embedding layer performs a lookup, not a calculation, so the returned embedding vector is identical to the matching row.
- Rows are counted from index zero, so token ID 3 selects the row at index 3, which is the fourth row.
- To debug a surprising lookup, trace the token ID, the zero-based row index, and the returned vector, and find the step where they stop matching the expected row.