Superposition in a Toy word2vec

Three attributes, two dimensions to put them in. A model can represent more features than it has dimensions, and this is what it costs.

Week 1 trained word2vec on 17 million words and the co-occurrence structure came from whatever people happened to write. This post does the opposite. I hand-build a tiny corpus so the structure is exactly what I want, then give the model fewer dimensions than it needs and watch what it does.

The idea comes from Anthropic’s Toy Models of Superposition. A model can represent more features than it has dimensions by packing feature directions into the same space, instead of giving each one its own axis. I wanted to see it happen.

The Setup

Three attributes, two values each:

Attribute	Values
Status	noble / common
Gender	masc / fem
Age	adult / young

Combined, they give eight words:

king     = noble  + masc + adult      queen    = noble  + fem + adult
man      = common + masc + adult      woman    = common + fem + adult
prince   = noble  + masc + young      princess = noble  + fem + young
boy      = common + masc + young      girl     = common + fem + young

Each of the six attribute values gets its own word, which I will call a pole. Each of the eight words above co-occurs with its three poles.

The corpus is built from a cycle. The six poles sit in a ring, arranged so that opposites are three steps apart:

noble — masc — adult — common — fem — young — noble

That ring is what spreads the poles out evenly. Each pole also gets a self-loop, and each corner word touches only its own three poles, once each. Sixty-six co-occurrence pairs in total.

import sys

sys.path.insert(0, "..")
import numpy as np

from itertools import combinations

POLES = ["noble", "common", "masc", "fem", "adult", "young"]
SYNONYMS = {
    "noble": ["royal", "crown", "regal"],  "common": ["peasant", "serf", "lowly"],
    "masc":  ["he", "him", "his"],         "fem":    ["she", "her", "hers"],
    "adult": ["grown", "elder", "mature"], "young":  ["child", "youth", "junior"],
}
CYCLE = ["noble", "masc", "adult", "common", "fem", "young"]   # opposites 3 apart
CORNERS = {
    "king":   ("noble", "masc", "adult"),  "queen":    ("noble", "fem", "adult"),
    "man":    ("common", "masc", "adult"), "woman":    ("common", "fem", "adult"),
    "prince": ("noble", "masc", "young"),  "princess": ("noble", "fem", "young"),
    "boy":    ("common", "masc", "young"), "girl":     ("common", "fem", "young"),
}


def sparse_corpus():
    """The pole cycle, self-loops, and each corner touching only its own poles."""
    corpus = [f"{CYCLE[i]} {CYCLE[(i + 1) % 6]}" for i in range(6)] * 2
    corpus += [f"{p} {p}" for p in POLES]
    for corner, poles in CORNERS.items():
        for p in poles:
            corpus += [f"{corner} {p}", f"{p} {corner}"]
    return corpus


print(f"{len(sparse_corpus())} co-occurrence pairs, "
      f"{len({w for s in sparse_corpus() for w in s.split()})} words")

66 co-occurrence pairs, 14 words

The model is skip-gram with negative sampling, the training method word2vec actually uses, written out in plain numpy. Every word gets a vector. For each word, pull it toward the words it appears next to and push it away from a random word it doesn’t.

Nothing here instructs the model to make noble and common opposite. That falls out of how often the words appear near each other.

