Lesson 5 - Guided Project: An Embedding Layer with Positions

Welcome to the Guided Project#

Across this module you took attention’s biggest blind spot apart and fixed it. You saw that self-attention is permutation-invariant and cannot tell word order on its own. You built a token embedding table that turns character ids into learnable vectors, with its forward gather and its scatter-add backward. Then you added position, first with the fixed sinusoidal construction and then with the learned positional embeddings that GPT actually uses. In this guided project you stop treating those as separate experiments and package them into one object: an EmbeddingLayer class that owns a token table and a position table, embeds a batch of token ids, and adds their positions in a single forward call.

This class is not throwaway demo code. It is literally the first thing the GPT does in Modules 6 and 7: raw token ids come in, and this layer turns them into the (B, T, C) tensor that every attention block and feed-forward network downstream consumes. Getting it clean and provably correct now means the rest of the model can assume its input is already position-aware. So the project ends where every serious layer should: a float64 gradient check that compares the analytic gradients to finite differences, and an explicit demonstration that the output shape is exactly what the SelfAttention head from Modules 2-3 expects.

By the end of this project, you will be able to:

• Structure the input stage of a transformer as a class holding a token table E and a learned position table P
• Implement forward(idx) so it gathers E[idx], adds P[:T], and caches what backward needs
• Implement backward(dout) so it scatter-adds dE with np.add.at and sums dP over the batch
• Gradient-check the whole class in float64, reporting max relative errors for E and P
• Confirm the (B, T, C) output plugs directly into the self-attention layers from earlier modules

This project assumes you followed the earlier lessons in this module, especially the token-embedding backward pass and the learned-positional-embedding lesson. You only need numpy (version 2.x is fine) — no PyTorch, no GPU, no API key.

Stage 1: The Class#

The input stage needs two learned tables. The token table E has shape (vocab_size, C): one C-dimensional vector for every character in the vocabulary, looked up by id. The position table P has shape (block_size, C): one C-dimensional vector for every position from 0 up to block_size - 1, looked up by slot. Both are parameters the model learns; both live on self. The class also holds a cache (what forward leaves for backward) and the gradient slots dE, dP.

Notice the three numbers __init__ commits to: vocab_size, block_size, and C. It does not commit to the batch size B or the actual sequence length T of any particular call — those arrive with the data. We use the canonical Lantern Bay corpus so vocab_size comes out to exactly 24, and we seed the weight draw so the whole class is reproducible. The weights start small (* 0.02, the standard transformer init).

import numpy as np

# ---------- corpus + tokenization (canonical) ----------
CORPUS = (
    "the lamp keeper walks the shore at dusk. "
    "the lamp keeper lights the lantern when the fog rolls in. "
    "the boats come home when the lantern glows. "
    "the boats wait outside the bay when the fog is thick. "
    "a small boat drifts near the rocks at dawn. "
    "the keeper rings the bell when a boat drifts near the rocks. "
    "the tide rises at dusk and falls at dawn. "
    "the gulls call over the bay when the tide is low. "
    "the keeper counts the boats and writes the count in a book. "
    "the lantern burns all night and rests at dawn. "
) * 20

chars = sorted(set(CORPUS))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = np.array(encode(CORPUS), dtype=np.int64)
print("vocab_size:", vocab_size, "| len(data):", len(data))


class EmbeddingLayer:
    """Token + learned-position embedding: the input stage of the GPT."""

    def __init__(self, vocab_size, block_size, C, seed=42):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.C = C
        rng = np.random.default_rng(seed)
        self.E = rng.standard_normal((vocab_size, C)) * 0.02   # token table
        self.P = rng.standard_normal((block_size, C)) * 0.02   # learned position table
        self.cache = None
        self.dE = self.dP = None


emb = EmbeddingLayer(vocab_size=vocab_size, block_size=32, C=16)
print("E shape (token table)   :", emb.E.shape)
print("P shape (position table):", emb.P.shape)
print("cache starts empty      :", emb.cache)

vocab_size: 24 | len(data): 10020
E shape (token table)   : (24, 16)
P shape (position table): (32, 16)
cache starts empty      : None

The corpus encodes to 10020 ids over a 24-character vocabulary, exactly as the module promised. The layer now holds a (24, 16) token table and a (32, 16) position table — enough to embed any sequence of up to 32 tokens drawn from the 24-character alphabet into a 16-dimensional space. Nothing has flowed through it yet; the cache is empty and the gradients are unset. The design mirrors the SelfAttention class you built in Module 2: parameters, a cache, and gradients, all living on self.

Stage 2: forward(idx)#

