melissa kate
<< Little Jane

What I Actually Learned Training a Tiny Transformer From Scratch

A journal of the real mechanics, not the textbook version.

---

The setup

---

Mechanics I now actually understand

1. The forward pass, in plain English

A token like "cold" is just an integer (say 47). The forward pass turns it into a guess about what comes next, via 7 stages:

  1. Text → blocks: conversations formatted as "A: hi\nB: yeah\n\n". The \n\n is the EOS signal.
  2. BPE tokenization: text becomes integer IDs based on a learned vocabulary.
  3. Batch sampling: random BLOCK_SIZE-length windows from the token stream.
  4. Forward pass: tokens → embeddings → attention → MLPs → logits per vocab token.
  5. Loss: cross-entropy between predicted next-token and actual next-token.
  6. Backward pass: AdamW nudges every parameter toward "less wrong."
  7. Sampling/generation: at inference, feed prompt → sample one token → append → repeat.

2. What attention actually does

Each token computes three projections of its vector: Query (Q), Key (K), Value (V).

It's a soft database lookup:

For each position i, the model:

  1. Dots Q_i with every earlier K_j → scores
  2. Softmaxes scores → weights that sum to 1
  3. Weighted sum of earlier V_j → blended context
  4. Adds that back to x_i via residual connection

Result: every token's representation now encodes "what I was, plus what I attended to in the past."

3. Why multi-head attention

Each head runs the Q/K/V mechanism independently. Different heads can learn different attention patterns:

Their outputs concatenate. More heads = more parallel patterns the model can track at once.

4. Why "deeper" matters

In a 6-layer model:

Each layer attends to the OUTPUT of the previous, so abstractions compound.

---

Lessons I had to learn the hard way

Lesson 1: More data is not always better

What I did: generated 1411 more multi-turn noir conversations via two LLMs talking to each other, retrained.

What happened: the new model was more noir-flavored but worse at carrying a conversation.

Why: the new data was stylistically narrow. The model has finite capacity — when you flood it with narrow content, it allocates more of that capacity to narrow patterns, AT THE COST of conversational primitives like "answer the question."

Generalized rule: distribution matters more than volume. Data diversity > data volume.

Lesson 2: Tokenization drifts with your data

The BPE tokenizer is trained on your corpus. Change the corpus, change the vocab.

When I added more noir data, the word "cold" became frequent enough to earn its own merge — went from 2 tokens (·co + ld) to 1 token (·cold). Old vocabulary words I cared about no longer existed identically.

The cost is hidden — you don't notice until you try to compare checkpoints and the token sequences don't line up.

Lesson 3: Attention sinks are a feature

When a head has no clear signal about what to attend to, it dumps attention onto the first token (position 0). This is the "attention sink" phenomenon.

It looks like a bug ("the model is bailing out!") but it's actually graceful uncertainty: instead of committing to attending to the wrong content, the head parks attention somewhere harmless.

The worse model lost this skill. Its deep layers had less sink behavior (0.52 vs 0.65 in layer 5). It became MORE confident, but confident in wrong patterns. On out-of-distribution inputs (casual prompts), the better model fell back gracefully; the worse one committed to nonsense.

Confidence ≠ correctness. Sometimes the right answer is to back off.

Lesson 4: The dropout / train-val inversion

Train loss was HIGHER than val loss the entire run. Looks wrong.

Why: dropout (0.1) is ON during training, OFF during eval. So train loss is computed on a deliberately handicapped model with 10% of activations zeroed. Val loss is the unhandicapped model. The gap is dropout, not generalization.

To compare: the "real" train loss (without dropout) is what to compare to val.

Lesson 5: Capacity vs data must scale together

The reason fine-tuning works so well: the Q/K/V projection matrices already know how to do pronoun resolution, subject tracking, topic continuity. Your tiny model has to learn all of that from 50K examples. A pretrained model inherits them.

Lesson 6: Visualizing attention is shockingly informative

Heatmaps of attention weights show you:

A bar chart of "what does the final token attend to" is even clearer — it directly shows what the model is reading to generate the next response.

Lesson 7: Catastrophic forgetting is real

What I did: full fine-tune of GPT-2 small (124M params) on my dialogue data, LR 5e-5, 3 epochs.

What happened: the model lost its English fluency. Plain prompt "The detective looked at" produced "I't to.B He still to." — base GPT-2 produces "him with a look of disbelief..."

Why: updating all 124M params with my narrow data overwrote the general English patterns. The model traded its language fluency for trying to memorize the noir format.

Diagnostic: running base GPT-2 and the fine-tuned version side-by-side on the same plain English prompt is the cleanest way to detect this — if base works and fine-tuned doesn't, you've forgotten.

