intermediate · ~20 min

Word Embeddings

Train tiny 2D word vectors from scratch and watch related words cluster together purely from co-occurrence statistics.

This module builds on Explainability Basics. Feel free to jump ahead anyway.

So far every input has been a number you could measure directly — a square footage, a pixel brightness, an (x, y) coordinate. Words aren't numbers. To feed "cat" or "bread" into a network, you need some way to turn a token into a vector first.

The naive approach — one-hot encoding, a vector of all zeros except a single 1 — technically works, but treats every pair of words as equally different: "cat" is just as far from "dog" as it is from "bread." A dense embedding instead learns a small vector per token — and the training signal below (predicting which words tend to appear near each other) naturally pulls related tokens' vectors closer together.

Beginner tip

Watch the points below during training: they start scattered randomly, and words from the same category visibly drift toward each other as training proceeds — nobody told the model which words are "animals" or "foods," it discovered the grouping purely from co-occurrence statistics.

🔍 Deep dive: Dot product as similarity

The training task here is genuinely simple: given two token vectors, predict whether they co-occur, using their dot product fed through a sigmoid. A large positive dot product means "these vectors point the same direction," which the loss function rewards for co-occurring pairs — so vectors that should be similar are gradually pushed toward pointing the same way. This is a tiny, literal version of the skip-gram idea behind word2vec.

Production note

Real embeddings (word2vec, GloVe, or a tf.keras.layers.Embedding layer trained end-to-end inside a larger model) use dozens to hundreds of dimensions, not 2 — 2D here is purely so the vectors can be plotted directly. Cosine similarity, used in the mini-project below, is the standard way to compare embeddings regardless of dimension, since it ignores vector length and only measures direction.

Playground

animalfoodcolor
catdogbirdfishapplebreadcheesesoupredbluegreenyellow
Step 0
steploss

Mini project

This toy vocabulary was only trained to cluster by category, not to learn fine-grained relationships within a category — so don't expect a "king − man + woman = queen"-style result here. What you will see: subtracting two same-category words nearly cancels out (their vectors are almost identical), so the result lands close to whichever third word you added.

+bread(cosine similarity 1.000)