Embeddings

in progressThis lesson is still being built - content, demos, and checkpoints may change before release.

Embeddings

Loading interactive demo…

Eight words, each a vector of two numbers. Every training step takes a pair the corpus put side by side and a stranger, scores each by a dot product with the centre word, and moves the three vectors so the pair points the same way and the stranger does not. Take the tour, or press and explore. Legend:

The panel beside the stage is the real EmbeddingTable class from the published library executing in your browser, with the current line highlighted.

The centre word of the triple.
Its neighbour, the word the corpus placed beside it.
The stranger, drawn from outside the pair.
The move a vector makes this step; the hollow point is the focused word.

pair cosine is the angle score of the current pair, loss the surprise at both scores, nearest the focused word's closest word by cosine after the run, pairs seen the triples finished. Tap a word to focus it.

The code panel shows bare executable code by default. The bar above it names the file (hover the name for the full published path); options holds the toggles: comments, the tracer instrumentation that produces this visualization, and verbose names for the terse identifiers (display only: the default is the byte-exact published file). The view switch adds plain-English narration: both interleaves it with the code, pseudo shows only the pseudocode. hide code collapses the panel; the rail it leaves behind brings it back.

Keys: Space play/pause · ←/→ step · Home/End jump · R reset · F fullscreen · G guided tour

A dot product is an angle. For vectors $a$ and $b$ ,

a \cdot b = \lVert a \rVert\, \lVert b \rVert \cos \theta,

so once lengths are divided out, similarity is the cosine of the angle between two words. The stage keeps two dimensions so the angle is visible; real tables use hundreds, and the same formula.

The objective. A training triple has a centre word $c$ , a word $o$ seen beside it, and a stranger $n$ drawn at random. With $\sigma$ the sigmoid, the model's scores are $p = \sigma(c \cdot o)$ and $q = \sigma(c \cdot n)$ , and the loss is the surprise at both:

L = -\ln p - \ln (1 - q).

It is small only when the pair scores high and the stranger low.

The gradients. Differentiating through the sigmoid gives

\frac{\partial L}{\partial c} = (p - 1)\, o + q\, n, \qquad \frac{\partial L}{\partial o} = (p - 1)\, c, \qquad \frac{\partial L}{\partial n} = q\, c,

and each vector steps against its own. The centre moves toward the neighbour and away from the stranger; the neighbour toward the centre; the stranger away from it. Every other word stays put.

Why the stranger. Without the $q$ term the loss is minimised by making every dot product large, which every vector can do by growing in one shared direction: all cosines go to 1 and nothing is learned. The stranger is what makes direction mean something. Edges: a zero vector has no cosine, which the library reports as 0; a word absent from the corpus moves only when it is drawn as a stranger; a rate near 1 overshoots and the pair swings past each other.

The idea that a word is characterised by the company it keeps is older than any neural network: linguists in the 1950s proposed that words used in the same contexts tend to mean similar things. Embeddings make that hypothesis a training objective. Each word owns a vector, pairs seen together are pulled to point the same way, and a random stranger is pushed away so the vectors cannot all collapse into one direction. The 2013 word2vec toolkit made this fast enough for whole corpora, and its vectors turned out to carry analogies as arithmetic.

1954Zellig Harris - States the distributional hypothesis: words that occur in the same contexts tend to have similar meanings.source
1950sJohn Rupert Firth - Popularises the slogan that a word is characterised by the company it keeps.source
2003Yoshua Bengio and colleagues - 'A Neural Probabilistic Language Model' learns word vectors jointly with a neural language model.source
2013Tomáš Mikolov and colleagues at Google - Word2vec, in two papers, trains skip-gram vectors with negative sampling far faster than earlier approaches; this stage runs its objective in two dimensions.source
2014Jeffrey Pennington, Richard Socher, and Christopher Manning - GloVe at Stanford fits vectors to global co-occurrence counts instead of one pair at a time.source

Used today

Language models - The first layer of a transformer is an embedding table; the vectors this stage learns are what self-attention reads.
Search and recommendation - Documents, products, and users are embedded so that nearest neighbours by cosine are the results.
Anything with tokens - Amino acids, code tokens, and graph nodes are embedded the same way: a lookup table trained by what appears near what.

rate0.4

steps12

corpus

pair cosine: 0.49
loss: 1.691
nearest to king: prince 0.70
pairs seen: 12

