L10.3 · LESSON 3 OF 6

A causal block

Assemble a tiny causal language-model experiment with a clear data split, baseline, training loop and evaluation.

THE IDEA

The causal block, traced

The model: look up a vector for each token and add a position vector; run one masked attention sublayer and add its output to the stream (residual); run an MLP and add that; finally score every vocabulary entry by multiplying with the transposed embedding table. Trace one tiny batch through the shapes and make sure no position can use a later token.

batch B2positions T8vocabulary V27width d16

Stage 1 of 5

▶ token ids (2, 8) integers in 0 … V−1

parameters: tokens 27·16 = 432, positions 8·16 = 128, attention 4·16² = 1,024, MLP 2·16·64 = 2,048 total = 3,632 (no biases or norms; output tied to the token table)

Step through the stages. Move the sliders to see how shapes and the parameter count change.

Predict first. In a shifted training window, what is the target for the first input token?

Chapter context

The capstone joins the earlier pieces into one inspectable experiment. The goal is to understand the model and its evidence, not to claim that a toy trained on a few paragraphs competes with a production assistant.

TRY IT YOURSELF

Part 1. Do it hands-on

Trace shapes and count parameters by hand, then write the masked attention, MLP and forward pass and prove that the future cannot leak.

batch B2positions T8vocabulary V27width d16

Stage 1 of 5

▶ token ids (2, 8) integers in 0 … V−1

parameters: tokens 27·16 = 432, positions 8·16 = 128, attention 4·16² = 1,024, MLP 2·16·64 = 2,048 total = 3,632 (no biases or norms; output tied to the token table)

logits shape for B = 2, T = 8, V = 27:

target for the first input of [5, 9, 2, 7] parameters for V 27, d 16, T 8, f 64

To do: A batch of B = 2 sequences, each T = 8 characters, vocabulary V = 27. The model outputs one score per vocabulary entry at every position. Type the three numbers of the logits shape
To do: A window of tokens [5, 9, 2, 7]: inputs are the first three, targets the last three. What is the target for the FIRST input token (5)?
To do: One block, V = 27, d = 16, T = 8, MLP width f = 64, no biases or norms, output layer tied to the embedding: token table (V·d) + position table (T·d) + attention (4·d²) + MLP (2·d·f). Total?

Part 2. Do it in code

init_params(V, d, T, f, seed) is provided: it returns a dict with E (V×d token table, also the output matrix), Pos (T×d), Wq, Wk, Wv, Wo (d×d), W1 (d×f) and W2 (f×d).
causal_mask(T): a boolean (T, T) array, True where position i may attend to position j (j ≤ i).
attention(x, p): x has shape (T, d). One head: q, k, v = x@Wq, x@Wk, x@Wv; scores = [email protected] / sqrt(d); masked positions get −inf before the softmax; return (weights @ v) @ Wo.
mlp(x, p): relu(x @ W1) @ W2.
forward(ids, p): ids is an int array (B, T). For each sequence: x = E[ids] + Pos[:T]; x = x + attention(x); x = x + mlp(x); logits = x @ E.T. Return logits of shape (B, T, V).
param_count(p): the total number of numbers in all arrays of p (each array counted once).

import numpy as np

def init_params(V, d, T, f, seed):
    rng = np.random.default_rng(seed)
    g = lambda *shape: rng.standard_normal(shape) * 0.1
    return {"E": g(V, d), "Pos": g(T, d), "Wq": g(d, d), "Wk": g(d, d), "Wv": g(d, d), "Wo": g(d, d), "W1": g(d, f), "W2": g(f, d)}

def causal_mask(T):
    pass

def attention(x, p):
    pass

def mlp(x, p):
    pass

def forward(ids, p):
    pass

def param_count(p):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

What is the target for the first input token in a shifted training window?

AThe last token of the sequence.BThe first input token itself.CThe next token in the same sequence: inputs are tokens 1…T, targets are tokens 2…T+1, and the mask hides later tokens.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a token sequence of length n and context length T, stride-one windows containing T inputs and their shifted targets require T+1 original tokens. The number of complete windows is max(0,n-T). Keep this counting convention distinct from nonoverlapping batching.

Worked example

A 400-token text is split into 320 training tokens and 80 validation tokens before windowing. With T=8, the training part contains 320−8=312 complete shifted windows, while validation contains 72. Constructing windows before the split would risk overlap across the boundary.

Reference code

The following example uses the Python standard library and does not call a remote service.

text='abcd'*100
train,validation=text[:320],text[320:]
context=8
windows=[train[i:i+context+1] for i in range(len(train)-context)]
assert len(text) == 400
assert len(windows) == 312
assert all(len(w)==9 for w in windows)
assert len(validation)-context == 72

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Toy output quality is not a benchmark claim. Splitting overlapping windows after construction leaks shared text. Save the tokenizer and architecture configuration, not only raw weights.