Embeddings

By the end of this lesson you will be able to explain what an embedding is, why it beats one-hot encoding, how similar items end up with similar vectors, and how to use nn.Embedding in PyTorch. Embeddings are the bridge between symbols (words, product IDs, PIN codes) and the numbers neural networks need.

Every transformer, recommender system and semantic search engine starts with an embedding layer. The next lesson, Attention & Transformers, builds directly on this one.

The problem with one-hot vectors

Networks need numbers. A simple way to turn a word into numbers is one-hot encoding: a vector as long as the vocabulary, all zeros except a 1 at the word's position. You met it for categories in Scaling & Encoding Features.

For deep learning it has two problems:

1. Size. With a vocabulary of 50,000 words, every word is a vector of 50,000 numbers, almost all zero.
2. No meaning. Every pair of one-hot vectors is equally far apart. "Chai" is as different from "coffee" as it is from "train".

Dense vectors with meaning

An embedding represents each item as a short, dense vector of, say, 50 or 300 numbers, where similar items get similar vectors.

Think of describing Indian dishes by a few traits: how spicy, how sweet, how liquid, how much it is a breakfast food. Masala dosa and idli would score alike; gulab jamun and rasgulla would score alike; and the two groups would be far apart. Each dish becomes a point in "taste space". An embedding is exactly that, except the network learns the traits itself, and they usually have no simple names.

To compare two vectors we use cosine similarity: 1 means they point the same way, 0 means unrelated, -1 means opposite.

Step 1: one-hot versus hand-made dense vectors.

import numpy as np

dishes = ["idli", "dosa", "gulab jamun", "rasgulla", "pani puri"]

def cosine(a, b):
    return a @ b / (np.linalg.norm(a) * np.linalg.norm(b))

one_hot = np.eye(len(dishes))
print("one-hot  idli vs dosa:", cosine(one_hot[0], one_hot[1]),
      "| idli vs rasgulla:", cosine(one_hot[0], one_hot[3]))

#                 spicy sweet liquid breakfast
dense = np.array([[0.2, 0.0, 0.1, 0.9],     # idli
                  [0.4, 0.0, 0.0, 0.8],     # dosa
                  [0.0, 1.0, 0.4, 0.0],     # gulab jamun
                  [0.0, 0.9, 0.6, 0.0],     # rasgulla
                  [0.8, 0.2, 0.7, 0.0]])    # pani puri
for i, j in [(0, 1), (2, 3), (0, 3), (4, 1)]:
    print(f"dense    {dishes[i]} vs {dishes[j]}: {cosine(dense[i], dense[j]):.2f}")

one-hot  idli vs dosa: 0.0 | idli vs rasgulla: 0.0
dense    idli vs dosa: 0.96
dense    gulab jamun vs rasgulla: 0.98
dense    idli vs rasgulla: 0.06
dense    pani puri vs dosa: 0.33

With one-hot vectors every pair scores 0. With the dense vectors, idli and dosa are very close, the two sweets are very close, and idli and rasgulla have almost nothing in common.

An embedding layer is a lookup table

An embedding layer stores one row of numbers per item, in a matrix of shape (number of items, embedding size). To embed item number 3, it simply returns row 3. That is all.

It is worth seeing that this lookup is the same as multiplying a one-hot vector by the matrix. So an embedding layer is a dense layer without bias, applied to one-hot inputs, but computed in a much cheaper way. Because it is just a layer of weights, it is trained by backpropagation along with the rest of the network.

Step 2: lookup equals one-hot times matrix.

import numpy as np

rng = np.random.default_rng(0)
vocab_size, dim = 6, 3
E = rng.normal(size=(vocab_size, dim)).round(2)    # the embedding table

word_ids = np.array([3, 0, 3, 5])                   # a short "sentence" of word numbers
looked_up = E[word_ids]                             # lookup: pick rows
one_hot = np.eye(vocab_size)[word_ids]
multiplied = one_hot @ E                            # the expensive way

print("table shape:", E.shape)
print("lookup result:\n", looked_up)
print("same as one-hot @ E?", np.allclose(looked_up, multiplied))

table shape: (6, 3)
lookup result:
 [[-1.27 -0.62  0.04]
 [ 0.13 -0.13  0.64]
 [-1.27 -0.62  0.04]
 [-0.73 -0.54 -0.32]]
same as one-hot @ E? True

The sentence of 4 word IDs became a (4, 3) matrix: one vector per word. The repeated word 3 got the same vector both times.

Step 3: meaning from context

Where do good embeddings come from? A key idea from linguistics is that words used in similar contexts tend to have similar meanings. "Chai" and "coffee" both appear near "drinks", "hot" and "morning".

Before neural networks, a classic recipe was: count which words appear near each other, weight the counts (here with PPMI, positive pointwise mutual information, which highlights pairs that occur together more than chance), then compress the table with SVD into a few dimensions. You can run it on a tiny corpus here.

import numpy as np

corpus = """priya drinks hot chai every morning
rohan drinks hot coffee every morning
meera drinks cold coffee after lunch
kabir drinks masala chai after lunch
priya eats spicy samosa with chai
rohan eats crispy samosa with coffee
meera eats sweet jalebi with chai
kabir eats sweet ladoo after lunch
the train leaves the station at night
the bus leaves the station at noon
the train reaches the platform late
the bus reaches the stop late
aarav boards the train at the station
isha boards the bus at the stop"""
sentences = [line.split() for line in corpus.splitlines()]
vocab = sorted({w for s in sentences for w in s})
index = {w: i for i, w in enumerate(vocab)}

