Model card: detailed v1.0 -> v1.1 weight-edit write-up

README.md

@@ -144,34 +144,104 @@ semantic-ID generative retrieval work.
 - **Versions in this repo:** `v1.1/` (recommended: `final/` plus the head edit below), `final/` (end of round 8), and the
   `pass7/` and `pass6/` checkpoints.
 
-## What's new in v1.1
+## What's new in v1.1: fixing repeats by editing the weights
 
-v1.0 (`final/`) sometimes listed the same song twice in one playlist (5% of tracks, e.g. *Leonard Cohen –
-Everybody Knows* three times). We fixed this by **editing the weights directly**, with no extra training:
+v1.0 (`final/`) sometimes listed the same song twice in one playlist. About **5% of tracks were repeats**, e.g.
+*Leonard Cohen – Everybody Knows* three times in "dark cabaret heartbreak". The usual fix is more fine-tuning. Instead
+we found the circuit responsible and **edited it directly: no gradient steps, no new data, about 0.03% of the weights**.
 
-1. **Collect:** 400 playlists the model wrote itself, with 239 repeated items and 3,692 fresh ones.
-2. **Correlate:** a linear probe finds where "about to repeat" is most visible. It peaks at layers 16–19.
-3. **Attribute causally:** gradient × activation (attribution patching) scores every MLP neuron and attention head by
-   how much it raises the log-probability of a *repeated* song address relative to a fresh one.
-4. **Intervene:** 21 candidate edits are scored on held-out playlists against random-unit controls. Switching off
-   "repeat-promoting" neurons or ablating the probe direction lowered repeats *and* normal picks alike. Boosting the four
-   strongest **repeat-suppressing heads** (L22 h8, L22 h7, L26 h7, L16 h14) was selective.
-5. **Apply:** multiply those heads' columns in `o_proj` by 2. That is 4 × 128 columns, about 0.03% of the weights.
+### 1. Collect the behaviour
+Greedy-decode 400 playlists from training prompts that the evaluation does not use. Label every list item as
+**repeat** (its 3-token song address already appeared earlier in the list) or **fresh**: 239 repeats and 3,692 fresh.
 
