Module 04: Embeddings

Introduction

An embedding turns a token ID into numbers a network can actually use. A token embedding is the learned vector stored in the table row selected by a token ID.

Why it matters for LLMs:

  • Everything downstream reads vectors: attention, feedforward layers, and output heads all operate on embeddings, never on raw IDs.
  • Lookup is the whole mechanism: one indexing operation converts a symbolic address into adjustable numeric information.
  • Learning lives in the rows: training never moves an ID — it reshapes the vectors the IDs point to, and useful geometry can emerge as a consequence.

What You’ll Learn

After this module, you can:

  • Explain why a token ID is a categorical address, not a measurement
  • Retrieve an embedding row by direct lookup and state the one-hot equivalence eₖᵀ E = E[k]
  • Trace [B, T] token IDs through one shared [V, D] table to [B, T, D] output vectors, including repeated IDs
  • Show that an optimizer step changes only the addressed rows while token IDs stay fixed
  • Inspect learned geometry with cosine similarity without mistaking it for literal meaning
  • Use nn.Embedding for the same lookup in PyTorch
  • Explain why token identity alone cannot encode sequence order

Prerequisites

This module requires familiarity with:

Orientation: From a Symbol to an Address

Module 03 showed how text is split into tokens and how each token receives an integer token ID. We only need that earlier story as a starting point: the ID is a label that lets us find the token again.

This lesson asks what happens next:

How can a symbolic word label become numeric information that a model can use and adjust?

Before any formula, keep four plain-language meanings in mind:

  • An address is a stable slot number used to find an item. A token ID is an address, not a measurement: the number’s size and its gap from another ID do not describe the tokens.
  • A table is an organized set of slots. Here, each vocabulary token owns one row in the embedding table.
  • A parameter is an adjustable number inside the model. Training can change a parameter when that number helps the model make better predictions.
  • A vector is an ordered list of numbers. In this lesson, a row’s two parameters form a vector that we can pass to the rest of the model or draw as a point.

That learned row vector is the token’s embedding. The address finds the row; the row supplies adjustable numeric information. We will first watch that lookup happen, then watch feedback change a row, and finally inspect the geometry that can emerge.

The False Number-Line Problem

To see why an address must not be treated like a measurement, consider this deliberately arbitrary assignment:

cat       → 42
airplane  → 43
dog       → 891

On a number line, cat appears much closer to airplane than to dog. That is a trap. The assignment says nothing about animals, vehicles, or similarity; a tokenizer could choose different addresses without changing the tokens.

TipPredict before you look

Which token appears numerically closest to cat? Does the tokenizer give us any reason to expect that pair to behave similarly in a model?

The false number line is the only place where this lesson uses arbitrary large addresses. From here onward, every required explanation and visual uses one small, fixed teaching world: eight tokens, IDs 0 through 7, and two learned values per row.

ImportantFirst principle: an ID only finds a row

The numerical value of a token ID carries no semantic meaning. Use it as an address for a learnable row instead of treating it as a measurement.

Table Lookup: An Address Chooses a Row

An embedding table is a table whose entries the model is allowed to adjust. In matrix notation, call the table E. It has V = 8 rows—one for each token in the toy vocabulary—and D = 2 numeric values in each row.

The ID does not undergo arithmetic. It selects a row:

token ID = address
table    = adjustable rows
E[k]     = row selected by address k

The values in a selected row are parameters: numbers that training may change. Follow one token from its label to its ID, highlighted row, and returned vector.

One-Hot as a Categorical Selector

A one-hot vector is another honest way to name one category. It is mostly zeros with one 1 at the chosen address. It does not claim that categories have a meaningful numeric order, and it does not create learned similarity by itself.

For a selected ID k, the one-hot selector is:

\mathbf{e}_k^{\top} := [0, \ldots, 0, \underbrace{1}_{\text{position }k}, 0, \ldots, 0].

Multiplying it by the table keeps exactly the selected row:

\mathbf{e}_k^{\top} E = E[k].

Predict: when the selected token changes, which one-hot position and which table row should change? Turn on Show math in the visual below to compare the sparse selector’s result with direct indexing. Direct indexing is still the mechanism to keep in mind: nn.Embedding retrieves E[k] without materializing a mostly-zero vector.

The highlighted row is the embedding returned by the lookup: a vector of learned parameters, not a reinterpretation of the ID. The selector explanation and the direct lookup therefore end at the same object, E[id].

