feat: V2 analysis — head patching, redundancy test, seed robustness, error bars

Author: diego-magana

.github/workflows/ci.yml

@@ -13,7 +13,15 @@ jobs:
       - name: Install dependencies
         run: |
           pip install torch --index-url https://download.pytorch.org/whl/cpu
-          pip install numpy matplotlib pytest
+          pip install numpy matplotlib pytest nbconvert jupyter
           pip install -e .
       - name: Run tests
-        run: pytest tests/ -v
+        run: pytest tests/ -v
+      - name: Smoke-execute the analysis notebooks
+        # These load the committed checkpoint (no training) and exercise the full
+        # analysis pipeline, so they guard the notebooks against code drift. The
+        # progression notebooks 01-03 train and are excluded for runtime.
+        run: |
+          jupyter nbconvert --to notebook --execute --inplace \
+            notebooks/progression/04_full_gpt.ipynb \
+            notebooks/05_attention_analysis.ipynb

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Diego Magana
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -9,22 +9,24 @@ The short version of what I found: in this 4-layer, 4-head character model, almo
 all of the predictive work runs through a single attention head — layer 1, head 1 —
 which copies the previous character into the position that predicts the next one.
 Two heads *look* like they do that job; only one of them actually matters, and it
-took intervening on the network to tell them apart.
+took intervening on the network to tell them apart. Retrain from a different seed
+and the same kind of head takes over — just not always at the same index.
 
 This is the third in a series after
 [micrograd](https://github.com/diego-magana/micrograd) (a scalar autograd engine)
 and [makemore](https://github.com/diego-magana/makemore) (n-gram to WaveNet
-character models). The activation-patching analysis here is the one I pointed
-forward to at the end of makemore.
+character models). The patching analysis here is the one I pointed forward to at the
+end of makemore.
 
 ---
 
 ## What I found
 
-I trained the full model to convergence and ran four analyses on it — two that just
-look at the network (attention patterns, the residual stream) and two that
-intervene on it (head ablation, activation patching). The interesting part is that
-looking and intervening disagree, and the disagreement is where the real result is.
+I trained the full model to convergence and ran six analyses on it — two that just
+look at the network (attention patterns, the residual stream) and four that
+intervene on it (head ablation, a pairwise-ablation redundancy test, residual
+patching, and head patching). The interesting part is that looking and intervening
+disagree, and the disagreement is where the real result is.
 
 **The model has two heads that look like previous-token heads.** Averaging attention
 over 512 held-out sequences, layer 0 is mostly diffuse, layer 1 leans entirely on
@@ -36,44 +38,75 @@ attention maps they read as the same kind of head.
 ![Mean attention maps](assets/attention_grid.png)
 
 **Only one of them matters.** When I zero each head and remeasure validation loss,
-L1 H1 costs ≈ 0.89 nats — seven times more than any other head. L0 H3, the *other*
-sharp previous-token head, costs ≈ 0.04. Nearly identical attention patterns, an
-order of magnitude apart in how much the model actually needs them. My read is
-redundancy: L1 H1 carries a stronger version of the same signal one layer later, so
-removing L0 H3 on its own barely registers. This is the part I'd point a reviewer
-at first — the attention maps alone would have told me these two heads were
+L1 H1 costs ≈ 0.89 ± 0.02 nats — seven times more than any other head and far
+outside the error bar. L0 H3, the *other* sharp previous-token head, costs
+≈ 0.04 ± 0.00. Nearly identical attention patterns, an order of magnitude apart in
+how much the model actually needs them. This is the part I'd point a reviewer at
+first: the attention maps alone would have told me these two heads were
 interchangeable, and they aren't close.
 
-**Activation patching shows where L1 H1's information goes.** I take a clean context,
-corrupt the character two positions before the end, and then patch the clean
-residual-stream activations back in one site at a time, measuring how much of the
-original prediction comes back. Through the embedding and block 0 the corrupted
-information sits at its own position. Then block 1 moves it forward: recovery jumps
-from ≈ 0.13 at the corrupted position to ≈ 0.75 at the prediction position, and the
-later blocks carry it the rest of the way. Block 1 is where L1 H1 lives, so this is
-the head's copy operation caught in the act.
+![Head ablation](assets/head_ablation.png)
+
+**The reason L0 H3 is free is redundancy — and I tested it instead of asserting it.**
+If L0 H3 is cheap to remove only because L1 H1 carries a stronger copy of the same
+signal downstream, then it should get expensive once L1 H1 is gone. It does: ablating
+L0 H3 on its own costs ≈ 0.04, but ablating it once L1 H1 is *already* ablated costs
+≈ 0.23 — almost six times as much. L0 H3 isn't dead weight; it's a backup
+previous-token head whose contribution is masked by the stronger one.
 