-| Metric (greedy, 110 prompts) | v1.0 `final/` | **v1.1** |
+### 2. Where is "about to repeat" visible? (correlational)
+At the position that predicts each item's first address token, take the residual stream at every layer. The
+difference of class means gives a direction, scored on a held-out half by Cohen's *d*:
+
+| layers | 1–9 | 10–15 | **16–19** | 20–21 | 22–28 |
+|---|---|---|---|---|---|
+| Cohen's *d* | 0.6–1.1 | 0.4–0.95 | **1.4–1.6** | 1.05–1.3 | 0.4–0.9 |
+
+The decision is most legible in the middle-late layers. Correlation is not causation, so we went one step further.
+
+### 3. Which units cause it? (attribution patching)
+For each list item, target = log p(the item's 3 address tokens). For every MLP neuron (`down_proj` input) and every
+attention head (its slice of the `o_proj` input), attribution = activation × ∂target/∂activation, summed over the
+3 predicting positions. Then:
+
+  repeat score = mean attribution on **repeated** items − mean attribution on **fresh** items
+
+The largest scores were *negative*, meaning units that push repeats down. These are **copy-suppression heads** that look
+back at songs already listed and lower their probability:
+
+| head | repeat score |
+|---|---|
+| L22 h8 | −0.068 |
+| L22 h7 | −0.046 |
+| L26 h7 | −0.042 |
+| L16 h14 | −0.040 |
+| (most repeat-promoting: L15 h6) | +0.028 |
+
+### 4. Test edits causally, against controls
+Each candidate edit was applied to a fresh copy of the weights and scored by teacher forcing on held-out playlists
+(eval prompts). Δ is the change in mean log-probability of repeated vs. fresh items; we want Δrepeat ≪ Δfresh ≈ 0.
+
+| edit | Δ repeat | Δ fresh | |
+|---|---|---|---|
+| 8 / 32 / 128 random neurons → 0 (control) | 0.00 / 0.00 / −0.03 | 0.00 / 0.00 / −0.01 | control is flat ✅ |
+| top 8 / 32 / 128 repeat-promoting neurons → 0 | −0.33 / −0.91 / −1.90 | −0.27 / −0.65 / −1.49 | hurts everything ❌ |
+| remove the probe direction (layers 15–19) | −0.41 | −0.41 | not selective ❌ |
+| boost top anti-repeat neurons ×1.5–2 | −0.06 to −0.21 | −0.03 to −0.11 | too weak |
+| top 2 / 4 repeat-promoting heads → 0 | −0.49 / −0.82 | −0.26 / −0.53 | partly selective |
+| **4 copy-suppression heads ×1.5** | **−0.68** | **−0.25** | selective ✅ |
+| **4 copy-suppression heads ×2** | **−1.64** | **−0.60** | **chosen** ✅ |
+| 4 copy-suppression heads ×3 | −3.81 | −1.51 | starts to damage |
+
+### 5. The edit
+```python
+import torch
+from transformers import AutoModelForCausalLM
+m = AutoModelForCausalLM.from_pretrained("asharrm/juke-1.7b-dev", subfolder="final", torch_dtype=torch.bfloat16)
+hd = m.config.head_dim                                   # 128
+with torch.no_grad():
+    for layer, head in [(22, 8), (22, 7), (26, 7), (16, 14)]:
+        m.model.layers[layer].self_attn.o_proj.weight[:, head * hd:(head + 1) * hd] *= 2.0
+m.save_pretrained("juke-1.7b-v1.1")
+```
+That is 4 × 128 columns of `o_proj`: 1.05M of 1.7B weights, about **0.03%**.
+
+### 6. Result (greedy decoding, 110 prompts, real generation)
+
+| metric | v1.0 `final/` | ×1.5 | **v1.1 (×2)** |
+|---|---|---|---|
+| repeated tracks inside a playlist | 5.0% | 2.9% | **1.8%** |
+| song addresses that exist | 99.9% | 100% | 99.9% |
+| song name matches its address | 96.6% | 97.1% | **98.0%** |
+| distinct songs across 40 unrelated prompts | 90.9% | 91.7% | 93.0% |
+| overlap with teacher playlists (trained prompts) | 7.2% | 5.9% | 5.6% |
+
+Before and after:
+
+| prompt | v1.0 | v1.1 |
 |---|---|---|
-| Repeated tracks inside a playlist | 5.0% | **1.8%** |
-| Song addresses that exist | 99.9% | 99.9% |
-| Song name matches its address | 96.6% | **98.0%** |
-| Distinct songs across 40 unrelated prompts | 90.9% | 93.0% |
-| Overlap with teacher playlists (trained prompts) | 7.2% | 5.6% |
-
-**Side effect:** when a genre has few memorized songs, v1.1 drifts to a neighbouring style more often instead of
-repeating. For example, "banda" gave cumbia and mambo.
-
-This follows ideas from recent interpretability work: *Neural Networks Are Already Clean* (Cammarata, Goodfire), the
-"Evil Qwen" single-neuron steering reproductions (HedwigAI), the copy-suppression heads catalogued by neuron-lens, and
-*Sparks of Weight-Space Programming* (Takehi et al., MIT CSAIL).
+| dark cabaret heartbreak | Leonard Cohen – Everybody Knows ×3 | The Dresden Dolls, Nellie McKay, Cold Cave, no repeats |
+| wistful indie rock from the 2000s | The Postal Service – Sleeping In ×3 | Jimmy Eat World, Bright Eyes, Death Cab for Cutie, Dashboard Confessional, The Shins… |
+| weird punk songs | The Adverts ×2 | Misfits, Bad Brains, Black Flag, Television, Dead Kennedys, The Undertones |
+
+**Side effect:** where the model knows few songs in a genre, it now drifts to a neighbouring style more often instead of
+repeating. For example, "banda" gave cumbia and mambo. The edit removes an obstacle; it cannot add songs the model never
+learned, which matches the weight-space-programming finding below.
+
+### Why this worked
+The training data never contained repeats, so the model learned a copy-suppression mechanism. It was simply too quiet,
+and greedy decoding sometimes let a strongly-recalled song win twice. Turning up a mechanism that already exists is a
+targeted fix. Switching off neurons that merely *correlate* with repeats damaged everything.
+
+Inspired by: *Neural Networks Are Already Clean* (Cammarata, Goodfire, 2026). The "Evil Qwen" single-neuron reproductions
+(HedwigAI, 2026), where the most active neurons were not the causal ones, so rank by causal attribution and check random
+controls. Copy-suppression heads catalogued in neuron-lens (stdlogits). *Sparks of Weight-Space Programming* (Takehi,
+Teoh, Zhang & Khattab, MIT CSAIL, 2026). Scripts: `repeat_direction.py`, `repeat_attrib.py`, `repeat_edit_test.py`,
+`repeat_apply.py` in the juke repo.
 
 ## Evaluation (v1.0 `final/`, greedy, 110 prompts; v1.1 numbers above)