Embedding Layers: Mapping Token IDs to Dense Vectors

Introduction

In our last lesson, we focused on tokenization, the process of converting raw text into a sequence of integer token IDs. For example, the sentence "Tokenization is fascinating!" became a sequence like [11241, 4712, 318, 12691, 0].

This raises a critical question: what do these numbers actually mean to a model? As they are, they're just arbitrary indices. The model has no inherent understanding that the token for "dog" is semantically closer to "cat" than it is to "car". For a neural network to perform meaningful operations, it needs a richer, more nuanced representation of these tokens.

This is where embeddings come in. This lesson will explain how an embedding layer maps each discrete token ID to a dense, multi-dimensional vector. These vectors, or embeddings, capture the semantic properties of tokens, placing them in a continuous space where similar concepts are closer together. This transformation is the second crucial step in preparing text for a language model.

From IDs to Meaningful Vectors

The core idea is to represent each token not as a single integer, but as a list of floating-point numbers—a vector. This vector lives in a high-dimensional "semantic space," and its position is learned by the model during training.

To build your intuition, let's start with a high-level conceptual overview. This video explains what embeddings are and how they differ from the tokens we discussed in the last lesson.

As the video explained, while a token is just an ID, an embedding is a dense vector designed to capture meaning. The model doesn't process the token IDs directly; it first converts them into these meaningful vectors. This is the job of the embedding layer.

The Embedding Layer: A Learnable Lookup Table

In PyTorch, this mapping is handled by the torch.nn.Embedding layer. The most effective way to think about this layer is as a learnable lookup table.

Imagine a large matrix where:

• The number of rows is equal to the size of your vocabulary (vocab_size).
• The number of columns is the size of the vector you want to represent each token with (embedding_dim). This dimension is a hyperparameter you choose (e.g., 768 for BERT, 4096 for Llama 2 7B).

When you pass a token ID (say, 42) to this layer, it simply returns the 42nd row of the matrix. This row is the embedding vector for that token.

How is this Implemented?

The "lookup table" mechanism is surprisingly simple. It's not a complex operation, but rather a direct and efficient memory access. To demystify this, let's look at how you could build an embedding layer from scratch.

As you can see, the core of the embedding layer is a single tensor of weights and an efficient indexing operation in the forward pass.

How Embeddings are Learned

A crucial point is that the values in this embedding matrix are not fixed; they are learned.

1. Initialization: The (vocab_size, embedding_dim) matrix is initialized with small random numbers. At this stage, the vectors are meaningless.
2. Training Loop: During training, a batch of text is tokenized and passed through the embedding layer to get the corresponding vectors. These vectors are then fed into the rest of the model (which we'll cover in future lessons) to predict the next token.
3. Backpropagation: The model calculates the error (loss) in its prediction. Using backpropagation, gradients are computed and flow back to all model parameters, including the embedding matrix.
4. Sparse Updates: A key efficiency is that gradients only update the rows (the embedding vectors) that were actually looked up in the forward pass. If your batch of text contained tokens [10, 58, 1004], only rows 10, 58, and 1004 of the embedding matrix will be adjusted.

Over many training iterations, the model learns to adjust the vectors such that words used in similar contexts are pushed closer together in the high-dimensional embedding space.

The following video provides a practical look at the torch.nn.Embedding layer and includes a fantastic visualization of how embeddings evolve during training.

Practical Considerations

When building a model, there are important trade-offs to consider, primarily related to the embedding dimension (embed_dim).

• Semantic Capacity: A larger dimension allows the model to capture more nuanced semantic relationships.
• Memory Cost: The memory required for the embedding table scales linearly with the embedding dimension. A large vocabulary combined with a large embedding dimension can result in a massive parameter matrix.
• Computational Cost: Subsequent layers of the model, like attention mechanisms, perform computations over these vectors. Larger vectors mean more computation and higher memory bandwidth requirements.

This resource provides a detailed breakdown of these trade-offs and shows typical embedding dimensions used in well-known models.

Final Code Example

Let's tie everything together with a complete PyTorch code snippet showing the full process from text to embeddings.

import torch
import torch.nn as nn
from transformers import AutoTokenizer




# 1. Tokenization (from our previous lesson)
# We use a pre-trained tokenizer to get the token IDs.
tokenizer = AutoTokenizer.from_pretrained("gpt2")
text = "Embeddings are vectors."
token_ids = tokenizer.encode(text, return_tensors="pt") 

print(f"Input text: '{text}'")
print(f"Token IDs: {token_ids}")
print(f"Shape of Token IDs tensor: {token_ids.shape}\n")



# token_ids will be a 2D tensor of shape [batch_size, sequence_length]
# In this case, it will be [1, 4]




# 2. Define the Embedding Layer
vocab_size = tokenizer.vocab_size  # For GPT-2, this is 50257
embedding_dim = 768              # A common choice for base-sized models

embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)




# 3. Map Token IDs to Dense Vectors
# This is the "lookup" step.
embeddings = embedding_layer(token_ids)

print(f"Shape of output embeddings tensor: {embeddings.shape}")
print(f"The model now has a {embedding_dim}-dimensional vector for each of the {token_ids.shape[1]} tokens.")

Notice the change in the tensor's shape. We started with a 2D tensor of integer IDs (batch_size, sequence_length) and ended with a 3D tensor of floating-point vectors (batch_size, sequence_length, embedding_dim). This 3D tensor is the standard input format for Transformer models.

Conclusion

In this lesson, we've bridged the gap between abstract token IDs and the meaningful numerical representations that neural networks require.

Key Takeaways:

• Embeddings are dense vectors that capture the semantic properties of tokens in a continuous space.
• An nn.Embedding layer functions as a lookup table, which is essentially a learnable matrix of shape (vocab_size, embedding_dim).
• A token's integer ID is used as an index to retrieve its corresponding dense vector from this matrix.
• These embedding vectors are initialized randomly and are learned via backpropagation during the model's training process, allowing them to capture complex linguistic relationships.
• The choice of embedding_dim is a critical hyperparameter that involves a trade-off between semantic capacity, memory usage, and computational cost.

Preview of the Next Lesson:

Now that we can convert text into a sequence of meaningful vectors, we're ready to build a model that can process them. In our next lesson, we will define autoregressive language modeling, the fundamental task of predicting the next token in a sequence given all the preceding tokens. This will set the stage for understanding how models like GPT generate text.