Row to Point: The Same Vector, Drawn Geometrically

A vector is an ordered list, so a two-value row can also be drawn as a coordinate pair. For the fixed toy world, cat (ID 0) starts with row [-0.72, +0.18]; those same values locate the point (-0.72, +0.18).

row [-0.72, +0.18] → vector → coordinates (-0.72, +0.18) → point

This is a change of view, not a new representation and not a projection. The two plotted axes are teaching coordinates chosen to make the mechanism visible. They are not literal features such as “animalness,” and the map is not a production model’s measurement.

Predict: before changing the stage, which numbers should appear in the table and at the point? The staged visual below begins with Lookup and Vector to point; the same two numbers should move from the highlighted row onto the axes.

TipStop point: lookup to point only

Use only Lookup → Vector to point for this first pass. The same visual also offers Training update; leave that state unselected until the Learning Feedback section below explains the adjustment.

Batch Lookup: One Table, Many Positions

The same table is reused when a model handles several sequences at once. A batch is a group of sequences; B counts the sequences and T counts the token positions in each sequence. If the table has V rows and D values per row, lookup changes the shape but not the table:

token IDs [B, T] → shared table E [V, D] → vectors [B, T, D]

The fixed toy batch is [[0, 1, 0], [3, 4, 7]]: two sequences of three IDs use the same [8, 2] table and produce [2, 3, 2] vectors. ID 0 appears twice in the first sequence, so both positions retrieve the identical cat row. Reusing a row is not copying a new parameter; it is asking the same table for the same address twice.

Predict: which two output positions should contain identical vectors, and which table row supplies them? Keep the answer in view as the batch shape visual introduces repeated-row identity.

Learning Feedback: Why a Row Changes

Lookup only retrieves values; learning gives those values a reason to move. During a training example, the model uses the selected vector to make a prediction. A loss is a numerical score for how far that prediction is from the desired result. The score sends feedback toward the parameters that participated in the prediction. Feedback loop at a glance:

input → prediction → feedback → small adjustment

Read that chain left to right: the input ID selects a row, the row helps the model make a prediction, comparing that prediction with a target produces feedback, and the feedback nudges the participating row values.

Two Module 02 ideas are enough here; this is a bridge back to autograd, not a second autograd lesson:

  • A gradient is a suggested direction and size of adjustment: it summarizes how the loss would change if a parameter moved. Gradient descent commonly steps opposite the gradient to reduce the loss.
  • An optimizer is the rule that applies that suggestion, often using a learning rate and other safeguards.

The causal story is therefore:

token → row lookup → prediction → loss/feedback → gradient → optimizer → row changes

Predict: if this example looked up ID 0, should the ID change, should every row change, or should the selected row’s values change? Return to the visual above and select Training update. It shows one small optimizer step: the selected row moves from before to after step, while the other seven rows stay visible but dimmed.

ImportantToken IDs do not learn. Embedding values do.

The ID remains an address. Training changes the parameters stored in table rows, which changes their vectors and therefore moves their points.

A real batch can touch several rows, and a repeated token can contribute more than once. The shared table is still one set of parameters; only rows reached by the example receive a nonzero update for that step.

Learned Geometry: Inspect, Don’t Read Meaning

One update moves one selected row and therefore one point. Across many batches, different rows can be selected on different steps, so their cumulative histories can reorganize the table. Useful neighborhoods or directions may emerge because they help reduce prediction loss; the model is not handed labels such as “animal” or “vehicle.”

Predict: before seeing the next visual, what should happen in each state? Before training shows the initial rows, During training shows one selective step for the current focus token, and After training shows a separately curated endpoint after many batches. The Compare token menu automatically excludes the current focus token.

Use this plot as inspection evidence, not literal meaning. A nearby point, direction, or analogy-like offset can be a hypothesis-generating probe to test on a trained model; it is not a human-labeled axis, a dictionary definition, or a universal rule. This 2D map is a teaching view of the fixed toy rows, not a production embedding visualization.

The Before training state is scattered. In During training, only the current focus token takes one illustrated gradient step; all seven other rows stay at their before coordinates. After training is a separate, curated endpoint after many batches in which different token rows were selectively updated on different steps. Its trails summarize cumulative histories rather than one shared update. The axes still have no simple labels such as “animalness”: real learned features are usually distributed across many coordinates.