The forward(idx) method takes a batch of integer token ids idx of shape (B, T) and returns the position-aware embeddings of shape (B, T, C). It does exactly two lookups and one add. The token lookup self.E[idx] is a gather: NumPy fancy-indexing an (B, T) id array into the (vocab_size, C) table produces an (B, T, C) array where entry [b, t] is row idx[b, t] of E. The position lookup self.P[:T] grabs the first T position vectors, shape (T, C), and adding it to the token embeddings broadcasts that same (T, C) slab across every sequence in the batch. Position t gets the same positional vector no matter which batch row it sits in.

forward caches idx and T — that is all backward needs, because the gradient routes entirely through which rows were looked up. Here it is, run on a real slice of the encoded corpus, decoded so you can see the actual characters going in.

import numpy as np

CORPUS = (
    "the lamp keeper walks the shore at dusk. "
    "the lamp keeper lights the lantern when the fog rolls in. "
    "the boats come home when the lantern glows. "
    "the boats wait outside the bay when the fog is thick. "
    "a small boat drifts near the rocks at dawn. "
    "the keeper rings the bell when a boat drifts near the rocks. "
    "the tide rises at dusk and falls at dawn. "
    "the gulls call over the bay when the tide is low. "
    "the keeper counts the boats and writes the count in a book. "
    "the lantern burns all night and rests at dawn. "
) * 20
chars = sorted(set(CORPUS))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
decode = lambda ids: "".join(itos[i] for i in ids)
data = np.array([stoi[c] for c in CORPUS], dtype=np.int64)


class EmbeddingLayer:
    """Token + learned-position embedding: the input stage of the GPT."""

    def __init__(self, vocab_size, block_size, C, seed=42):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.C = C
        rng = np.random.default_rng(seed)
        self.E = rng.standard_normal((vocab_size, C)) * 0.02
        self.P = rng.standard_normal((block_size, C)) * 0.02
        self.cache = None
        self.dE = self.dP = None

    def forward(self, idx):
        B, T = idx.shape
        assert T <= self.block_size, "sequence longer than block_size"
        token_emb = self.E[idx]          # (B, T, C) gather
        pos_emb = self.P[:T]             # (T, C)
        out = token_emb + pos_emb        # (B, T, C) broadcast over batch
        self.cache = (idx, T)
        return out


emb = EmbeddingLayer(vocab_size=vocab_size, block_size=32, C=16)
T = 16
idx = data[:T].reshape(1, T)             # (B=1, T=16) real token ids
out = emb.forward(idx)
print("idx shape :", idx.shape)
print("decoded   :", repr(decode(idx[0].tolist())))
print("out shape :", out.shape, "= (B, T, C)")
print("cached T  :", emb.cache[1])

idx shape : (1, 16)
decoded   : 'the lamp keeper '
out shape : (1, 16, 16) = (B, T, C)
cached T  : 16

The first 16 characters of the Lantern Bay corpus, 'the lamp keeper ', become a (1, 16, 16) tensor: one 16-dimensional vector per character, each already carrying its position. The two 'e' characters at positions 2 and 13 gather the same row of E (both are token id for 'e') but add different rows of P, so their final vectors differ — which is the entire point of positional embeddings, and exactly the order information plain attention lacked. The cache now holds idx and T = 16, the only two things backward will read.

Stage 3: backward(dout)#

The backward(dout) method takes the gradient of the loss with respect to the output, dout of shape (B, T, C), and produces the gradients of the two tables. Both are the reverse of a lookup, and each has a characteristic shape.

Because forward gathered rows of E (possibly the same row many times — every 'e' in the batch pulled the same row), backward must scatter-add: every place a row was used sends its slice of dout back to that row, and repeated uses accumulate. Plain self.dE[idx] += dout silently drops duplicate indices, so we use np.add.at(self.dE, idx, dout), which is the buffered, duplicate-safe scatter-add. The result dE has the same shape as E, (vocab_size, C), with a nonzero row only for characters that actually appeared.

The position gradient is simpler. P[:T] was added to every sequence in the batch, so its gradient is dout summed over the batch axis, landing in rows 0..T-1; positions T..block_size-1 were never used this step and stay zero. There is deliberately no dX: the inputs are integer token ids, not continuous values, so there is nothing to differentiate with respect to the input. backward returns None and stores dE and dP on self for the optimizer.

    def backward(self, dout):
        idx, T = self.cache
        self.dE = np.zeros_like(self.E)
        np.add.at(self.dE, idx, dout)              # scatter-add over token ids
        self.dP = np.zeros_like(self.P)
        self.dP[:T] = dout.sum(axis=0)             # sum over batch into used rows
        # no dX: inputs are integer ids, not differentiable
        return None

With that method added to the class, run a backward pass using a pretend upstream gradient shaped like the output:

np.random.seed(42)
dout = np.random.randn(*out.shape)       # pretend gradient from the layer above
emb.backward(dout)
print("dout shape:", dout.shape)
print("dE shape  :", emb.dE.shape, "= E shape")
print("dP shape  :", emb.dP.shape, "= P shape")
print("dP rows T..end are zero:", np.allclose(emb.dP[T:], 0.0))

dout shape: (1, 16, 16)
dE shape  : (24, 16) = E shape
dP shape  : (32, 16) = P shape
dP rows T..end are zero: True

The shapes match their tables exactly — dE is (24, 16) like E, dP is (32, 16) like P — and rows 16..31 of dP are all zero because this sequence was only 16 tokens long, so positions 16 through 31 received no gradient. That zeroing is not a special case to remember; it falls straight out of self.dP[:T] = ... leaving the untouched rows at their initialized zero.

Stage 4: Verify#

A backward pass you cannot verify is a backward pass you cannot trust. We prove this one two ways: a gradient check (do the analytic gradients match finite differences?) and a connection check (does the output slot into the attention layers from earlier modules?).

Gradient Check#

The idea is the definition of a derivative. Nudge one entry of a table by $\pm \epsilon$, measure how much a scalar loss changes, and divide — that finite difference must match the analytic gradient:

We use a simple scalar loss $L = \sum (\text{out} \odot G)$ for a fixed random matrix $G$, because then the analytic gradient flowing into backward is just dout = G. We run the check in float64 (finite differences are numerically delicate; single precision would swamp the signal with rounding noise) and report the max relative error across every entry of E and P. The token ids come from the real corpus so that duplicate indices genuinely exercise the np.add.at scatter.

def scalar_loss(out, G):
    return np.sum(out * G)

def rel_err(a, b):
    return np.max(np.abs(a - b) / np.maximum(1e-12, np.abs(a) + np.abs(b)))

def numeric_grad(layer, idx, table, G, eps=1e-6):
    """Central-difference gradient of scalar_loss w.r.t. every entry of `table`."""
    g = np.zeros_like(table)
    it = np.nditer(table, flags=['multi_index'])
    while not it.finished:
        k = it.multi_index
        old = table[k]
        table[k] = old + eps; lp = scalar_loss(layer.forward(idx), G)
        table[k] = old - eps; lm = scalar_loss(layer.forward(idx), G)
        table[k] = old
        g[k] = (lp - lm) / (2 * eps)
        it.iternext()
    return g

B, T, C = 2, 5, 8
lay = EmbeddingLayer(vocab_size=vocab_size, block_size=32, C=C)
idx = data[:B * T].reshape(B, T)          # real ids, shape (2, 5)
G = np.random.default_rng(0).standard_normal((B, T, C))   # fixed upstream gradient

lay.forward(idx)
lay.backward(G)                           # analytic dE, dP
num_E = numeric_grad(lay, idx, lay.E, G)
num_P = numeric_grad(lay, idx, lay.P, G)

print("max rel err  dE:", rel_err(lay.dE, num_E))
print("max rel err  dP:", rel_err(lay.dP, num_P))

max rel err  dE: 2.4590216652835383e-10
max rel err  dP: 4.654007604730881e-11

Both max relative errors are around $10^{-10}$ — far under the 1e-5 bar we set for “the backward pass is correct.” They are tiny (much smaller than the ~1e-6 you saw for attention) because an embedding layer is a linear lookup with no softmax or nonlinearity to accumulate finite-difference noise; the only arithmetic is a gather and an add. That dE passes is the real trophy: it means the np.add.at scatter correctly accumulated the many duplicate token ids in the batch. Swap it for a plain += and this number would explode.

It Connects to Attention#