counts = np.zeros((len(vocab), len(vocab)))
for s in sentences:
    for i, w in enumerate(s):
        for j in range(max(0, i - 2), min(len(s), i + 3)):   # 2 words either side
            if i != j:
                counts[index[w], index[s[j]]] += 1

total = counts.sum()
p_word = counts.sum(axis=1, keepdims=True) / total
with np.errstate(divide="ignore"):
    pmi = np.log((counts / total) / (p_word * p_word.T))
ppmi = np.maximum(pmi, 0)

U, S, Vt = np.linalg.svd(ppmi)
emb = U[:, :5] * S[:5]                                 # 5-dimensional embeddings
emb /= np.linalg.norm(emb, axis=1, keepdims=True)

print("vocabulary size:", len(vocab))
for word in ["chai", "priya", "samosa"]:
    sims = emb @ emb[index[word]]
    nearest = [vocab[k] for k in np.argsort(-sims)[1:4]]
    print(f"{word:<7} nearest: {nearest}")

vocabulary size: 37
chai    nearest: ['coffee', 'every', 'priya']
priya   nearest: ['rohan', 'samosa', 'spicy']
samosa  nearest: ['spicy', 'crispy', 'with']

Nobody told the program that chai and coffee are drinks, or that Priya and Rohan are people. Their vectors ended up close because they appear in similar places. With millions of sentences instead of fourteen, this works remarkably well.

word2vec (Mikolov and colleagues, 2013) learned such vectors with a small neural network trained to predict nearby words, and it made embeddings famous. In modern models, embeddings are simply the first layer of the network, learned together with the main task.

Embeddings in PyTorch

nn.Embedding(num_embeddings, embedding_dim) is the lookup table. Its input is a tensor of int64 indices of any shape, and it adds one dimension of size embedding_dim at the end.

import torch
from torch import nn

torch.manual_seed(0)
emb = nn.Embedding(num_embeddings=1000, embedding_dim=8, padding_idx=0)
batch = torch.tensor([[12, 45, 7, 0, 0],          # sentence 1, padded with 0s
                      [3, 99, 512, 18, 2]])       # sentence 2
vectors = emb(batch)
print("ids:", tuple(batch.shape), "-> vectors:", tuple(vectors.shape))
print("padding row is zero:", torch.all(emb.weight[0] == 0).item())
print("parameters:", emb.weight.numel())

# a tiny text classifier: average the word vectors, then a linear layer
mask = (batch != 0).unsqueeze(-1).float()
sentence_vectors = (vectors * mask).sum(dim=1) / mask.sum(dim=1)
classifier = nn.Linear(8, 2)
print("sentence vectors:", tuple(sentence_vectors.shape), "-> logits:", tuple(classifier(sentence_vectors).shape))

ids: (2, 5) -> vectors: (2, 5, 8)
padding row is zero: True
parameters: 8000
sentence vectors: (2, 8) -> logits: (2, 2)

Sentences of different lengths are padded to the same length with a special index, here 0. padding_idx=0 keeps that row at zero and never trains it, and the mask makes sure padding does not count in the average. Averaging word vectors throws away word order, but it is a surprisingly strong baseline for text classification.

Embeddings beyond words

Embeddings work for any large set of categories:

• Products and users in a recommender: an online grocery store in Bengaluru can learn that customers who buy idli batter also buy coconut chutney (see Recommender System Basics).
• Categorical columns in tables, such as PIN codes or store IDs with thousands of values.
• Whole sentences, images or documents, produced by a trained model, used for semantic search and retrieval. The Embeddings & Semantic Search lesson shows how that is used with language models.

The embedding size is a hyperparameter. Small sets of categories need only a few dimensions; large vocabularies in language models use hundreds or thousands.

Common mistakes

• Passing float tensors to nn.Embedding. Indices must be integers (int64).
• Index out of range. Every ID must be less than num_embeddings; reserve extra rows for "unknown" and "padding".
• Letting padding affect the result. Use padding_idx and masks when averaging.
• Comparing embeddings with plain distance when their lengths differ. Cosine similarity compares direction and is usually what you want.
• Assuming embeddings are unbiased. They learn whatever associations the training text contains, including social stereotypes. See Responsible Deep Learning.

Key points

• One-hot vectors are huge and treat all items as equally different.
• An embedding maps each item to a short dense vector, so similar items get similar vectors.
• An embedding layer is a trainable lookup table, equivalent to one-hot times a weight matrix.
• Words used in similar contexts get similar vectors; word2vec made this popular, and modern models learn embeddings as their first layer.
• nn.Embedding takes integer IDs and returns vectors; use padding_idx and masks for padded sequences.

import numpy as np

dishes = ["idli", "dosa", "gulab jamun", "rasgulla", "pani puri", "upma"]
vectors = np.array([[0.2, 0.0, 0.1, 0.9], [0.4, 0.0, 0.0, 0.8], [0.0, 1.0, 0.4, 0.0],
                    [0.0, 0.9, 0.6, 0.0], [0.8, 0.2, 0.7, 0.0], [0.3, 0.0, 0.1, 0.9]])
unit = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)

def most_similar(name):
    i = dishes.index(name)
    sims = unit @ unit[i]
    sims[i] = -1                       # ignore the dish itself
    best = int(sims.argmax())
    return dishes[best], round(float(sims[best]), 3)

for name in ["upma", "pani puri"]:
    print(name, "->", most_similar(name))

upma -> ('idli', 0.995)
pani puri -> ('rasgulla', 0.513)