Lesson 8: LoRA preserves fluency by construction

What LoRA does: freezes the base model entirely, adds small low-rank adapter matrices on top of the attention projections (~0.5% of total params trainable).

Why it can't catastrophically forget: you can't overwrite what you can't update. The base model's English is preserved by definition.

What I observed: LoRA samples DID stay closer to coherent English than the full fine-tune. Real words, occasional phrases ("Don't know," "Be of use"). But still not good enough — see next lesson.

Lesson 9: Pretrained model alignment with your task matters more than I expected

The kicker: LoRA-adapted GPT-2 was STILL worse than my from-scratch tiny model.

Why: GPT-2 small was trained in 2019 on web articles, code, forum posts. Its deep instincts are "after a colon comes more text," "be long, be articulate, be formal." My task is "be Han: 4 words, dry as a bone." Those instincts ACTIVELY FIGHT the task. The LoRA adapters spend their tiny capacity trying to undo GPT-2's verbosity rather than learning noir.

The from-scratch model didn't have that problem — it learned terse dialogue as its native distribution from the start. There were no "be verbose" weights to override.

The real rule: "fine-tune a pretrained model" works only when the pretrained model is already in your task's neighborhood. For:

For a narrow stylistic task, a tiny from-scratch model can absolutely beat a poorly-aligned pretrained model. The pretrained advantage only kicks in when the base is task-aligned.

Lesson 10: Right base + LoRA still isn't enough at this data scale

After GPT-2 failed (wrong base), I tried DialoGPT-small (GPT-2 architecture pretrained on Reddit dialogues — supposedly the "right" base) with LoRA.

Result: also word salad. "see lady? she getting in hang? This'm her in place too though" — coherent vocab, no grammar, no terseness.

Why: DialoGPT was trained on Reddit dialogues, which are RAMBLY and casual — almost the opposite of noir terseness. So while DialoGPT nominally "knows dialogue," it doesn't know MY KIND of dialogue. The 117M base + 0.5% LoRA capacity wasn't enough to pull it from Reddit-rambling to Bogart-tight on ~50K of my pairs.

The triangulated conclusion across three experiments:

ApproachResult
From-scratch 12MStylistically nailed, conversationally limited
GPT-2 + LoRA (wrong base)Word salad
DialoGPT + LoRA (closer base but still small + wrong-flavor)Word salad

The data is the bottleneck, not the model. For my goal (terse noir character that converses), I need either:

Lesson 11: Cognitive friction in data collection produces better data AND more of it

The discovery: I expected writing noir replies for noir scenes (midnight at the docks, etc.) to be easiest. It wasn't. Replying as a 1940s detective to a backyard barbecue or juice cleanse turned out to be way more fun to write — and I produced more pairs, faster.

Why this works mechanically:

Why the training data is better too:

Practical implication: when designing collection prompts, prefer settings that create comedic or cognitive asymmetry with the character's voice. The asymmetry is the engine — engagement AND signal both come from the same source.

This was a real reason the earlier "more noir data" attempts crowded out conversational primitives. The data was internally consistent (noir voice + noir setting), but the model never learned to apply voice OUTSIDE its training distribution.

---

My decision

I'm sticking with the from-scratch path and treating data collection as the long-term project. Reasons:

Production answer for "I want to chat with Han right now": Claude API via /han. The trained tiny model is the learning artifact, not the product.

---

The numbers from my run

Best from-scratch model

Worse (overspecialized) model

What the canonical model produced

how are you?       → It's fine, I guess.
what's up?         → Nothing much. You?
cold out tonight   → If you say so
heard about Walter? → I'd rather hear the department.

Stylistically nailed, conversationally limited.

---

The honest assessment

What this experiment did well:

What it can't do, and won't with more data alone:

The architectural ceiling is real. Going from "stylistic mimic" to "actually conversational" requires either dramatically more data + a bigger model OR fine-tuning a pretrained model that already has language fluency baked in.

---

What I'd do differently if starting over

  1. Mix registers in training data. Don't let any one style dominate.
  2. Generate dialogue with Claude API, not local 14B models. The local models are too dumb to maintain real coherence; their "noir" was performative, not conversational.
  3. Visualize attention earlier. I should have looked at attention weights after each training run, not at the end.
  4. Pick a task-aligned base model BEFORE fine-tuning. GPT-2 was a terrible base for terse dialogue. Should have started with DialoGPT or an instruction-tuned model.
  5. Use LoRA from the start, never full fine-tune. Catastrophic forgetting is the default, not the exception.
  6. Save checkpoints more aggressively. Compare attention patterns across iterations, not just final models.

---

What I still don't fully understand

---

Things I'm proud of

You don't get to feel any of that from reading the GPT-2 paper.