Build an LLMMiddleLesson 385 min read

Step 2 · Embeddings and position

Turning token ids into vectors, and telling the model what order they came in. Two lookup tables, one subtle idea.

🧪 Lab · 15 min 📝 Lesson test

Lesson in motion

In 60 seconds

Step 2 · Embeddings and position

Turning token ids into vectors, and telling the model what order they came in. Two lookup tables, one subtle idea.

1/4

In simple words

Each word-chunk gets a list of numbers, like coordinates on a map. Chunks that mean similar things end up near each other on the map. The model learns the map by itself.

Token id 15496 means nothing arithmetically — it is just a slot number. The first thing the model does is look up a learned vector for it.

Embeddings, in three linespython

import torch, torch.nn as nn

vocab_size, d_model = 50257, 768
tok_emb = nn.Embedding(vocab_size, d_model)   # a (50257, 768) lookup table

ids = torch.tensor([[15496, 995, 0]])         # batch 1, sequence 3
x = tok_emb(ids)                              # -> (1, 3, 768)

That is it. nn.Embedding is a matrix where row i is the vector for token i, and the lookup is differentiable, so training moves those rows around until similar tokens sit near each other.

The order problem

Attention, which you will meet next module, has no idea about order. Feed it "dog bites man" and "man bites dog" and — without position information — it sees the same bag of vectors. We must inject position explicitly.

Method	How it works	Used by
Learned absolute	A second lookup table indexed by position 0,1,2…	GPT-2, BERT
Sinusoidal	Fixed sine and cosine waves of different frequencies	Original transformer paper
RoPE (rotary)	Rotate the query and key vectors by an angle proportional to position	Llama, Mistral, Qwen, most modern models
ALiBi	Add a distance penalty directly to attention scores	Some long-context models

Why RoPE won

Absolute position embeddings teach the model "this is slot 7". RoPE instead makes the attention score between two tokens depend on how far apart they are, which is what actually matters in language. It also extrapolates far better to sequences longer than anything seen in training.

Rotary position embeddingspython

def rope_frequencies(head_dim, seq_len, base=10000.0, device="cpu"):
    # one angular frequency per pair of dimensions
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    pos = torch.arange(seq_len, device=device).float()
    ang = torch.outer(pos, inv)                 # (seq_len, head_dim/2)
    return torch.cos(ang), torch.sin(ang)

def apply_rope(x, cos, sin):
    # x: (batch, heads, seq, head_dim) -- rotate each adjacent pair of dims
    x1, x2 = x[..., 0::2], x[..., 1::2]
    cos, sin = cos[None, None], sin[None, None]
    out1 = x1 * cos - x2 * sin
    out2 = x1 * sin + x2 * cos
    return torch.stack((out1, out2), dim=-1).flatten(-2)

Do this

Read that geometrically. Each pair of dimensions is treated as a point on a circle and spun by an angle proportional to its position. When you later take a dot product between a query at position m and a key at position n, the result depends on m − n — relative distance falls out of the maths for free.

What embeddings are not

They are not a dictionary of meanings. A token vector is a starting point that every later layer rewrites in context.
The famous "king − man + woman ≈ queen" arithmetic came from older word-vector models. Transformer token embeddings are much less interpretable in isolation.
Embedding and output layers are often tied — the same matrix, transposed, produces logits. It saves parameters and usually helps.

Watch and read more

**Attention in transformers, visually**3Blue1Brown · video

RoPE explainedYouTube search

Lab

Position encoding, and the failure when you remove it.

~15 min

The problem

Take a working tiny transformer. Delete the position encoding and retrain. Measure the loss gap and inspect what the model can and cannot now do. Then implement RoPE and compare extrapolation beyond the trained length.

Starter codepython

# Prove attention is permutation-invariant without positions
x = torch.randn(1, 5, 64)
perm = torch.randperm(5)
assert torch.allclose(attn(x).sum(1), attn(x[:, perm]).sum(1), atol=1e-5)

You are done when

Measured loss gap with and without positions
You demonstrated permutation invariance
RoPE extrapolated further than learned positions

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Rescaling RoPE frequencies extends context with no retraining. Why is this insufficient, and what does the model actually lack?Reveal

Rescaling makes long positions representable — the angles no longer alias — but the model has never seen a dependency spanning that distance, so it has learned no attention pattern that uses one. You get fluent output over a long input and poor use of it. What is missing is training signal at length: continued pretraining on genuinely long documents where the answer depends on distant context.

Please sign in to continue.

Questions people ask

How big should d_model be?

It scales with the model. Roughly: 128–384 for a tiny teaching model, 768 for GPT-2 small, 4096 for a 7B model, 8192+ at frontier scale. It should be divisible by the number of attention heads.

Are these the same as sentence embeddings for search?

Different things with the same name. Retrieval embeddings represent a whole passage as one vector, produced by a model trained for similarity. Token embeddings are per-token inputs to a generative model.

Why does the embedding table dominate small models?

With a 50k vocabulary and d_model 768, that table alone is 38M parameters. In a 100M-parameter model it is a third of everything. This is why tiny models often use smaller vocabularies.

Can I add RoPE to a model trained with learned positions?

Not without substantial retraining. Position encoding is baked into what every attention head learned. Conversions exist and are all approximate.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.