Measuring a Relationship: Cosine Similarity

Now that vectors are visible as points and directions, we can ask a purposeful question: how aligned are the two selected vectors? Cosine similarity measures the angle between them while ignoring overall length:

\operatorname{cos}(\mathbf{u}, \mathbf{v}) = \frac{\mathbf{u}^{\top}\mathbf{v}}{\lVert\mathbf{u}\rVert\,\lVert\mathbf{v}\rVert}.

The current Focus token, Compare token, and Training state controls drive this measurement:

Cosine is a measurement of one learned geometric relationship—not a definition of meaning, an antonym detector, or proof that every model must organize these tokens the same way.

Implementation Bridge: nn.Embedding Does the Same Lookup

The visual mechanism has a direct PyTorch name: nn.Embedding(vocab_size, embedding_dim). This module stores the table as a trainable parameter. Calling it with integer IDs asks for rows by address; it does not first build one-hot vectors.

For B sequences of T IDs, the shape contract is:

token IDs          [B, T]
embedding table    [V, D]
output vectors     [B, T, D]

The executable example uses the same eight-token, two-dimensional table and the same shared batch ([[0, 1, 0], [3, 4, 7]]) as the Batch Lookup visual, plus a real optimizer step on ID 0 that mirrors the Learning Feedback visual. It inspects the output shape, the repeated-row identity, and exactly which row changes. The default padding_idx=None leaves row 0 trainable, so it remains the cat row in this teaching world.

import torch
import torch.nn as nn

torch.manual_seed(4)

# Same toy world as the diagrams: V = 8 tokens and D = 2 values per row.
initial_rows = torch.tensor([
    [-0.72,  0.18],  # 0 cat
    [ 0.64, -0.62],  # 1 dog
    [-0.12,  0.78],  # 2 kitten
    [ 0.18, -0.08],  # 3 airplane
    [-0.82, -0.32],  # 4 bicycle
    [ 0.38,  0.70],  # 5 king
    [-0.46,  0.44],  # 6 queen
    [ 0.82,  0.12],  # 7 code
], dtype=torch.float32)

embedding = nn.Embedding(num_embeddings=8, embedding_dim=2)
with torch.no_grad():
    embedding.weight.copy_(initial_rows)

ids = torch.tensor([
    [0, 1, 0],  # same shared batch as the Batch Lookup visual; ID 0 repeats
    [3, 4, 7],
])
vectors = embedding(ids)
print("token ID shape:", list(ids.shape))
print("output shape:  ", list(vectors.shape))
# Both positions of ID 0 retrieve the same cat row, as the batch visual shows.
print("repeated ID 0 identical:", bool((vectors[0, 0] == vectors[0, 2]).all()))

# A tiny downstream objective touches only row 0, mirroring the
# Learning Feedback visual: loss sends feedback to the participating row.
optimizer = torch.optim.SGD(embedding.parameters(), lr=0.25)
target = torch.tensor([0.88, 0.98])
before = embedding.weight.detach().clone()
loss = (embedding(torch.tensor([0]))[0] - target).pow(2).sum()

optimizer.zero_grad()
loss.backward()
optimizer.step()

after = embedding.weight.detach().clone()
changed_rows = torch.where((after - before).abs().sum(dim=1) > 0)[0].tolist()
print("row 0 before:", [round(value, 2) for value in before[0].tolist()])
print("row 0 after: ", [round(value, 2) for value in after[0].tolist()])
print("changed rows:", changed_rows)
token ID shape: [2, 3]
output shape:   [2, 3, 2]
repeated ID 0 identical: True
row 0 before: [-0.72, 0.18]
row 0 after:  [0.08, 0.58]
changed rows: [0]

The input IDs remain integers. embedding.weight is the learnable matrix. Both [0, 0] and [0, 2] positions returned the same row because they share ID 0. The optimizer changed only row 0 because that was the row used by the loss; the other seven rows stayed fixed, exactly as the feedback visual’s selective update shows.

One Problem Remains: Order

Token lookup depends on token identity, not sequence position. The same address retrieves the same row wherever the token appears:

cat at position 2  → E[0]
cat at position 9  → E[0]

same token
same token embedding
different sequence position