def train(corpus, dim, seed=9, window=2, lr0=0.01, epochs=2000):
    vocab = sorted({w for s in corpus for w in s.split()})
    index = {w: i for i, w in enumerate(vocab)}

    pairs = [(c, x) for s in corpus
             for ids in [[index[w] for w in s.split()]]
             for i, c in enumerate(ids) for j, x in enumerate(ids)
             if 0 < abs(i - j) <= window]

    counts = np.array([sum(s.split().count(w) for s in corpus) for w in vocab], float)
    unigram = counts ** 0.75
    unigram /= unigram.sum()

    sigmoid = lambda x: 1 / (1 + np.exp(-np.clip(x, -50, 50)))
    rng = np.random.RandomState(seed)
    W_in = rng.uniform(-0.5 / dim, 0.5 / dim, (len(vocab), dim))
    W_out = rng.uniform(-0.5 / dim, 0.5 / dim, (len(vocab), dim))

    for epoch in range(epochs):
        lr = lr0 * (1 - epoch / epochs)
        for centre, context in pairs:
            negatives = rng.choice(len(vocab), size=1, p=unigram)
            negatives = negatives[negatives != context]
            targets = np.concatenate(([context], negatives))
            labels = np.r_[1, [0] * (len(targets) - 1)]
            error = sigmoid(W_in[centre] @ W_out[targets].T) - labels
            W_out[targets] -= lr * np.outer(error, W_in[centre])
            W_in[centre] -= lr * (error @ W_out[targets])

    return {w: v for w, v in zip(vocab, W_in - W_in.mean(0))}

Three Dimensions

Three attributes suggests three dimensions, one per attribute. Train it that way first and that is what you get: the three attribute directions come out near perpendicular, nothing has to share, and every analogy lands on the right word.

def pole(vectors, name):
    """The dense corpus has no token named after its pole, so average its synonyms."""
    if name in vectors:
        return vectors[name]
    return np.mean([vectors[t] for t in SYNONYMS[name]], axis=0)


def axis(vectors, a, b):
    """The direction an attribute points, from one pole to its opposite."""
    d = pole(vectors, a) - pole(vectors, b)
    return d / np.linalg.norm(d)


AXES = {"status": ("noble", "common"), "gender": ("masc", "fem"), "age": ("adult", "young")}

vectors_3d = train(sparse_corpus(), dim=3)
directions = {name: axis(vectors_3d, a, b) for name, (a, b) in AXES.items()}

names = list(AXES)
for i, a in enumerate(names):
    for b in names[i + 1:]:
        angle = np.degrees(np.arccos(np.clip(directions[a] @ directions[b], -1, 1)))
        print(f"{a:>7} to {b:<7} {angle:6.1f} degrees apart")

 status to gender    80.1 degrees apart
 status to age       99.1 degrees apart
 gender to age       80.3 degrees apart

Near enough to right angles. Each attribute has an axis to itself.

Two Dimensions

Now take one away. Same eight words, same three attributes, one fewer dimension to put them in.

vectors = train(sparse_corpus(), dim=2)

print("pole angles, measured from the x axis\n")
for p in sorted(POLES, key=lambda p: np.degrees(np.arctan2(*vectors[p][::-1])) % 360):
    print(f"  {p:<8} {np.degrees(np.arctan2(*vectors[p][::-1])) % 360:7.1f}")

print()
for name, (a, b) in AXES.items():
    v, w = vectors[a], vectors[b]
    opposite = np.degrees(np.arccos(np.clip(v @ w / (np.linalg.norm(v) * np.linalg.norm(w)), -1, 1)))
    print(f"  {a:>6} to {b:<8} {opposite:6.1f} degrees")

pole angles, measured from the x axis

  noble       11.4
  masc        71.2
  adult      133.1
  common     192.5
  fem        251.4
  young      312.7

   noble to common    178.9 degrees
    masc to fem       179.8 degrees
   adult to young     179.6 degrees

The six poles have spread into a hexagon. Each attribute is still an opposite pair, near enough antipodal. But the three axes now sit 60 degrees apart rather than 90.

They cannot be perpendicular any more, so they share the plane. Three attributes represented in a space with room for two, paid for by giving up the right angles. That’s superposition.

import matplotlib.pyplot as plt
from _style import COLOURS, MUTED, figure, style_axes

fig, ax = figure(width=6.6, height=6.0)

ring = np.array([vectors[p] for p in CYCLE + [CYCLE[0]]])
ax.plot(ring[:, 0], ring[:, 1], color="#cfd6de", linewidth=1.2, linestyle="--", zorder=1)

