AIM-03-002-P · Primer — Token Embeddings

The idea in one sentence

A token embedding is a learned dense vector — a single row of a big lookup matrix — that stands in for a discrete token id so the network can do arithmetic on words instead of just comparing IDs for equality.

Why it exists

Neural networks only understand continuous vectors, but text starts life as a sequence of arbitrary integers (token ids) from a tokenizer's vocabulary — and integers 3796 and 3797 being adjacent tells you nothing about whether the words they represent are related. It's like a library that assigns every book a random accession number: number 3796 could be "Moby Dick" and 3797 could be "Quantum Field Theory," with the numeric closeness meaning nothing. An embedding replaces that arbitrary numbering with a GPS coordinate in a learned "meaning space" — books on similar topics end up near each other, and the network can measure that nearness with ordinary dot products and distances.

How it works

Concretely, an embedding layer is just a matrix \(E \in \mathbb{R}^{V \times d}\), where \(V\) is the vocabulary size and \(d\) is the model's hidden dimension. Row \(i\) of \(E\) is token \(i\)'s embedding. Looking up a token is mathematically a one-hot matrix multiply, \(e = E^\top x\) where \(x\) is a one-hot vector with a 1 at position \(i\), but every real implementation (PyTorch's nn.Embedding, the input layer of every Hugging Face model) skips the wasted multiply-by-zero work and just gathers row \(i\) directly — an \(O(d)\) array copy, not an \(O(V d)\) matmul. GPT-2 small, for example, uses \(V = 50257\) and \(d = 768\), so \(E\) is a 50,257 × 768 matrix (about 38.6M numbers) that starts out randomly initialized and gets tuned by ordinary backprop alongside every other weight in the model, driven purely by the language-modeling loss — there is no separate "meaning" objective, similarity in embedding space is an emergent side effect of predicting the next token well.

Here's a worked example with a toy vocabulary of 6 tokens and \(d = 4\): \(E = \begin{bmatrix} 0.10 & -0.20 & 0.05 & 0.30 \\ 0.44 & 0.12 & -0.08 & 0.19 \\ -0.15 & 0.33 & 0.27 & -0.05 \\ 0.08 & -0.02 & 0.41 & 0.22 \\ 0.19 & 0.19 & -0.19 & 0.05 \\ 0.00 & 0.00 & 0.00 & 0.00 \end{bmatrix}\) with rows 0–5 meaning "the," "cat," "sat," "on," "mat," "\(\langle\text{pad}\rangle\)." The sentence "cat sat" tokenizes to ids \([1, 2]\), and lookup simply reads off rows 1 and 2: \([0.44, 0.12, -0.08, 0.19]\) and \([-0.15, 0.33, 0.27, -0.05]\) — no computation, just indexing.

This same row-gather mechanic explains why mean pooling matters in practice: a real BPE tokenizer splits an out-of-vocabulary word like "unhappiness" into multiple subword pieces (say "un", "happi", "ness"), each with its own row in \(E\) — there is no single embedding for the whole word. A common way to get one representative vector anyway — for a quick similarity check, a bag-of-embeddings classifier, or probing a model's vocabulary knowledge — is to average the looked-up rows elementwise: \[\bar{e} = \frac{1}{|S|}\sum_{i \in S} E_i\] where \(S\) is the set of token ids and \(E_i\) is row \(i\). For "un", "happi", "ness" with (toy) embeddings \([1,0]\), \([0.6,0.4]\), \([0.2,0.8]\), the pooled vector is \(\left(\frac{1+0.6+0.2}{3}, \frac{0+0.4+0.8}{3}\right) = (0.6, 0.4)\) — a single point standing in for the whole word, computed with nothing fancier than a sum and a divide.

See it — Token ids gather rows from the embedding matrix, which are then mean-pooled into one vector.

flowchart TD
  Ids["Token ids<br/>[1, 2]"] --> Matrix["Embedding<br/>matrix E<br/>(V rows x d cols)"]
  Matrix --> Row1["Row 1<br/>(cat)"]
  Matrix --> Row2["Row 2<br/>(sat)"]
  Row1 --> Pool["Mean pool<br/>elementwise avg"]
  Row2 --> Pool
  Pool --> Out["Pooled<br/>vector"]

The one gotcha

The embedding lookup is context-free: the token id for "bank" gets the exact same row whether the sentence is about rivers or money — every occurrence of a given token id maps to the identical vector straight out of \(E\). Disambiguation only happens after the embedding layer, once attention and the transformer's other layers mix that vector with its neighbors. It's easy to conflate "the embedding" (context-free, one row per id) with "the contextual representation" (a different vector per occurrence, produced by later layers) — they are not the same object, and confusing them leads people to expect a static embedding table to already encode word-sense disambiguation that it structurally cannot.

The one-line takeaway

Token id to vector is a matrix row lookup.

Where you'll meet it

PyTorch nn.Embedding; Hugging Face model.get_input_embeddings() on any AutoModel; tiktoken produces the ids that feed the lookup (it does not itself hold embeddings); the OpenAI and Anthropic embeddings-style APIs expose a pooled sentence-level version of this same idea.

Named concepts: token embeddings