Predict: if both occurrences get the same vector, what additional signal could tell the model where each occurrence sits? Token embeddings alone do not resolve order. Position needs separate machinery, which we will preview only as optional context after the core mental model; Module 05 supplies the attention mechanics later.

Mental Model

Keep this one-screen pipeline:

  • ID = address
  • row = parameters
  • vector = embedding
  • training = changes rows
  • geometry = consequence of learning

Check Your Understanding

Answer these without formulas.

  1. What does a token ID represent?
    Answer An address used to select a vocabulary entry and an embedding-table row.
  2. What is an embedding?
    Answer The learned vector stored in the selected row.
  3. What changes during training?
    Answer The values in embedding rows and other model parameters change; token IDs do not.
  4. Why can related tokens develop useful geometric relationships?
    Answer Prediction loss sends gradients that move their vectors into relationships that help the model perform its task.
  5. What does nn.Embedding do when it receives integer IDs?
    Answer It retrieves the addressed rows directly as vectors, without building one-hot selectors; a training step then changes only the rows its loss used.
  6. Why does position require additional machinery?
    Answer The same token retrieves the same token embedding wherever it appears in a sequence.

Exercises

Exercise 1: Batch Shapes and Repeated Rows

Rebuild the toy [8, 2] table and run the shared batch through it. Confirm the [B, T] to [B, T, D] shape contract and that both positions of the repeated ID retrieve the same row.

import torch
import torch.nn as nn

torch.manual_seed(4)

# Same toy world as the diagrams: V = 8 tokens and D = 2 values per row.
initial_rows = torch.tensor([
    [-0.72,  0.18],  # 0 cat
    [ 0.64, -0.62],  # 1 dog
    [-0.12,  0.78],  # 2 kitten
    [ 0.18, -0.08],  # 3 airplane
    [-0.82, -0.32],  # 4 bicycle
    [ 0.38,  0.70],  # 5 king
    [-0.46,  0.44],  # 6 queen
    [ 0.82,  0.12],  # 7 code
], dtype=torch.float32)

embedding = nn.Embedding(num_embeddings=8, embedding_dim=2)
with torch.no_grad():
    embedding.weight.copy_(initial_rows)

ids = torch.tensor([
    [0, 1, 0],  # same shared batch as the Batch Lookup visual; ID 0 repeats
    [3, 4, 7],
])
vectors = embedding(ids)
print("token ID shape:", list(ids.shape))
print("output shape:  ", list(vectors.shape))
print("repeated ID 0 identical:", bool((vectors[0, 0] == vectors[0, 2]).all()))
token ID shape: [2, 3]
output shape:   [2, 3, 2]
repeated ID 0 identical: True

Exercise 2: Only the Addressed Row Changes

Run a tiny optimizer step against one row and check exactly which rows moved. The loss below uses only the kitten row; change the target or the addressed ID and re-run to confirm the update stays selective.

# Fresh table from Exercise 1's initial_rows, so earlier cells cannot leak in.
embedding = nn.Embedding(num_embeddings=8, embedding_dim=2)
with torch.no_grad():
    embedding.weight.copy_(initial_rows)

# A tiny downstream objective touches only row 2 (kitten), mirroring the
# Learning Feedback visual: loss sends feedback to the participating row.
optimizer = torch.optim.SGD(embedding.parameters(), lr=0.25)
target = torch.tensor([0.50, -0.50])
before = embedding.weight.detach().clone()

# Your implementation here: write a loss that uses only the kitten row,
# e.g. loss = (embedding(torch.tensor([2]))[0] - target).pow(2).sum()
loss = (embedding(torch.tensor([2]))[0] - target).pow(2).sum()

optimizer.zero_grad()
loss.backward()
optimizer.step()

after = embedding.weight.detach().clone()
changed_rows = torch.where((after - before).abs().sum(dim=1) > 0)[0].tolist()
print("changed rows:", changed_rows)
assert changed_rows == [2], "only the addressed row should change"
print("Only the addressed row changed; the other seven rows stayed fixed.")
changed rows: [2]
Only the addressed row changed; the other seven rows stayed fixed.

Exercise 3: What Cosine Similarity Proves

Compute the cosine similarity between the starting cat and dog rows, then read the number for exactly what it is.

# Cosine similarity between the cat and dog rows of the starting table.
cat = initial_rows[0]
dog = initial_rows[1]
cos = torch.nn.functional.cosine_similarity(cat, dog, dim=0)
print(f"cosine(cat, dog) = {float(cos):.3f}")
cosine(cat, dog) = -0.866

