A Starting Loss Twice as Bad as Guessing: From Shakespeare to the 2.49 Baseline

The previous part took backward() apart down to the bottom. This one starts over on a new dataset — and from here to the end of stage one there is a single goal: write a GPT from scratch.

The sixth episode (WaveNet) is skipped. It turns the previous MLP into a hierarchy and the loss falls further, but it does so with parts that have already been taken apart — so, on value for time, straight on to the seventh. Which is why the source list goes from Ep.5 to Ep.7; that gap is deliberate.

New data, and a different way to split it

The vocabulary is sorted by ASCII: 0 is newline, 1 is space, then punctuation, upper case, lower case. The only digit anywhere in Shakespeare is 3 — a character-level vocabulary is determined entirely by the data, and anything absent from it raises KeyError on sight.

n = int(0.9 * len(data))
train_data = data[:n]      # first 90%: 1,003,854
val_data   = data[n:]      # last 10%: 111,540

⚠️ Split by position, never shuffle. Part 12 set the rule as “split by name, not by sample” — shuffling was safe there because each name is an independent item. Here the text is continuous, and shuffling would leave validation sentences continuous with training context: immediate leakage.

A token is something you define

A token is an entry in the model’s vocabulary, corresponding to one integer ID. It is the smallest unit the model can handle.

A neural network only computes numbers, and text is a string — tokens are the bricks of that bridge:

text  →  [tokenizer]  →  integers  →  [model]  →  integer  →  [decode]  →  text
'hi'  →               →  [46, 47]  →           →  ...

🔑 The stoi / itos written back in parts 7 to 22 were already the simplest possible tokenizer, just not called one at the time. The model knows nothing but integers, start to finish.

Three granularities

The same sentence, The cats are unhappy:

Subword splits are decided by frequency: a common word is one token, a rare one is broken into a few pieces, and a rarer one all the way down to bytes — so OOV never happens.

❗ Why sequence length matters so much

Attention costs O(T²). For the same passage of text:

A factor of 16. That is why nobody trains a real LLM at character level, and why context windows are finite — a 128K window costs 256 times what 8K does (16²).

BPE is derived from statistics, not designed

① start with a vocabulary of all 256 byte values
② count which adjacent pair occurs most often (say e and s)
③ merge it into a new token (es) and add it to the vocabulary
④ repeat until the vocabulary reaches its target size (say 50k)

The whole process comes out of the data, with no linguistic knowledge required — point the same algorithm at Chinese, code or DNA and it runs unchanged.

Does it merge across words? Yes. GPT attaches the space to the following word, so " have" and "have" are two different IDs, and " of the" is often a single token — BPE only knows frequency, not words. It does not run away, though: most tokenizers pre-split with a regular expression and never cross punctuation or newlines, and the vocabulary size is capped.

Two practical consequences:

• Numbers split unpredictably (1234 might become 123 + 4) — one of the reasons early LLMs were bad at arithmetic; Llama 3 and Qwen later mandated digit-by-digit splitting
• The same word has several IDs: hello, " hello" and Hello are three unrelated symbols to the model — one extra space in a prompt can change the output

🔑 A tokenizer is welded to the model’s weights

It is derived from a corpus, but trained entirely separately from the model: the tokenizer is trained once and frozen, then used to cut text for training the model. It has no gradients and never updates.

Swap the tokenizer and a word’s ID changes, the embedding table is read at the wrong row, and the whole model is scrap. Which is why tokenizer.json has to come down with the weights.

★ This is the real reason for choosing Qwen

The biggest trap in borrowing a tokenizer is language match: for the same Chinese text, Qwen (150k vocabulary, heavily Chinese corpus) may need a third of the tokens Llama 2 does (32k).

🔑 Not “it feels better at Chinese”, but: half the tokens for the same text ⟹ twice the context and half the cost. And by the O(T²) arithmetic above, halving the tokens cuts attention’s cost to a quarter.

It cuts the other way too: a small model with a large vocabulary is badly proportioned — a 10M parameter model on a 50k vocabulary spends 50000 × 768 = 38.4M on the embedding table alone, which dwarfs the rest.

One chunk holds T samples

x = train_data[:8]          # input
y = train_data[1:9]         # answers = x shifted by one

