Embedding spaces

From an integer to a direction in space

The embedding matrix has one row per vocabulary entry and one column per dimension — for GPT-2 that is 50,257 rows by 768 columns. Tokenizing gives you row indices; the embedding layer looks up those rows. That is the whole operation.

What makes it interesting is that the rows are learned. Nothing forces 'king' and 'queen' near each other at initialisation — they start random. Training on prediction pushes tokens that appear in similar contexts toward similar positions, and semantic structure appears as a side effect of getting the next word right.

The four things called 'embeddings'

Cosine similarity, and why not Euclidean distance

Cosine similarity is the dot product of two vectors divided by their lengths — the cosine of the angle between them. It runs from 1 (same direction) through 0 (orthogonal, unrelated) to −1 (opposite).

Direction is used rather than distance because vector magnitude in these spaces tracks token frequency far more than meaning. Two words can be semantically identical while one is far longer simply because it appeared more often. Normalising away the length removes that nuisance.

There is a neat identity worth knowing: cosine similarity on mean-centred vectors is exactly the Pearson correlation between them. So 'these two embeddings are similar' and 'these two vectors correlate' are the same statement, which is why so much interpretability analysis reduces to correlation matrices.

import numpy as np

def cosine(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

a = np.random.randn(300)
b = a + 0.5 * np.random.randn(300)

print(cosine(a, b))                                   # high: b is mostly a
print(cosine(a, 100 * a))                             # 1.0 — scale is irrelevant

# Mean-centre first and cosine IS Pearson correlation.
ac, bc = a - a.mean(), b - b.mean()
print(cosine(ac, bc), np.corrcoef(a, b)[0, 1])        # identical

# Batched: normalise the rows, then one matmul gives every pair.
def cosine_matrix(X):
    Xn = X / np.linalg.norm(X, axis=1, keepdims=True)
    return Xn @ Xn.T

Note: Normalise once and multiply once. Looping over pairs is the standard beginner performance mistake here.

Contextual embeddings, and unembedding

In a transformer the input embedding is only the starting point. Each layer adds to it, so the representation of a token at layer 12 has absorbed information from the surrounding tokens. That evolving vector is what the residual stream carries.

At the output the process runs backwards. The final hidden state is multiplied by an unembedding matrix — often the transpose of the input embedding, a trick called weight tying — to produce one score per vocabulary entry. Softmax turns those scores into a distribution over next tokens.

You can apply the unembedding to any layer's hidden state, not just the last. Doing so shows what the model would have predicted had it stopped there, which is the whole idea behind the logit lens later in the curriculum.

Position embeddings

Attention has no inherent sense of order — it computes a weighted sum, and sums do not care about arrangement. Without positional information, 'the dog bit the man' and 'the man bit the dog' would be identical inputs.

The fix is to add a position-dependent vector to each token embedding before the first layer. GPT-2 learns these vectors, one per position up to the context length, which is why its context length is a hard architectural limit. Sinusoidal schemes compute them instead, and modern models mostly use rotary embeddings, which rotate the query and key vectors by an angle proportional to position — encoding relative distance rather than absolute index and generalising better beyond the training length.

Inspecting learned position embeddings is instructive: nearby positions have high cosine similarity and it decays smoothly with distance, so the model has discovered a notion of 'close by' without ever being told to.

Training embeddings yourself

Building this end to end is what makes it concrete. You need four pieces: a data loader that yields (context, target) pairs from a tokenized corpus; a model that embeds the context and projects to vocabulary scores; a loss that measures how wrong those scores are; and a training loop.

The loss is cross-entropy on the next token, which is where the semantics come from. To lower it, the model must place tokens that predict similar continuations near each other — so the geometry is not an objective, it is a consequence.

import torch, torch.nn as nn
from torch.utils.data import DataLoader, Dataset

class NextTokenData(Dataset):
    """Fixed-length windows; the target is the window shifted by one."""
    def __init__(self, ids, block=16):
        self.ids, self.block = torch.tensor(ids), block

    def __len__(self):
        return len(self.ids) - self.block - 1

    def __getitem__(self, i):
        return self.ids[i : i + self.block], self.ids[i + 1 : i + self.block + 1]

class EmbeddingLM(nn.Module):
    def __init__(self, vocab, dim=64):
        super().__init__()
        self.emb = nn.Embedding(vocab, dim)
        self.out = nn.Linear(dim, vocab)

    def forward(self, x):
        return self.out(self.emb(x))          # [B, T, vocab]

model = EmbeddingLM(vocab_size)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

loader = DataLoader(NextTokenData(train_ids), batch_size=32, shuffle=True)

for epoch in range(5):
    for x, y in loader:
        logits = model(x)
        # Flatten to [B*T, vocab] and [B*T] — cross-entropy wants 2-D input.
        loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    print(epoch, loss.item())

# Snapshot model.emb.weight each epoch to watch the geometry form:
# similarities among related tokens rise while unrelated pairs stay flat.

Note: CrossEntropyLoss expects raw logits, not softmax output — it applies log-softmax internally. Passing softmax probabilities is a silent, common bug that just trains badly.

What stability tells you

Retrain with a new seed and measure the correlation between the two similarity matrices. Frequent tokens land in highly reproducible relative positions; rare tokens move a great deal, because a handful of occurrences is not enough evidence to place them.

That gives you a practical reliability check: any interpretability claim about a rare token, a specific dimension, or a single neuron should be repeated across seeds before you believe it. Much of the observational work later in this curriculum lives or dies on exactly this discipline.