word2vec from Scratch

Week 1 of the residency. How to turn free text into numbers a model can use, starting from the feature engineering problem that makes you need it.

Week 1 didn’t start with word2vec. It started with a boring problem. word2vec was where we ended up when the boring problem stopped working.

The problem: predict someone’s income from what you know about them.

A neural network takes numbers in and gives a number out. The job is the caption under that diagram. Turn what you know about a person into numbers.

Some of it’s easy.

Gender is a flag. Age is already a number. Height in centimetres. Weight in kilograms. Four features, no thought required.

Then education. The easy run stops.

Encoding Education

There are four options on that slide.

Ordinal. None is 0, high school 1, bachelor 2, master 3, PhD 4. This encodes the ordering. It also tells the model that a PhD is four times a high school diploma. Nobody believes that.

Years of schooling. 0, 12, 15, 17, 21. Better, because the spacing means something. But you had to know the answer. You supplied the domain knowledge, not the model.

One-hot. Five slots, one of them set to 1. It makes no claims about ordering or distance. The cost is a slot per category. It also makes every pair equally far apart, so a master’s and a PhD are as different as a PhD and no schooling.

Learned embeddings. Give each level a short list of numbers. Start them random. Let training move them. No assumptions, no ordering imposed. The model decides what the numbers mean.

The fourth option is what this week is about. At this scale it looks like overkill. There are five categories. One-hot would do.

Encoding City

City has the same problem as education. No ordering, and a thousand categories instead of five. One-hot is now a thousand slots of mostly zeros.

There is a way out.

Break the category into features that describe it. A city is its population, its average age, its average income, its crime rate. Four numbers instead of a thousand. Each one means something. Two similar cities end up with similar numbers.

This is good feature engineering. A lot of real models are built this way. It also has a ceiling, and the next feature is above it.

Free Text

Now the form has a box that says tell us about yourself. Someone has written a paragraph.

There is no decomposition to reach for. Cities have populations. Paragraphs do not have an agreed set of four numbers. The categories are unbounded, because anyone can write anything. Every trick from the last three slides is gone.

The red question on that slide is the subject of week 1:

How can we produce meaningful numbers from free text?

The Distributional Hypothesis

What is the meaning of bardiwac?

He handed her her glass of bardiwac.

Beef dishes are made to complement the bardiwacs.

Nigel staggered to his feet, face flushed from too much bardiwac.

I dined off bread and cheese and this excellent bardiwac.

You have never seen that word before. After four sentences you know it is a drink. Probably a red wine. Probably alcoholic.

Nobody defined it for you. You got it from the company it keeps.

This is the distributional hypothesis. A word can be described by the words that appear around it. If that holds, you can compute a word’s meaning from a large pile of ordinary text. Nobody has to label anything.

The learned-embedding option is still the plan. What changed is that there is now a way to train it without labels.

The word2vec Recipe

Three steps. The second one is the one people skip.

• Slide a five-word window over the text
• Hide the middle word
• Predict it from the other four

That gives you an enormous number of training examples from raw text, with no annotation.

Step two is the important bit. “For dinner we served dark red ___ with steak” does not have one right answer. It could be bardiwac, wine, merlot or malbec. So the model is not trained to output one word. It is trained to output a score for every word in the vocabulary.

That is the part that makes it work. Words that fit the same gaps end up with similar scores, and similar scores pull their vectors together.

There are two ways to arrange that prediction.

CBOW takes the context and predicts the middle word. Four words go in. Each is looked up in a table of vectors. The four vectors are averaged into one. That one vector is used to score every word in the vocabulary.

Skip-gram runs it the other way. One word in, predict each of its neighbours.

I trained both.

One thing to notice: that sentence, Hello my name is Bes. It comes back in week 3 going through a transformer, in week 5 being transcribed by Whisper, and in week 6 being scored by a reward model. It is the course’s running example. By the end you have seen the same five words pass through six architectures.

Negative Sampling

Step two is expensive.

