Embedding Matrix

The embedding matrix E is a learned table with shape [V, D]: V rows (one per vocabulary entry, usually) and D columns (the embedding dimension). A token ID selects a row; that row is the token's embedding.

The transformer needs numbers it can do arithmetic on, and token IDs are not those numbers (see Token Embeddings). A table solves three problems at once.

• Fast access: a token ID is a ready-made index. Reading row i needs no search.
• Learnable: every number in the table is a parameter. Training can nudge a token's row to make the model's predictions better.
• Rich representation: a row of hundreds of numbers can hold far more than a single integer.

A token ID works like a row ID. The same number is the key into two different tables. (Small IDs below are illustrative.)

• The vocabulary is a dictionary-like mapping ("cat" → 1), usually a hash map.
• The embedding matrix is not a dictionary. It is a dense array indexed by position, E[1], and stores no token strings.
• They line up only through the shared number. Change the tokenizer and every ID points at the wrong row.

Is the vector stored inside the token ID?

No. For token ID 1 (cat), the matrix row is a short list of learned numbers (illustrative, 4 values):

The ID is the address; the four numbers are the vector. Here the dimension is 4. A real matrix with 50,000 tokens and 4,096 dimensions has shape 50,000 × 4,096: 50,000 rows of 4,096 numbers each. See Vector Dimensions for a slider that shows how the vector grows, and How Embeddings Are Learned for where these numbers come from.

What the matrix is not

• Not a dictionary of definitions. A row is a learned numeric pattern, not a stored meaning.
• Not context-aware. A row does not depend on the other tokens. Context is added later (see Contextual Embeddings).
• Not the whole model. Knowledge also lives in the transformer weights. The matrix is one set of learned parameters among several.

Use the matrix from the introduction (V = 5, D = 3). All numbers are invented for teaching.

Lookup equals a one-hot product

Frameworks index the row directly instead of multiplying, because the one-hot vector is almost all zeros.

Why the gradient is sparse

Row $j$ contributes to the loss $\mathcal{L}$ only through the positions where the token ID equals $j$:

If token $j$ is absent from the batch, the sum is empty and the row gets zero gradient for that step.

Parameter count and the output side

Many language models reuse $\mathbf{E}$ to turn the final hidden vector $\mathbf{h}$ into scores over the vocabulary ("weight tying"): $\text{logits} = \mathbf{h}\,\mathbf{E}^{\top}$. GPT-2 does this.

Type words from the toy vocabulary (the cat dog bank river). The lab shows their IDs, highlights the rows that get read, and stacks them into the output matrix.

Part A — pure Python:

vocab = {"the": 0, "cat": 1, "dog": 2, "bank": 3, "river": 4}     # Table 1: token -> ID (illustrative)

E = [                                                              # Table 2: ID -> vector (illustrative)
    [ 0.2,  0.5, -0.1],
    [ 0.8, -0.3,  0.4],
    [-0.2,  0.7,  0.1],
    [ 0.1,  0.6, -0.5],
    [ 0.0,  0.4, -0.6],
]

ids = [vocab[w] for w in "the cat".split()]
print(ids)                                    # [0, 1]

X = [E[i] for i in ids]                       # the lookup
print(X)                                      # [[0.2, 0.5, -0.1], [0.8, -0.3, 0.4]]

# same thing as a one-hot product
one_hot = [0, 1, 0, 0, 0]
row = [sum(one_hot[r] * E[r][c] for r in range(5)) for c in range(3)]
print(row)                                    # [0.8, -0.3, 0.4]

# one gradient-descent step on row 1
grad, lr = [0.5, -1.0, 0.2], 0.1
E[1] = [round(E[1][c] - lr * grad[c], 3) for c in range(3)]
print(E[1])                                   # [0.75, -0.2, 0.38]

Part B — PyTorch:

import torch
import torch.nn as nn

emb = nn.Embedding(num_embeddings=5, embedding_dim=3)     # E has shape [5, 3], random start
ids = torch.tensor([0, 1])

x = emb(ids)
print(x.shape)                                             # torch.Size([2, 3])
print(torch.equal(x, emb.weight[ids]))                     # True  -> a lookup is row indexing

one_hot = torch.nn.functional.one_hot(ids, num_classes=5).float()
print(torch.allclose(one_hot @ emb.weight, x))             # True  -> same as a one-hot product

emb(torch.tensor([1, 1])).sum().backward()                 # use token 1 twice
used = (emb.weight.grad.abs().sum(dim=1) > 0).nonzero().flatten().tolist()
print(used)                                                # [1]  -> only row 1 received gradient

• Memory and size: V × D parameters. Doubling the vocabulary doubles this table.
• Adding tokens: after adding tokens to a tokenizer, call model.resize_token_embeddings(len(tokenizer)) so each new ID has a row. New rows start untrained and need fine-tuning to become useful.
• Weight tying: reusing E as the output projection saves V × D parameters.
• Fine-tuning choices: some recipes freeze E to save memory or keep the original vocabulary behaviour; others train it, especially when new tokens are added.
• Swapping tokenizers: not possible without retraining or remapping, because rows are aligned to the original vocabulary.

🚫 Common Misconceptions to Correct

• Misconception: the vocabulary and the embedding matrix are one thing.
  Reality: The vocabulary belongs to the tokenizer; the matrix belongs to the model. The ID joins them.
• Misconception: the lookup learns.
  Reality: The lookup only reads. Learning happens when the optimizer updates the rows during training.
• Misconception: the matrix is the model's whole knowledge.
  Reality: The transformer weights hold a large part of it too.

📌 Summary — Everything in Six Bullets

• E is a learned table of shape [V, D]: one row per vocabulary entry, usually.
• The token ID is the row index. It is a key, not a meaning.
• Two tables, one ID. The vocabulary maps tokens to IDs; E maps IDs to vectors.
• Lookup reads, training writes. Only rows used in a batch receive gradient.
• Geometry carries meaning. Closeness of vectors matters; closeness of row numbers does not.
• Pipeline: Text → Token ID → Embedding matrix row → Vector → Transformer.