Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
413 changes: 413 additions & 0 deletions educational_material/embeddings/embeddings.py

Large diffs are not rendered by default.

Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
82 changes: 82 additions & 0 deletions educational_material/embeddings/report.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
# Do directions in embedding space carry meaning?

Experiments on the token embeddings of Bengio (Q5), GPT-2 (Q6) and SmolLM2 (Q7), to pick examples for the course. All numbers come from [`results.md`](results.md), which is generated by [`embeddings.py`](embeddings.py) along with the figures in [`plots/`](plots/).

## In short

- **GPT-2 and SmolLM2 both show it clearly.** Their nearest neighbours make sense, the classic analogies come out top-1 in both, and a "gender" or "plural" direction learned from a few word pairs correctly orders every held-out pair. GPT-2 is a little cleaner on most tests.
- **The Bengio model shows nothing.** Its nearest neighbours look random (king → stup, glancing, form, dad…), and every analogy and direction test is at chance level. It was trained on about 200k words of fiction, in which "queen" appears 4 times, "aunt" 9 times and "cat" 10 times. That is too little text for meaning to settle into the embeddings. Use it in Q5 as the "before" picture rather than retraining it (see [Bengio: retrain or not?](#bengio-retrain-or-not)).

## Shortlist for the course, best first

| # | Example | Models | Why it works for teaching | Figure |
|---|---|---|---|---|
| 1 | **A gender direction**, learned from man→woman, he→she, boy→girl, father→mother and son→daughter, then applied to new pairs | GPT-2, SmolLM2 | Every held-out pair lands the right way round (king/queen, uncle/aunt, actor/actress, brother/sister, husband/wife). The arrows are nearly parallel, which is the idea in a single picture. The same plot as Bengio's shows the contrast. | [`direction_view_gender.png`](plots/direction_view_gender.png) |
| 2 | **king − man + woman = queen**, plus **uncle − man + woman = aunt** and **actor − man + woman = actress** | GPT-2, SmolLM2 | The famous example, top-1 in both models, and the two extra pairs show it isn't a one-off. Note that "queen" is already king's second-nearest neighbour in GPT-2, so the offset only has to nudge it to first place. | table in `results.md` |
| 3 | **Paris − France + Germany = Berlin** (country → capital) | GPT-2, SmolLM2 | This is the relation where the direction does the real work: a capital is rarely among its country's plain nearest neighbours (60% top-5 in GPT-2, 10% in SmolLM2), but adding the offset finds it (91% / 58% top-1). The PCA plot recreates the well-known word2vec figure (Mikolov et al., 2013). | [`pca_country_capital.png`](plots/pca_country_capital.png) |
| 4 | **The plural direction and number words** | GPT-2 (best); SmolLM2 | Held-out plurals all fall on the plural side. As a bonus, GPT-2 puts the number words in order along the plural direction: one, single, two, three, several, many, hundred. That is a nice "it learned this without being told" moment. | [`direction_plural.png`](plots/direction_plural.png) |
| 5 | **A bias hook**: occupations projected onto the gender direction | GPT-2 (clearest); SmolLM2 | In GPT-2, nurse, dancer and secretary sit towards the female end, and engineer, soldier and doctor towards the male end. SmolLM2 is similar but noisier: secretary sits in the middle and programmer lands on the female side. It's a short, real illustration that embeddings absorb the associations in their training text. | [`direction_gender.png`](plots/direction_gender.png) (bottom row of each panel) |
| 6 | **Nearest neighbours**: Monday → the other weekdays, red → colours, three → numbers, walked → walk, stepped, wandered | GPT-2, SmolLM2; Bengio as the contrast | The quickest intuition-builder: similar words sit close together. Bengio's random lists beside them make the Q5 → Q6 jump obvious. | table in `results.md` |
| 7 | **pizza − Italy + Japan = sushi** | GPT-2, SmolLM2 | Memorable and top-1 in both, but it's cherry-picked: country → food is unreliable overall (10–20% top-1). Use it as a light aside, not as evidence. | table in `results.md` |

**Suggested allocation:** in Q5, show Bengio's random neighbours and say why they're random. In Q6, use examples 1, 2, 4 and 6 with GPT-2. In Q7, show the same plots side by side with SmolLM2, then add example 3 (and 5 if there's room for the bias discussion). One honest point for Q7: SmolLM2 is newer and better at chat, yet its input embeddings are no cleaner than GPT-2's. The embedding table is only the first layer, and most of what a modern model "knows" is built up in the layers after it.

## What worked in each model

**Analogies:** share of all a − b + c combinations within each relation where the expected word is ranked first, after centring. GPT-2 and SmolLM2 have every word; Bengio misses most country, capital and food words, but scores 0% on everything it does have.

| relation | GPT-2 | SmolLM2 | plain nearest-neighbour baseline (top-5, GPT-2 / SmolLM2) |
|---|---|---|---|
| plural | 100% | 100% | 100% / 100% |
| past tense | 94% | 80% | 100% / 100% |
| comparative | 100% | 88% | 75% / 62% |
| gender | 89% | 70% | 100% / 83% |
| country → nationality | 100% | 100% | 100% / 100% |
| country → capital | 91% | 58% | 60% / 10% |
| number → ordinal | 55% | 45% | 0% / 0% |
| country → food/drink | 10% | 20% | 0% / 0% |

The baseline column is worth knowing about when choosing examples. For plurals, tenses and nationalities, the answer is usually among the word's plain nearest neighbours already: "dogs" is right next to "dog". Those analogies work, but they don't prove much about directions. Country → capital and number → ordinal are where the offset clearly earns its place.

**Directions** (built from five pairs, tested on held-out pairs): GPT-2 and SmolLM2 order every held-out pair correctly for plural, gender and country → capital. Bengio scores 2/5 for plural and 1/4 for gender, which is chance.

## A concept is a direction, not an axis

Each word is a point in a space with 64 dimensions (Bengio), 768 (GPT-2) or 576 (SmolLM2). The gender direction above is not one of those dimensions. It is the average of woman − man, she − he, girl − boy and so on, scaled to length 1, so it points diagonally through the space and uses nearly every coordinate a little:

| gender direction | Bengio | GPT-2 | SmolLM2 |
|---|---|---|---|
| dimensions | 64 | 768 | 576 |
| coordinates near zero (abs < 0.01) | 11% | 27% | 17% |
| coordinates needed for 90% of its length | 31 (48%) | 305 (40%) | 259 (45%) |
| largest single coordinate | 0.33 | 0.13 | 0.14 |
| each coordinate if spread perfectly evenly | 0.125 | 0.036 | 0.042 |

The plural and capital directions look the same (see `results.md`). Training gives no reason for the axes to line up with concepts, so there is no "gender coordinate". Meaning lives in combinations of many coordinates. Measuring a word against a concept is just a dot product: the word's unit vector · the direction gives one number, how far along that concept the word sits. That number is the x-axis in the direction-view plots. A 768-dimensional space has room for very many such directions at once, each nearly at right angles to the others.

For students, one sentence covers it: a word is a point in 768-D space, and a concept is a direction through that space, not one of its axes.

## Method notes

- **Whole-word tokens only.** Each word is looked up as a single token: "Ġking" (a leading space) in GPT-2 and SmolLM2, which share this convention, and "▁king" in Bengio's SentencePiece vocabulary. Bengio's text was lowercased, so it uses "paris" rather than "Paris". Neighbour searches consider whole-word tokens only, and exclude the input words and their case variants (" King" after " king").
- **Coverage:** GPT-2 has all 137 test words, and SmolLM2 has all but "tacos" and "vodka". Bengio lacks 29, mostly countries, capitals and foods, plus words like actress, prince and faster.
- **Centring** (subtracting the mean embedding before normalising) makes little difference to these results. It even makes SmolLM2's neighbour lists slightly worse, letting in tokens like "the", "a" and "in". The course can skip centring without losing anything. The plots use centred vectors; `results.md` has both versions.
- **Two kinds of 2D plot.** A plain PCA of the word pairs works for countries and capitals, but for gender it mostly separates groups of words (royalty vs family) and hides the arrows in GPT-2. The "direction view" plots fix that: x is the direction learned from *other* pairs, and y is the main remaining variation (PCA). A random embedding can't fake that view, which is why Bengio's arrows scatter while the other two models' point right. All the PCA and direction-view plots are in `plots/`, and the plain PCA ones are probably only worth using for countries.

## Bengio: retrain or not?

**Recommendation: don't retrain.** The untrained-looking result is itself a good teaching point: an embedding only means something if the model has seen each word in enough contexts. It's also authentic, since word2vec's famous results came from billions of words. With 4 sightings of "queen", there's nothing to learn.

If you do want to show the effect of data volume with the same model, here is my estimate (not run). I timed the model's training steps on this Mac's GPU:

- **text8** (about 17M words, so about 20M BPE tokens) takes about 12 minutes per epoch at the current batch size of 64. The current 20 epochs would take about 4 hours. With batch 1024 and a higher learning rate, it's about 3 minutes per epoch, so 5–10 epochs would take 15–30 minutes, plus a few minutes to retrain the BPE vocabulary and build the dataset.
- gensim isn't installed, but text8 can also be downloaded directly (a 31MB zip from mattmahoney.net), so a retrain wouldn't need a new dependency.
- It would very likely give sensible neighbours and some working analogies. However, the current learning rate (1e-4) is also low, so a retrain changes two things at once, and wouldn't be a clean "more data only" comparison.

## Rerunning

```bash
poetry run python -m educational_material.embeddings.embeddings
```

This takes about 20 seconds once the models are cached. GPT-2 and SmolLM2 download into `$HF_HOME` (default `/tmp/hf`) on the first run. It regenerates `results.md` and `plots/`, but not this report.
Loading
Loading