word2vec

You can learn a lot about a stranger from their friends. If two people keep turning up with the same crowd, they probably have something in common. Linguists put it the same way: "You shall know a word by the company it keeps" (J. R. Firth, 1957). "Coffee" and "tea" both show up next to "cup", "hot" and "morning", so they must be related.

word2vec turns that idea into a game. Every word gets a spot on a map of meaning (a short list of numbers, its vector). The game: looking only at a word's spot, guess which words sit near it in real sentences. Every wrong guess nudges the spots a little. After millions of nudges, words that keep the same company end up in the same neighbourhood, because that's the only way to guess well for all of them.

Step 1, make guessing pairs. Take the sentence "river bank fish" and a window of 1 (look one word left and right). Each word, the center, is paired with each neighbour, its context:

(river, bank), (bank, river), (bank, fish), (fish, bank)

No human labels anything: the text itself says which words appeared together.

Step 2, score a pair. Put everything in 2 dimensions. The center word has vector v = (1, 0), its true context word has u = (0, 1), and one random "noise" word (a word we pretend is a context, to have something to push away) has uₙ = (0, −1). The score of a pair is their dot product (multiply matching positions, add): v·u = 1·0 + 0·1 = 0, and v·uₙ = 0. Zero means "no opinion".

Step 3, turn scores into probabilities with the sigmoid σ(x), which squashes any number into 0 to 1 (σ(0) = 0.5, big positive → near 1, big negative → near 0). The model says the true pair is 50% likely and the noise pair is 50% likely. That's bad: we want ~100% and ~0%.

