Step 2 — Break your model on purpose

Step 2 — Break your model on purpose 🔬¶

Goal: understand what every part of your Step 1 model actually does — by removing it, one piece at a time, and watching exactly what breaks.

Why this step matters¶

In Step 1 you wrote a tiny transformer by hand and finished when you could explain every line out loud. This step is where you prove those explanations to yourself.

You can’t truly understand a part by reading about it. You understand it by deleting it and seeing what falls over. Remove the causal mask and the model learns to cheat. Remove the residual connections and a deep model refuses to learn at all. Each break turns a line you believe matters into a line you know matters.

This takes one afternoon and is worth more than a month of tutorials. It’s also the step people skip, because it feels like going backwards. It isn’t. Every hour here saves you a day of confused debugging in Step 9.

The shape of this chapter 🗺️¶

Five experiments, against the exact model from Step 1 — the same Head, Block, Config, and get_batch. Each one breaks a different part and leaves a different fingerprint:

The rest of this chapter is the five experiments, in the same order.

1 · Remove the causal mask 🔓¶

This is the most important bug in the entire field. In Head.forward, Step 1 had this line — it hides the future before the softmax:

wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # <-- Step 1: hide the future
wei = F.softmax(wei, dim=-1)

Break it by deleting that one line:

wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
# wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # <-- BROKEN: future visible
wei = F.softmax(wei, dim=-1)

Now every position can peek at the answer sitting one step ahead:

What you’ll see — the loss looks incredible, then the samples are garbage:

iter    0: train 4.71
iter  200: train 0.08     <- too good to be true
iter  400: train 0.02
sample: े ैं ् प ्ा ...    <- pure noise

2 · Normalize after, not before 🔄¶

Step 1 used pre-norm: the LayerNorm sits inside the x + (...), so a clean copy of x runs untouched from input to output.

def forward(self, x):
    x = x + self.sa(self.ln1(x))   # Step 1: pre-norm
    x = x + self.ff(self.ln2(x))
    return x

Break it by switching to post-norm — the original 2017 design, where the norm happens after the add:

def forward(self, x):
    x = self.ln1(x + self.sa(x))   # post-norm
    x = self.ln2(x + self.ff(x))
    return x

What you’ll see — training is jumpier, and it gets clearly worse as you stack more blocks (try n_layer = 8).

3 · Remove the residual connections 🛣️¶

Now remove the x + — the “add the input back” from Step 1’s block:

def forward(self, x):
    x = self.sa(self.ln1(x))   # BROKEN: no residual (dropped the `x +`)
    x = self.ff(self.ln2(x))
    return x

Set n_layer = 10 so the model is genuinely deep, and train.

What you’ll see — the loss barely moves. A deep model refuses to learn:

iter    0: train 4.71
iter 1000: train 4.55
iter 3000: train 4.51     <- barely moves

4 · Learning rate 10× too high 🌡️¶

One number, in Config:

class Config:
    ...
    learning_rate = 3e-4   # Step 1: a calm, sensible value

Break it:

class Config:
    ...
    learning_rate = 3e-3   # BROKEN: 10× too high

What you’ll see — instead of falling smoothly, the loss jumps up, or turns into NaN and never recovers:

iter    0: train 4.71
iter  100: train 3.10
iter  200: train 7.88     <- a loss spike
iter  300: train nan

5 · Overfit on purpose 🧠¶

Keep the model exactly as Step 1, but feed it only a tiny slice of text and train for a long time. Point the data at, say, the first 100 KB and raise max_iters.

Watch both losses — Step 1’s estimate_loss already prints train and val:

iter    0: train 4.70  val 4.70
iter 2000: train 0.40  val 3.10     <- the gap opens
iter 5000: train 0.05  val 3.60     <- train falls, val CLIMBS
sample: (a verse copied word-for-word from your file)

Training loss near zero while validation loss rises is the classic fingerprint:

What you should see ▶️¶

Put side by side, the five failures each have a distinct signature. Learning to read these fingerprints is the real skill of this chapter:

• 🔓 Loss too low + output noise → leakage (you’re seeing the answer).

• 🔄 Loss jumpy, worse when deep → a normalization or depth problem.

• 🛣️ Loss flat on a deep model → the gradient can’t reach the early layers.

• 🌡️ Loss spikes up or NaN → learning rate too high.

• 🧠 Train down, val up → memorisation.

Where people usually get stuck¶

Changing two things at once. Then a result appears and you can’t say which change caused it. One change, one run, one note.

Forgetting to revert — especially the causal mask. Before every new experiment, restore your known-good Step 1 model. (This is exactly why the flag harness above is worth the ten minutes to set up.)

Not writing it down. The five notes are the output of this step — not the training runs. Skip the notes and you did the experiments but kept none of the understanding.

You are ready to move on when¶

You have five short notes, one per experiment, in your own words — and your model is back to the correct Step 1 version.

A good test: hand the notes to another engineer. If they learn something from your five sentences, you did this properly. Then head to Step 3, where we finally fix the real problem — how we chop Sanskrit and Urdu into pieces.