for p in POLES:
    x, y = vectors[p]
    ax.scatter(x, y, s=46, color=MUTED, zorder=3)
    ax.annotate(p, (x, y), xytext=(6, 4), textcoords="offset points", fontsize=9, color=MUTED)

for corner in CORNERS:
    x, y = vectors[corner]
    faded = np.linalg.norm(vectors[corner]) < 0.5
    ax.scatter(x, y, s=54, color=COLOURS[1] if faded else COLOURS[0], zorder=4)
    ax.annotate(corner, (x, y), xytext=(7, -10 if faded else 5),
                textcoords="offset points", fontsize=9,
                color=COLOURS[1] if faded else COLOURS[0])

ax.add_patch(plt.Circle((0, 0), 0.22, fill=False, linestyle=":", color=COLOURS[1], linewidth=1.2))
ax.set_aspect("equal")
style_axes(ax, grid=None)
ax.set_xticks([]); ax.set_yticks([])
for s in ax.spines.values():
    s.set_visible(False)
fig.tight_layout()

Figure 1: Three attributes in two dimensions. The six poles spread into a hexagon, and two of the eight corner words have been cancelled to the origin.

Two Words Disappear

Look at the middle of that plot. queen and boy are sitting on the origin, on top of each other, despite sharing no attributes at all.

They have not drifted there. They have been cancelled.

for corner in sorted(CORNERS, key=lambda c: np.linalg.norm(vectors[c])):
    print(f"  {corner:<9} {np.linalg.norm(vectors[corner]):.3f}")

  boy       0.054
  queen     0.061
  prince    1.549
  man       1.589
  princess  1.594
  king      1.642
  woman     1.648
  girl      1.671

Six words sit out at about 1.6. Two are at 0.05.

A corner word is the sum of its three poles. queen is noble + fem + adult, and in the hexagon those three sit at 11°, 251° and 133°. Each pair is 120° apart.

Three vectors of equal length at 120° sum to zero, so queen has nowhere to go but the centre. boy is common + masc + young, the same tripod rotated.

Those are the only two corners this happens to, and it’s not luck. Go around the pole cycle taking every other pole and you get exactly two triples, {noble, adult, fem} and {common, masc, young}. Every other corner takes three poles spanning a 120° arc, which add up constructively instead.

So you can predict which two words will vanish by looking at the cycle, before you train anything. I didn’t expect that. I assumed interference would show up as general noise. Instead it picked two specific words and deleted them.

The Analogies Still Pass

Take the predicted point, look up its nearest word, and every analogy comes back correct.

TESTS = [("king", "man", "woman", "queen"), ("king", "prince", "princess", "queen"),
         ("man", "king", "prince", "boy"), ("king", "queen", "woman", "man"),
         ("prince", "boy", "girl", "princess")]

print(f"{'analogy':<32}{'nearest':>10}{'error':>8}{'margin':>9}")
for a, b, c, expected in TESTS:
    target = vectors[a] - vectors[b] + vectors[c]
    ranked = sorted(((np.linalg.norm(vectors[w] - target), w) for w in CORNERS))
    (err, nearest), (runner_up, _) = ranked[0], ranked[1]
    flag = "" if nearest == expected else "  <- wrong"
    print(f"{a + ' - ' + b + ' + ' + c:<32}{nearest:>10}{err:>8.3f}{runner_up - err:>9.3f}{flag}")

analogy                            nearest   error   margin
king - man + woman                   queen   0.069    0.031
king - prince + princess             queen   0.082    0.001
man - king + prince                    boy   0.120    0.003
king - queen + woman                   man   0.069    1.511
prince - boy + girl               princess   0.182    1.315

Every one is right. Look at the margins.

The bottom two beat their runner-up by well over a point. They are in no danger. The middle two win by 0.001 and 0.003. A model doesn’t get an answer right by that margin. A coin lands on its edge by that margin.