Correct gradients are half the job; the other half is that this layer produces exactly what the rest of the model expects to eat. Modules 2 and 3 built self-attention around a single convention: activations flow as (B, T, C). The whole reason we built the embedding layer is to manufacture that tensor from raw ids. So the final check feeds the embedding output straight into the module’s self-attention forward pass and confirms the shapes line up with no glue code in between.

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def self_attention(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    hs = Q.shape[-1]
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(hs)
    weights = softmax(scores, axis=-1)
    return weights @ V

X = lay.forward(idx)                       # (B, T, C) input stage output
rng = np.random.default_rng(42)
hs = 4
Wq = rng.standard_normal((C, hs)) * 0.1
Wk = rng.standard_normal((C, hs)) * 0.1
Wv = rng.standard_normal((C, hs)) * 0.1
att = self_attention(X, Wq, Wk, Wv)
print("embedding out :", X.shape, "-> attention consumes it")
print("attention out :", att.shape, "= (B, T, hs)")

embedding out : (2, 5, 8) -> attention consumes it
attention out : (2, 5, 4) = (B, T, hs)

The embedding layer hands attention a (2, 5, 8) = (B, T, C) tensor, attention projects it through the (8, 4) weight matrices, and out comes (2, 5, 4) = (B, T, hs) — no reshaping, no adapters. That is the connection that matters: the input stage you just packaged is the literal first call inside the GPT, and its output type is precisely the input type of everything downstream. Run this whole lesson a second time and, because every random draw is seeded, you get byte-for-byte identical shapes and errors — this course has no nondeterminism to hide behind.

Practice Exercises#

Try these before checking the hints. They reuse the EmbeddingLayer class and the setup you built above.

Exercise 1: Prove the Duplicate-Safe Scatter Matters#

Replace np.add.at(self.dE, idx, dout) in backward with the naive self.dE[idx] += dout and rerun the Stage 4 gradient check. Watch the dE relative error blow up while dP stays tiny, then explain why only dE breaks.

# Your code here: swap in dE[idx] += dout, re-run numeric_grad for E and P

Exercise 2: Freeze the Position Table#

Positional embeddings are learnable, but you can also freeze them. Add a flag so backward still computes dE but leaves dP at zero, then confirm dE is unchanged from the full backward pass.

# Your code here: add freeze_positions and skip the dP update when it is set

Exercise 3: A Longer Sequence Fills More of dP#

In Stage 3 the sequence had T = 16, so rows 16..31 of dP were zero. Run forward and backward with T = 32 (the full block_size) and confirm that every row of dP is now nonzero.

# Your code here: use idx = data[:32].reshape(1, 32), run forward/backward, inspect dP

Summary#

You folded a whole module into one reusable, provably correct component — the input stage every later module builds on. Let’s review what you built.

Key Concepts#

The Class Design

• __init__(vocab_size, block_size, C) stores a token table E of shape (vocab_size, C) and a learned position table P of shape (block_size, C), both seeded and initialized small (* 0.02)
• The layer commits only to vocab_size, block_size, and C, never to B or the actual T, so the same object embeds any batch of any length up to block_size
• Three kinds of state live on self: parameters (E, P), a forward cache, and gradients (dE, dP)

forward and backward

• forward(idx) gathers token_emb = E[idx], slices pos_emb = P[:T], and returns out = token_emb + pos_emb of shape (B, T, C), caching only idx and T
• backward(dout) scatter-adds dE with np.add.at (duplicate-safe) and sums dP = dout.sum(axis=0) into rows [:T]; there is no dX because the inputs are integer ids
• A gradient always matches the shape of its table: dE like E at (24, C), dP like P at (32, C)

Proving Correctness

• A float64 central-difference gradient check gave max relative errors of 2.5e-10 (dE) and 4.7e-11 (dP) — far under the 1e-5 bar, and much smaller than attention’s because the layer is purely linear
• The (B, T, C) output plugged straight into the Module 2-3 self_attention forward pass, producing (B, T, hs) with no glue code
• Everything is seeded and byte-for-byte reproducible on a second run

Why This Matters#

This EmbeddingLayer is the front door of the GPT. When you assemble the full model in Modules 6 and 7, the very first line of its forward pass is x = self.embed.forward(idx) — raw token ids in, a position-aware (B, T, C) tensor out — and everything after it (the causal attention blocks, the feed-forward networks, the final projection back to vocabulary logits) assumes that tensor already knows both what each token is and where it sits. Because you gradient-checked both tables here, the training loop never has to wonder whether the input stage is correct; it inherits a component that is already known to be right.

More broadly, you now understand the two ideas that make sequence modeling possible at all. A token embedding is a learned dictionary from discrete symbols to continuous vectors, trained by a scatter-add that routes each row’s gradient back to exactly the tokens that used it. A positional embedding is what rescues order from attention’s permutation-invariance, added in as a second learned table indexed by slot instead of symbol. Every large language model, however enormous, opens with exactly this: look up the token, add its position, and hand the result to attention.

Continue Building Your Skills#

You now have a trustworthy input stage — a class that turns bare token ids into the position-aware vectors attention needs. In the next module, The Transformer Block, you will stack the pieces you have gradient-checked into the repeating unit that gives the architecture its name. You will wrap your SelfAttention head and a small position-wise feed-forward network in residual connections and layer normalization, so a token’s representation flows forward largely unchanged while each sub-layer adds a refinement on top. That block is the thing a GPT repeats n_layer times; the embedding layer you just built feeds the first one, and its output feeds the next, all the way up the stack. Keep this class close — Module 6 opens by calling its forward on a real batch of Lantern Bay text.