Word Embeddings: Token IDs Become Dense Vectors

Token IDs are arbitrary integers. The number 47458 means whatever the BPE merge that produced it happened to mean — there is no semantic structure in the integer itself. The model's job, before any attention happens, is to turn each ID into a dense vector that does have semantic structure.

This is a learned lookup. The embedding matrix has shape (vocab_size × d_model). Indexing it by token ID returns a row of size d_model — that row is the token's initial representation. During training, gradients flow back to those rows, nudging them so that tokens used in similar contexts end up with similar vectors.

The result is a high-dimensional space where geometry encodes meaning. "Cat" and "dog" end up close. "King" and "queen" end up close. "Run" and "running" end up close. This is the substrate that all subsequent layers refine — they don't operate on words, they operate on points in this space.

Code

Embedding matrix in PyTorch·python

import torch
import torch.nn as nn

vocab_size, d_model = 50_257, 768   # GPT-2 base
emb = nn.Embedding(vocab_size, d_model)

# emb.weight has shape (50257, 768) and is learned
ids = torch.tensor([464, 2368, 3290])     # "The cat sat"
vectors = emb(ids)                        # (3, 768)

# Vectors start random. Training reshapes the matrix so that
# semantically similar tokens have similar rows.

What 'similar' means geometrically·python

import torch.nn.functional as F

# After training, you can probe what's near what
def neighbors(token, k=5):
    target = emb.weight[token_id(token)]     # (d_model,)
    sims = F.cosine_similarity(emb.weight, target.unsqueeze(0))
    return sims.topk(k).indices.tolist()

# In a trained Transformer, neighbors('cat') typically include
# 'dog', 'kitten', 'cats', 'feline', 'pet' — in some order.
# In a randomly-initialized model, you'd get random neighbors.

Exercise

Load gpt2 and extract its input embedding matrix (model.transformer.wte.weight). Pick a target word like 'computer.' Find its 10 nearest neighbors by cosine similarity in the embedding space. Are the neighbors mostly synonyms, related concepts, morphological variants, or noise? What does that tell you about what the embedding has captured?

Word Embeddings: Token IDs Become Dense Vectors

Code

External links

Exercise

Progress

Comments 0