-![Activation patching recovery](assets/activation_patching.png)
+**Residual patching shows where the information flows.** I take a clean context,
+corrupt the character two positions before the end, and patch the clean
+residual-stream activations back in one site at a time, measuring how much of the
+original prediction comes back (as a contrastive logit difference, on the 189 of 256
+examples where the corruption actually flipped the prediction). Through the embedding
+and block 0 the corrupted information sits at its own position. Then block 1 moves it
+forward: recovery jumps from ≈ 0.24 at the corrupted position to ≈ 0.74 at the
+prediction position, and the later blocks carry it the rest of the way.
+
+![Residual patching recovery](assets/activation_patching.png)
+
+**Head patching pins it on L1 H1 directly.** Residual patching localizes the move to
+block 1, but block 1 has four heads and an FFN — it can't tell me *which* head does
+the copying. So I patch one level down: splice a single head's clean output into the
+corrupted run and leave everything else corrupted. Restoring L1 H1 alone recovers
+≈ 0.51 ± 0.03 of the prediction — five times the next head. One head out of sixteen
+accounts for roughly half the effect by itself. It isn't all of it, and it shouldn't
+be, since the previous-token signal also leaks through other paths; but I'm no longer
+inferring L1 H1 from "block 1 matters" plus "L1 H1 ablates hardest," I'm routing its
+clean output into a broken run and watching the prediction come back.
+
+![Head patching](assets/head_patching.png)
 
 **The prediction sharpens steadily with depth.** Reading each layer's residual stream
 through the final unembedding (the logit lens), top-1 next-character accuracy climbs
 0.06 → 0.12 → 0.21 → 0.34 → 0.51 across the four blocks, with the last block doing
-the most. One wrinkle I didn't expect: the embedding-level logit-lens cross-entropy
-is *worse* in this converged model than it was in an under-trained one, even though
-embedding-level accuracy is unchanged. The embeddings specialize to feed the deeper
-layers rather than to be read out directly, so the lens — which uses the final
-unembedding — reads them less faithfully. A reminder that the logit lens is a lower
-bound on what a layer knows, not a decoder of it.
-
-Three alternative readings of the result — and the limits of each method — are at 
-the end of the notebook.
+the most. The logit lens is a lower bound, not a decoder: it reads intermediate
+streams through the *final* unembedding, so it under-reads any feature a layer has
+computed but not yet rotated into the output basis. The embedding-level readout
+already lands at 6%, a weak unigram prior baked into the static embeddings before any
+attention runs.
+
+**Is "L1 H1" an accident of initialization?** Partly. I retrained the same
+architecture from four more seeds (at a reduced 3k-step budget — the specialization
+is unambiguous well before convergence) and recorded, for each, the single most
+causally important head. Across all five seeds it is *always* a sharp previous-token
+head (prev 0.90–0.97) sitting in layer 0 or 1 — but which head it lands in moves
+around (L1 H1 here, L0 H0 / L0 H1 / L0 H3 elsewhere). So the mechanism replicates and
+the head index doesn't. "L1 H1 carries the prediction" is true of *this* model; the
+durable statement is "training reliably builds a previous-token head, and the model
+leans on it."
+
+What this doesn't show, and the three things I'd want before calling it a circuit
+(path patching, the QK/OV weight-space read, and anything non-local — `block_size=32`
+and 0.21M parameters leave no room for induction heads, so this is the one mechanism
+a tiny local model *can* have) are spelled out at the end of
 [`notebooks/05_attention_analysis.ipynb`](notebooks/05_attention_analysis.ipynb).
 
 ---
 
 ## What it builds
 
-The four analyses run on a model I build up one ingredient at a time in
+The analyses run on a model I build up one ingredient at a time in
 [`notebooks/progression/`](notebooks/progression). To keep the comparison about the
 *architecture* rather than about training length, I hold the three transformer
 stages at a fixed 5,000-step budget:
