Embedding layers

Learning what "similar" means

Think about a supermarket checkout. Nobody tags a tent and a sleeping bag as "camping gear." Yet the same shoppers often buy them on the same trip. Watch enough baskets and the pattern shows up by itself.

An embedding table can learn that pattern. The vectors start random and mean nothing. They gain meaning from the training task. Here's a classic recipe, used by word2vec and many recommenders:

1. Positive pairs. Take two products from the same basket. The right answer is "bought together."
2. Negative pairs. Take two products at random. The right answer is "not bought together." A random pair is usually a good guess for that.
3. Score each pair. Multiply the two vectors number by number, and add up the results. This is the dot product, from Phase 2.
4. Adjust. Training pushes the score up for positive pairs and down for negative pairs.

Here's the score as a formula. Product A's vector holds the numbers $a_1, a_2, \ldots, a_n$. Product B's vector holds $b_1, b_2, \ldots, b_n$. The letter $n$ is how many numbers each vector has.

In plain words: pair up the two vectors' numbers, multiply each pair, and add. Vectors that point the same way get a big score. Vectors that point in different directions get a small one. Longer vectors push the score further, too.

print(1 * 3 + 2 * 4)

A quick example with 2 numbers per vector. Say the tent is [0.9, 0.1] and the sleeping bag is [0.8, 0.3]. The score is 0.9 × 0.8 + 0.1 × 0.3 = 0.72 + 0.03 = 0.75. If the kettle is [0.1, 0.9], the tent and kettle score 0.09 + 0.09 = 0.18. The worked example below does the same sums in Python.

names = ["tent", "sleeping bag", "kettle", "mug"]
table = [[0.9, 0.1],
         [0.8, 0.3],
         [0.1, 0.9],
         [0.2, 0.7]]

one_hot = [1, 0, 0, 0]                    # ID 0, the tent
by_multiplying = [sum(o * row[j] for o, row in zip(one_hot, table)) for j in range(2)]
print("lookup:", table[0])
print("one-hot times table:", by_multiplying)

def score(a, b):
    return round(sum(x * y for x, y in zip(a, b)), 2)

for i in [1, 2, 3]:
    print("tent and", names[i], "->", score(table[0], table[i]))

Training on shopping baskets

The next example runs the whole recipe. There are fifteen products in three departments: outdoor, kitchen and office. There are 400 made-up shopping baskets. Here's what the code does.

• Builds baskets. Each basket picks one department and three products from it. About one basket in five also gets a random extra item, as real shoppers do.
• Lists positive pairs. Every pair of products that share a basket.
• Trains for 300 steps. Each step scores 128 positive pairs and 128 random pairs, then updates the table. Each product gets 8 numbers. Adam is an optimizer: the rule that updates the weights.
• Finds neighbors. It compares every product with every other using cosine similarity, from Phase 1. It prints the three closest products to the tent, the kettle and the planner.
• Draws a picture. Each vector has 8 numbers, too many to plot. PCA, from Phase 2, squashes them down to 2. Each dot's color shows its department.

import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

products = ["tent", "sleeping bag", "headlamp", "trail shoes", "rain jacket",
            "coffee mug", "kettle", "tea towels", "cutting board", "candle",
            "notebook", "desk lamp", "pen set", "planner", "laptop sleeve"]
department = [0] * 5 + [1] * 5 + [2] * 5            # outdoor, kitchen, office (the model never sees this)
rng = np.random.default_rng(0)
baskets = []
for _ in range(400):
    d = rng.integers(3)
    items = set(rng.choice([i for i in range(15) if department[i] == d], size=3, replace=False).tolist())
    if rng.random() < 0.2:
        items.add(int(rng.integers(15)))             # sometimes a random extra item
    baskets.append(sorted(items))
pairs = torch.tensor([(a, b) for bk in baskets for a in bk for b in bk if a != b])

torch.manual_seed(0)
emb = nn.Embedding(15, 8)
opt = torch.optim.Adam(emb.parameters(), lr=0.05)
for step in range(300):
    pos = pairs[torch.randint(0, len(pairs), (128,))]          # bought together
    neg = torch.randint(0, 15, (128, 2))                       # random pairs
    a, b = torch.cat([pos[:, 0], neg[:, 0]]), torch.cat([pos[:, 1], neg[:, 1]])
    target = torch.cat([torch.ones(128), torch.zeros(128)])
    score = (emb(a) * emb(b)).sum(dim=1)                       # dot product of the two vectors
    loss = nn.functional.binary_cross_entropy_with_logits(score, target)
    opt.zero_grad()
    loss.backward()
    opt.step()

E = emb.weight.detach().numpy()
En = E / np.linalg.norm(E, axis=1, keepdims=True)
sim = En @ En.T                                                # cosine similarity between all products
for i in [0, 6, 13]:
    nearest = np.argsort(-sim[i])[1:4]
    print(f"{products[i]:<10} → {[products[j] for j in nearest]}")

xy = (E - E.mean(0)) @ np.linalg.svd(E - E.mean(0))[2][:2].T    # PCA to 2-D for a picture
plt.scatter(xy[:, 0], xy[:, 1], c=department, cmap="viridis", s=60)
for (x, y), name in zip(xy, products):
    plt.annotate(name, (x, y), fontsize=8)

The department labels are only used to build the baskets and color the plot. The model never sees them. Yet its vectors sort themselves into three groups, purely from what's bought together.

That's the same thing that makes "king" and "queen" neighbors in word embeddings. Items used in similar contexts get similar vectors.

Note: the neighbors use cosine similarity, not the raw dot product. Cosine compares only the direction of two vectors, so a vector's length can't sway the ranking.

Embeddings inside language models

A language model reads tokens: words or pieces of words, as you saw in Phase 1. Each token has an ID. So the model's very first layer is an embedding table.

• Token embeddings. The first layer is an nn.Embedding(vocab_size, d_model). The vocab_size is how many different tokens the model knows. The d_model is how many numbers each token gets, called the model width. Llama 3 8B knows about 128,000 tokens, with 4,096 numbers each. Rows times columns gives about 525 million weights, in this one table alone.
• Position embeddings. A second table (or a formula) adds where each token sits in the text. You'll meet these in Phase 5.
• Tied weights. The model's last layer gives a score to every token in the vocabulary, to pick the next one. Many models reuse the token embedding table, turned on its side, as that last layer. GPT-2 does this. It saves weights, and it links "reading" a token to "writing" it.

print(50_000 * 768)

How this relates to embedding APIs

In Phase 1 you used embedding APIs. They return one vector for a whole text. That vector comes from a model's later layers, after the model has read the whole text. Often the model was trained specifically so that similar texts get similar vectors.

An nn.Embedding table is simpler. It holds one vector per ID. It learns only from the task you train it on. The geometry is the same: close vectors mean similar things. The scope is different.

For marketers: product, customer and campaign embeddings power "customers who bought this also bought", lookalike audiences and content recommendations. Train vectors on behavior (co-purchases, co-views, sessions), and similarity search over them does the rest.