Both of the close ones are aiming at queen or boy, the two words sitting on the origin a fraction apart. Anything pointing at one of them is choosing between two words that are, for practical purposes, in the same place.

And the coin does land the other way. Re-run the same corpus at other seeds and the close analogies come back as the wrong word.

smallest corner	analogies misread		smallest corner	analogies misread
0.054 (seed 9)	0		0.451 (seed 16)	2
0.200 (seed 21)	2		0.507 (seed 1)	2
0.284 (seed 13)	0		0.539 (seed 10)	1
0.413 (seed 6)	2		0.543 (seed 18)	3
0.441 (seed 2)	2		0.545 (seed 20)	1
0.442 (seed 23)	1		0.554 (seed 15)	3

Seed 9 is the run above. It has the worst collapse of the twelve, with queen and boy annihilated several times harder than anywhere else, and zero misreads. The two mildest collapses have three misreads each.

Whichever way that relationship runs, it is not “worse embedding, worse score”. Whether the analogies pass tells you very little about whether the embedding is any good.

The three-dimensional model makes the same point from the other side. There the errors are larger, and every answer is nearest its target by a mile. The 2D model is better at hitting the target and worse at being understood, because in 2D the target has someone else standing on it.

The Corpus Decides

None of this works on a corpus you would get by accident.

My first attempt was dense. Each pole carried three interchangeable synonym tokens, every corner word co-occurred with all of them, and I added extra pairs between the synonyms to hold each pole together. Ninety co-occurrence pairs against the sixty-six above.

It worked some of the time and not others. On four seeds out of ten, two or three of the attribute directions ended up on the same line.

def dense_corpus():
    """Every corner co-occurs with all three tokens of each of its poles, plus struts."""
    corpus = [f"{corner} {tok}" for corner, poles in CORNERS.items()
              for p in poles for tok in SYNONYMS[p]]
    corpus += [f"{a} {b}" for toks in SYNONYMS.values() for a, b in combinations(toks, 2)]
    return corpus


dense = train(dense_corpus(), dim=2)
dense_dirs = {name: axis(dense, a, b) for name, (a, b) in AXES.items()}

print(f"{'':>18}{'sparse':>9}{'dense':>9}")
names = list(AXES)
for i, a in enumerate(names):
    for b in names[i + 1:]:
        sp = np.degrees(np.arccos(np.clip(axis(vectors, *AXES[a]) @ axis(vectors, *AXES[b]), -1, 1)))
        dn = np.degrees(np.arccos(np.clip(dense_dirs[a] @ dense_dirs[b], -1, 1)))
        print(f"{a + ' to ' + b:>18}{sp:>8.0f}°{dn:>8.0f}°")

                     sparse    dense
  status to gender      59°     179°
     status to age     121°       2°
     gender to age      62°     176°

Same model, same attributes, same seed. Only the corpus changes. In the dense column the three attribute directions have collapsed onto one line: 179° and 2° are both “the same axis”, just pointing opposite ways.

What fixed it was giving the model less to work with, not more. The extra edges were the problem. Every corner pulling on every synonym of every pole gives the poles a reason to bunch together, and once two of them bunch there’s nothing holding the third direction apart.

I assumed this was a space problem that a third dimension would fix. It’s not. At three dimensions, with a whole dimension per attribute going spare, the dense corpus still sits about 27° off orthogonal on average against the sparse corpus’s 12°, usually by fusing two attributes onto one axis.

Conclusion

• A model can represent more features than it has dimensions, by giving up the right angles between them
• The cost is interference, and here it is not spread evenly as noise
• It picked two specific words and cancelled them to zero, predictably, from the structure of the corpus alone
• The analogy tests still pass, two of them by 0.001
• Passing the test tells you almost nothing about whether the space is any good

That last point is the one I would take anywhere else. The arithmetic still moves in the right direction, the errors stay small, and the analogies still come back correct. Superposition does not announce itself.

The code is on GitHub, and the idea is from Anthropic’s Toy Models of Superposition.