@@ -106,15 +139,17 @@ reference run to within ≈ 0.01 nats.
 ```bash
 pip install -e .                       # editable install; pulls torch, numpy, matplotlib
 python train_gpt.py                    # reproduce assets/gpt.pth (~15 min on CPU, 30k steps)
+python seed_sweep.py --seed 0          # reproduce one row of the seed-robustness check
 jupyter lab notebooks/                 # run the progression, then 05_attention_analysis
 ```
 
-The tests cover the package's correctness invariants — including the one the whole
-patching analysis rests on, that splicing a run's own clean activations back in
-changes nothing:
+The tests cover the package's correctness invariants — including the two the patching
+analyses rest on, that splicing a run's own clean activations (residual *or* a single
+head's output) back in changes nothing:
 
 ```bash
-pytest                                 # 13 tests, runnable from any directory
+pytest                                 # 16 tests, runnable from any directory
+pytest -m "not slow"                   # 10 fast tests; skips the training-based ones
 ```
 
 I commit `assets/gpt.pth` (the trained 30k model) so the analysis notebook runs on
@@ -132,17 +167,18 @@ gpt/
 │   ├── data.py              char tokenizer, train/val split, reproducible batching
 │   ├── layers.py            Head, MultiHeadAttention, FeedForward, Block + ActivationCache
 │   ├── train.py             training loop, isolated eval, generation, seeding
-│   └── analysis.py          attention stats, ablation, activation patching, logit lens
+│   └── analysis.py          attention stats, ablation + redundancy, residual & head patching, logit lens
 ├── models/
 │   ├── bigram.py            the baseline
 │   └── gpt.py               GPT + GPTConfig, instrumented forward, checkpointing
 ├── notebooks/
 │   ├── 05_attention_analysis.ipynb     ← the analysis (start here)
 │   └── progression/         01 bigram → 02 single head → 03 multihead+FFN → 04 full GPT
-├── assets/                  gpt.pth, loss history, generated figures
+├── assets/                  gpt.pth, loss history, analysis_summary + seed_robustness JSON, figures
 ├── data/                    input.txt (Tiny Shakespeare)