To turn raw scores into probabilities you need a softmax. Exponentiate each score, then divide by the total so they add up to one. The dividing is the problem. The total is over the whole vocabulary. Here that is 71,290 words, after dropping everything seen fewer than five times.

So every training step touches 71,290 words to learn one thing.

Negative sampling changes the question. Instead of asking which of these 71,290 words goes here, ask is this pair real, or did I make it up. Once for the true neighbour, then five more times for words picked at random. Six comparisons instead of 71,290.

Where the random words come from matters:

• Pick uniformly and nearly every one is a rare word the model already scores near zero. It learns nothing.
• Pick by raw frequency and nearly all of them are the.
• Mikolov’s papers raise the frequencies to the power 3/4, which sits between the two.

Training It

The rest of this page trains that model. It is small enough to finish while the page builds: two million characters of text, a fiftieth of the real run, 64 numbers per word.

import collections

import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download

REVISION = "9113ea48905b4f7178b333919ed5ec1a474561d7"
path = hf_hub_download("roshbeed/ai-residency-blog-data", "text8/text8-2m.txt",
                       repo_type="dataset", revision=REVISION)

words = open(path).read().split()
counts = collections.Counter(words)
vocab = [w for w, c in counts.most_common() if c >= 10]
stoi = {w: i for i, w in enumerate(vocab)}
ids = np.array([stoi[w] for w in words if w in stoi], dtype=np.int64)

print(f"{len(words):,} tokens, {len(counts):,} distinct, {len(vocab):,} kept")

336,026 tokens, 26,147 distinct, 3,704 kept

One more step before training. the appears about 70,000 times and sits next to everything, so it tells you nothing about its neighbours. Common words get thrown away, and the more common a word is the more often it goes.

frequency = np.bincount(ids, minlength=len(vocab)).astype(np.float64)
frequency /= frequency.sum()

t = 1e-3
keep_probability = np.minimum(1.0, np.sqrt(t / frequency))
rng = np.random.default_rng(0)
kept = ids[rng.random(len(ids)) < keep_probability[ids]]

print(f"'the' survives with probability {keep_probability[stoi['the']]:.3f}, "
      f"'philosophy' with {keep_probability[stoi['philosophy']]:.3f}")

WINDOW = 5
centres, contexts = [], []
for offset in range(1, WINDOW + 1):
    centres.append(kept[offset:]);  contexts.append(kept[:-offset])
    centres.append(kept[:-offset]); contexts.append(kept[offset:])

centre = torch.from_numpy(np.concatenate(centres))
context = torch.from_numpy(np.concatenate(contexts))
noise = torch.from_numpy(frequency ** 0.75 / (frequency ** 0.75).sum())
print(f"{len(centre):,} training pairs")

'the' survives with probability 0.115, 'philosophy' with 1.000
1,839,810 training pairs

There are two tables of vectors. One holds a word’s vector when it’s the centre of a window. The other holds it when it’s somebody’s neighbour. Only the first is kept at the end. The second exists to give the first something to be scored against.

The loss below pushes the true pair up and five invented pairs down.

DIMENSIONS, K = 64, 5
generator = torch.Generator().manual_seed(0)

centre_vectors = (torch.randn(len(vocab), DIMENSIONS, generator=generator) * 0.01).requires_grad_()
context_vectors = torch.zeros(len(vocab), DIMENSIONS, requires_grad=True)
optimiser = torch.optim.Adam([centre_vectors, context_vectors], lr=2e-3)


def loss_on(centre_batch, context_batch):
    """One real pair scored up, K invented pairs scored down."""
    v = centre_vectors[centre_batch]
    real = F.logsigmoid((v * context_vectors[context_batch]).sum(-1))
    fake_ids = torch.multinomial(noise, len(centre_batch) * K,
                                 replacement=True, generator=generator)
    fake = context_vectors[fake_ids.view(len(centre_batch), K)]
    invented = F.logsigmoid(-(fake @ v.unsqueeze(-1)).squeeze(-1)).sum(-1)
    return -(real + invented).mean()


