Word embeddings

1. Intuition

With TF-IDF, every word is its own axis, so all words are equally unrelated. Embeddings instead squeeze each word into a short dense vector (50 to 1,000 numbers) where direction carries meaning.

How do we find those vectors without labels? Self-supervision. Word2Vec slides a window over text and plays a prediction game:

• Skip-gram: given the center word, predict the words around it.
• CBOW (continuous bag of words): given the surrounding words, predict the center word.

To win, the model must give words that share contexts similar vectors, because similar vectors make similar predictions. "Cat" and "kitten" both appear near "pet," "fur," and "vet," so they end up close together.

Static embeddings give "bank" one vector whether it means a riverbank or a financial bank. Contextual embeddings from transformers compute a different vector for each occurrence, using the surrounding sentence. That is what modern NLP uses.

2. Visualization

Synthetic corpus where each group of words appears with its own context words. A count-based method (co-occurrence, positive PMI, then SVD to 2 dimensions) recovers the three groups: cosine similarity is 0.997 for cat and dog but −0.997 for cat and king. Real embeddings are learned from billions of words and have hundreds of dimensions.

3. The math

Symbols

Embedding lookup

A token ID $i$ selects row $i$ of $E$. Equivalently, multiply a one-hot vector by $E$: $e = \text{onehot}(i)^\top E$. Gradients during training update only the rows that were used.

Skip-gram with negative sampling

For each (center word $w$, true context word $c$) pair, plus $k$ random "negative" words $n_1..n_k$, maximize

That is binary classification: push the dot product up for real pairs and down for random pairs. Words that share contexts get pushed toward the same context vectors, so their embeddings become similar.

Count-based alternative

Build a word-word co-occurrence matrix, convert counts to positive pointwise mutual information, $\text{PPMI}(w, c) = \max\left(0, \ln\frac{P(w, c)}{P(w)P(c)}\right)$, and reduce dimensions with SVD. Word2Vec has been shown to implicitly factorize a matrix closely related to this one; the script uses this method because it needs no deep-learning library.

Worked example: similarity and analogy

Toy 3D embeddings: king $= [0.9, 0.8, 0.1]$, queen $= [0.9, 0.1, 0.8]$, man $= [0.1, 0.9, 0.0]$, woman $= [0.1, 0.1, 0.9]$.

Analogy. king − man + woman $= [0.9 - 0.1 + 0.1,\; 0.8 - 0.9 + 0.1,\; 0.1 - 0.0 + 0.9] = [0.9, 0.0, 1.0]$. Cosine with queen: dot $= 0.81 + 0 + 0.8 = 1.61$; norms $\sqrt{0.81 + 1} = 1.345$ and $\sqrt{0.81 + 0.01 + 0.64} = 1.208$; cosine $= 1.61 / 1.625 = 0.991$. The nearest word to the analogy vector is "queen." (These numbers are hand-picked to illustrate the idea; real analogies hold only approximately.)

4. Implementation

Build it from scratch

We train skip-gram with negative sampling in about 50 lines of NumPy, on ten hand-written sentences about cats, dogs, mats, and rugs (synthetic, repeated 20 times). Two matrices are learned: input vectors $E$ and output vectors $U$, both $[V, d]$. $E$ is the embedding table you keep.

Training pairs from a window. Every word within 2 positions of a center word becomes a (center, context) pair. These pairs are the only labels; no human labelled anything.

def skipgram_pairs(docs, vocab, window=2):
    pairs = []
    for doc in docs:
        for i, center in enumerate(doc):
            for j in range(max(0, i - window), min(len(doc), i + window + 1)):
                if i != j:
                    pairs.append((vocab[center], vocab[doc[j]]))
    return np.array(pairs)

Predict: the sentence "the cat sat on the mat" has 6 words. How many pairs does it give with window 2? (Edge words have fewer neighbors.)

Where negatives come from. Negatives are drawn from word frequencies raised to the 3/4 power, as in the original Word2Vec. The power flattens the distribution so "the" is not drawn almost every time.

def noise_dist(docs, vocab):
    counts = np.zeros(len(vocab))
    for doc in docs:
        for t in doc:
            counts[vocab[t]] += 1
    p = counts ** 0.75                        # flattens: rare words sampled more than their frequency
    return p / p.sum()

Predict: one word appears 81 times and another once. With the 3/4 power, how many times more likely is the frequent word to be drawn? (81 is 3⁴.)

One gradient step. The loss for a pair is $-\log\sigma(e_w \cdot u_c) - \sum_j \log\sigma(-e_w \cdot u_{n_j})$. Its derivative with respect to a score $s$ is $\sigma(s) - 1$ for the true pair and $\sigma(s)$ for each negative. Each error multiplies the other side's vector to give that side's gradient.

def sgd_step(E, U, w, c, negs, lr):
    e = E[w].copy()                           # [d]
    g_pos = sigmoid(e @ U[c]) - 1.0           # d(loss)/d(score) for the true pair: pushes score up
    g_neg = sigmoid(U[negs] @ e)              # [k] for the negatives: pushes scores down
    E[w] -= lr * (g_pos * U[c] + g_neg @ U[negs])
    U[c] -= lr * g_pos * e
    np.subtract.at(U, negs, lr * g_neg[:, None] * e)   # safe if a negative repeats
    return -np.log(sigmoid(e @ U[c]) + 1e-9) - np.log(1 - g_neg + 1e-9).sum()

Predict: only rows w, c, and the negatives change. With $V = 24$, $d = 16$, $k = 5$, how many of the $2 \times 24 \times 16 = 768$ parameters can one step touch at most?

The loop. Shuffle the pairs each epoch and draw $k = 5$ fresh negatives per pair.

