diff --git a/educational_material/embeddings/embeddings.py b/educational_material/embeddings/embeddings.py new file mode 100644 index 0000000..fa5ce59 --- /dev/null +++ b/educational_material/embeddings/embeddings.py @@ -0,0 +1,413 @@ +"""Do directions in embedding space carry meaning? + +Experiments on the token embeddings of the course's three neural models: +Bengio (Q5), GPT-2 (Q6) and SmolLM2 (Q7). Writes every table to results.md and +the figures to plots/, next to this file. + + poetry run python -m educational_material.embeddings.embeddings + +GPT-2 and SmolLM2 download from Hugging Face on first run (into $HF_HOME, default /tmp/hf). +""" +import itertools +import os +from dataclasses import dataclass, field +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +HERE = Path(__file__).resolve().parent +STORAGE = HERE.parent.parent / "evaluation_function" / "models" / "storage" +PLOTS = HERE / "plots" + +FROM_COLOUR, TO_COLOUR, INK, MUTED, SURFACE = "#2a78d6", "#eb6834", "#0b0b0b", "#52514e", "#fcfcfb" + +NEIGHBOUR_WORDS = ["king", "woman", "dog", "house", "red", "happy", "walked", "three", "France", "Monday", "coffee"] + +# (x, y) pairs where y = x + relation. Analogies take every ordered pair of pairs: +# y2 ≈ x2 - x1 + y1, e.g. queen ≈ king - man + woman. +RELATIONS = { + "plural": [("cat", "cats"), ("dog", "dogs"), ("car", "cars"), ("house", "houses"), ("book", "books"), + ("tree", "trees"), ("bird", "birds"), ("hand", "hands"), ("girl", "girls"), ("door", "doors")], + "past tense": [("walk", "walked"), ("play", "played"), ("look", "looked"), ("open", "opened"), ("go", "went"), + ("see", "saw"), ("run", "ran"), ("eat", "ate"), ("take", "took"), ("come", "came")], + "comparative": [("big", "bigger"), ("small", "smaller"), ("fast", "faster"), ("old", "older"), + ("young", "younger"), ("tall", "taller"), ("good", "better"), ("bad", "worse")], + "gender": [("man", "woman"), ("king", "queen"), ("uncle", "aunt"), ("actor", "actress"), ("boy", "girl"), + ("brother", "sister"), ("father", "mother"), ("son", "daughter"), ("husband", "wife"), + ("he", "she"), ("prince", "princess"), ("Mr", "Mrs")], + "number → ordinal": [("one", "first"), ("two", "second"), ("three", "third"), ("four", "fourth"), + ("five", "fifth"), ("six", "sixth"), ("seven", "seventh"), ("ten", "tenth")], + "country → capital": [("France", "Paris"), ("Germany", "Berlin"), ("Italy", "Rome"), ("Japan", "Tokyo"), + ("Spain", "Madrid"), ("England", "London"), ("Russia", "Moscow"), ("China", "Beijing"), + ("Egypt", "Cairo"), ("Greece", "Athens")], + "country → nationality": [("France", "French"), ("Germany", "German"), ("Italy", "Italian"), + ("Japan", "Japanese"), ("Spain", "Spanish"), ("China", "Chinese"), + ("Russia", "Russian"), ("Greece", "Greek"), ("England", "English")], + "country → food/drink": [("Italy", "pizza"), ("Japan", "sushi"), ("Germany", "beer"), ("France", "wine"), + ("Mexico", "tacos"), ("India", "curry"), ("Russia", "vodka")], +} + +SHOWCASE = [("king", "man", "woman", "queen"), ("uncle", "man", "woman", "aunt"), + ("actor", "man", "woman", "actress"), ("cats", "cat", "dog", "dogs"), + ("walked", "walk", "play", "played"), ("bigger", "big", "small", "smaller"), + ("Paris", "France", "Germany", "Berlin"), ("French", "France", "Italy", "Italian"), + ("pizza", "Italy", "Japan", "sushi"), ("first", "one", "two", "second")] + +# Directions are averaged over the first pairs and tested on held-out ones. +DIRECTIONS = { + "plural": dict( + train=RELATIONS["plural"][:5], + test=RELATIONS["plural"][5:], + probes=["one", "single", "two", "three", "several", "many", "hundred"], + ), + "gender": dict( + train=[("man", "woman"), ("he", "she"), ("boy", "girl"), ("father", "mother"), ("son", "daughter")], + test=[("king", "queen"), ("uncle", "aunt"), ("actor", "actress"), ("brother", "sister"), + ("husband", "wife")], + probes=["nurse", "doctor", "teacher", "engineer", "secretary", "soldier", "dancer", "programmer"], + ), + "country → capital": dict( + train=RELATIONS["country → capital"][:5], + test=RELATIONS["country → capital"][5:], + probes=[], + ), +} + +PCA_SETS = { + "gender": [("man", "woman"), ("king", "queen"), ("uncle", "aunt"), ("boy", "girl"), ("brother", "sister")], + "plural": [("cat", "cats"), ("dog", "dogs"), ("car", "cars"), ("house", "houses"), ("book", "books")], + "country → capital": [("France", "Paris"), ("Germany", "Berlin"), ("Japan", "Tokyo"), ("Spain", "Madrid"), + ("Russia", "Moscow")], +} + + +def normalise(m): + return m / np.linalg.norm(m, axis=-1, keepdims=True).clip(1e-8) + + +@dataclass +class Space: + """A model's token embedding matrix, restricted to whole-word tokens for lookups.""" + + name: str + vectors: np.ndarray + pieces: list + prefix: str # marks a word-initial token: "▁" (SentencePiece) or "Ġ" (byte-level BPE) + lowercase: bool # Bengio's corpus was lowercased + unit: dict = field(init=False) + + def __post_init__(self): + self.ids = {p: i for i, p in enumerate(self.pieces)} + n = len(self.prefix) + self.word_ids = np.array([i for i, p in enumerate(self.pieces) + if p.startswith(self.prefix) and p[n:].isascii() and p[n:].isalpha()]) + self.unit = {False: normalise(self.vectors), True: normalise(self.vectors - self.vectors.mean(0))} + self.word_unit = {c: m[self.word_ids] for c, m in self.unit.items()} + + def id(self, word): + return self.ids.get(self.prefix + (word.lower() if self.lowercase else word)) + + def has(self, *words): + return all(self.id(w) is not None for w in words) + + def word(self, i): + return self.pieces[i][len(self.prefix):] + + def vec(self, word, centred): + return self.unit[centred][self.id(word)] + + def nearest(self, target, exclude, k, centred): + sims = self.word_unit[centred] @ normalise(target) + skip = {w.lower() for w in exclude} + top = np.argpartition(-sims, k + 30)[:k + 30] + out = [] + for j in top[np.argsort(-sims[top])]: + w = self.word(self.word_ids[j]) + if w.lower() not in skip: + out.append(w) + skip.add(w.lower()) # drop case variants: " King" after " king" + if len(out) == k: + break + return out + + def neighbours(self, word, k=8, centred=True): + return self.nearest(self.vec(word, centred), [word], k, centred) + + def analogy(self, a, b, c, k=5, centred=True): + """a - b + c, e.g. king - man + woman.""" + target = self.vec(a, centred) - self.vec(b, centred) + self.vec(c, centred) + return self.nearest(target, [a, b, c], k, centred) + + def direction(self, pairs, centred=True): + pairs = [p for p in pairs if self.has(*p)] + return normalise(np.mean([self.vec(y, centred) - self.vec(x, centred) for x, y in pairs], axis=0)) + + +def load_bengio(): + import sentencepiece as spm + import torch + + weights = torch.load(STORAGE / "bengio_model.pt", map_location="cpu")["emb.weight"].numpy() + sp = spm.SentencePieceProcessor(model_file=str(STORAGE / "bpe.model")) + pieces = [sp.id_to_piece(i) for i in range(sp.get_piece_size())] + return Space("Bengio", weights, pieces, prefix="▁", lowercase=True) + + +def load_hf(name, model_id): + from transformers import AutoModelForCausalLM, AutoTokenizer + + cache = os.environ.get("HF_HOME", "/tmp/hf") + model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=cache) + tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache) + weights = model.get_input_embeddings().weight.detach().float().numpy() + pieces = tokenizer.convert_ids_to_tokens(list(range(weights.shape[0]))) + return Space(name, weights, pieces, prefix="Ġ", lowercase=False) + + +# ---------------------------------------------------------------- experiments + +def vocab_table(spaces): + words = sorted({w for pairs in RELATIONS.values() for p in pairs for w in p} | set(NEIGHBOUR_WORDS), + key=str.lower) + lines = ["| model | vocab | whole-word tokens | test words found | missing |", "|---|---|---|---|---|"] + for s in spaces: + missing = [w for w in words if not s.has(w)] + lines.append(f"| {s.name} | {len(s.pieces)} × {s.vectors.shape[1]} | {len(s.word_ids)} | " + f"{len(words) - len(missing)}/{len(words)} | {', '.join(missing) or '-'} |") + return lines + + +def neighbour_tables(spaces): + lines = [] + for centred in (True, False): + lines += [f"\n### Nearest neighbours ({'centred' if centred else 'not centred'})\n", + "| word | " + " | ".join(s.name for s in spaces) + " |", "|---" * (len(spaces) + 1) + "|"] + for w in NEIGHBOUR_WORDS: + cells = [", ".join(s.neighbours(w, centred=centred)) if s.has(w) else "-" for s in spaces] + lines.append(f"| {w} | " + " | ".join(cells) + " |") + return lines + + +def analogy_scores(space, pairs, centred): + """Top-1 and top-5 hits over all ordered pairs of pairs, plus a baseline that ignores the offset.""" + pairs = [p for p in pairs if space.has(*p)] + top1 = top5 = base5 = n = 0 + for (x1, y1), (x2, y2) in itertools.permutations(pairs, 2): + if len({x1, y1, x2, y2}) < 4: + continue + guesses = space.analogy(x2, x1, y1, centred=centred) + top1 += guesses[0] == y2 + top5 += y2 in guesses + base5 += y2 in space.neighbours(x2, k=5, centred=centred) + n += 1 + return top1, top5, base5, n + + +def analogy_tables(spaces): + lines = [] + for centred in (True, False): + lines += [f"\n### Analogy accuracy ({'centred' if centred else 'not centred'})\n", + "Each cell: top-1 % / top-5 % (baseline top-5 %: nearest neighbours of x2 alone, no offset)" + " [number of analogies].\n", + "| relation | " + " | ".join(s.name for s in spaces) + " |", "|---" * (len(spaces) + 1) + "|"] + for rel, pairs in RELATIONS.items(): + cells = [] + for s in spaces: + t1, t5, b5, n = analogy_scores(s, pairs, centred) + cells.append(f"{100 * t1 / n:.0f} / {100 * t5 / n:.0f} ({100 * b5 / n:.0f}) [{n}]" if n else "-") + lines.append(f"| {rel} | " + " | ".join(cells) + " |") + lines += ["\n### Showcase analogies (centred, top 5; ✓ = expected answer ranked first)\n", + "| analogy | expected | " + " | ".join(s.name for s in spaces) + " |", + "|---" * (len(spaces) + 2) + "|"] + for a, b, c, want in SHOWCASE: + cells = [] + for s in spaces: + if not s.has(a, b, c, want): + cells.append("-") + continue + got = s.analogy(a, b, c) + cells.append(("✓ " if got[0].lower() == want.lower() else "") + ", ".join(got)) + lines.append(f"| {a} - {b} + {c} | {want} | " + " | ".join(cells) + " |") + return lines + + +def direction_tables(spaces): + lines = [] + for name, d in DIRECTIONS.items(): + lines += [f"\n### Direction: {name}\n", + f"Built from {', '.join(f'{y} - {x}' for x, y in d['train'])}. " + "Held-out: share of pairs where the second word projects further along the direction.\n", + "| model | centring | held-out pairs ordered | probe projections (high → low) |", "|---|---|---|---|"] + for s in spaces: + for centred in (True, False): + u = s.direction(d["train"], centred) + test = [p for p in d["test"] if s.has(*p)] + ok = sum(s.vec(y, centred) @ u > s.vec(x, centred) @ u for x, y in test) + probes = sorted(((s.vec(w, centred) @ u, w) for w in d["probes"] if s.has(w)), reverse=True) + lines.append(f"| {s.name} | {'centred' if centred else 'raw'} | {ok}/{len(test)} | " + + (", ".join(f"{w} {v:+.2f}" for v, w in probes) or "-") + " |") + return lines + + +def spread_table(spaces): + """How many embedding coordinates each learned direction uses (it is not one axis).""" + lines = ["\n### How spread out the directions are\n", + "Each unit-length direction (centred) over the model's coordinates. If one coordinate carried the concept, " + "it would be near 1 and the rest near 0.\n", + "| direction | model | dims | coords near zero (abs < 0.01) | coords for 90% of length | largest coord " + "| even spread would be |", "|---|---|---|---|---|---|---|"] + for name, d in DIRECTIONS.items(): + for s in spaces: + a = np.sort(np.abs(s.direction(d["train"])))[::-1] + k90 = int(np.searchsorted(np.cumsum(a ** 2), 0.9)) + 1 + lines.append(f"| {name} | {s.name} | {len(a)} | {(a < 0.01).mean():.0%} | {k90} ({k90 / len(a):.0%}) " + f"| {a[0]:.2f} | {1 / np.sqrt(len(a)):.3f} |") + return lines + + +# ---------------------------------------------------------------------- plots + +def style(ax): + ax.set_facecolor(SURFACE) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + for side in ("left", "bottom"): + ax.spines[side].set_color(MUTED) + ax.tick_params(colors=MUTED, labelsize=8) + + +def draw_pairs(ax, space, pairs, P): + for i, (x0, y0) in enumerate(P[0::2]): + x1, y1 = P[2 * i + 1] + ax.annotate("", (x1, y1), (x0, y0), arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.2)) + colours = [FROM_COLOUR, TO_COLOUR] * len(pairs) + for (px, py), word, colour, dy in zip(P, [w for p in pairs for w in p], colours, label_offsets(P)): + ax.plot(px, py, "o", ms=8, color=colour, mec=SURFACE, mew=2) + ax.annotate(space.word(space.id(word)), (px, py), (0, dy), textcoords="offset points", ha="center", + va="center", fontsize=9, color=INK) + + +def label_offsets(P): + """Put a label above its point, or below if that would collide with a label already placed.""" + span = np.ptp(P, axis=0).clip(1e-9) + placed, offsets = [], [] + for point in P / span: + for dy in (10, -11, 21, -22): + spot = point + [0, dy / 300] + if all(abs(spot[0] - q[0]) > 0.12 or abs(spot[1] - q[1]) > 0.05 for q in placed): + break + placed.append(spot) + offsets.append(dy) + return offsets + + +def pair_panels(spaces, title, filename, panel): + """One panel per model, side by side; `panel` returns the 2D points for a model's pairs, or None.""" + fig, axes = plt.subplots(1, len(spaces), figsize=(4.2 * len(spaces), 4), facecolor=SURFACE) + for ax, s in zip(axes, spaces): + style(ax) + ax.set_title(s.name, color=INK, fontsize=11) + ax.set_yticks([]) + drawn = panel(ax, s) + if drawn is None: + ax.set_xticks([]) + ax.text(0.5, 0.5, "too few of these words\nin the vocabulary", transform=ax.transAxes, ha="center", + va="center", color=MUTED, fontsize=9) + else: + draw_pairs(ax, s, *drawn) + ax.margins(0.18) + fig.suptitle(title, color=INK, fontsize=12) + fig.tight_layout() + fig.savefig(PLOTS / filename, dpi=150, facecolor=SURFACE) + plt.close(fig) + + +def plot_pca(spaces, name, pairs): + """Plain 2D PCA of the pairs' embeddings, oriented so the mean arrow points right.""" + def panel(ax, s): + found = [p for p in pairs if s.has(*p)] + if len(found) < 3: + return None + X = np.array([s.vec(w, True) for p in found for w in p]) + X -= X.mean(0) + P = X @ np.linalg.svd(X, full_matrices=False)[2][:2].T + if (P[1::2] - P[0::2]).mean(0)[0] < 0: + P[:, 0] *= -1 + ax.set_xticks([]) + return found, P + + pair_panels(spaces, f"{name}: 2D PCA of the word embeddings", f"pca_{slug(name)}.png", panel) + + +def plot_direction_view(spaces, name, d): + """Held-out pairs placed along a direction learned from other pairs (x) against the main remaining + variation (y, PCA). Arrows pointing right mean the direction generalises to new words.""" + def panel(ax, s): + test = [p for p in d["test"] if s.has(*p)] + if len(test) < 3 or not any(s.has(*p) for p in d["train"]): + return None + u = s.direction(d["train"]) + X = np.array([s.vec(w, True) for p in test for w in p]) + x = X @ u + rest = X - np.outer(x, u) + rest -= rest.mean(0) + y = rest @ np.linalg.svd(rest, full_matrices=False)[2][0] + ax.axvline(0, color=MUTED, lw=0.6, ls=":") + ax.set_xlabel(f"along the {name} direction", color=MUTED, fontsize=9) + return test, np.column_stack([x, y]) + + pair_panels(spaces, f"{name}: held-out pairs along the direction learned from other pairs", + f"direction_view_{slug(name)}.png", panel) + + +def plot_direction(spaces, name, d, labels): + """Each word's position along a direction: held-out pairs on two rows, probe words on a third.""" + fig, axes = plt.subplots(len(spaces), 1, figsize=(9, 2.6 * len(spaces)), facecolor=SURFACE) + rows = [labels[0], labels[1], "probe words"] + for ax, s in zip(axes, spaces): + u = s.direction(d["train"]) + test = [p for p in d["test"] if s.has(*p)] + groups = [[x for x, _ in test], [y for _, y in test], [w for w in d["probes"] if s.has(w)]] + style(ax) + for row, (words, colour) in enumerate(zip(groups, (FROM_COLOUR, TO_COLOUR, MUTED))): + proj = [s.vec(w, True) @ u for w in words] + ax.plot(proj, [row] * len(proj), "o", ms=8, color=colour, mec=SURFACE, mew=2) + for i, (p, w) in enumerate(sorted(zip(proj, words))): + ax.annotate(s.word(s.id(w)), (p, row), (0, (8, -14, 18)[i % 3]), textcoords="offset points", + ha="center", fontsize=8, color=INK) + ax.set_yticks(range(3), rows, fontsize=9, color=INK) + ax.set_ylim(-0.9, 3.0) + ax.set_title(s.name, loc="left", color=INK, fontsize=10) + axes[-1].set_xlabel(f"projection onto the {name} direction (cosine)", color=MUTED, fontsize=9) + fig.tight_layout() + fig.savefig(PLOTS / f"direction_{slug(name)}.png", dpi=150, facecolor=SURFACE) + plt.close(fig) + + +def slug(text): + return "".join(c if c.isalnum() else "_" for c in text.lower()).strip("_").replace("___", "_") + + +def main(): + spaces = [load_bengio(), load_hf("GPT-2", "gpt2"), load_hf("SmolLM2", "HuggingFaceTB/SmolLM2-135M-Instruct")] + + lines = ["# Embedding experiments: raw results", "", + "Generated by `embeddings.py`; see `report.md` for the interpretation. Vectors are unit-normalised; " + "'centred' means the mean embedding was subtracted first. Neighbour searches only consider whole-word " + "tokens and exclude the input words.", "", "### Vocabulary coverage", ""] + for table in (vocab_table, neighbour_tables, analogy_tables, direction_tables, spread_table): + lines += table(spaces) + (HERE / "results.md").write_text("\n".join(lines) + "\n") + + PLOTS.mkdir(exist_ok=True) + for name, pairs in PCA_SETS.items(): + plot_pca(spaces, name, pairs) + plot_direction_view(spaces, name, DIRECTIONS[name]) + plot_direction(spaces, "plural", DIRECTIONS["plural"], ("singular", "plural")) + plot_direction(spaces, "gender", DIRECTIONS["gender"], ("male", "female")) + print(f"Wrote {HERE / 'results.md'} and {PLOTS}/") + + +if __name__ == "__main__": + main() diff --git a/educational_material/embeddings/plots/direction_gender.png b/educational_material/embeddings/plots/direction_gender.png new file mode 100644 index 0000000..8998d81 Binary files /dev/null and b/educational_material/embeddings/plots/direction_gender.png differ diff --git a/educational_material/embeddings/plots/direction_plural.png b/educational_material/embeddings/plots/direction_plural.png new file mode 100644 index 0000000..9ea3a2c Binary files /dev/null and b/educational_material/embeddings/plots/direction_plural.png differ diff --git a/educational_material/embeddings/plots/direction_view_country_capital.png b/educational_material/embeddings/plots/direction_view_country_capital.png new file mode 100644 index 0000000..21609ae Binary files /dev/null and b/educational_material/embeddings/plots/direction_view_country_capital.png differ diff --git a/educational_material/embeddings/plots/direction_view_gender.png b/educational_material/embeddings/plots/direction_view_gender.png new file mode 100644 index 0000000..11312a8 Binary files /dev/null and b/educational_material/embeddings/plots/direction_view_gender.png differ diff --git a/educational_material/embeddings/plots/direction_view_plural.png b/educational_material/embeddings/plots/direction_view_plural.png new file mode 100644 index 0000000..35665ed Binary files /dev/null and b/educational_material/embeddings/plots/direction_view_plural.png differ diff --git a/educational_material/embeddings/plots/pca_country_capital.png b/educational_material/embeddings/plots/pca_country_capital.png new file mode 100644 index 0000000..47c9016 Binary files /dev/null and b/educational_material/embeddings/plots/pca_country_capital.png differ diff --git a/educational_material/embeddings/plots/pca_gender.png b/educational_material/embeddings/plots/pca_gender.png new file mode 100644 index 0000000..53bcf3f Binary files /dev/null and b/educational_material/embeddings/plots/pca_gender.png differ diff --git a/educational_material/embeddings/plots/pca_plural.png b/educational_material/embeddings/plots/pca_plural.png new file mode 100644 index 0000000..20f6aaf Binary files /dev/null and b/educational_material/embeddings/plots/pca_plural.png differ diff --git a/educational_material/embeddings/report.md b/educational_material/embeddings/report.md new file mode 100644 index 0000000..c4b7fc8 --- /dev/null +++ b/educational_material/embeddings/report.md @@ -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. diff --git a/educational_material/embeddings/results.md b/educational_material/embeddings/results.md new file mode 100644 index 0000000..23b5845 --- /dev/null +++ b/educational_material/embeddings/results.md @@ -0,0 +1,143 @@ +# Embedding experiments: raw results + +Generated by `embeddings.py`; see `report.md` for the interpretation. Vectors are unit-normalised; 'centred' means the mean embedding was subtracted first. Neighbour searches only consider whole-word tokens and exclude the input words. + +### Vocabulary coverage + +| model | vocab | whole-word tokens | test words found | missing | +|---|---|---|---|---| +| Bengio | 8000 × 64 | 5813 | 108/137 | actress, Athens, Beijing, Berlin, Cairo, Chinese, curry, Egypt, England, faster, Germany, Greece, India, Italian, London, Madrid, Mexico, Moscow, pizza, prince, princess, Russia, Russian, sixth, Spain, sushi, tacos, tenth, vodka | +| GPT-2 | 50257 × 768 | 32064 | 137/137 | - | +| SmolLM2 | 49152 × 576 | 31586 | 135/137 | tacos, vodka | + +### Nearest neighbours (centred) + +| word | Bengio | GPT-2 | SmolLM2 | +|---|---|---|---| +| king | stup, glancing, form, dad, kate, we, with, forbid | kings, queen, prince, monarch, royal, emperor, ruler, lord | kings, Queen, emperor, John, royal, Prince, president, in | +| woman | winston, picked, piece, crime, coughlin, rode, divided, presented | women, girl, female, man, lady, person, mother, wife | women, man, girl, men, person, female, lady, people | +| dog | felix, weep, fuzzy, tri, thank, are, hub, roberta | dogs, canine, puppy, cat, pet, animal, horse, pets | dogs, cat, pet, horse, the, a, D, in | +| house | envy, cover, wally, girl, drinks, missed, laughing, robards | houses, home, mansion, apartment, homes, family, household, room | houses, home, car, a, household, in, family, the | +| red | flesh, plastic, rubber, column, collect, running, roll, rot | blue, yellow, green, white, black, purple, pink, orange | black, white, blue, green, yellow, B, re, a | +| happy | redhead, screaming, english, curtain, courage, portra, sorrow, carried | unhappy, happier, pleased, glad, thrilled, proud, satisfied, delighted | happiness, good, is, in, P, a, L, positive | +| walked | notice, cool, determined, purse, turn, enjoyed, guns, testament | walk, walks, walking, stood, stepped, wandered, went, climbed | walk, walks, walking, stood, stepped, marched, traveled, went | +| three | brood, lat, instead, eyel, this, front, eng, appre | four, two, five, six, seven, nine, eight, several | four, two, five, six, seven, eight, nine, several | +| France | myn, alice, if, all, sixteen, bunch, bigger, wondering | French, Germany, Europe, Spain, Italy, Paris, Belgium, England | French, Germany, Europe, Spain, England, Russia, Italy, China | +| Monday | cash, brian, entertain, behaved, slide, visited, lazy, belonged | Tuesday, Thursday, Wednesday, Friday, Sunday, Saturday, week, yesterday | Tuesday, Wednesday, Thursday, Friday, Saturday, Sunday, January, July | +| coffee | itself, nephew, reuveni, interview, million, instance, admired, shades | espresso, tea, Starbucks, caffeine, coff, beer, chocolate, cafe | tea, beer, food, water, chocolate, wine, re, the | + +### Nearest neighbours (not centred) + +| word | Bengio | GPT-2 | SmolLM2 | +|---|---|---|---| +| king | stup, glancing, form, dad, we, kate, with, bul | kings, queen, prince, monarch, emperor, ruler, royal, lord | kings, queen, emperor, prince, kingdom, monarch, president, Pharaoh | +| woman | winston, picked, piece, crime, coughlin, rode, divided, presented | women, girl, female, lady, man, wife, WOM, mother | women, girl, lady, female, females, girls, ladies, man | +| dog | weep, felix, fuzzy, tri, thank, are, hub, roberta | dogs, canine, puppy, pet, cat, animal, pets, puppies | dogs, cat, pet, puppy, cats, canine, horse, pets | +| house | envy, cover, wally, drinks, girl, missed, laughing, robards | houses, mansion, home, apartment, homes, household, kitchen, housing | houses, household, home, homes, housing, kitchen, room, host | +| red | flesh, rubber, plastic, column, collect, running, roll, rot | blue, yellow, green, purple, white, pink, black, orange | blue, white, black, green, yellow, purple, pink, brown | +| happy | redhead, screaming, english, curtain, courage, portra, sorrow, carried | happier, unhappy, pleased, glad, thrilled, happiest, delighted, joyful | unhappy, happier, happiness, happily, satisfied, pleased, joyful, delighted | +| walked | notice, cool, determined, purse, turn, enjoyed, guns, testament | walks, walk, walking, wandered, stood, stepped, marched, climbed | walks, walk, Walking, wandered, stroll, climbed, traversed, marched | +| three | brood, lat, instead, eyel, this, front, eng, appre | four, two, five, six, seven, nine, eight, several | four, two, five, six, seven, eight, several, nine | +| France | myn, alice, all, if, sixteen, bigger, bunch, wondering | French, Germany, Spain, Italy, Europe, Belgium, Paris, Denmark | Spain, Germany, Italy, Belgium, Russia, England, French, Europe | +| Monday | cash, brian, entertain, behaved, slide, visited, lazy, belonged | Tuesday, Thursday, Wednesday, Friday, Sunday, Saturday, Mondays, Tues | Tuesday, Wednesday, Thursday, Saturday, Friday, Sunday, Sundays, weekend | +| coffee | itself, nephew, reuveni, interview, million, instance, admired, shades | espresso, tea, caffeine, Starbucks, coff, beer, cocoa, cafe | tea, caffeine, beer, cocoa, chocolate, coff, cereal, pizza | + +### Analogy accuracy (centred) + +Each cell: top-1 % / top-5 % (baseline top-5 %: nearest neighbours of x2 alone, no offset) [number of analogies]. + +| relation | Bengio | GPT-2 | SmolLM2 | +|---|---|---|---| +| plural | 0 / 0 (0) [90] | 100 / 100 (100) [90] | 100 / 100 (100) [90] | +| past tense | 0 / 1 (0) [90] | 94 / 100 (100) [90] | 80 / 100 (100) [90] | +| comparative | 0 / 0 (0) [42] | 100 / 100 (75) [56] | 88 / 96 (62) [56] | +| gender | 0 / 0 (0) [90] | 89 / 100 (100) [132] | 70 / 94 (83) [132] | +| number → ordinal | 0 / 0 (0) [30] | 55 / 79 (0) [56] | 45 / 80 (0) [56] | +| country → capital | 0 / 0 (0) [6] | 91 / 100 (60) [90] | 58 / 92 (10) [90] | +| country → nationality | 0 / 0 (0) [2] | 100 / 100 (100) [72] | 100 / 100 (100) [72] | +| country → food/drink | - | 10 / 43 (0) [42] | 20 / 20 (0) [20] | + +### Analogy accuracy (not centred) + +Each cell: top-1 % / top-5 % (baseline top-5 %: nearest neighbours of x2 alone, no offset) [number of analogies]. + +| relation | Bengio | GPT-2 | SmolLM2 | +|---|---|---|---| +| plural | 0 / 0 (0) [90] | 100 / 100 (100) [90] | 97 / 97 (100) [90] | +| past tense | 0 / 1 (0) [90] | 94 / 100 (100) [90] | 73 / 100 (100) [90] | +| comparative | 0 / 0 (0) [42] | 100 / 100 (75) [56] | 96 / 100 (88) [56] | +| gender | 0 / 0 (0) [90] | 86 / 100 (100) [132] | 66 / 97 (100) [132] | +| number → ordinal | 0 / 0 (0) [30] | 57 / 82 (0) [56] | 68 / 82 (0) [56] | +| country → capital | 0 / 0 (0) [6] | 90 / 100 (60) [90] | 74 / 88 (10) [90] | +| country → nationality | 0 / 0 (0) [2] | 100 / 100 (89) [72] | 99 / 100 (44) [72] | +| country → food/drink | - | 10 / 38 (0) [42] | 10 / 25 (0) [20] | + +### Showcase analogies (centred, top 5; ✓ = expected answer ranked first) + +| analogy | expected | Bengio | GPT-2 | SmolLM2 | +|---|---|---|---|---| +| king - man + woman | queen | gasp, there, rent, slim, tract | ✓ queen, princess, kings, girl, queens | ✓ queen, kings, women, female, emperor | +| uncle - man + woman | aunt | fil, deserted, grad, cared, tract | ✓ aunt, niece, grandmother, girlfriend, mother | ✓ aunt, grandmother, nephew, grandfather, cousin | +| actor - man + woman | actress | - | ✓ actress, actresses, actors, comedian, singer | ✓ actress, actors, singer, singers, dancer | +| cats - cat + dog | dogs | roberta, class, fraud, affair, weep | ✓ dogs, puppies, pets, canine, animals | ✓ dogs, horses, animals, birds, pets | +| walked - walk + play | played | lists, eleven, chill, tick, mon | ✓ played, plays, playing, went, had | ✓ played, plays, playing, players, had | +| bigger - big + small | smaller | six, concept, offic, doorway, someone | ✓ smaller, larger, tiny, greater, large | ✓ smaller, larger, tiny, greater, large | +| Paris - France + Germany | Berlin | - | ✓ Berlin, German, Munich, Hamburg, London | ✓ Berlin, German, London, Chicago, Boston | +| French - France + Italy | Italian | - | ✓ Italian, Italians, Spanish, German, Greek | ✓ Italian, German, Greek, Japanese, Spanish | +| pizza - Italy + Japan | sushi | - | ✓ sushi, Japanese, burger, pizz, hamb | ✓ sushi, coffee, tea, restaurant, dinner | +| first - one + two | second | glanced, their, kitt, rubbing, sixty | three, second, four, five, six | ✓ second, three, four, five, initial | + +### Direction: plural + +Built from cats - cat, dogs - dog, cars - car, houses - house, books - book. Held-out: share of pairs where the second word projects further along the direction. + +| model | centring | held-out pairs ordered | probe projections (high → low) | +|---|---|---|---| +| Bengio | centred | 2/5 | hundred +0.42, two +0.04, three +0.03, many +0.00, single -0.01, one -0.03, several -0.23 | +| Bengio | raw | 2/5 | hundred +0.42, two +0.04, three +0.04, many +0.00, single -0.00, one -0.03, several -0.22 | +| GPT-2 | centred | 5/5 | hundred -0.09, many -0.12, several -0.15, three -0.18, two -0.19, single -0.22, one -0.24 | +| GPT-2 | raw | 5/5 | hundred +0.01, many +0.01, several -0.00, three -0.05, two -0.06, single -0.10, one -0.12 | +| SmolLM2 | centred | 5/5 | hundred -0.20, many -0.27, several -0.27, two -0.27, three -0.29, single -0.30, one -0.30 | +| SmolLM2 | raw | 5/5 | hundred +0.16, several -0.05, single -0.10, many -0.14, three -0.16, two -0.18, one -0.27 | + +### Direction: gender + +Built from woman - man, she - he, girl - boy, mother - father, daughter - son. Held-out: share of pairs where the second word projects further along the direction. + +| model | centring | held-out pairs ordered | probe projections (high → low) | +|---|---|---|---| +| Bengio | centred | 1/4 | nurse +0.19, soldier +0.12, engineer +0.08, doctor -0.01, teacher -0.10, secretary -0.14 | +| Bengio | raw | 1/4 | nurse +0.19, soldier +0.12, engineer +0.08, doctor -0.01, teacher -0.10, secretary -0.14 | +| GPT-2 | centred | 5/5 | nurse +0.22, dancer +0.19, secretary +0.10, teacher +0.04, programmer +0.04, doctor -0.00, soldier -0.03, engineer -0.05 | +| GPT-2 | raw | 5/5 | nurse +0.23, dancer +0.19, secretary +0.12, teacher +0.08, programmer +0.07, doctor +0.04, soldier +0.01, engineer -0.00 | +| SmolLM2 | centred | 5/5 | dancer +0.15, nurse +0.10, programmer +0.04, secretary -0.01, teacher -0.07, soldier -0.08, doctor -0.09, engineer -0.13 | +| SmolLM2 | raw | 5/5 | dancer +0.33, nurse +0.32, programmer +0.27, secretary +0.23, soldier +0.17, teacher +0.17, doctor +0.16, engineer +0.15 | + +### Direction: country → capital + +Built from Paris - France, Berlin - Germany, Rome - Italy, Tokyo - Japan, Madrid - Spain. Held-out: share of pairs where the second word projects further along the direction. + +| model | centring | held-out pairs ordered | probe projections (high → low) | +|---|---|---|---| +| Bengio | centred | 0/0 | - | +| Bengio | raw | 0/0 | - | +| GPT-2 | centred | 5/5 | - | +| GPT-2 | raw | 5/5 | - | +| SmolLM2 | centred | 5/5 | - | +| SmolLM2 | raw | 5/5 | - | + +### How spread out the directions are + +Each unit-length direction (centred) over the model's coordinates. If one coordinate carried the concept, it would be near 1 and the rest near 0. + +| direction | model | dims | coords near zero (abs < 0.01) | coords for 90% of length | largest coord | even spread would be | +|---|---|---|---|---|---|---| +| plural | Bengio | 64 | 8% | 29 (45%) | 0.28 | 0.125 | +| plural | GPT-2 | 768 | 24% | 300 (39%) | 0.28 | 0.036 | +| plural | SmolLM2 | 576 | 19% | 251 (44%) | 0.15 | 0.042 | +| gender | Bengio | 64 | 11% | 31 (48%) | 0.33 | 0.125 | +| gender | GPT-2 | 768 | 27% | 305 (40%) | 0.13 | 0.036 | +| gender | SmolLM2 | 576 | 17% | 259 (45%) | 0.14 | 0.042 | +| country → capital | Bengio | 64 | 2% | 29 (45%) | 0.37 | 0.125 | +| country → capital | GPT-2 | 768 | 26% | 321 (42%) | 0.13 | 0.036 | +| country → capital | SmolLM2 | 576 | 20% | 246 (43%) | 0.12 | 0.042 |