Embeddings from scratch

The previous two posts got us as far as integer ids. A word-level tokenizer turns "Shall I hear more" into [8761, 5032, 4705, 6430] (Why word-level tokenizers break), and BPE does the same with subword pieces so nothing is ever out-of-vocabulary (Building BPE from scratch). That is enough to represent the text, but it leaves the model with a slightly absurd situation: id 5032 is no closer to 5033 than to 9000. The integers are names, not meaning. Hand a network those raw ids and there is no useful relationship for it to generalize from.

This is where embeddings come in. Give every id a vector, then learn those vectors rather than assigning them by hand. The training signal is simple and self-supervised: use a token to predict its neighbors, or use the neighbors to predict the token. Tokens that occur in similar contexts get pushed toward similar predictions, so their vectors tend to end up near one another. That is the word2vec idea (Mikolov et al., 2013), and it is small enough to build from scratch.

What is an embedding, really?

Before training one, I find it useful to make the word vector less mysterious. An embedding is just a list of numbers that describes something along a set of axes. If we choose the axes ourselves, a bear might be big, furry, wild, and not very royal; a dog might share the furriness while scoring lower on wildness. Similar things get similar lists. That is the whole intuition - the interesting part is how a model finds useful axes without us naming them.

For a more abstract version, imagine axes for royalty, gender, and age. king is roughly royal, male, adult; queen is royal, female, adult; man and woman share the gender direction without the royalty. The famous analogy then becomes ordinary arithmetic: king − man removes the male component, and adding woman puts the female component back, landing near queen. In other words, we are estimating the queen's features from the other three vectors. The figure below lets you switch into arithmetic mode and try it.

Loading feature vectors…

So the goal is not to hand-design vectors like these. It is to learn them from text. That leaves the first mechanical question: where do the vectors actually live, and how do you look one up?

The lookup table

Start with the mechanism, before any learning. An embedding layer is a matrix W of shape (vocab_size, embedding_dim) - one row per token. "Embedding token i" means taking row i of W.

You can write that as a matrix multiply: encode id i as a one-hot vector and compute one_hot @ W. That multiplication only selects the row where the 1 is, so the detour is wasted work. nn.Embedding skips it and indexes the row directly; it is the row-select, nothing more.

Loading embedding table…

The rows above are hand-picked to make the point - king and queen sit close, apple sits apart. In a fresh model W is random, and the rows mean nothing. The whole job of training is to move those rows into positions where distance encodes something. For a gentle, picture-first tour of the same idea, Jay Alammar's The Illustrated Word2vec is excellent.

Where the signal comes from

There are no labels here - just a stream of tokens. The trick of word2vec is to manufacture a supervised task out of raw text: for each position, treat the token as the target and the tokens around it as the context. Slide a fixed window across the corpus and every position yields one (target, context) example. With window_size=2 the context is the two tokens on each side:

def create_context_windows(token_ids, window_size=2):
    windows = []
    for i in range(window_size, len(token_ids) - window_size):
        target = token_ids[i]
        context = token_ids[i - window_size:i] + token_ids[i + 1:i + window_size + 1]
        windows.append((target, context))
    return windows

Loading context window…

There is one problem with feeding this stream straight in. A handful of tokens - newlines, punctuation, the, and, of - dominate the counts (the same Zipfian skew from the tokenization posts), and their sheer frequency drowns out the signal from rarer content words. Word2vec's fix is subsampling: drop a token with a probability that grows with its frequency, so common tokens survive less often (Mikolov et al., 2013 - Distributed Representations):

subsample_t = 1e-4
 
def keep_prob(token_id):
    freq = token_counts[token_id] / n_tokens
    return min(1.0, np.sqrt(subsample_t / freq))

A rare token has freq near subsample_t, so keep_prob is close to 1 and it is almost always kept; a token like the has a large freq, so the square root pulls its keep probability far below 1. On Karpathy's Tiny Shakespeare, encoded with GPT-2's tiktoken tokenizer and remapped to the 11,706 ids that actually occur, this turns ~300k tokens into roughly 128k training tokens. Those "words" are GPT-2 subword pieces, not whole words.

Skip-gram vs CBOW

There are two ways to point the prediction. Skip-gram predicts each context token from the target; CBOW (continuous bag-of-words) predicts the target from the averaged context. Both optimize a softmax over the whole vocabulary. For skip-gram, the probability of a context token $w_c$ given target $w_t$ is:

$P(w_c \mid w_t) = \frac{\exp\!\left(v_{w_c}^\top v_{w_t}\right)}{\sum_{w \in V} \exp\!\left(v_w^\top v_{w_t}\right)}$

Training maximizes that probability for the (target, context) pairs the corpus actually contains, which is the same as minimizing cross-entropy of the logits against the true context ids. Skip-gram is one embedding lookup followed by one linear layer; the complete pair of classes is in the collapsible below.

CBOW has the same two layers, but embeds all four context tokens and averages them before the linear layer. mean(dim=1) collapses the window into a single vector, so the order of the context tokens is thrown away (hence "bag of words"). The complete pair of classes is in the collapsible below.

One mechanic in the skip-gram training loop is worth slowing down on. The model produces one logit vector per target, but a target has four context tokens to predict. The loop expands the single prediction into four independent classifications with repeat_interleave:

pred = skipgram(x)
 
# One target predicts every context slot: tile the logits so each
# (target, context_token) pair becomes an independent classification.
_, n_context = y.shape
pred = pred.repeat_interleave(n_context, dim=0)
y = y.flatten()
 
loss = loss_fn(pred, y)

So a batch of 128 targets becomes 512 (target, context) classification problems, each pushing the same target logits toward a different neighbor. That is the standard skip-gram objective - "predict each surrounding word independently" - realized with a tile-and-flatten instead of a loop.

These are first-iteration implementations, built for clarity. For the full mathematical derivation of both objectives, Xin Rong's word2vec Parameter Learning Explained works through every gradient; Chris McCormick's skip-gram tutorial is the friendlier walk-through.

Training, and an honest result

Both models use a full softmax: the linear layer produces one logit for each of the 11,706 vocabulary ids, and cross-entropy normalizes over all of them. There is no negative sampling here; this is the plain, expensive version. Skip-gram runs 60 epochs, CBOW 80, both at lr=1e-3.

Loading loss curves…

The two curves tell different stories. Skip-gram falls from ~9.0 to ~5.65 and flattens: it is solving four hard predictions per target and never gets very confident. CBOW drops from ~8.8 to ~0.26. That looks better, but it is mostly a warning sign. Averaging the context gives CBOW an easier target and enough capacity to memorize this small corpus rather than generalize.

The neighbors confirm the suspicion. Ranking every content token by cosine similarity on the 128k training tokens gives semantically thin results. A few real skip-gram examples:

There are faint hits - you finds thee and grace, lord finds gracious - but the top scores hover around 0.3, and most neighbors are noise or broken subword fragments like entle and ifled. That is what raw, from-scratch word2vec on this corpus honestly produces: the famous clean analogies come from models trained on billions of tokens, not one megabyte of Shakespeare.

Seeing the space

Even weak embeddings have structure worth looking at, but 128 dimensions do not fit on a screen. The map below projects the top-60 content tokens down to 2D with t-SNE, with skip-gram on one side and CBOW on the other. Hover a point for its nearest neighbors.

Loading embedding map…

I used t-SNE rather than PCA because this view is about local neighborhoods, not the directions of greatest global variance. That is also what cosine similarity reads from the embeddings, so the projection is at least showing the kind of structure we are asking about.

For a cleaner read, here is a pairwise cosine-similarity heatmap over a curated 16-word Shakespeare vocabulary (king, queen, lord, love, death, ...), for each model. Bright cells are pairs the model placed close together:

Loading cosine matrix…

What we skipped, and where this goes

Two useful pieces were left out on purpose.

Negative sampling. The full softmax normalizes over all 11,706 ids on every example - fine here, but hopeless at real vocabulary sizes. Word2vec's speed trick replaces it with a binary problem: push the true pair together and a few randomly sampled negative pairs apart. That turns an 11,706-way softmax into a handful of logistic updates (Mikolov et al., 2013 - Distributed Representations; McCormick's negative sampling tutorial walks through it). That is the single change that made word2vec trainable at scale.

The king − man + woman ≈ queen analogy. It does hold on large models, but the usual demo quietly excludes the three input words when looking for the nearest vector. Without that exclusion the answer is often just king again. The effect is real, but the evaluation detail matters. Later count-based methods like GloVe reach similar geometry from co-occurrence statistics directly.

The deeper limitation is that these vectors are trained in isolation, on a proxy task, and then frozen. The next direction in this series drops the proxy: put an nn.Embedding at the front of a GPT and let the language-modeling loss shape it end-to-end, so the embeddings are learned for the task that actually uses them - and on enough text for the geometry to mean something.