-├── tests/                   test_smoke.py, test_analysis.py
-└── train_gpt.py
+├── tests/                   test_smoke.py (fast invariants), test_analysis.py (slow integration)
+├── train_gpt.py             reproduce the 30k analysis checkpoint
+└── seed_sweep.py            reproduce the seed-robustness check
 ```
 
 ---
@@ -153,9 +189,29 @@ A few places where the decision that mattered wasn't the obvious one:
 
 - **I loop over a `ModuleList` instead of `nn.Sequential` for the blocks.** The
   source stacks them in `Sequential`. I loop explicitly so the forward pass can
-  capture each layer's attention and residual stream, route head-ablation flags to
-  the right layer, and overwrite activations for patching. `Sequential` hides the
-  loop and forbids exactly the per-layer access the whole analysis needs.
+  capture each layer's attention, residual stream, and per-head outputs, route
+  head-ablation and head-patching to the right layer, and overwrite activations for
+  patching. `Sequential` hides the loop and forbids exactly the per-layer access the
+  whole analysis needs.
+- **Two patching primitives, because they answer different questions.** Residual
+  patching overwrites the residual *sum* at a (layer, position) — good for localizing
+  *where* information is, blind to *which head* put it there. Head patching overwrites
+  a single head's output, which is what isolates L1 H1 rather than just block 1. Both
+  write out of place (a patch goes into a cloned tensor so a corrupted run can't
+  clobber the cached clean activations), and both no-op invariants — splicing a run's
+  own clean residual, or its own clean head output, back in — are unit tests. Without
+  those invariants every recovery number is suspect, so I pinned them down rather than
+  trusting them.
+- **Patching uses a contrastive logit-difference metric and a flip-based selection.**
+  Recovery is measured as the clean-vs-corrupt logit gap rather than one token's
+  probability, which is less sensitive to overall distribution shifts, and I keep only
+  the examples where the corruption actually flipped the top-1 prediction (189 of 256)
+  — a selection I can state plainly instead of a hand-tuned threshold.
+- **Every intervention reports an error bar.** Ablation deltas use the standard error
+  of the *paired* per-batch difference (ablated minus intact on the same batch, which
+  cancels the base-loss noise); patching recoveries report the standard error across
+  examples. It's what lets me say the 7× ablation gap and the 5× head-patching gap are
+  real and not sampling.
 - **Evaluation draws from its own isolated RNG.** In the source, training and
   evaluation sample from the same global generator, so the trained weights quietly
   depend on how often you evaluate. I re-seed a separate generator inside
@@ -170,15 +226,6 @@ A few places where the decision that mattered wasn't the obvious one:
   of the query·key dot product grows with the *head* dimension, so that's what the
   scaling has to cancel. Get it wrong and nothing crashes — softmax just goes peaky
   at initialization and training stalls, which is the worst kind of bug.
-- **Patching writes out of place, and the no-op case is tested.** A patch writes
-  into a cloned tensor so a corrupted run can't overwrite the cached clean
-  activations, and a test asserts that patching a run's own clean activations back
-  in changes the output by nothing. Without that invariant every recovery number is
-  suspect, so I pinned it down with a test rather than trusting it.
-- **Recovery ratios are gated on a usable denominator.** Patching divides by the
-  clean-minus-corrupt metric gap; I drop the examples where the corruption barely
-  moved the prediction (about 150 of 192 survive) so the ratio doesn't blow up on
-  noise.
 - **The blocks are pre-norm.** LayerNorm sits inside each residual branch and leaves
   the skip path clean, so gradients reach the early layers — which is why notebook 03
   (multi-head + FFN with no residuals) is the harder model to optimize despite being
@@ -191,6 +238,7 @@ A few places where the decision that mattered wasn't the obvious one:
 The architecture and training recipe follow Andrej Karpathy's *"Let's build GPT:
 from scratch, in code, spelled out"* and
 [nanoGPT](https://github.com/karpathy/nanoGPT); the corpus is Tiny Shakespeare. What
-I added is the interpretability instrumentation (activation cache, head-ablation and
-activation-patching APIs), the reproducibility engineering (isolated evaluation,
-self-describing checkpoints), and the four-part analysis this README leads with.
+I added is the interpretability instrumentation (activation cache, head-ablation,
+residual- and head-patching APIs), the reproducibility engineering (isolated
+evaluation, self-describing checkpoints, seeded interventions with error bars), and
+the six-part analysis this README leads with.

assets/activation_patching.png

@@ -0,0 +1,72 @@
+[
+ {
+  "seed": 0,
+  "steps": 3000,
+  "base_loss": 1.893,
+  "top_head": [
+   0,
+   1
+  ],
+  "top_delta": 0.4933,
+  "top_prev": 0.973,
+  "top_entropy": 0.137,
+  "top_label": "previous-token",
+  "top_is_prev_token_L01": true
+ },
+ {
+  "seed": 1,
+  "steps": 3000,
+  "base_loss": 1.8948,
+  "top_head": [
+   0,
+   3
+  ],
+  "top_delta": 0.5447,
+  "top_prev": 0.967,
+  "top_entropy": 0.159,
+  "top_label": "previous-token",
+  "top_is_prev_token_L01": true
+ },
+ {
+  "seed": 2,
+  "steps": 3000,
+  "base_loss": 1.9034,
+  "top_head": [
+   0,
+   0
+  ],
+  "top_delta": 0.1232,
+  "top_prev": 0.896,
+  "top_entropy": 0.46,
+  "top_label": "previous-token",
+  "top_is_prev_token_L01": true
+ },
+ {
+  "seed": 3,
+  "steps": 3000,
+  "base_loss": 1.8821,
+  "top_head": [
+   0,
+   1
+  ],
+  "top_delta": 0.2665,
+  "top_prev": 0.943,
+  "top_entropy": 0.272,
+  "top_label": "previous-token",
+  "top_is_prev_token_L01": true
+ },
+ {
+  "seed": 1337,
+  "steps": 3000,
+  "base_loss": 1.8903,
+  "top_head": [
+   1,
+   1
+  ],
+  "top_delta": 0.43,
+  "top_prev": 0.941,
+  "top_entropy": 0.23,
+  "top_label": "previous-token",
+  "top_is_prev_token_L01": true
+ }
+]