Embedding lookup tables

A raw ID, such as a word's position in the vocabulary or a user's number, tells a model nothing about similarity. ID 7 is not "closer" to ID 8 than to ID 300. An embedding fixes that. It maps each ID to a vector, and training adjusts those vectors until their geometry reflects something useful.

Walking through a lookup

import torch
import torch.nn as nn

torch.manual_seed(0)
emb = nn.Embedding(num_embeddings=10_000, embedding_dim=128)
ids = torch.tensor([3, 7, 3])
vecs = emb(ids)
print(vecs.shape)
print(torch.equal(vecs[0], vecs[2]))
print(vecs.requires_grad)

Output:

torch.Size([3, 128])
True
True

The table has 10,000 rows of 128 numbers each. The lookup of [3, 7, 3] returns three vectors. Positions 0 and 2 both asked for ID 3, so they are identical. The output requires gradients, so training can move rows 3 and 7.

Only the rows that were looked up receive gradient. A user who never appears in a batch gets no update from that batch.

Why it works

The table starts with small random numbers. The meaning comes from the loss. If two IDs appear in similar contexts, or are clicked together, a good loss pushes their rows toward each other. Nothing in the table is hand-labeled.

The same idea applies in several places:

• Word tokens in language models.
• Users and items in recommenders. See recommender systems.
• Categorical features in tabular networks.
• Patch IDs in vision transformers.

Choosing the size

The embedding dimension is a trade-off. More dimensions can hold more nuance, but they cost memory and search time. Gains often flatten out at a few hundred dimensions for retrieval. When you know the dimension you will serve, train at that size, rather than shrinking later.

If you normalize each vector to length 1, the dot product equals the cosine similarity. Many vector indexes assume that, so decide early.

Common mistakes

• Feeding an ID as a plain number. A model that gets 7 as a scalar assumes 7 lies between 6 and 8. Use an embedding instead.
• Ignoring unseen IDs. A new user or item has an untrained row. Give it a warm start from content, or a shared default.
• Using a huge table with few examples per ID. Most rows never train well. Hash or share rows when the ID space is huge.

Check yourself

1. If a batch contains only ID 3, why does row 7 get no gradient? Answer: the lookup only copies the rows it touched, so gradients flow only to those rows.
2. What shape comes out for 3 IDs and a dimension of 128? Answer: (3, 128).

Key takeaways

• An embedding layer is a learned table; a lookup returns one row per ID.
• Geometry comes from the training loss, not from hand-set values.
• Only looked-up rows get gradient, so rare IDs learn slowly.
• Normalize vectors when you want dot product to mean cosine.

Further learning

Next, see the classic way to learn word vectors from context in word2vec.