LM3-T1 — From Token IDs to Embedding Vectors

Learning objectives

After this lesson, you should be able to:

• Describe an embedding table as a learned matrix of shape V × d and the lookup as row selection.
• Show that multiplying a one-hot vector by the table gives the same row as indexing.
• Compute dot products and cosine similarity between embedding vectors and interpret them.
• State how the table is trained and compute its parameter count.

The idea: give every token a location

Think of every token as a point in a space with d dimensions, for example d = 768. The embedding table stores one point per vocabulary entry. When the model reads the token with ID 1, it fetches row 1 of the table. That row, a vector of d numbers, is what flows through the rest of the network.

Why vectors?

A vector can express degrees of similarity that an ID cannot. Two tokens with similar roles (“robot” and “arm”, both things that act) can end up with vectors that point in similar directions, while an unrelated token (“box”) points elsewhere. Operations such as dot products then measure how alike two tokens are. The model is free to arrange the points however helps it predict the next token; nobody designs the axes. A single coordinate rarely has a human-readable meaning, but directions and distances often carry structure.

Where the numbers come from

The table is a matrix of learned parameters. It starts as small random numbers and is changed by gradient descent along with every other weight (Module 5). A token that is used in similar contexts receives similar gradient updates and drifts to a similar place. That is the practical form of the old idea that a word is characterised by the company it keeps.

Lookup, not computation

The embedding step does no arithmetic on its input: it copies a row. This is why it is fast, and why the gradient with respect to the table is sparse: in each step only the rows of the tokens that actually appeared receive a non-zero gradient.

What embeddings do not contain

The row for “box” is the same wherever the token appears. It knows nothing about its neighbours or its position. Context-dependence comes later: attention (Lesson LM3-T3) mixes the vectors of neighbouring tokens, and position information is added separately (Lesson LM3-T2).

The mathematics: lookup, one-hot products and similarity

Let E be the embedding matrix with V rows (one per token) and d columns. For a sequence of token IDs t₁, …, t_n the output is the matrix X whose row i is row t_i of E:

V

vocabulary size: number of rows of E.

d

embedding width: number of columns; the model’s hidden size.

t_i

the integer ID of the i-th token, 0 ≤ t_i < V.

X

the matrix of embeddings of the sequence; shape (n, d), or (B, n, d) for a batch of B sequences.

One-hot view. Let o_i be the one-hot vector of length V with a 1 at position t_i and 0 elsewhere. Then

because every other term of the product is multiplied by 0. Frameworks index instead of multiplying: the result is identical and needs d copies rather than V · d multiplications.

Similarity. Two embedding vectors a and b have a dot product and a cosine similarity:

a_k, b_k

component k of the two vectors, k = 1, …, d.

‖a‖

the length of a: the square root of the sum of its squared components.

cos(a, b)

between −1 and 1: 1 means the same direction, 0 unrelated (orthogonal), −1 opposite.

Attention (Lesson LM3-T3) uses the raw dot product, which also grows with vector length; cosine similarity removes the length and is the usual way to inspect embeddings.

Size. The table has V · d parameters. GPT-2 small has V = 50,257 and d = 768, so 50,257 × 768 = 38,597,376 parameters: roughly 31% of its 124 million. In float32 that is about 154 MB for the table.

A worked example: four tokens with three features

Use the toy table from the figure, with token IDs robot = 0, arm = 1, box = 2, tool = 3:

Dot products. robot · arm = 1.0 × 0.9 + 0.0 × 0.1 + 1.0 × 0.8 = 1.7. robot · box = 0 + 0 + 0 = 0. robot · tool = −1.0 + 0 + 0.5 = −0.5.

Cosine. ‖robot‖ = √2 ≈ 1.414, ‖arm‖ = √(0.81 + 0.01 + 0.64) = √1.46 ≈ 1.208, so cos(robot, arm) = 1.7 / (1.414 × 1.208) ≈ 0.995, almost the same direction. cos(robot, box) = 0 (orthogonal) and cos(robot, tool) = −0.5 / (1.414 × 1.118) ≈ −0.316 (a mild opposition). In this invented table “robot” and “arm” are close, “box” is unrelated.

