The Lookup Table That Learns: Token Embeddings and the Order-Blind Trap

Where we left off: a grid of meaningless numbers

In the previous post I turned a short story into batches of token IDs, shape [8, 4] — 8 windows, 4 token IDs each. That tensor is the input to everything that follows. There's just one problem with it: the numbers inside are meaningless as quantities.

The tokenizer assigned " tea" the ID 8887 and " you" the ID 345. Those aren't measurements. 8887 does not mean tea is bigger, heavier, or 25× more anything than you. It's a row number in a dictionary — an address, not a value. Feed it to a network as-is and the maths would treat 8887 as vastly larger than 345, which is nonsense.

This post is chapter 2's finale: how those flat, meaningless IDs become the actual tensor that chapter 3's attention consumes. It comes down to two lookup tables and one addition.

An embedding is just a lookup table

The fix is to stop describing a token with one number and start describing it with a small list of numbers — a vector with room to carry meaning. Collect one such vector per vocabulary word and you have a table: one row per word. That table is the embedding layer.

import torch
# toy: 6 words, 3 numbers each
embedding_layer = torch.nn.Embedding(6, 3)
embedding_layer(torch.tensor([3]))   # -> row 3, a single vector. No math, just a fetch.

The word lookup is exact. Looking up token ID 3 means "go to row 3, hand me what's there." No multiplication, no transformation — a row fetch. One ID returns one row (a vector); a batch of IDs returns a stack of rows (a matrix). The table itself never changes; you just index into it.

Now the real sizes. GPT-2's vocabulary is 50,257 tokens, and I gave each one a 256-number row:

token_embedding_layer = torch.nn.Embedding(50257, 256)
token_embedding_layer.weight.numel()   # 12,865,792 numbers to learn
token_embeddings = token_embedding_layer(inputs)   # inputs [8,4]  ->  [8,4,256]

So the [8, 4] grid of IDs becomes a [8, 4, 256] cube: the same 32 tokens, each now backed by 256 numbers instead of 1. That 256 is simply "how many numbers describe the meaning of one token." It's a knob — GPT-2 small uses 768, GPT-3 uses 12,288; I used 256 to keep it light.

The word "learns" is doing real work

Two things about that table are easy to miss. First, right now its 12.8 million numbers are random noise. nn.Embedding initialises them randomly, so at this moment " tea" and " war" have equally meaningless vectors. They only become meaningful during training (chapter 5), when backpropagation nudges every row a little at a time. The embedding table is not a preprocessing step — it is part of the model, and those 12.8M parameters are learned alongside everything else.

Second, why a lookup and not the textbook "one-hot vector times a weight matrix"? They're mathematically identical — but one-hot would mean multiplying by a 50,257-wide vector that is all zeros except a single 1. The lookup skips all of it and jumps straight to the one row that survives. Same result, none of the wasted multiplications.

The order-blind trap

Here's the catch that the finale exists to fix. A lookup table returns a row based only on the ID — nothing else. So the same token ID always returns the same row, no matter where it appears:

a = token_embedding_layer(torch.tensor([345]))
b = token_embedding_layer(torch.tensor([345]))
torch.equal(a, b)   # True

Now think about "the cat sat on the mat." Both instances of "the" share one ID, so the lookup hands back byte-for-byte identical vectors for each. The embedding has no idea that one came first and the other came later. To the model, the sentence is an unordered bag of tokens — it cannot tell "the cat sat on the mat" from "mat the on sat cat the." And word order plainly carries meaning: "dog bites man" is not "man bites dog."

This is not a bug in the code. It is a property of what an embedding is: a function of the token alone. Position has to be injected separately.

The fix: a second table, indexed by slot

The trick is almost cheeky in its simplicity. Build a second embedding table — but index it by slot number (position in the window) instead of by token ID:

context_length = 4
pos_embedding_layer = torch.nn.Embedding(context_length, 256)
pos_embeddings = pos_embedding_layer(torch.arange(4))   # slots 0,1,2,3  ->  [4,256]

This table has one row per position: a vector that means "I'm in slot 0," another for slot 1, and so on. Then you simply add the position vector onto the token vector. The proof is the cleanest demo in the whole chapter — feed the same token four times and watch position break the tie:

Before the addition, all four rows are identical — order is invisible. After adding the four slot vectors, all four rows differ, because each slot contributed a different position signal. The model can finally tell positions apart.

Two tables, one sum — and chapter 2 is done

Zoom back out to the full batch and one detail makes it click. The token embeddings are [8, 4, 256] (8 sentences), but the positional embeddings are only [4, 256] (4 slots, no sentence dimension). Adding them just works, via broadcasting:

token_embeddings.shape   # [8, 4, 256]
pos_embeddings.shape     # [4, 256]
input_embeddings = token_embeddings + pos_embeddings   # [8, 4, 256]

The same four position vectors are added to every sentence in the batch — which is exactly right: slot 0 means "slot 0" regardless of which sentence you're in. One small [4, 256] table serves all 8 rows.

And that's the end of chapter 2. The whole journey — raw text → token IDs → sliding-window batches → token embeddings → positional embeddings → input_embeddings — lands on a single [8, 4, 256] tensor. It's just two lookup tables and an addition, but it is the exact input the attention mechanism consumes next. The lookup table that carries meaning, plus a second one that carries order: that's what a token becomes before a transformer ever looks at it.