The BPE tokenizer hands me integers. For the sentence I love AI I get [15, 824, 391]. Most explanations stop there and jump straight to transformer diagrams. I did that for months — treated nn.Embedding as a black box and moved on. Then I built AL-1 in pure NumPy and had to implement the layer myself. The punchline is almost insultingly simple: it is a lookup table. Row $i$ is the vector for token ID $i$. The interesting part is why a lookup table is the right move, what those rows become after training, and how gradients rewrite only the rows you actually used.
This post is the bridge between “text as numbers” and “text the model can reason over.” Get it wrong and token ID $824$ is just a bigger scalar than $391$, which would imply love > AI in some numeric sense — nonsense. I need each ID to become a vector I can do geometry on, where distance means something and cat lands near dog but far from car. A scalar ID cannot do that. A learned row in a matrix can.
flowchart LR
TXT["'I love AI'"] --> TOK["BPE tokenizer"]
TOK --> IDS["IDs [15, 824, 391]"]
IDS --> EMB["Embedding lookup"]
EMB --> VEC["Vectors (3, d_model)"]
VEC --> TR["Transformer blocks"]
Why can’t the model use token IDs directly?
A token ID is a name, not a quantity. The tokenizer assigned love → 824 and AI → 391 by merge order during BPE training, not by meaning. If I feed raw integers into a matrix multiply, the network reads $824$ as “4.7× bigger than $175$” and starts fitting spurious magnitude relationships. The IDs carry zero semantic structure — cat and kitten could be $82$ and $49{,}001$, as far apart as any two random numbers.
What I want instead: each token becomes a vector in $\mathbb{R}^{d}$, where $d$ is the model width. Distances should mean something. Co-occurring tokens should end up near each other after training. You cannot get that from a scalar, so you map every ID to its own learned vector and let gradient descent arrange them.
What is a token embedding?
An embedding is a learned vector attached to each token in the vocabulary. Stack all of them and you get the embedding matrix $W$:
Here $V$ is vocabulary size and $d$ is d_model. Row $i$ is the vector for token ID $i$. In GPT-2 small:
Parameter count:
Almost 39 million parameters before attention. That is the single largest weight in a small GPT, and via weight tying it is reused as the output projection — the lm_head that maps hidden states back to vocabulary logits. Same shape $V \times d$, same meaning: “what a token means going in” tied to “how likely it is coming out.”
Take a toy vocabulary of three tokens with $d = 4$:
Row $0 \to$ dog. Row $1 \to$ cat. Row $2 \to$ car. Rows are tokens, columns are features. Nothing more.
How does the embedding lookup work?
Here is the entire layer from zyn/nn.py, give or take:
class Embedding:
def __init__(self, vocab_size, d_model, std=0.02):
w = np.random.randn(vocab_size, d_model) * std
self.weight = Tensor(w) # trainable
def forward(self, indices): # indices: (B, T) int array
return self.weight.gather(indices, dim=0) # (B, T, d_model)
gather is fancy indexing — weight[indices]. Given a tiny matrix:
And indices:
The output stacks rows $2$, then $0$, then $1$:
No multiplication. No activation. Input shape is $(B, T)$ — batch of $B$ sequences, $T$ tokens each. Output shape is $(B, T, d)$. Every integer becomes a $d$-dimensional vector, and that tensor is what the first transformer block consumes. For a real batch of shape $(32, 1024)$ with $d = 768$, the output is $(32, 1024, 768)$ floats.
Why multiply by 0.02 at initialization?
np.random.randn gives unit-variance Gaussians — values around $\pm 1$, occasionally $\pm 3$. Those are far too large to start. A row like $[1.5, -0.8, 2.1, \dots]$ fed into the residual stream blows up activations and the first few gradient steps thrash. Scaling by $\text{std} = 0.02$ shrinks everything:
A fresh row looks like $[0.012, -0.017, 0.006, \dots]$ — tiny, near the origin, all tokens starting roughly equal. That $0.02$ is the GPT-2 initializer. Training nudges rows apart slowly instead of starting from chaos.
How is lookup the same as a one-hot matrix multiply?
The equivalence that finally made embeddings click for me: a row lookup is exactly a one-hot vector times the matrix. Represent token $2$ as the one-hot row:
Multiply by a $3 \times 2$ matrix:
First product entry:
Second product entry:
Result:
That is precisely row $2$ of $W$. The zeros kill every other row; the single $1$ selects one. The embedding layer is a linear layer whose input happens to be one-hot. For $V = 50257$ you would multiply a 50k-wide vector that is all zeros except one slot. gather skips the charade and grabs the row directly. Same math, none of the floating-point busywork.
How do embeddings learn meaning?
The rows start as random noise. Meaning shows up through gradients. Say the model reads I love and predicts the next token. The loss $L$ flows back into the embedding matrix as $\frac{\partial L}{\partial W}$ — but only the rows that were actually used get a nonzero gradient. If the sequence touched token IDs $15$, $824$, and $391$, then rows $15$, $824$, and $391$ update. The other $50254$ rows sit untouched this step. That sparsity is why embedding training is cheap despite the parameter count.
Over millions of steps the geometry organizes itself. Tokens appearing in similar contexts get pulled together:
Nobody programmed dog ≈ cat. The objective — predict the next token — makes co-occurring tokens converge as a side effect.
flowchart TB
subgraph space["Embedding space (schematic)"]
dog["dog"] --- cat["cat"]
cat --- lion["lion"]
car["car"] --- truck["truck"]
end
Why does d_model size matter?
With $d = 2$ each token has two numbers to encode everything — meaning, part of speech, tense, sentiment. Not enough room; tokens collide. Bump to $d = 768$ and each vector has 768 directions to spread information across. Wider embeddings hold more nuance, at the cost of more parameters and compute. GPT-2 small uses $768$; GPT-2 XL uses $1600$. It is a capacity dial, not a free lunch — double $d$ and you double this 39M-parameter layer.
Where does this fit in the AL-1 build?
Pipeline so far: autograd engine → gradient checker → BPE tokenizer → token embeddings → positional encoding → attention.
The tokenizer hands me IDs. The embedding layer turns them into the $(B, T, d)$ tensor that flows through every block. But this lookup is order-blind. The vector for cat is identical whether cat is the first token or the hundredth — gather does not know or care about position. A transformer with only token embeddings sees a bag of words. Positional encoding fixes that by adding a second matrix that injects “where” on top of “what.”
I implemented this by hand mostly to kill the mystery. nn.Embedding is one line and it hides the fact that the biggest weight in your model is a dictionary you are slowly teaching to mean something. Once you have watched the rows drift from noise into structure, the rest of the transformer stops feeling like magic and starts feeling like consequences.
Full worked example: I love AI with IDs [15, 824, 391]
I will walk through the entire embedding layer for I love AI — lookup, one-hot equivalence, loss gradient, and row update — with every arithmetic step on its own line. I use a toy width $d = 4$ so the numbers stay readable; GPT-2 small uses $d = 768$, but the mechanics are identical.
Step 1 — Token IDs from BPE
The BPE tokenizer maps:
Sequence:
Shape $(T) = (3)$ for one sentence. Batched with $B = 1$, shape is $(1, 3)$.
Step 2 — Rows in the embedding matrix
Pull the three relevant rows from $W \in \mathbb{R}^{50257 \times 4}$ (toy $d = 4$):
Step 3 — Embedding lookup via gather
gather stacks the rows in index order:
Shape $(T, d) = (3, 4)$. With batch dimension: $(1, 3, 4)$. Row $0$ of $X$ is the embedding for I. Row $1$ is love. Row $2$ is AI. This $(B, T, d)$ tensor is what the first transformer block receives.
Step 4 — One-hot equivalence for token 15 (I)
A one-hot vector for token ID $15$ in a vocabulary of size $V = 50257$ has a $1$ at index $15$ and $0$ everywhere else:
Multiply by $W$:
Only $j = 15$ contributes:
So:
Same result as gather. For love (ID $824$):
For AI (ID $391$):
Stacking all three one-hot rows and multiplying is the same as gather on $[15, 824, 391]$.
Step 5 — Loss gradient flowing back
Suppose the sentence passes through the transformer and the loss $L$ produces upstream gradients at each embedding position. After backprop through later layers, the gradient at position $0$ (I) is:
At position $1$ (love):
At position $2$ (AI):
Because $X[t] = W[\text{indices}[t]]$, the chain rule gives $\frac{\partial L}{\partial W[k]} = g_t$ when $\text{indices}[t] = k$, and $\mathbf{0}$ otherwise. Only three rows receive signal:
Every other row of $\frac{\partial L}{\partial W}$ is the zero vector. That is the sparsity I mentioned — one training step touches three rows out of $50257$.
Step 6 — Gradient matrix (schematic)
The full $\frac{\partial L}{\partial W}$ is mostly zeros. The three nonzero rows:
Step 7 — Row update with learning rate η = 0.1
SGD update: $W \leftarrow W - \eta \, \frac{\partial L}{\partial W}$. Only the used rows change.
Row 15 (I), component by component:
Updated row:
Row 824 (love), component by component:
Updated row:
Row 391 (AI), component by component:
Updated row:
Step 8 — What changed and what did not
After one SGD step on I love AI:
Rows $0$ through $14$, $16$ through $390$, $392$ through $823$, and $825$ through $50256$ are unchanged. Repeat this across millions of sentences and the geometry of the table drifts from random noise into structure — tokens that share contexts get pulled together, one sparse row update at a time.
If the same token ID appears twice in one sequence, its row gradient is the sum of the per-position gradients. That is how repeated tokens accumulate more signal in a single step.
Embedding closeness is useful for semantic search, retrieval, clustering, and recommendations. But a static token row cannot distinguish two meanings of the same word — bank as a river edge vs a financial institution share one row until the transformer layers above create contextual representations. That is a different problem, solved downstream. Here, at the embedding layer, the job is simple: turn each integer from BPE into a vector the rest of the model can work with. Everything else is gradient descent rearranging rows you actually used.
Support the writing
If this post helped, a coffee keeps the deep dives, paper picks, and news digests coming.
Buy Me a Coffee