def train(docs, vocab, d=16, window=2, k=5, epochs=60, lr=0.05, seed=42):
    rng = np.random.default_rng(seed)
    V = len(vocab)
    E = rng.normal(0, 0.1, (V, d))
    U = rng.normal(0, 0.1, (V, d))
    pairs, p = skipgram_pairs(docs, vocab, window), noise_dist(docs, vocab)
    for epoch in range(epochs):
        loss = 0.0
        for w, c in pairs[rng.permutation(len(pairs))]:
            loss += sgd_step(E, U, w, c, rng.choice(V, size=k, p=p), lr)
        if epoch in (0, epochs - 1):
            print(f"epoch {epoch + 1:3d}: mean loss per pair = {loss / len(pairs):.3f}")
    return E

Predict: "mat" and "rug" never appear in the same sentence. Will they end up close? Why?

Subwords for unseen words (the fastText idea). A static model has no vector for a word it never saw. fastText represents a word as the sum of vectors for its character n-grams (lengths 3 to 6, with < and > marking the word's start and end), so "kittenish" borrows from "kitten."

def char_ngrams(word, n_min=3, n_max=6):
    w = f"<{word}>"                           # boundary markers make prefixes and suffixes distinct
    return {w} | {w[i:i + n] for n in range(n_min, n_max + 1) for i in range(len(w) - n + 1)}

Predict: "" has 8 characters. How many n-grams of length 3 to 6 does it have, plus the whole word?

Use the library

In practice you train or load embeddings with a library (pip install gensim). sg=1 selects skip-gram, negative=5 sets $k$, and gensim adds what our loop skips: subsampling of frequent words, a learning rate that decays, and multithreaded C code.

from gensim.models import Word2Vec, FastText
model = Word2Vec(sentences=tokenized_sentences, vector_size=100, window=5, sg=1, negative=5, min_count=5)
model.wv.most_similar("kitten")
ft = FastText(sentences=tokenized_sentences, vector_size=100, window=5, min_n=3, max_n=6)
ft.wv["kittenish"]          # works even if the word was never seen

In a neural network the embedding layer is the same $[V, d]$ table, trained with the rest of the model. A lookup is row indexing:

V, d = 10_000, 128
E = np.random.default_rng(0).normal(0, 0.1, (V, d))   # [V, d], learned during training
vectors = E[np.array([12, 845, 3])]                   # [3, d]

For most applications today, use a pretrained contextual encoder instead (pip install sentence-transformers):

from sentence_transformers import SentenceTransformer
encoder = SentenceTransformer("all-MiniLM-L6-v2")
vecs = encoder.encode(["my cat sleeps", "a kitten naps"], normalize_embeddings=True)
print(vecs[0] @ vecs[1])     # high similarity despite no shared words

Run it

From guide/ (takes a few seconds):

uv run code/09-classical-nlp/word2vec_scratch.py

Key output:

V = 24 words, 2800 training pairs per epoch
epoch   1: mean loss per pair = 2.590
epoch  60: mean loss per pair = 1.485
E shape: (24, 16)
nearest to cat : [('dog', 0.988), ('mouse', 0.831), ('mat', 0.569)]
nearest to mat : [('rug', 0.993), ('slept', 0.785), ('sat', 0.775)]
'kitten' has 19 n-grams; shares 14 with 'kittens', e.g. ['<ki', '<kit', '<kitt', '<kitte']
'kittenish' (never seen) still shares 14 n-grams with 'kitten'

"mat" and "rug" end up with cosine 0.993 without ever sharing a sentence, because both follow "on the." That is the distributional hypothesis, produced by 60 epochs of a binary classifier. The loss stays well above zero because the negatives are random and the corpus is tiny.

The figure script (count-based PPMI + SVD embeddings and the TF-IDF comparison) is code/09-classical-nlp/nlp.py.

5. Engineering

Uses. Features for downstream models; semantic search and RAG (embed documents and queries, retrieve by cosine similarity); clustering and deduplication; recommendation (user and item embeddings).

Choosing an embedding model. Match the domain and language; check the maximum input length; benchmark on your retrieval or similarity task rather than trusting leaderboards; consider dimension (storage and speed) and whether vectors are normalized.

Bias. Embeddings learned from web text encode social biases present in that text (for example, gender associations with occupations). Audit before using them in decisions about people.

Watch out

Failure modes. Out-of-vocabulary words with static embeddings (subword models fix this); ambiguous words averaged into one vector; mixing vectors from different embedding models in one index; re-embedding documents with a new model but not the queries (or vice versa).

Common mistakes

• Reading 2D projections of embeddings as exact geometry.
• Comparing unnormalized vectors with dot products when the model expects cosine.
• Assuming analogies always work; they are approximate and cherry-picked in demos.

6. Knowledge check

Practice exercise. An embedding matrix has a vocabulary of 50,000 tokens and dimension 768. How many parameters does it have, and what is the shape of the output for a batch of 8 sequences of 128 tokens?

Implementation challenge. Turn code/09-classical-nlp/word2vec_scratch.py into CBOW: average the input vectors of the context words within the window, and predict the center word with the same negative-sampling loss. Compare the nearest neighbors of "cat" and "mat" with skip-gram's.

Summary

• Embeddings are dense vectors where words with similar contexts are close, learned by self-supervised prediction (skip-gram, CBOW) or count factorization.
• An embedding layer is a lookup table of shape $[V, d]$.
• Cosine similarity compares embeddings; analogies are a fun but approximate property.
• Static embeddings give one vector per word; contextual (transformer) embeddings give one per occurrence.

Next: Named entity recognition

Related: From text to vectors · Self-attention · Model card: embedding model

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Word Embeddings — Word2Vec from Scratch, GloVe, FastText, and Subword Embeddings.

Interview angle