with torch.no_grad():
    baseline = loss_on(centre[:8192], context[:8192]).item()
print(f"loss before training: {baseline:.3f}, and (1 + {K}) * ln 2 = {(1 + K) * np.log(2):.3f}")

loss before training: 4.159, and (1 + 5) * ln 2 = 4.159

Those two numbers matching is a useful check. A model that knows nothing is guessing on six yes/no questions. Each costs ln 2. If a run starts anywhere else, the setup is wrong, not the model.

EPOCHS, BATCH = 25, 8192
history = []

for epoch in range(EPOCHS):
    order = torch.randperm(len(centre), generator=generator)
    total = steps = 0
    for i in range(0, len(order) - BATCH, BATCH):
        batch = order[i:i + BATCH]
        loss = loss_on(centre[batch], context[batch])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
        total += loss.item()
        steps += 1
    history.append(total / steps)

embeddings = F.normalize(centre_vectors.detach(), dim=1)


def nearest(word, k=6):
    similarity = embeddings @ embeddings[stoi[word]]
    similarity[stoi[word]] = -1
    return [vocab[i] for i in similarity.topk(k).indices.tolist()]


print(f"loss {baseline:.3f} -> {history[-1]:.3f}\n")
for word in ("king", "france", "computer", "three", "war", "music"):
    print(f"{word:10} {', '.join(nearest(word))}")

loss 4.159 -> 2.213

king       alexander, darius, philip, macedon, plutarch, kings
france     germany, spain, morocco, italy, brazil, portugal
computer   software, windows, hardware, extension, animation, processing
three      two, four, five, one, eight, zero
war        civil, defeat, loyalists, battle, occupation, generals
music      dance, folk, literature, radio, opera, publishing

That’s the answer to the red question. Every word now has 64 numbers attached. The numbers put related words near each other. Nothing was labelled to make it happen.

three sits with the other small numbers because it turns up where they turn up.

To turn a paragraph into numbers, look up each word and average. Crude, and enough to close the loop we opened with:

The free-text box goes in. Every feature on that form is now a number.

The full run uses fifty times this much text. It scores at or above the published word2vec numbers on the same data:

benchmark	untrained	skip-gram	CBOW	published
WordSim-353	0.028	0.726	0.730	0.68
SimLex-999	0.008	0.297	0.290	0.30
Google analogies	0.000	0.464	0.408	0.38

The untrained column is not decoration. A checkpoint in this project once reported a loss of 10.02. That reads like a trained model. An untrained one reports 11.09.

Checking the Vectors

The famous claim is that directions in this space mean something. Take king, subtract man, add woman, and you should land near queen.

query	top 3
bigger - big + small	smaller, larger, large
walked - walk + run	ran, running, runs
paris - france + italy	bologna, turin, rome
king - man + woman	throne, ermengarde, anjou

Three work. The famous one does not.

The reason is about the data, not the method. This text is 17 million words of Wikipedia. king appears mostly in lists of monarchs and succession prose. So its vector leans towards monarchy and lineage, not towards male. Push it along the gender direction and you land on the nearest female role in that same context. queen comes fifteenth.

The gender direction is there. father : mother :: son : daughter lands first, and so do three other family analogies. It is king in particular that is not mostly about being a man here.

Conclusion

Week 1 starts with a feature engineering problem and ends with a way to solve it.

• Some features are already numbers
• Some categories can be decomposed into numbers by hand
• Free text can be neither, so you learn the numbers instead
• The training signal is free, because the text supplies its own labels

The second project of week 1 is the income model with a real dataset. Predict how many upvotes a Hacker News post gets, from its title, its timestamp, its link and its author. A number to predict, some easy features, and one free-text field.

These vectors fill that field in. That post is here.

The window that makes this work is also what limits it. CBOW and skip-gram only ever see a few words either side. Replace that fixed window with attention over the whole sequence and you get BERT and GPT. That is where week 3 goes.

The full project, with the sweep and the deployed API, is on GitHub.