input [18]                           → predict 47
input [18, 47]                       → predict 56
…
input [18, 47, 56, 57, 58, 1, 15, 47] → predict 58

A chunk of length 9 yields 8 training samples. Two benefits: one forward pass produces 8 predictions at once, and the model sees context lengths from 1 to 8 — at inference it starts from nothing, so it has to handle short contexts.

⚠️ This differs from part 9. There block_size=3 meant “always look at 3”; here it is “be able to handle 1 through 8”.

The two dimensions (B, T)

Every step is 32 × 8 = 256 predictions. Like a test paper of 32 pages with 8 fill-in-the-blanks each, every page drawn from a different chapter, all marked in one pass.

⚠️ But those 256 are not fully independent samples — the 8 positions of one sequence come from the same stretch of the original. For a bigram it makes no difference (it only sees the current character), but with attention it does: position 3 must not peek at position 5’s answer. That is the entire reason the lower-triangular mask exists later on.

A lookup is a fully connected layer

class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

super().__init__() lets nn.Module set up its registry first; omit it and you get an AttributeError, with no exceptions. self.xxx = nn.Embedding(...) registers itself as a parameter automatically — which is why there is no more hand-written parameters() and zip(parameters, grads) the way part 22 needed.

The two 65s in nn.Embedding(65, 65) mean different things

A trainable table of 65 rows by 65 columns, random to begin with:

The first 65 is “there are 65 characters” (rows); the second is “how long a vector each character becomes” (columns). They coincide here and mean entirely different things.

❗ Why the second one is forced to be 65: a bigram has only this table and no other layer, so what comes out of the lookup has to be the scores over all 65 candidates already. Pick 32 and there would be 32 scores, leaving 33 characters unpredictable.

Later it gets split apart: nn.Embedding(65, 32) + nn.Linear(32, 65), and those 32 dimensions in the middle are the room attention works in. Normally the embedding dimension is a free hyperparameter and far smaller than the vocabulary (GPT-2: 50k vocabulary, 768 embedding) — part 9 made the case that an embedding’s value is compression, and compression is what forces semantic structure out. A bigram compresses nothing, so it learns no semantics.

one-hot × W = a lookup

Draw the model as a network: 65 input nodes on the left (one-hot), 65 outputs on the right, fully connected in between with W as the weights. When the current character is 'h' only that one cell is 1, so only the 65 edges leaving 'h' matter, and the output is exactly W’s row 46.

🔑 nn.Embedding is a fully connected layer with no bias, skipping the step of building a one-hot and multiplying. Part 9 covered one-hot’s geometry (27 orthogonal dimensions, compression as a necessity); this is its algebraic face.

(B, T, C)

logits = self.token_embedding_table(idx)    # (B,T) → (B,T,C)

🔑 From here to reading Qwen’s source, these three shapes are everywhere: (B,T) for token IDs, (B,T,C) after embedding, (B,T,T) for attention weights.

❗ Those 65 numbers do not “represent one character” — they are the scores of 65 candidates, the j-th corresponding to itos[j]. It is not saying “I am b”; it is scoring all 65 at once.

⚠️ The whole forward pass has no multiplication, no addition and no activation function. 256 character indices go in, 256 rows of scores come out — that is all of a bigram’s computation.

Why flattening with view does not misalign

F.cross_entropy only eats two dimensions, so logits.view(B*T, C) is needed. Why does C not get scrambled?

Part 10 covered it: memory holds one run of numbers, and B, T, C are a reading convention; row-major order makes the last dimension’s stride 1, so a position’s 65 numbers are contiguous. view(B*T, C) cuts every 65, which is exactly one whole position; row i is position (b, t) with i = b×T + t. targets flattens in the same order, so they correspond by construction.

⚠️ The dangerous version: view(C, B*T) — right shape, no error, entirely wrong results. Transposing needs .transpose(). This is the second instance of part 21’s “neither errors nor is obvious”: there it was fancy indexing, here it is dimension order. Both slip past the shape check.

🔑 Shapes never change behind your back in PyTorch. Every reshaping has a line of view / reshape / permute doing it.

generate: all there is to autoregression

for _ in range(max_new_tokens):
    logits, loss = self(idx)            # (B,T,C)
    logits = logits[:, -1, :]           # keep only the last position (B,C)
    probs = F.softmax(logits, dim=-1)
    idx_next = torch.multinomial(probs, num_samples=1)
    idx = torch.cat((idx, idx_next), dim=1)