embedding-table.ts

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Sigmoid } from "./activation";/** One training triple: a centre word, a word seen beside it, and a word that was not. */export interface ContextTriple {  center: number;  context: number;  negative: number;}/** The smallest probability the loss will take a logarithm of. */const FLOOR = 1e-12;/** * Word embeddings learned by skip-gram with one negative sample: each * word owns a vector, and a training triple pulls the centre and its * true neighbour together while pushing the centre and a random * stranger apart. Similarity is the cosine of the angle between two * vectors, so a dot product is the whole story. Vectors start uniform * in [-1, 1] from the seed alone. * * Time: O(dim) per triple. Space: O(vocab · dim). */export class EmbeddingTable {  private readonly vectors: number[][];  private readonly sigmoid = new Sigmoid();  constructor(    readonly vocab: readonly string[],    readonly dim: number,    readonly learningRate: number,    seed = 1,    private readonly tracer?: Tracer,  ) {    if (vocab.length < 2) throw new RangeError("a vocabulary needs at least two words");    if (!Number.isInteger(dim) || dim < 1) throw new RangeError("the dimension must be a whole number of at least one");    if (!(learningRate > 0)) throw new RangeError("the learning rate must be positive");    const random = new Mulberry32(seed);    this.vectors = vocab.map(() => Array.from({ length: dim }, () => random.between(-1, 1)));  }  /** A copy of one word's vector. */  vector(index: number): number[] {    const v = this.vectors[index];    if (!v) throw new RangeError(`no word at index ${index}`);    return [...v];  }  /** Cosine similarity of two words' vectors; zero when either has no length. */  similarity(a: number, b: number): number {    const va = this.vector(a);    const vb = this.vector(b);    const dot = va.reduce((s, x, i) => s + x * vb[i]!, 0);    const na = Math.hypot(...va);    const nb = Math.hypot(...vb);    return na === 0 || nb === 0 ? 0 : dot / (na * nb);  }  /** The word most similar to another, and how similar. */  nearest(index: number): { index: number; cosine: number } {    let best = { index: -1, cosine: -Infinity };    for (let i = 0; i < this.vocab.length; i++) {      if (i === index) continue;      const cosine = this.similarity(index, i);      if (cosine > best.cosine) best = { index: i, cosine };    }    return best;  }  /** Trains on each triple in turn and returns the loss per triple. */  train(triples: readonly ContextTriple[]): number[] {    const losses: number[] = [];    for (let step = 0; step < triples.length; step++) {      const { center, context, negative } = triples[step]!;      const c = this.vector(center);      const centerVector = this.vector(center);      const o = this.vector(context);      const contextVector = this.vector(context);      const n = this.vector(negative);      const negativeVector = this.vector(negative);      // ▸ three words become three vectors: the centre, its neighbour, a stranger      this.tracer?.emit("lookup", { step, center, context, negative, centerWord: this.vocab[center], contextWord: this.vocab[context], negativeWord: this.vocab[negative] });      const scores: number[] = [];      for (const [kind, partner] of [["context", o], ["negative", n]] as const) {      for (const [kind, partner] of [["context", contextVector], ["negative", negativeVector]] as const) {        const dot = c.reduce((s, x, i) => s + x * partner[i]!, 0);        const dot = centerVector.reduce((s, x, i) => s + x * partner[i]!, 0);        const probability = this.sigmoid.apply(dot);        scores.push(probability);        // ▸ a dot product squashed to a probability: how alike are these two        this.tracer?.emit("score", { step, kind, dot, probability });      }      const p = scores[0]!;      const q = scores[1]!;      const loss = -Math.log(Math.max(p, FLOOR)) - Math.log(Math.max(1 - q, FLOOR));      losses.push(loss);      // ▸ surprise that the neighbour scored low or the stranger high      this.tracer?.emit("loss", { step, loss });      const gradients: [number, number[]][] = [        [center, c.map((_, i) => (p - 1) * o[i]! + q * n[i]!)],        [center, centerVector.map((_, i) => (p - 1) * contextVector[i]! + q * negativeVector[i]!)],        [context, c.map((x) => (p - 1) * x)],        [context, centerVector.map((x) => (p - 1) * x)],        [negative, c.map((x) => q * x)],        [negative, centerVector.map((x) => q * x)],      ];      for (const [index, gradient] of gradients) {        const before = this.vector(index);        const after = before.map((x, i) => x - this.learningRate * gradient[i]!);        this.vectors[index] = after;        // ▸ each vector steps against its share of the surprise        this.tracer?.emit("move", { step, word: this.vocab[index], index, gradient, before, after });      }    }    return losses;  }}

loss =: 1.691

12 steps done with the loss at 1.691; the focused word's nearest word is now prince at a cosine of 0.70.

import { Mulberry32 } from "../random";import type { Tracer } from "../trace";import { Sigmoid } from "./activation";/** One training triple: a centre word, a word seen beside it, and a word that was not. */export interface ContextTriple { center: number; context: number; negative: number;}/** The smallest probability the loss will take a logarithm of. */const FLOOR = 1e-12;/** * Word embeddings learned by skip-gram with one negative sample: each * word owns a vector, and a training triple pulls the centre and its * true neighbour together while pushing the centre and a random * stranger apart. Similarity is the cosine of the angle between two * vectors, so a dot product is the whole story. Vectors start uniform * in [-1, 1] from the seed alone. * * Time: O(dim) per triple. Space: O(vocab · dim). */export class EmbeddingTable { private readonly vectors: number[][]; private readonly sigmoid = new Sigmoid(); constructor( readonly vocab: readonly string[], readonly dim: number, readonly learningRate: number, seed = 1, private readonly tracer?: Tracer, ) { if (vocab.length < 2) throw new RangeError("a vocabulary needs at least two words"); if (!Number.isInteger(dim) || dim < 1) throw new RangeError("the dimension must be a whole number of at least one"); if (!(learningRate > 0)) throw new RangeError("the learning rate must be positive"); const random = new Mulberry32(seed); this.vectors = vocab.map(() => Array.from({ length: dim }, () => random.between(-1, 1))); } /** A copy of one word's vector. */ vector(index: number): number[] { const v = this.vectors[index]; if (!v) throw new RangeError(`no word at index ${index}`); return [...v]; } /** Cosine similarity of two words' vectors; zero when either has no length. */ similarity(a: number, b: number): number { const va = this.vector(a); const vb = this.vector(b); const dot = va.reduce((s, x, i) => s + x * vb[i]!, 0); const na = Math.hypot(...va); const nb = Math.hypot(...vb); return na === 0 || nb === 0 ? 0 : dot / (na * nb); } /** The word most similar to another, and how similar. */ nearest(index: number): { index: number; cosine: number } { let best = { index: -1, cosine: -Infinity }; for (let i = 0; i < this.vocab.length; i++) { if (i === index) continue; const cosine = this.similarity(index, i); if (cosine > best.cosine) best = { index: i, cosine }; } return best; } /** Trains on each triple in turn and returns the loss per triple. */ train(triples: readonly ContextTriple[]): number[] { const losses: number[] = []; for (let step = 0; step < triples.length; step++) { const { center, context, negative } = triples[step]!; const c = this.vector(center); const centerVector = this.vector(center); const o = this.vector(context); const contextVector = this.vector(context); const n = this.vector(negative); const negativeVector = this.vector(negative); // ▸ three words become three vectors: the centre, its neighbour, a stranger this.tracer?.emit("lookup", { step, center, context, negative, centerWord: this.vocab[center], contextWord: this.vocab[context], negativeWord: this.vocab[negative] }); const scores: number[] = []; for (const [kind, partner] of [["context", o], ["negative", n]] as const) { for (const [kind, partner] of [["context", contextVector], ["negative", negativeVector]] as const) { const dot = c.reduce((s, x, i) => s + x * partner[i]!, 0); const dot = centerVector.reduce((s, x, i) => s + x * partner[i]!, 0); const probability = this.sigmoid.apply(dot); scores.push(probability); // ▸ a dot product squashed to a probability: how alike are these two this.tracer?.emit("score", { step, kind, dot, probability }); } const p = scores[0]!; const q = scores[1]!; const loss = -Math.log(Math.max(p, FLOOR)) - Math.log(Math.max(1 - q, FLOOR)); losses.push(loss); // ▸ surprise that the neighbour scored low or the stranger high this.tracer?.emit("loss", { step, loss }); const gradients: [number, number[]][] = [ [center, c.map((_, i) => (p - 1) * o[i]! + q * n[i]!)], [center, centerVector.map((_, i) => (p - 1) * contextVector[i]! + q * negativeVector[i]!)], [context, c.map((x) => (p - 1) * x)], [context, centerVector.map((x) => (p - 1) * x)], [negative, c.map((x) => q * x)], [negative, centerVector.map((x) => q * x)], ]; for (const [index, gradient] of gradients) { const before = this.vector(index); const after = before.map((x, i) => x - this.learningRate * gradient[i]!); this.vectors[index] = after; // ▸ each vector steps against its share of the surprise this.tracer?.emit("move", { step, word: this.vocab[index], index, gradient, before, after }); } } return losses; }}