The sequence [0, 1, 2] (robot, arm, box) becomes the 3 × 3 matrix whose rows are the first three rows of the table. A batch of two such sequences has shape (2, 3, 3). Nothing about order or neighbours is in the matrix yet.

The same thing in PyTorch

import torch
import torch.nn as nn

table = torch.tensor([[ 1.0, 0.0, 1.0],
                      [ 0.9, 0.1, 0.8],
                      [ 0.0, 1.0, 0.0],
                      [-1.0, 0.0, 0.5]])          # V = 4, d = 3

emb = nn.Embedding(num_embeddings=4, embedding_dim=3)
with torch.no_grad():
    emb.weight.copy_(table)                       # use our toy numbers instead of random ones

ids = torch.tensor([0, 1, 2])
x = emb(ids)                                      # lookup: shape (3, 3)
print(tuple(x.shape))                             # (3, 3)

one_hot = torch.nn.functional.one_hot(ids, num_classes=4).float()
print(torch.allclose(x, one_hot @ table))         # True: one-hot times the table equals the lookup

cos = torch.nn.functional.cosine_similarity
print(round(cos(table[0], table[1], dim=0).item(), 3))   # 0.995
print(round(cos(table[0], table[3], dim=0).item(), 3))   # -0.316
print(sum(p.numel() for p in emb.parameters()))   # 12  (V * d)

nn.Embedding holds the table as emb.weight, shape (V, d), and is called with an integer tensor of any shape; the output has one extra trailing dimension of size d.

Engineering implications and limitations

IDs must be integers (long tensors). Passing floats raises an error. A frequent bug is an ID that is out of range (≥ V), which gives an index error on CPU and a cryptic device-side assertion on GPU.

Initialisation matters little but not nothing. The book initialises from a normal distribution; any small random values break symmetry. All-zero embeddings would give every token the same gradient.

The embedding table is often tied to the output layer. Some models reuse the transposed table for the final projection to vocabulary logits (weight tying), saving V · d parameters and often helping small models. GPT-2 does this; the book’s version keeps them separate for clarity and mentions the choice.

Rare tokens get poor embeddings. A token that appears twice in the training data hardly moves from its random start. This is one reason subword tokenizers avoid very rare tokens, and a reason that glitchy, almost-unused tokens can behave strangely.

Do not over-read geometry. Nearest-neighbour lists and 2-D projections of embeddings are useful for debugging, but a single dimension is not “the gender axis”, and cosine similarity in one layer says nothing about what the full model does.

Memory. The table is a large share of a small model’s parameters. Training updates it with the optimizer state (Adam keeps two extra values per parameter), so the training memory for the table is about four times its size.

For robots. When mixing text with other signals (Physical AI and VLA tracks), non-text inputs such as image patches or joint readings are also mapped to vectors of the same width d and treated as tokens, so the idea of “everything becomes a row of a matrix” carries over; see the Vision-Language-Action track for how that is done.

Common misconceptions

Engineering takeaways

1. An embedding table E of shape (V, d) maps token IDs to d-dimensional vectors; the lookup selects rows and equals a one-hot product.
2. The sequence of n tokens becomes a matrix of shape (n, d), or (B, n, d) for a batch.
3. Dot products and cosine similarity measure how alike two vectors are; the table is learned, so similarity emerges from training.
4. The table has V × d parameters: 38,597,376 for GPT-2 small, about a third of the model.
5. Embeddings carry no context or position by themselves; attention and position encodings supply both.

References and further reading

See Chapter 2 of the book (creating token embeddings) and the PyTorch documentation for torch.nn.Embedding.

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: why attention needs position information and how it is added to the embeddings.

Related content

Nothing below is required to finish this lesson.

• Before
  LM2-T6: Batching, Data Loaders and Data Splits

  Group examples into batches, shuffle correctly, and split data so validation loss is honest.

• Lab
  LM3-L1: Compute Attention Step by Step

  Compute scores, softmax weights and context vectors for a short sentence by hand-written code.