Step 4, measure the mistake as a loss (a single number that's big when the guesses are bad): −ln σ(0) − ln σ(−0) = 0.693 + 0.693 = 1.386.

Step 5, nudge. Move v a little toward u and away from uₙ, and move u and uₙ a little toward or away from v. With a step size (learning rate) of 0.1 we get v = (1, 0.1), u = (0.05, 1), uₙ = (−0.05, −1). Now v·u = 0.15 and v·uₙ = −0.15: the true pair scores higher, the noise pair lower, and the loss drops to 1.242. Word2vec is this nudge, repeated millions of times.

flowchart LR
  S["river bank fish"] --> W[Slide a window<br/>over each word]
  W --> P["(bank, river)<br/>(bank, fish) ..."]
  P --> T[Raise the score of<br/>each true pair]
  N[Draw k random<br/>noise words] --> T2[Lower the score of<br/>each noise pair]
  T --> U[Nudge the vectors<br/>in each pair]
  T2 --> U
  U -->|next pair| P

In code: skipgram_pairs slides the window over a sentence and returns every (center, context) pair; build_corpus writes the lesson's synthetic sentences for it to slide over.

Every word has two vectors: vᵥ when it's the center and uᵥ when it's a context. For a center c, its true context o and k noise words n₁ … nₖ:

On the example: L = −log σ(0) − log σ(−0) = 0.693 + 0.693 = 1.386; after the nudge, −log σ(0.15) − log σ(0.15) = 0.621 + 0.621 = 1.242.

In Python:

import math
def sigma(x):
    # σ(x) = 1 / (1 + e^(-x))
    return 1 / (1 + math.exp(-x))
def dot(a, b):
    return sum(a_i * b_i for a_i, b_i in zip(a, b))
def L(v_c, u_o, noise):
    # -log σ(u_o · v_c)
    return (-math.log(sigma(dot(u_o, v_c)))
            # - Σ_i log σ(-u_ni · v_c)
            - sum(math.log(sigma(-dot(u_n, v_c))) for u_n in noise))
round(L((1, 0), (0, 1), [(0, -1)]), 3)  # → 1.386
# after the nudge
round(L((1, 0.1), (0.05, 1), [(-0.05, -1)]), 3)  # → 1.242

The nudge follows the gradient: for each vector, the direction in which a small change would increase the loss fastest. We step the opposite way. For this loss the gradients are short enough to derive by hand (sgns_grads), using the fact that the slope of −log σ(x) is σ(x) − 1:

On the example: g_o = 0.5 − 1 = −0.5 and g₁ = 0.5, so ∂L/∂vc = −0.5·(0, 1) + 0.5·(0, −1) = (0, −1). Stepping against it with learning rate 0.1: vc = (1, 0) − 0.1·(0, −1) = (1, 0.1).

In Python:

import math
def sigma(x):
    return 1 / (1 + math.exp(-x))
v_c, u_o, u_n = (1, 0), (0, 1), (0, -1)
# σ(u_o · v_c) - 1
g_o = sigma(sum(u * v for u, v in zip(u_o, v_c))) - 1
# σ(u_n1 · v_c)
g_1 = sigma(sum(u * v for u, v in zip(u_n, v_c)))
g_o, g_1  # → (-0.5, 0.5)
# ∂L/∂v_c = g_o u_o + Σ_i g_i u_ni
grad = [g_o * o + g_1 * n for o, n in zip(u_o, u_n)]
grad  # → [0.0, -1.0]
# step against it, learning rate 0.1
[v - 0.1 * g for v, g in zip(v_c, grad)]  # → [1.0, 0.1]

Noise words are drawn in proportion to their count raised to the 3/4 power:

On an example: counts 1 and 16 become 1^0.75 = 1 and 16^0.75 = 8, so the shares are 1/9 and 8/9 (0.111 and 0.889) instead of 1/17 and 16/17 (0.059 and 0.941).

In Python:

counts = [1, 16]
# count(w)^0.75
weights = [count ** 0.75 for count in counts]
weights  # → [1.0, 8.0]
# divide by Σ over w′ so the shares add to 1
[round(w / sum(weights), 3) for w in weights]  # → [0.111, 0.889]
# without the 3/4 power
[round(c / sum(counts), 3) for c in counts]  # → [0.059, 0.941]

In code: sgns_loss computes L for one center, one context and k noise words, sgns_update takes one step against the gradient, and noise_distribution builds P(w). train_sgns runs the whole loop in mini-batches and returns a WordVectors table; gradient_check confirms the hand-derived gradients against a numerical estimate.

Why it matters a full softmax over the vocabulary (primer.ml.attention explains softmax) would score every word, 100,000 or more, for every training pair. Negative sampling scores k + 1. That trick is what made training on billions of words practical in 2013, and the same pull-together, push-apart shape reappears in every modern embedding model (primer.ml.embeddings.contrastive).

Everyday picture on a street map, "walk two blocks east" is the same walk wherever you start. word2vec's map works the same way: the walk from "man" to "woman" is roughly the same walk as from "king" to "queen", because in both cases the contexts change from he/his to she/her.

Tiny example with made-up 2-D vectors: man = (1, 0), woman = (1, 1), king = (3, 0). Then king − man + woman = (3 − 1 + 1, 0 − 0 + 1) = (3, 1), exactly where queen = (3, 1) would be.

On the example: v = (3, 1) and queen = (3, 1) point the same way, so the cosine is 1, and queen wins.

In Python:

import math
def cos(a, b):
    dot = sum(a_i * b_i for a_i, b_i in zip(a, b))
    return dot / (math.sqrt(sum(a_i ** 2 for a_i in a)) * math.sqrt(sum(b_i ** 2 for b_i in b)))
v = {"man": (1, 0), "woman": (1, 1), "king": (3, 0), "queen": (3, 1)}
a, b, c = "king", "man", "woman"
# v_a - v_b + v_c
target = [x_a - x_b + x_c for x_a, x_b, x_c in zip(v[a], v[b], v[c])]
target  # → [3, 1]
# w ∉ {a, b, c}
candidates = [w for w in v if w not in {a, b, c}]
# arg max of the cosine
w_hat = max(candidates, key=lambda w: cos(v[w], target))
w_hat, round(cos(v[w_hat], target), 2)  # → ('queen', 1.0)

This lesson trains on a small synthetic corpus built from three independent attributes (male/female, royal/common, adult/child), so the effect appears in seconds on a laptop.

In code: analogy returns the word nearest v_a − v_b + v_c, skipping the three inputs; nearest lists a word's closest neighbours by cosine, and nearest_to_vector does the same for any vector.

Why it matters this was the first clear evidence that learned vectors capture meaning as geometry. It's why "embedding" became the default way to represent text, and why vector arithmetic (averaging, subtracting, finding nearest neighbours) works on meaning.

Everyday picture instead of playing the guessing game, keep a tally sheet: for every pair of words, count how often they appear near each other. Then ask which pairs turn up together more than chance would predict.

Tiny example two words that each appear only with themselves: counts [[2, 0], [0, 2]]. Out of 4 sightings, each word appears half the time (P = 0.5), and each diagonal pair also appears half the time. Chance would predict 0.5 · 0.5 = 0.25, but we see 0.5, twice as often, so the score is log(0.5 / 0.25) = log 2 = 0.693. The off-diagonal pairs never co-occur, and their score is clipped to 0.

On the example: PMI = log(0.5 / (0.5 · 0.5)) = log 2 = 0.693 on the diagonal; the off-diagonal pairs have P(w, c) = 0, so PMI = −∞ and PPMI = 0.

In Python:

import math
# co-occurrence counts
X = [[2, 0], [0, 2]]
# 4 sightings
total = sum(sum(row) for row in X)
# P(w): share of each row
P_w = [sum(row) / total for row in X]
# P(c): share of each column
P_c = [sum(col) / total for col in zip(*X)]
def ppmi(w, c):
    P_wc = X[w][c] / total
    # log 0 = -∞
    pmi = math.log(P_wc / (P_w[w] * P_c[c])) if P_wc else -math.inf
    # PPMI = max(PMI, 0)
    return max(pmi, 0.0)
[[round(ppmi(w, c), 3) for c in range(2)] for w in range(2)]  # → [[0.693, 0.0], [0.0, 0.693]]

The PPMI table has one row per word, as many columns as the vocabulary, and is mostly zeros. SVD (singular value decomposition, see primer.notation) compresses it into a few columns that keep its main patterns; those few numbers per row are the word's embedding.

flowchart LR
  T[Text] --> C[Count co-occurrences<br/>in a window]
  C --> P[PPMI: keep pairs that meet<br/>more often than chance]
  P --> S[SVD: compress each row<br/>to d numbers]
  S --> E[Word vectors]

In code: cooccurrence counts the pairs in a window, ppmi rescales the table by chance and clips negatives to 0, and ppmi_svd_embeddings runs all three steps and keeps d columns of the SVD.

GloVe (Pennington et al., 2014) is the best-known counting method. It fits vectors so that each dot product predicts the log of the co-occurrence count:

On an example: if "ice" and "cold" appear together 20 times, training nudges the vectors until w_ice · w̃_cold + b_ice + b̃_cold ≈ log 20 = 3.0. With 2-D vectors w_ice = (1, 1) and w̃_cold = (1, 0.5) and offsets 0.25 each, the left side is 1.5 + 0.25 + 0.25 = 2.0, still 1.0 short, so training keeps pushing the two vectors to agree more.

In Python:

import math
# log X_ij: the target for 20 co-occurrences
round(math.log(20), 1)  # → 3.0
w_ice, w_cold_tilde = (1, 1), (1, 0.5)
b_ice, b_cold_tilde = 0.25, 0.25
# w_i · w̃_j + b_i + b̃_j
left = sum(a * b for a, b in zip(w_ice, w_cold_tilde)) + b_ice + b_cold_tilde
# the left side, and how far it still falls short
left, round(math.log(20) - left, 1)  # → (2.0, 1.0)

Why it matters counting and predicting are two routes to the same geometry. That's reassuring: embeddings aren't magic, they're compressed co-occurrence statistics.

Everyday picture a phone book lists one address per name. If two different people are called "Bank", the book can only print one address, somewhere between the two. A static embedding table has the same problem: "bank" gets one vector for both the riverbank and the bank account.

Tiny example after training on this lesson's corpus, the single "bank" vector has cosine about 0.52 with river words and 0.50 with money words: it sits halfway between.

The fix is context. A transformer doesn't look words up once and stop. It computes a fresh vector for every word in every sentence by attending over the neighbours (primer.ml.attention): the new vector is a weighted average of the sentence's vectors,

On the example: in "river bank fish", "bank" blends in "river" and "fish", and its cosine with the river words climbs from 0.52 to about 0.83. The same move by hand, in 2-D where the first axis is "river" and the second is "money": river = (1, 0), bank = (1, 1), fish = (1, 0), and α = (0.25, 0.5, 0.25). Then bank′ = 0.25·(1, 0) + 0.5·(1, 1) + 0.25·(1, 0) = (1, 0.5), and its cosine with the river axis climbs from 0.71 to 0.89.

In Python:

import math
# river, bank, fish: axis 1 is "river", axis 2 "money"
x = [(1, 0), (1, 1), (1, 0)]
# attention weights, Σ_j α_j = 1
alpha = [0.25, 0.5, 0.25]
# Σ_j α_j x_j
bank_new = [sum(a_j * x_j[i] for a_j, x_j in zip(alpha, x)) for i in range(2)]
bank_new  # → [1.0, 0.5]
def cos_with_river(v):
    # cosine with (1, 0)
    return v[0] / math.sqrt(v[0] ** 2 + v[1] ** 2)
round(cos_with_river(x[1]), 2), round(cos_with_river(bank_new), 2)  # → (0.71, 0.89)

flowchart LR
  subgraph Static["Static (word2vec, GloVe)"]
    B1["bank"] --> V1["one vector,<br/>both senses blended"]
  end
  subgraph Contextual["Contextual (transformer)"]
    R["river bank fish"] --> A1[attention mixes<br/>in neighbours] --> V2["bank ≈ riverside"]
    M["bank loan cash"] --> A2[attention mixes<br/>in neighbours] --> V3["bank ≈ finance"]
  end

Why it matters this is why BERT (2018) and every embedding model since are built on transformers. Search, RAG and classification all need "Python the language" and "python the snake" to be different points on the map.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.