02 · Embeddings 🧭

4 min read

A word's integer ID is just a name tag — it carries no meaning. Embeddings turn those IDs into rich vectors where similar words live near each other.

1. The 'Why'

After tokenization (sub-module 01) you have integer IDs like cat → 3, dog → 4. But those numbers are arbitrary labels: the fact that dog is 4 and cat is 3 doesn't mean dog is "one more than" cat in any meaningful sense. Feeding raw IDs to a network implies false numeric relationships. The obvious fix — one-hot encoding, a vector that's all zeros except a single 1 at the word's index — removes the false ordering, but it's disastrous in its own way: with a 50,000-word vocabulary, every word becomes a 50,000-dimensional vector that's 99.998% zeros, and crucially, every word is equally distant from every other. One-hot vectors know that "cat" and "dog" are different, but not that they're both animals and more alike than "cat" and "democracy."

Embeddings solve this elegantly. Instead of a sparse one-hot, each token maps to a short, dense vector of learnable numbers — say 100 or 300 dimensions. These vectors start random and are trained along with the rest of the model, so the network learns to place words with similar meanings or uses near each other in the vector space. The famous result is that embeddings capture semantic structure: related words cluster, and directions in the space can encode relationships. nn.Embedding implements this as a simple, efficient lookup table, and it's the standard first layer of virtually every NLP model — the piece that finally turns symbols into something a network can reason about.

2. Core Concepts

nn.Embedding(num_embeddings, embedding_dim). A learnable lookup table with one row per vocabulary token. num_embeddings = vocabulary size; embedding_dim = the length of each vector. Given integer IDs, it returns the corresponding rows.

It's a lookup, not a matmul. Conceptually an embedding is a one-hot times a weight matrix, but PyTorch implements it as a direct row lookup — far faster and more memory-efficient. You pass integer IDs, not one-hot vectors.

Learnable by default. The embedding table is a parameter (requires_grad=True), so it's trained by backprop like any weight. The network shapes the vector space to serve the task.

padding_idx. Tell the layer which ID is <pad> so its embedding stays all-zeros and receives no gradient — keeping padding truly inert.

Pretrained embeddings. You can initialize the table with pretrained vectors (GloVe, word2vec) instead of random — a form of transfer learning (Module 07) for text.

3. Code in Action

The embedding lookup

import torch
import torch.nn as nn

vocab_size, dim = 20, 8
embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=dim)

ids = torch.tensor([[3, 4, 6, 0],      # a padded batch of token IDs
                    [7, 5, 2, 0]])      # (batch=2, seq_len=4)
vectors = embedding(ids)                # look up each ID's row
print(vectors.shape)                    # torch.Size([2, 4, 8]) -> (batch, seq, dim)

Keeping padding inert with `padding_idx`

import torch
import torch.nn as nn

emb = nn.Embedding(20, 8, padding_idx=0)   # ID 0 (<pad>) -> fixed zero vector
print(emb(torch.tensor([0]))[0][:4])       # tensor([0., 0., 0., 0.]) -> all zeros
# Its row won't be updated during training, so padding never influences learning.

Embedding as the first layer of a text model

import torch
import torch.nn as nn

class TextClassifier(nn.Module):
    def __init__(self, vocab_size, dim, num_classes):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim, padding_idx=0)  # IDs -> vectors
        self.fc = nn.Linear(dim, num_classes)                     # simple classifier head

    def forward(self, ids):
        vecs = self.embed(ids)          # (batch, seq, dim)
        pooled = vecs.mean(dim=1)       # average word vectors -> (batch, dim)
        return self.fc(pooled)          # -> logits

model = TextClassifier(vocab_size=20, dim=8, num_classes=2)
print(model(torch.tensor([[3, 4, 6, 0]])).shape)   # torch.Size([1, 2])

4. Common Pitfalls

Pitfall 1 — Passing one-hot vectors instead of IDs. nn.Embedding expects a LongTensor of integer indices, not one-hot floats. Handing it one-hots (or float IDs) throws a type/shape error. Feed raw integer token IDs straight from numericalization.

Pitfall 2 — Vocabulary/embedding size mismatch. num_embeddings must be at least as large as your largest token ID + 1. If an ID exceeds the table size, you get an index-out-of-range error at runtime. Size the embedding to the full vocabulary (including specials).

Pitfall 3 — Ignoring padding_idx. Without it, the <pad> token gets a trainable, non-zero embedding that leaks into pooled representations and wastes capacity. Set padding_idx so padding contributes nothing.

5. Further Reading & Watch List

⬅️ Prev: 01 · Text Data & Tokenization · ➡️ Next: 03 · Recurrent Layers & LSTMs