Embeddings

Problem

Text arrives as symbols: "cat", "dog", "car". Symbols have no arithmetic. You can't add "cat" to "dog", and the number you'd get by assigning each word an ID (cat = 17, dog = 18, car = 19) is worse than useless, because it claims that cat and dog are close and that car is just one step further, for no reason except alphabetical luck. A model fed raw IDs would try to learn from that false geometry.

The alternative that avoids false geometry, giving each word its own on/off switch (a "one-hot" vector with 6 slots for this page's tiny vocabulary, 50,257 for GPT-2), treats every pair of words as equally unrelated and wastes huge amounts of memory on zeros. What we want instead is a short dense vector for each symbol in which similar symbols get similar vectors, and we want the model to find those vectors itself.

Intuition

Think of a spreadsheet with one row per word and a handful of columns. Looking a word up means reading its row. That's the whole mechanism. The rows start out random, and training nudges them: every time a word's row helps the model make a better prediction, the row moves a little in the direction that helped. Words that appear in the same kinds of sentences get nudged in the same directions, so their rows drift together.

Linguists put it as "you shall know a word by the company it keeps". "Cat" and "dog" are both followed by "eats" and "chases", so after training their rows are close. "Car" is followed by "drives" and "needs", so its row lands elsewhere. Nobody wrote down that cats and dogs are animals; the geometry fell out of the statistics.

Highlighted rows are the ones this sentence reads. Nothing is computed: each output row is a copy of a table row, so the layer costs memory and index arithmetic, not multiplications.

Mechanics

Lookup equals one-hot times matrix

Let E be the embedding matrix with V rows (the vocabulary) and d columns (the embedding dimension). For a token with index i, its one-hot vector e_i has a single 1 in position i, and

embed(i) = e_i^T E = E[i, :]    shapes: [1, V] @ [V, d] → [1, d]

The matrix product with a one-hot row just picks out row i, so a linear layer without a bias fed one-hot inputs is an embedding layer, and nn.Embedding is that layer without the wasted multiplications by zero. For a batch of index tensors of any shape (*), the output is (*, d). A batch of 2 sentences of 6 tokens becomes a [2, 6, 3] tensor.

What the gradient does

Because the forward pass copies rows, the backward pass adds into rows. If upstream gradient G_t (shape [d]) arrives at the output for position t whose token is i_t, then

∂L/∂E[j, :] = Σ_{t : it = j} G_t

Rows of tokens that do not appear in the batch get zero gradient, and a token appearing three times gets the sum of three gradients. That's why the gradient of a 50,000-row table is mostly zeros, why sparse=True exists (the gradient is stored as just the touched rows, but only some optimizers accept it), and why scale_grad_by_freq=True divides each row's gradient by how often its token occurs in the mini-batch, so that frequent words don't get most of the movement. padding_idx blocks the gradient for one row, which is how padding tokens stay at the all-zero vector they're initialized to.

Learning geometry: word2vec

Nothing above says why similar words end up close. That comes from the training task. The skip-gram model of word2vec (Mikolov et al., 2013) gives every word two vectors, v_w as a center word and u_w as a context word, and trains them so that a center word predicts the words within a window around it:

p(o | c) = exp(u_o · v_c) / Σ_w exp(u_w · v_c),   L = −Σ_(c,o) log p(o | c)

This is softmax regression (see logistic regression and loss functions) where the "features" are the learned vector v_c. The gradient on v_c is Σ_w (p(w|c) − 1[w = o]) u_w, and the gradient on each u_w is (p(w|c) − 1[w = o]) v_c. Two words that are followed by the same neighbours receive nearly the same pushes, so their vectors converge. In practice the full softmax over a vocabulary of a million words is too slow, and the paper replaces it with hierarchical softmax or negative sampling (the follow-up paper), which updates only a few sampled wrong answers per step. The objective is the same idea.

Each point is the learned center vector of one word, in two dimensions, after training on 78 (center, context) pairs from 13 three-word sentences. Real embeddings use 100 to 12,288 dimensions, so they can keep many such groupings apart at once.

Static versus contextual

An embedding table gives one vector per word, whatever sentence it's in, so "bank" gets the same row in "river bank" and "bank loan". These are static embeddings (word2vec, GloVe, fastText). Modern language models use the table only as the first layer: transformer layers then mix each token's vector with its neighbours', so the vector at the top of the stack is contextual: it depends on the whole sentence (ELMo, BERT, GPT). A sentence-embedding model pools those contextual vectors into one vector per sentence or document, which is what retrieval systems search.

Scaling and weight tying

Two details from the original Transformer: the embedding output is multiplied by √d_model, so its scale is comparable to the added position signal, and the input table is shared with the final output projection (pre-softmax linear layer). Sharing works because the output layer computes a dot product between the final hidden vector and one vector per word, which is the same shape as the input table (Press and Wolf, 2017). For GPT-2 small (V = 50,257, d = 768) the table has 38,597,376 parameters, so tying saves 38,597,376 of an untied 163,037,184, leaving 124,439,808.

Worked example

Implementation

import torch
import torch.nn as nn

emb = nn.Embedding(num_embeddings=6, embedding_dim=3, padding_idx=0)
emb.weight[0]                         # tensor([0., 0., 0.]), zeroed at construction
ids = torch.tensor([[1, 2, 3, 0, 0, 0]])
out = emb(ids)                        # shape [1, 6, 3]

# the same thing as a matrix product with one-hot rows
onehot = torch.nn.functional.one_hot(ids, num_classes=6).float()   # [1, 6, 6]
torch.allclose(onehot @ emb.weight, out)                           # True

out.sum().backward()
emb.weight.grad[0]                    # zeros: padding_idx gets no gradient
emb.weight.grad[1]                    # tensor([1., 1., 1.])

# weight tying for a language model head
head = nn.Linear(3, 6, bias=False)    # weight shape [6, 3], same as the table
head.weight = emb.weight              # one Parameter, used twice

# start from existing vectors, frozen by default
pretrained = nn.Embedding.from_pretrained(torch.randn(6, 3), freeze=True)

Skip-gram with a full softmax, in plain Python:

import math, random

def skipgram(sentences, dim=2, window=2, lr=0.05, epochs=300, seed=0):
    vocab = sorted({w for s in sentences for w in s})
    ix = {w: i for i, w in enumerate(vocab)}
    rng = random.Random(seed)
    V = [[rng.gauss(0, 0.5) for _ in range(dim)] for _ in vocab]   # center vectors
    U = [[rng.gauss(0, 0.5) for _ in range(dim)] for _ in vocab]   # context vectors
    pairs = [(ix[s[i]], ix[s[j]]) for s in sentences for i in range(len(s))
             for j in range(max(0, i - window), min(len(s), i + window + 1)) if j != i]
    for _ in range(epochs):
        for c, o in pairs:
            logits = [sum(a * b for a, b in zip(V[c], u)) for u in U]
            top = max(logits)
            exps = [math.exp(x - top) for x in logits]
            probs = [e / sum(exps) for e in exps]
            grad_v = [0.0] * dim
            for k, u in enumerate(U):
                g = probs[k] - (k == o)
                for d in range(dim):
                    grad_v[d] += g * u[d]
                    u[d] -= lr * g * V[c][d]
            for d in range(dim):
                V[c][d] -= lr * grad_v[d]
    return dict(zip(vocab, V))

Tradeoffs

• Memory is the cost. V × d parameters, plus the optimizer state for each. A 50,257 × 768 float32 table is 154 MB, and Adam keeps two more copies. Recommendation models with hundreds of millions of IDs spend most of their memory here, which is why they shard tables and use sparse=True gradients with optimizers that support them.
• Out-of-vocabulary words. A word-level table has no row for an unseen word. Subword tokenization (BPE and relatives) fixes this by making the vocabulary pieces of words; fastText instead sums character n-gram vectors.
• Dimension. Too small and distinct meanings collide; too large and rarely seen rows stay close to their random start. Typical sizes run from 64 for small models to over 10,000 for the largest language models.
• Static versus contextual. Static vectors are cheap, interpretable and reusable across tasks, but conflate senses of a word. Contextual vectors resolve that at the cost of a full forward pass of a large network.
• Weight tying saves a whole table and often helps perplexity in small models, but ties the input and output geometries together, which large models sometimes prefer to keep separate.
• Pitfalls. Indices must be in [0, V): an out-of-range ID raises an error on CPU and can silently misbehave or crash on GPU. Learned geometry also reproduces the biases in the training text; "nearest neighbour" is not "true meaning".
• Compared with one-hot. One-hot is exact and needs no training but has V dimensions and no similarity structure; an embedding is a learned low-rank factorization of that input layer.

Check yourself

Connections

References

• Mikolov, Chen, Corrado and Dean. Efficient Estimation of Word Representations in Vector Space. 2013 (arXiv 1301.3781).
• Mikolov, Sutskever, Chen, Corrado and Dean. Distributed Representations of Words and Phrases and their Compositionality. NeurIPS, 2013 (arXiv 1310.4546; negative sampling).
• Pennington, Socher and Manning. GloVe: Global Vectors for Word Representation. EMNLP, 2014.
• Press and Wolf. Using the Output Embedding to Improve Language Models. EACL, 2017 (arXiv 1608.05859).
• Peters et al. Deep Contextualized Word Representations. NAACL, 2018 (arXiv 1802.05365; ELMo).
• Devlin, Chang, Lee and Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. 2018 (arXiv 1810.04805).
• PyTorch v2.15.0: torch/nn/modules/sparse.py and F.embedding.