Predict, draw, append, predict again. ChatGPT emitting one character at a time is running this loop.

❗ Why only the last position

Suppose generation has reached \nHe, so the output is (1, 3, 65):

The characters the first T−1 predictions correspond to are already in the sequence. Training uses every position (each has an answer to check against); generation uses only the last.

multinomial draws by probability rather than taking the maximum — part 12’s sampling section covered it, and argmax sends the output into loops; temperature was covered in part 8 (divide the logits by T before softmax). Neither is repeated here.

❗ Every round recomputes the whole sequence and uses only the last position

For a bigram this is pure waste — the first T−1 positions’ predictions are computed and discarded. But the shape of the code is there for attention: once attention is added, the last position pools everything before it, and this function does not change by a line. (Real inference solves the recomputation with a KV cache.)

The starting point is idx = torch.zeros((1, 1), dtype=torch.long), that is [[0]] — index 0 is a newline, so it means “start a new paragraph”. Replace it with encode('ROMEO:\n') and the model carries on with Romeo’s lines.

🔑 Which is what a prompt is: encode your words into idx and run generate.

★ A starting loss twice as bad as guessing

starting loss 4.87        uniform guessing: ln 65 = 4.174

Converting 4.87 back, the model gives the correct answer an average probability of e⁻⁴·⁸⁷ = 0.77%, while guessing uniformly with eyes closed gives 1/65 = 1.54%.

Exactly half. Straight after initialisation the model is twice as bad as doing nothing at all — it is confidently wrong.

Part 14 diagnosed this: the initial weights are too large, the logits’ standard deviation is stretched, and softmax piles the probability onto the wrong few. Part 10’s “starting loss of 17 instead of 3.3” is a heavier case of the same illness. Nothing is fixed here — a few hundred steps bring it down on its own.

💡 This ruler works every time: uniform guessing costs ln(vocabulary size). For 27 characters that is 3.2958 (part 10’s ruler); for 65 it is 4.174. A starting loss above it means the initialisation is wrong.

The 2.49 baseline

A single batch’s loss is noisy, so take the mean over 200 batches and report train and val separately:

@torch.no_grad()
def estimate_loss():
    ...
    model.eval()
    # draw eval_iters batches, take each one's loss, average
    model.train()

This is part 22’s split_loss() upgraded — there the whole dataset went in at once, here it is 200 batches, which is what makes it fit in memory when the data is large. model.eval() / model.train() do nothing for a bigram (no BN, no Dropout), but they go in now because evaluation will have to switch Dropout off later.

step    0: train 4.7305, val 4.7241
step  300: train 2.8110, val 2.8249
step  600: train 2.5434, val 2.5682
step  900: train 2.4932, val 2.5088
…
step 2700: train 2.4738, val 2.4911

Three things:

① It flattens by step 900. The remaining 1800 steps buy 0.02 — this is the bigram’s ceiling, and seeing only the previous character, no amount of further training changes that.

② train ≈ val, 0.02 apart. With 4,225 parameters against a million characters there is no capacity to memorise. Part 12’s test applies unchanged: a small gap says “it can be made bigger”, and an absolute value of 2.49 says “it should be”.

③ That 2.49 is the control group for everything that follows. With attention it reaches about 2.4; with the whole Transformer built, 1.48.

What it generates already has shape: capitals mostly at line starts, punctuation at word ends, plausible word lengths, the “character name, colon, newline” script format, alternating vowels and consonants. KIN is one G away from KING.

But not one real word — seeing only the previous character, at hou it sees only u, and cannot get to house.

The optimiser: the first ready-made one

optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)

The p.data += -lr * grad hand-written in part 22 is the simplest optimiser there is (SGD), and optimizer.step() does it in one line. AdamW does three more things:

Which is why the learning rate’s magnitude differs too: SGD commonly uses 0.1, the Adam family 1e-3 to 3e-4 — it rescales the step internally.

⚠️ optimizer.zero_grad() cannot be skipped — PyTorch accumulates gradients. Part 6’s ten-seed comparison measured what that costs: the version that never zeroed ranged over twenty-six orders of magnitude. All those += in the previous part are the same accumulation semantics.

What is being set up