This number measures the angle between two curated teaching rows: it tells you how aligned they start out, which is all a similarity readout ever claims. It does not prove cat and dog mean similar things — these rows are hand-picked starting points, and only training feedback plus task behavior could turn a neighborhood into evidence.

Optional Chapter: Positional Representations and Production Scale

The required lesson ends with the mental model above and the unresolved order problem. Everything in this chapter is optional: it previews ways to add position and a few engineering choices without making the token-lookup story depend on attention.

Four ways to add position

All four methods answer the same question—how can two occurrences of one token carry different positional information? They contribute that information in different places:

  • Learned absolute positions store one trainable vector per position. LearnedPositionalEmbedding looks up the row for position t and can add it to the token vector.
  • Fixed sinusoidal positions compute sine and cosine waves at several rates. SinusoidalPositionalEmbedding adds a repeatable, non-trainable pattern so each absolute position has a distinct signal.
  • RoPE (rotary position embedding) rotates paired features by an angle tied to position. RotaryPositionalEmbedding and apply_rotary_pos_emb apply that rotation later to the query/key representations, so relative offsets can affect their relationship.
  • ALiBi (Attention with Linear Biases) leaves token vectors alone and contributes a distance-dependent penalty where attention scores are compared. alibi_slopes supplies head-specific slopes and alibi_bias builds the corresponding biases.

These are four different homes for position, not four required implementations to memorize now. The optional position visual asks one focused question at a time; defer score calculations and the rest of attention vocabulary to Module 05.

Predict: if cat at positions 2 and 9 retrieves the same E[0], which part of each method below supplies the difference? Recall the order problem above — cat at position 2 → E[0] and cat at position 9 → E[0] — and use the Position method control to ask one focused question at a time. This preview is optional: the required mental model stands without it, and score details wait for Module 05.

The visual keeps each method’s embeddings.py API name and shape contract next to its mechanism. TransformerEmbedding combines only the learned or sinusoidal input variant with dropout; RoPE and ALiBi act later, where Module 05 compares queries, keys, and scores.

Optional production context

A token table contains V × D parameters, so vocabulary size and vector width affect memory. A larger vocabulary can shorten token sequences but enlarges the table. Some models tie the input table to output-vocabulary weights; padding_idx and sparse-gradient settings are implementation choices rather than changes to the direct-lookup idea. In a production model, D is high-dimensional and inspection usually uses projections or similarity tools; it is not a list of D human-labeled semantic axes.

Common Pitfalls

When working with embeddings, watch out for these issues:

  1. Treating ID size or ID gaps as meaningful quantities: the numbers are addresses; arithmetic on them manipulates slots, not tokens.
  2. Calling the ID itself an embedding: the embedding is the selected row vector, not the integer that found it.
  3. Assuming every table row changes on every optimizer step: only rows the loss used receive gradients; the rest stay fixed.
  4. Reading a toy two-dimensional map as a literal or universal map of meaning: the 2D coordinates are teaching aids, and similarity readouts are inspection evidence, not proofs.

Going Deeper

Core Papers:

Practical Resources:

  • PyTorch nn.Embedding — The direct-lookup API used in the bridge above, with its padding and sparsity options.

Summary

Key takeaways:

  1. ID = address — A token ID is a categorical slot number; its size and its gaps from other IDs carry no meaning.
  2. Row = parameters — Each ID selects one row of the shared [V, D] table, whose values training may adjust.
  3. Vector = embedding — The selected row’s values leave the table as a vector the rest of the model can use. eₖᵀ E = E[k] states the same lookup as a one-hot equivalence, while nn.Embedding indexes the row directly.
  4. Training = changes rows — Loss sends gradients to participating rows only, so an optimizer step moves the addressed rows while token IDs stay fixed; [B, T] IDs return [B, T, D] vectors with repeated IDs retrieving identical rows.
  5. Geometry = consequence of learning — Cosine neighborhoods are inspection evidence and hypothesis-generating probes, not literal or universal meanings. Token identity alone also leaves order unresolved, which is why position needs separate machinery.

What’s Next

Module 05: Attention starts with these embedding vectors and shows how tokens compare and combine information. It also supplies the machinery needed to understand why modern position methods act on attention rather than on token IDs.