| # Lab 8 (capstone): build your own edit from the model's own loops |
|
|
| **Big thing:** With no hand-written texts, the unedited model's own looping answers give a loop direction close to the agent's. A weight edit built from them ("well before the loop", recipe B) reaches `58.2%` accuracy on 256 unseen practice problems, against `61.5%` for the released edit. |
|
|
| **Why it matters:** In Lab 3, we rebuilt the released edit from 8 short math scratchpads that the AI agent wrote and labeled by hand. Here you extract the direction and the 64 neurons per layer from the base model's failed answers on 128 practice problems, check your direction with four diagnostics before editing any weights, and see why the comparison window inside each answer changes unseen accuracy by `15.2` percentage points. |
|
|
| **Time and hardware:** About 5 minutes on 1 NVIDIA A100 GPU to extract the edits, 1 minute on CPU to build the edited models, and 30 minutes on 2 A100 GPUs (one run per GPU) to test both on `dev256`. |
|
|
| **Theory guide:** Sections 9 and 13.3. |
|
|
| ## Before you run: predict |
|
|
| Write down a short guess for each question before you run the script. |
|
|
| 1. For each looping answer, we compare a 36-token window (a token is a word piece) inside the loop against an earlier 36-token window from the same answer. Recipe A ("just before the loop") ends 32 tokens before verbatim repetition begins. Recipe B ("well before the loop") sits halfway between the start of the answer and the start of the loop. Which recipe will produce a stronger edit on unseen problems, and why? |
| 2. The AI agent needed only 4 pairs of hand-written texts to find its direction. Cosine similarity measures how closely two directions align. With the model's own noisy, long answers, roughly how many random looping answers must we average before the direction $u$ settles down (say, cosine similarity above `0.95` with the direction from a held-out set of loops, kept separate from those used to estimate the direction)? Will recipe A need more answers than recipe B? |
| 3. In Lab 3, neuron 11 of layer 19 had a score nearly five times larger than any other neuron on the agent's 8 short texts. When we score all 3,584 neurons on 225 long math answers from the model itself, will neuron 11 of layer 19 still rank first? |
|
|
| ## Run it |
|
|
| Connect to `sampler`. The `labs/lab8_capstone.py` script reads the base model's answers in `runs/build/base`. Later on this page, `labs/lab4_endings.py` also reads the full answer texts from the `dev/base` and `dev/p50c` runs, which Lab 4 makes. If the machine lacks any of these three runs, make the missing ones first with `python scripts/run_jobs.py plans/sampler.json --only build/base dev/base dev/p50c --gpus 5-7`. That command skips any run that already finished. With all three missing, it takes about 45 minutes, one run per GPU. |
|
|
| Run the four main commands from `labs/lab8_capstone.py`: |
|
|
| ```bash |
| ssh sampler |
| . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code |
| nvidia-smi |
| CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py # about 5 minutes on one GPU |
| CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants configs/capstone.json # CPU, about a minute |
| python scripts/run_jobs.py plans/extra.json --gpus 6,7 --only 'dev/cap*' # about 30 minutes, one run per GPU |
| python scripts/summarize.py /nvme/wsp-lab/runs # writes /nvme/wsp-lab/summary.md |
| ``` |
|
|
| `scripts/run_jobs.py` skips any shard with a `.done.json` marker. To generate `dev/capA-a6` and `dev/capB-a6` fresh instead of keeping the finished runs, move `/nvme/wsp-lab/runs/dev/capA-a6` and `/nvme/wsp-lab/runs/dev/capB-a6` aside first. |
|
|
| What each step writes: |
|
|
| - `labs/lab8_capstone.py` reads the base model's answers in `/nvme/wsp-lab/runs/build/base` (`build128`, 128 practice problems never used for testing). It writes `results/lab8/capstone.json`, `results/lab8/example_answer.json`, `configs/capstone.json` (defining `capA-a6` and `capB-a6`), and three figures: `results/figures/lab8_loop_onset.png`, `results/figures/lab8_data_needed.png`, and `results/figures/lab8_one_answer.png`. |
| - `scripts/build_models.py --variants configs/capstone.json` writes two edited model folders in `/nvme/wsp-lab/models/`: `qwen35-capA-a6` and `qwen35-capB-a6`. |
| - `scripts/run_jobs.py plans/extra.json --gpus 6,7 --only 'dev/cap*'` generates 2 answers per problem across the 256 unseen problems of `dev256` (`512` answers per model; `dev/capA-a6` took `31.8` minutes and `dev/capB-a6` took `26.1` minutes) and writes `merged.jsonl` and `summary.json` under `/nvme/wsp-lab/runs/dev/capA-a6` and `/nvme/wsp-lab/runs/dev/capB-a6`. |
| - `scripts/summarize.py` updates the summary table. |
|
|
| Two more scripts fill in the rest of the section 3 table. `labs/lab4_endings.py` sorts answers by how they ended ([Lab 4](lab4_dose.md) explains the four groups), and `labs/lab7_side_effects.py` measures side effects on WikiText-2 and the agent's 8 texts ([Lab 7](lab7_side_effects.md)). Run both on `sampler`: |
|
|
| ```bash |
| CUDA_VISIBLE_DEVICES="" python labs/lab4_endings.py dev/base dev/p50c dev/capA-a6 dev/capB-a6 --out lab8/endings # CPU, a few seconds per run |
| CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --models p50c l15only capA-a6 capB-a6 # under a minute on one GPU |
| ``` |
|
|
| `--out lab8/endings` writes `results/lab8/endings.json`, and the side-effects run writes `results/lab7/try-models-p50c-l15only-capA-a6-capB-a6/`. |
|
|
| Passing any non-default flag to `labs/lab8_capstone.py` (such as `--max-samples 16`, `--min-before 2048`, or `--alpha 8`) writes to `results/lab8/try-<flags>/` (`capstone.json`, `example_answer.json`, `variant.json`, `plan.json`, and three figures) without touching the main files or `configs/capstone.json`. Full workflow for an experiment: |
|
|
| ```bash |
| CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --max-samples 16 # about 5 minutes on one GPU |
| CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants results/lab8/try-max-samples-16/variant.json # CPU |
| python scripts/run_jobs.py results/lab8/try-max-samples-16/plan.json --gpus 0,1 # about 30 minutes, one run per GPU |
| CUDA_VISIBLE_DEVICES="" python labs/lab4_endings.py dev/capA-max-samples-16 dev/capB-max-samples-16 --out lab8/try-max-samples-16/endings |
| CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --models capA-max-samples-16 capB-max-samples-16 capA-alpha-8 capB-alpha-8 --try-name lab8-tries |
| ``` |
|
|
| (`--alpha 8` works the same way in `try-alpha-8`.) We ran both `lab8_capstone.py` tries side by side on GPUs 6 and 7, then built both models and ran each plan on two GPUs. The `dev256` runs took `28 to 33` minutes each, except `dev/capA-max-samples-16`, which took `43.5` minutes because most answers ran to the 32,768-token cap (mean `28,525` tokens). The Lab 7 check of all four models took `33` seconds (`results/lab7/try-lab8-tries/`). |
|
|
| ## What you should see |
|
|
| ### 1. Loops found in `build128` and the extracted directions |
|
|
| Running `labs/lab8_capstone.py` prints the loop counts, the percentiles of loop period and onset token, the diagnostics for both recipes, and the onset profile: |
|
|
| ```text |
| 256 answers in /nvme/wsp-lab/runs/build/base: 230 hit the token cap, 225 of those end in a verbatim loop |
| using 225 loops (loop starts at least 512 answer tokens in) |
| loop period, 10/25/50/75/90th percentile: [18.0, 27.0, 44.0, 78.0, 144.0] tokens |
| loop starts at answer token, same percentiles: [2813.4, 4021.0, 5822.0, 8013.0, 10953.8] |
| recipe A (just before the loop), layer 15: cos with released u +0.722; released neurons picked 38/64; split-half cos 0.978; agent's 4 pairs, looping minus productive along u: [0.78, 0.88, 1.69, 1.16] |
| answers needed: [(1, 0.6), (2, 0.628), (4, 0.712), (8, 0.841), (16, 0.919), (32, 0.952), (64, 0.97), (112, 0.978)] |
| recipe A (just before the loop), layer 19: cos with released u +0.724; released neurons picked 28/64; split-half cos 0.978; agent's 4 pairs, looping minus productive along u: [1.15, 1.27, 2.4, 1.73] |
| recipe B (well before the loop), layer 15: cos with released u +0.782; released neurons picked 42/64; split-half cos 0.988; agent's 4 pairs, looping minus productive along u: [0.9, 1.04, 1.7, 1.25] |
| answers needed: [(1, 0.676), (2, 0.781), (4, 0.853), (8, 0.912), (16, 0.958), (32, 0.97), (64, 0.984), (112, 0.988)] |
| recipe B (well before the loop), layer 19: cos with released u +0.773; released neurons picked 28/64; split-half cos 0.989; agent's 4 pairs, looping minus productive along u: [1.32, 1.5, 2.35, 1.83] |
| cos between recipe A's u and recipe B's u: {'15': 0.951, '19': 0.959} |
| ``` |
|
|
| The loop finder works on tokens: an answer counts if its tokens end in one block repeated back to back at least 3 times over at least 64 tokens. It finds `225` of the `230` cut answers; the summary table's word-level check counts `222` of `230` (`96.5%`), a 3-answer difference. In `results/lab8/capstone.json`, all `225` loops start at least 512 tokens into the answer, so all `225` are used. The loop period (the length in tokens of one repeated block) has 10th, 25th, 50th, 75th, and 90th percentiles of `18, 27, 44, 78, 144` tokens. The answer token where the loop starts has percentiles of `2,813; 4,021; 5,822; 8,013; 10,954`. |
|
|
| How the two window recipes compare against the released edit (`p50c`) and each other: |
|
|
| | Recipe (comparison window) and layer | Cosine with released $u$ | Released neurons picked (of 64) | Split-half cosine (stability) | Gap size ($\|\overline{\Delta h}\|$) | Agent's 4 pairs: looping minus productive along $u$ | Top neuron (score, next score) | |
| |---|---|---|---|---|---|---| |
| | Recipe A (just before the loop), layer 15 | `0.722` | `38 / 64` | `0.978` | `0.689` | `+0.78, +0.88, +1.69, +1.16` | `46` (`0.0120`, next `294` at `0.0108`) | |
| | Recipe A (just before the loop), layer 19 | `0.724` | `28 / 64` | `0.978` | `1.037` | `+1.15, +1.27, +2.40, +1.73` | `11` (`0.121`, next `171` at `0.024`) | |
| | Recipe B (well before the loop), layer 15 | `0.782` | `42 / 64` | `0.988` | `1.418` | `+0.90, +1.04, +1.70, +1.25` | `46` (`0.0261`, next `0` at `0.0187`) | |
| | Recipe B (well before the loop), layer 19 | `0.773` | `28 / 64` | `0.989` | `2.234` | `+1.32, +1.50, +2.35, +1.83` | `11` (`0.210`, next `1` at `0.045`) | |
|
|
| Four patterns in these numbers: |
|
|
| - **Agreement with each other and with the released edit:** Both recipes give nearly the same unit direction (a length-1 vector in the 1,024-number stream; cosine similarity is the dot product of two unit vectors, where `1` means the same direction and `0` means at right angles). Their mutual cosine `cos_A_B` is `0.951` at layer 15 and `0.959` at layer 19, yet recipe B has more than double the raw gap size (the length of the average stream difference) and sits closer to the released edit's direction. The table above lists the gap sizes and cosines for each layer. |
| - **Top neuron:** At layer 19, neuron `11` ranks first in both recipes with a score nearly five times the next neuron's. The table above lists both scores for each recipe. The AI agent switched off this same neuron in its 4th forward pass (request file `fwd_5` in `external/lm_understand-logs/exp8_math/runs/p50c/requests/`; an earlier failed request did not count), taking 32 practice problems from 3 to 5 correct (`+6.25` points, standard error `4.35`, which its notes call "weak positive, not confirmed"). At layer 15, both recipes' top three neurons (`46, 294, 0`) are among Lab 3's top four (`294, 371, 0, 46`). |
| - **Transfer to the agent's 8 hand-written texts:** Both recipes were built only from the base model's long math loops, yet projecting the last 36 tokens of all 4 hand-written pairs (`labs/data/astra_pairs.json`) onto $u$ still separates the looping text from the productive text in every pair. The table above lists the gaps along the direction for each pair. |
| - **Data needed:** At layer 15, the average cosine similarity between a direction from $N$ random answers and a direction from a held-out half (`113` answers) climbs fast: |
|
|
| | Answers used ($N$) | `1` | `2` | `4` | `8` | `16` | `32` | `64` | `112` | |
| |---|---:|---:|---:|---:|---:|---:|---:|---:| |
| | Recipe A (just before the loop), layer 15 | `0.60` | `0.63` | `0.71` | `0.84` | `0.92` | `0.95` | `0.97` | `0.978` | |
| | Recipe B (well before the loop), layer 15 | `0.68` | `0.78` | `0.85` | `0.91` | `0.96` | `0.97` | `0.98` | `0.988` | |
|
|
|  |
|
|
| ### 2. When the loop starts in the residual stream |
|
|
| Aligning all `225` loops at the token where verbatim repetition begins (`token 0`) and averaging across answers shows when the residual stream and the 64 released neurons start drifting: |
|
|
| | Quantity (mean over 225 loops) | `2,048 to 1,024` tokens before | `512 to 128` tokens before | Last `128` tokens before (`-128..0`) | First `128` tokens inside (`0..128`) | `128 to 256` tokens inside | |
| |---|---:|---:|---:|---:|---:| |
| | Stream along released $u$, layer 15 | `1.14` | `1.39` | `1.58` | `1.95` | `2.20` | |
| | Stream along released $u$, layer 19 | `2.27` | `2.66` | `2.95` | `3.50` | `3.96` | |
| | Trigger $s(t)$ (released 64 neurons), layer 15 | `0.135` | `0.173` | `0.206` | `0.288` | `0.366` | |
| | Trigger $s(t)$ (released 64 neurons), layer 19 | `0.183` | `0.237` | `0.277` | `0.376` | `0.512` | |
|
|
| Both the stream projection along $u$ and the 64 neurons' push $s(t) = u \cdot (W_S a_S(t))$ (the Lab 6 trigger measuring how hard those 64 neurons push along $u$ at token $t$) rise hundreds of tokens before repetition starts, then jump at token `0`. |
|
|
|  |
|
|
| In `results/lab8/example_answer.json` and `lab8_one_answer.png`, the median-onset answer (`math_train_algebra_1625`, sample `1`) runs to `32,768` tokens, and its loop starts at answer token `5,822` with period `51` tokens. The three shaded 36-token windows show where each recipe reads the stream: window B covers tokens `2,893 to 2,929`, window A covers `5,754 to 5,790`, and the loop window covers `5,924 to 5,960` (starting two periods after token `5,822`). |
|
|
|  |
|
|
| ### 3. Test results on the 256 unseen problems (`dev256`) |
|
|
| The table below combines `dev256` results (`256` unseen problems, `2` answers each = `512` answers per run), answer endings from `results/lab8/endings.json`, and side effects from `results/lab7/try-models-p50c-l15only-capA-a6-capB-a6/side_effects.json`. From Lab 7, perplexity ($e$ to the mean negative log-probability of the true next token; base model `20.009` on WikiText-2, lower is better) and KL divergence ($\text{KL}(\text{base} \parallel \text{edited})$, how far next-token probabilities shift from the unedited model, where `0` means identical) measure how much an edit disturbs ordinary text and the agent's 8 texts: |
|
|
| | Model / edit | Accuracy % | Cut at cap % | Mean tokens | Gain vs unedited (pp), rescued / lost | Correct when finished % | Endings: cut / stopped / escaped loop (correct) / normal (correct) | Perplexity change | KL: WikiText / productive / looping | |
| |---|---:|---:|---:|---|---:|---|---:|---| |
| | `dev/base` (unedited) | `12.1 +/- 1.6` | `88.7` | `29,683` | baseline | `100.0` | `453 (4) / 0 / 3 (2) / 56 (56)` | `0%` (`20.009`) | `0 / 0 / 0` | |
| | `capA-a6` (just before the loop, recipe A) | `43.0 +/- 2.6` | `52.5` | `22,699` | `+30.9 +/- 2.3`, `95 / 1` | `89.3` | `237 (3) / 1 (0) / 141 (94) / 133 (123)` | `+0.51%` | `0.0178 / 0.0103 / 0.0162` | |
| | `capB-a6` (well before the loop, recipe B) | `58.2 +/- 2.9` | `30.7` | `19,899` | `+46.1 +/- 2.6`, `113 / 0` | `83.7` | `108 (0) / 0 / 145 (67) / 259 (231)` | `+0.92%` | `0.0271 / 0.0117 / 0.0208` | |
| | `p50c` (released edit, alpha 6) | `61.5 +/- 2.8` | `19.7` | `16,104` | `+49.4 +/- 2.6`, `125 / 0` | `76.6` | `72 (0) / 0 / 96 (29) / 344 (286)` | `+0.79%` | `0.0333 / 0.0097 / 0.0241` | |
|
|
| Paired per-problem differences across the 256 problems: |
|
|
| - `capB-a6` beats `capA-a6` by `+15.2 +/- 2.2` percentage points. |
| - `capB-a6` trails the released edit `p50c` by `-3.3 +/- 1.5` percentage points. |
|
|
| An answer closed normally if it wrote `</think>` with no repeated snippet in the 6,000 characters before it ([Lab 4](lab4_dose.md)). Both capstone edits escape a running loop about equally often (`141` for recipe A, `145` for recipe B), while `capB-a6` nearly doubles normal endings (`259` vs `133`) and more than halves answers trapped at the cap (`108` vs `237`). A likely reading, though these counts do not prove it: recipe B also seems to stop more loops from starting. Rescued / lost in the table counts problems, relative to the unedited model. |
|
|
| ## Why it works |
|
|
| ### 1. From the model's own failures to a rank-one weight edit |
|
|
| A language model reads and writes one token (word piece) at a time. At each token, it updates a 1,024-number residual stream across 24 layers (blocks). Each block has a token mixer and an MLP (multi-layer perceptron) whose output matrix `down_proj` has 3,584 columns, one per neuron. In Lab 3, each of the agent's 4 text pairs shared the same prompt and opening lines, so subtracting the productive text from the looping text canceled the problem topic. |
|
|
| In `runs/build/base`, each answer cut at the 32,768-token cap starts with productive reasoning before falling into a verbatim loop (the same token block repeating to the cap). Comparing a window inside the loop against an earlier window from the same answer via teacher forcing (feeding the saved tokens through the model in one pass and recording activations with PyTorch hooks) cancels the problem topic within each answer. |
|
|
| ```mermaid |
| flowchart LR |
| S["Start of answer<br/>(token n0)"] --> WB["Well before the loop (recipe B)<br/>(36 tokens at midpoint of n0..onset)"] |
| WB --> WA["Just before the loop (recipe A)<br/>(36 tokens ending 32 tokens before onset)"] |
| WA --> ON["Loop starts<br/>(token onset, period p)"] |
| ON --> WL["Loop window: inside the loop<br/>(36 tokens starting at onset + 2p)"] |
| ``` |
|
|
| For every looping answer whose loop starts at least `512` tokens into the answer (`on - n0 >= 512`), `labs/lab8_capstone.py` slices three 36-token windows (`WIN = 36`): |
|
|
| - **Inside the loop (`loop`):** 36 tokens starting two periods after loop onset (`on + 2 * p`), once verbatim repetition is underway. |
| - **Just before the loop (recipe A):** 36 tokens ending `GAP = 32` tokens before loop onset (`on - 68` to `on - 32`). |
| - **Well before the loop (recipe B):** 36 tokens centered halfway between the start of the answer and loop onset (`n0 + (on - n0) // 2 - 18` to `+ 18`). |
|
|
| Because the model is causal (each token reads only earlier tokens), the script truncates each forward pass right after the needed windows (`end = max(w["loop"][1], min(on + 256, len(ids)))`) instead of running all `32,768` tokens. That finishes all 225 answers in about 5 minutes on one GPU. |
|
|
| At layers $L \in \{15, 19\}$, the remaining steps follow Lab 3: |
|
|
| 1. **Direction $u$:** For each answer $i$ and recipe $k \in \{A, B\}$, subtract the mean residual stream leaving block $L$ over window $k$ from the mean over the loop window: $\Delta h_{i, L}^{(k)} = \bar{h}_{i, L}^{\text{loop}} - \bar{h}_{i, L}^{(k)}$. Average $\Delta h_{i, L}^{(k)}$ across all 225 answers and scale to unit length: |
| $$u_L^{(k)} = \frac{\overline{\Delta h}_L^{(k)}}{\|\overline{\Delta h}_L^{(k)}\|}$$ |
| 2. **Neuron scores:** Over the last `AWIN = 24` tokens of each window, find the mean activation difference $\overline{\Delta a}_L^{(k)}[j]$ (loop window minus comparison window) for each of the 3,584 MLP neurons. Multiply by how much column $j$ of `down_proj` writes along $u_L^{(k)}$: |
| $$\text{score}_j = \left(u_L^{(k)} \cdot W_{\text{down}, L}[:, j]\right) \times \overline{\Delta a}_L^{(k)}[j]$$ |
| 3. **Rank-one edit on the top 64 columns:** Sort the 3,584 neurons by $\text{score}_j$ descending, keep the top 64 indices $S$, and apply the rank-one update (one column vector $u$ times one row vector $u^T W_S$) $W' = W - \alpha \, u \, (u^T W_S)$ with edit strength $\alpha = 6$. This scales those 64 neurons' push along $u$ by $1 - \alpha = -5$. |
|
|
| ```mermaid |
| flowchart LR |
| R["225 looping answers<br/>from build128 (runs/build/base)"] --> W["Take 36-token windows:<br/>loop minus window A or B"] |
| W --> U["Average stream difference<br/>across answers -> unit vector u"] |
| U --> N["Score 3,584 neurons:<br/>(u . W[:, j]) * mean(da_24)<br/>Keep top 64 at alpha = 6"] |
| N --> T["Build edited models and test<br/>on 256 unseen problems (dev256)"] |
| ``` |
|
|
| ### 2. Why "well before the loop" beats "just before the loop" |
|
|
| Recipe A and recipe B have cosine similarity `0.951` at layer 15 and `0.959` at layer 19, and both pick `28 of 64` released neurons at layer 19 (`38` vs `42` at layer 15). Yet recipe B beats recipe A on `dev256` by `15.2` percentage points (`58.2%` vs `43.0%`). |
|
|
| In `lab8_loop_onset.png` and Lab 6 Part 3, hundreds of tokens before exact repetition starts at token `0`, the model is already re-checking the same numbers and drifting toward the loop: |
|
|
| - At layer 15, the stream projection along the released $u$ climbs from `1.14` (`2,048 to 1,024` tokens before onset) to `1.58` in the last `128` tokens before onset, then jumps to `2.20` inside the loop. |
| - In Lab 6 Part 3, the 64 neurons' push $s_{15}(t)$ nearly doubled from `0.096` in the early half of the pre-loop answer to `0.178` in the last 512 tokens before onset, then jumped to `0.424` inside the loop. |
| |
| Window A (`-68` to `-32` tokens before onset) sits inside that late drift zone, where the stream has already moved about halfway to the loop. When recipe A computes `stream(in the loop) - stream(just before the loop)`, it subtracts out the early drift and keeps mostly the final jump into verbatim repetition. Three results reflect this: |
| |
| 1. **The raw gap vector is less than half as long:** Recipe A's $\|\overline{\Delta h}\|$ is `0.689` vs `1.418` for recipe B at layer 15 (`1.037` vs `2.234` at layer 19). Because part of the loop signal cancels out, background token differences between the two windows make up a larger share of the normalized unit vector $u_A$, lowering its cosine with the released $u$ (`0.722` vs `0.782`) and its layer-15 neuron overlap (`38` vs `42`). |
| 2. **Smaller neuron activation gaps:** Because loop neurons are already firing harder in the last 68 pre-loop tokens, subtracting window A's activations shrinks $\overline{\Delta a}_{24}$ on the neurons that build up the drift. For example, the `score` of neuron 11 at layer 19 drops from `0.210` in B to `0.121` in A, and neuron 46 at layer 15 drops from `0.0261` to `0.0120`. |
| 3. **Fewer loops versus more escapes on `dev256`:** Both edits escape a loop about equally often (`141` for recipe A, `145` for recipe B), while recipe B nearly doubles normal endings (`259` vs `133`). A likely reading, which the endings counts fit but do not prove, is that recipe B's direction includes the pre-loop drift, so its edit pushes back before repetition starts, whereas recipe A seems to act mainly once a loop is underway. |
| |
| ### 3. What split-half stability and the data-needed curve measure |
| |
| To check that the direction reflects a shared property across problems instead of sample noise, the script runs 20 random split-half checks. Each trial shuffles the 225 loops into two non-overlapping halves of `112` and `113` answers ($H_1$ and $H_2$), computes a unit direction from each half, and takes their cosine similarity. Recipe B's halves agree at `0.988 +/- 0.003` (layer 15) and `0.989 +/- 0.003` (layer 19), and recipe A's agree at `0.978` at both layers. Different loops point along nearly the same direction in the 1,024-dimensional stream. |
| |
| To measure how many answers are needed, the script draws $N \in \{1, 2, 4, 8, 16, 32, 64, 112\}$ answers from $H_1$ and measures the cosine of their direction with the full held-out half $H_2$: |
| |
| ```mermaid |
| xychart-beta |
| title "Layer 15: recipe B (upper line) vs recipe A (lower line)" |
| x-axis "Answers used" ["1", "2", "4", "8", "16", "32", "64", "112"] |
| y-axis "Cosine with held-out half" 0.5 --> 1.0 |
| line [0.600, 0.628, 0.712, 0.841, 0.919, 0.952, 0.970, 0.978] |
| line [0.676, 0.781, 0.853, 0.912, 0.958, 0.970, 0.984, 0.988] |
| ``` |
| |
| At $N = 1$ (one looping answer), recipe B already has a `0.68` cosine with the held-out half. By $N = 16$, it reaches `0.958` (`0.919` for recipe A), and past `32` answers the curve is nearly flat (`0.970` to `0.988`). In `lab8_data_needed.png`, the dashed curve (cosine with the released $u$) levels off near `0.78` for recipe B and `0.72` for recipe A. Averaging more answers reduces sampling noise around the natural-loop direction, while the gap between `0.78` and `1.00` is a systematic difference between windows of long natural answers and the agent's 4 hand-written pairs. |
| |
| Because each `build128` problem has two answers, splitting by answer can put two answers to the same problem in opposite halves. To check whether that inflates agreement, we reran with `--split-by problem`, keeping both answers of each problem in the same half: |
| |
| ```bash |
| CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --split-by problem # writes results/lab8/try-split-by-problem/ |
| ``` |
| |
| The 225 loops come from 120 problems (105 with two loops), giving halves of 108 to 117 loops. The split-half cosines barely change: `0.977 / 0.978` for recipe A at layers 15 / 19 (vs `0.978 / 0.978` by answer) and `0.988 / 0.987` for recipe B (vs `0.988 / 0.989`). Both data-needed curves shift by at most `0.068` at $N = 1$. From $N = 16$ on, held-out agreement shifts by under `0.006` and released-$u$ agreement by under `0.009`. When the half that the $N$ answers come from holds fewer than 112 loops, the final point uses all of them (`108` to `112` in `answers_used_min_max` in `capstone.json`). The edits themselves are unchanged because the split affects only these two checks. |
| |
| ### 4. Why we keep `build128` separate from `dev256` and use paired differences |
| |
| Section 13.3 of the theory guide states a basic evaluation rule: test an edit on problems not used to build it. Here `build128` (`128` problems, `256` answers) is used only to extract $u$ and the 64 neurons, and `dev256` (`256` separate problems, `512` answers) is used only for testing. |
| |
| On `dev256`, `capB-a6` scores `58.2% +/- 2.9%` and `p50c` scores `61.5% +/- 2.8%`, where the standard error (SE) is the typical random error from testing a finite sample of problems (a gap of 2 SE or more is unlikely to be luck). Judging from the standalone SEs (`2.9%` and `2.8%`), the `3.3`-point gap looks like it could be chance variation. Because both models run on the same 256 problems, ranging from easy GSM8K arithmetic to hard competition algebra, problem-to-problem difficulty inflates both standalone SEs. Taking the paired difference on each problem $p$ (the fraction of sample answers solved by model 1 minus the fraction solved by model 2 on problem $p$) and computing the SE across those 256 differences cancels out shared problem difficulty: |
| |
| - `capB-a6` minus `p50c`: `-3.3 +/- 1.5` percentage points (more than 2 SE below `p50c`, so `capB-a6` is close to the released edit and measurably a little weaker). |
| - `capB-a6` minus `capA-a6`: `+15.2 +/- 2.2` percentage points (nearly seven standard errors). |
| |
| ## Experiments: sample count and edit strength |
| |
| We ran two experiment flags on `sampler` to test how the capstone edit behaves with fewer loops (`--max-samples 16`) or a stronger edit (`--alpha 8`). |
| |
| ### 1. How the directions and weight changes compare |
| |
| - `--max-samples 16` (`results/lab8/try-max-samples-16/`) keeps the first `16` usable loops in file order (`used[:16]`), a file prefix instead of a random draw across all 128 problems. Because `build128` lists its 64 GSM8K problems before its 64 MATH problems, all 16 loops come from `9` GSM8K problems (`gsm8k_train_3030`, `3613`, `4033`, `1318`, `4077`, `6248`, `6350`, `6176`, `7193`). With 16 loops, the data-needed curve stops at `8` answers, and the split-half cosine compares two random halves of `8` loops each (vs `112` and `113` in the main run). |
| - `--alpha 8` (`results/lab8/try-alpha-8/`) uses all `225` loops, so it keeps the main run's unit directions $u$ and 64 neurons and raises edit strength from $\alpha = 6$ to $\alpha = 8$. The relative weight change (Frobenius norm (the overall size) of $W' - W$ divided by Frobenius norm of $W$, layer 15 / layer 19) is `0.078 / 0.075` for `capA-max-samples-16`, `0.097 / 0.092` for `capB-max-samples-16`, `0.141 / 0.129` for `capA-alpha-8`, and `0.149 / 0.147` for `capB-alpha-8`. |
| |
| | Edit variant | Loops used | Alpha | Cosine with released $u$ (L15 / L19) | Split-half cosine (L15 / L19) | Released neurons picked (L15 / L19) | |
| |---|---:|---:|---|---|---| |
| | `capA-a6` (just before, main) | `225` | `6` | `0.722 / 0.724` | `0.978 / 0.978` | `38 / 28` | |
| | `capA-max-samples-16` (just before, 16 loops) | `16` | `6` | `0.574 / 0.596` | `0.632 / 0.622` | `27 / 21` | |
| | `capB-a6` (well before, main) | `225` | `6` | `0.782 / 0.773` | `0.988 / 0.989` | `42 / 28` | |
| | `capB-max-samples-16` (well before, 16 loops) | `16` | `6` | `0.687 / 0.682` | `0.828 / 0.864` | `34 / 20` | |
| |
| With 16 loops from 9 GSM8K problems, the directions of recipes A and B still agree (`0.873` cosine at layer 15 and `0.902` at layer 19, vs `0.951` and `0.959` on 225 loops). Neuron 11 of layer 19 still ranks first in both recipes. Both directions also separate the agent's 8 texts in all 4 pairs (`+0.69, +0.83, +1.67, +1.10` for recipe B at layer 15). |
| |
| ### 2. Evaluation on `dev256`, endings, and breakdown by source |
| |
| | Edit variant | Accuracy % | Cut % | Mean tokens | Correct when finished % | Gain vs unedited (pp), rescued / lost | Paired change vs same recipe at alpha 6 with all 225 loops | Perplexity change | KL: WikiText / productive / looping | |
| |---|---:|---:|---:|---:|---|---|---:|---| |
| | `capA-a6` (just before, main) | `43.0 +/- 2.6` | `52.5` | `22,699` | `89.3` | `+30.9 +/- 2.3`, `95 / 1` | baseline for recipe A | `+0.51%` | `0.0178 / 0.0103 / 0.0162` | |
| | `capA-alpha-8` (just before, alpha 8) | `42.8 +/- 2.8` | `53.5` | `22,772` | `89.9` | `+30.7 +/- 2.5`, `87 / 4` | `-0.2 +/- 2.3` | `+1.61%` | `0.0298 / 0.0183 / 0.0265` | |
| | `capA-max-samples-16` (just before, 16 loops) | `17.6 +/- 1.9` | `82.6` | `28,525` | `96.6` | `+5.5 +/- 1.9`, `34 / 17` | `-25.4 +/- 2.3` | `+0.10%` | `0.0081 / 0.0063 / 0.0086` | |
| | `capB-a6` (well before, main) | `58.2 +/- 2.9` | `30.7` | `19,899` | `83.7` | `+46.1 +/- 2.6`, `113 / 0` | baseline for recipe B | `+0.92%` | `0.0271 / 0.0117 / 0.0208` | |
| | `capB-alpha-8` (well before, alpha 8) | `54.5 +/- 2.9` | `34.0` | `20,870` | `81.7` | `+42.4 +/- 2.6`, `107 / 1` | `-3.7 +/- 1.6` | `+2.76%` | `0.0468 / 0.0206 / 0.0340` | |
| | `capB-max-samples-16` (well before, 16 loops) | `53.7 +/- 2.8` | `39.6` | `20,725` | `88.7` | `+41.6 +/- 2.6`, `110 / 0` | `-4.5 +/- 1.8` | `+0.06%` | `0.0204 / 0.0102 / 0.0173` | |
| | `p50c` (released edit, alpha 6) | `61.5 +/- 2.8` | `19.7` | `16,104` | `76.6` | `+49.4 +/- 2.6`, `125 / 0` | (alpha 8 vs 6: `-0.8 +/- 1.6`) | `+0.79%` | `0.0333 / 0.0097 / 0.0241` | |
| |
| Answer endings in `results/lab8/try-alpha-8/endings.json` and `results/lab8/try-max-samples-16/endings.json` (out of 512 answers, with correct counts in parentheses): |
| |
| - `capA-alpha-8`: `242 (5)` cut while thinking, `0` stopped while thinking, `130 (87)` escaped a loop, `140 (127)` closed normally (vs `capA-a6`: `237 (3)` cut, `1 (0)` stopped, `141 (94)` escaped a loop, `133 (123)` closed normally). |
| - `capB-alpha-8`: `133 (1)` cut while thinking, `1 (0)` stopped while thinking, `96 (44)` escaped a loop, `282 (234)` closed normally (vs `capB-a6`: `108 (0)` cut, `0` stopped, `145 (67)` escaped a loop, `259 (231)` closed normally). |
| - `capA-max-samples-16`: `422 (4)` cut while thinking (`96.0%` of its cut answers end in repetition, `94.8%` in a verbatim loop), `0` stopped, `38 (34)` escaped a loop, `52 (52)` closed normally. |
| - `capB-max-samples-16`: `150 (1)` cut while thinking (`76.8%` of its cut answers end in repetition, `36.9%` in a verbatim loop), `0` stopped, `183 (107)` escaped a loop, `179 (167)` closed normally. |
| |
| Breakdown by problem source on `dev256` (`128` GSM8K and `128` MATH problems, `256` answers each): |
| |
| | Edit variant | GSM8K accuracy % / cut % | MATH accuracy % / cut % | |
| |---|---|---| |
| | `capA-a6` (just before, main) | `55.5 / 37.5` | `30.5 / 67.6` | |
| | `capA-alpha-8` (just before, alpha 8) | `51.2 / 41.8` | `34.4 / 65.2` | |
| | `capA-max-samples-16` (just before, 16 loops) | `23.8 / 76.2` | `11.3 / 89.1` | |
| | `capB-a6` (well before, main) | `71.5 / 13.3` | `44.9 / 48.0` | |
| | `capB-alpha-8` (well before, alpha 8) | `68.0 / 15.6` | `41.0 / 52.3` | |
| | `capB-max-samples-16` (well before, 16 loops) | `69.1 / 20.3` | `38.3 / 59.0` | |
| | `p50c` (released edit, alpha 6) | `74.6 / 4.7` | `48.4 / 34.8` | |
| |
| Using 16 loops instead of 225 costs recipe A `-31.6 +/- 3.4` points on GSM8K and `-19.1 +/- 2.9` on MATH (paired vs `capA-a6`), and costs recipe B `-2.3 +/- 2.7` on GSM8K and `-6.6 +/- 2.5` on MATH (paired vs `capB-a6`). All 16 loops came from GSM8K, which may explain why recipe B loses slightly more on MATH, though the difference between sources (about 4 points) is within noise. |
| |
| Two lessons from these runs: |
| |
| 1. **Turning up $\alpha$ did not help either recipe.** Raising $\alpha$ from `6` to `8` leaves `capA` flat (`-0.2 +/- 2.3` points) and drops `capB` by `-3.7 +/- 1.6` points, while WikiText-2 perplexity damage grows from `+0.51%` and `+0.92%` at $\alpha = 6$ to `+1.61%` and `+2.76%` at $\alpha = 8$. Even the released edit gains nothing at $\alpha = 8$ (`-0.8 +/- 1.6` points, perplexity `+3.33%`). Raising $\alpha$ scales the whole update, including any part of $u$ that points the wrong way. With one stronger setting per recipe, we cannot tell whether a partly wrong direction or a strength already past its peak explains the drop. |
| 2. **A few loops are enough for recipe B ("well before the loop"), but not for recipe A ("just before the loop").** From 16 loops (9 problems), recipe B still reaches `53.7%` accuracy (`+41.6` points over the unedited model), `4.5 +/- 1.8` points below 225-loop `capB-a6` (`58.2%`). It reaches that accuracy even though its direction is noisier and sits further from the released $u$. In a split-half check, two halves of 8 loops agree at cosine `0.828` and `0.864`, against `0.988` and `0.989` with halves of 112 and 113 loops. The `capB-max-samples-16` row in the direction table above gives the cosines with the released direction for both layers. Recipe A on the same 16 loops collapses to `17.6%` (`-25.4 +/- 2.3` points vs `capA-a6`), and `82.6%` of its answers are still cut. Its direction is the noisiest of all: two halves of 8 loops agree at only `0.632` and `0.622` at layers 15 and 19, and the direction table above lists its lower cosines with the released $u$. Its WikiText KL is `0.0081` (vs `0.0204` for `capB-max-samples-16` and `0.0178` for `capA-a6`), but [Lab 4](lab4_dose.md) and [Lab 7](lab7_side_effects.md) suggest that KL size does not explain the accuracy gap. The random-direction control `randdir` (same 64 neurons, same size of weight change, random direction) has WikiText KL `0.0046` and gains `+2.3` points. Applying the agent's direction to all 3,584 neurons (`all2`) at $\alpha = 2$ has almost the same KL as `capA-max-samples-16` (`0.0078` vs `0.0081`) and gains `+37.3` points. This suggests that, in this setup, the gain depends on where the change points instead of how far it shifts WikiText predictions. Recipe B on 16 loops beats recipe A on 225 loops by `+10.7 +/- 2.1` points (paired), suggesting that window choice matters more than loop count here. One caveat: `--max-samples 16` takes the first 16 loops in file order (all from 9 GSM8K problems). In `data_needed` (`results/lab8/capstone.json`), random 16-loop draws agreed better with the released $u$ at layers `15 / 19` (`0.666 / 0.683` for A and `0.757 / 0.746` for B, vs `0.574 / 0.596` and `0.687 / 0.682` on the first 16), so that narrow sample may explain part of the drop. We did not build or test edits from random 16-loop sets. Data quality, window placement, and sample count trade off against one another. |
| |
| ## Check yourself |
| |
| 1. Why does `labs/lab8_capstone.py` take both the loop window and the baseline window (window A or window B) from the same answer, instead of averaging all looping answers and subtracting the average of the 26 non-looping answers in `build128`? |
| <details><summary>Answer</summary> |
| |
| The 26 finished answers come from 18 problems, and 215 of the 225 loops come from other problems where neither answer finished. Subtracting the finished answers' average would mix problem differences (topic, numbers, problem type) into $u$ along with the loop signal. Taking both windows from the same answer cancels the prompt and topic within each example, as the shared opening did in Lab 3. |
| </details> |
| |
| 2. At layer 15, recipe A and recipe B have a cosine similarity of `0.951`, yet recipe B scores `58.2%` on `dev256` and recipe A scores `43.0%`. What happens in the residual stream during the last few hundred tokens before a loop that makes window A a worse baseline than window B? |
| <details><summary>Answer</summary> |
| |
| The residual stream along $u$ and the 64 neurons' push $s(t)$ are already climbing hundreds of tokens before the verbatim loop starts (`stream15` rises from `1.14` well before the loop to `1.58` in the last 128 tokens before onset, and `2.20` inside). Window A sits 32 to 68 tokens before loop onset, so subtracting it removes much of the loop buildup, halves the gap norm $\|\overline{\Delta h}\|$ (`0.689` vs `1.418`), and likely yields an edit that reacts mainly after a loop starts instead of stopping the drift into a loop. |
| </details> |
| |
| 3. When you run `labs/lab8_capstone.py --max-samples 16`, recipe B keeps most of its accuracy (`53.7%` vs `58.2%`), while recipe A drops from `43.0%` to `17.6%`. Which diagnostic numbers from the 16-loop extraction and the main run's `data_needed` table warn you before running `dev256` that recipe A will suffer much more than recipe B, and why did the first 16 loops fall below the random-16 curve? |
| <details><summary>Answer</summary> |
| |
| First, `--max-samples 16` splits the 16 loops into two halves of `8` loops. Recipe B's two halves still agree at cosine `0.828` (layer 15) and `0.864` (layer 19), with released-$u$ cosines `0.687` and `0.682`. In recipe A, subtracting window A cancels half the loop signal (`gap_size` is half as large), so sample noise dominates: its two halves agree at `0.632` and `0.622`, and its released-$u$ cosines drop to `0.574` and `0.596`. Second, `--max-samples 16` takes `used[:16]` in file order (all from 9 GSM8K problems) instead of 16 random loops across all 128 problems; in `data_needed`, random 16-loop sets had released-$u$ cosines of `0.666 / 0.683` (recipe A) and `0.757 / 0.746` (recipe B). |
| </details> |
| |
| 4. Why is the standard error on the paired difference between `capB-a6` and `p50c` (`-3.3 +/- 1.5` percentage points) nearly half as large as the standalone standard error on either run (`58.2 +/- 2.9%` and `61.5 +/- 2.8%`)? |
| <details><summary>Answer</summary> |
| |
| The standalone standard error (`2.8%` or `2.9%`) includes the wide difficulty variation across the 256 problems in `dev256`. Because both edits are tested on the same 256 problems, pairing the scores problem by problem subtracts each problem's baseline difficulty and leaves only the variance of the difference between the two edits. |
| </details> |
| |
| ## Try this next |
| |
| 1. **Restrict to loops that start late (`--min-before 2048`).** We did not run this experiment, so write down your prediction first. Out of the `225` verbatim loops in `build128`, `218` start at least `2,048` tokens into the answer (10th percentile `2,813` tokens). We expect `218` loops to be used and the directions $u$ and neuron overlaps to barely move, because Lab 6 Part 3 and the onset profile show the same gradual pre-loop drift regardless of onset position. Run the extraction (about 5 minutes on one GPU): |
| ```bash |
| CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --min-before 2048 |
| ``` |
| Compare `cos_with_released_u`, `overlap_with_released_64`, and `split_half_cos_mean_sd` in `results/lab8/try-min-before-2048/capstone.json` against the default `--min-before 512` run. |
| 2. **Move window A further back from the loop onset (`GAP = 32` to `GAP = 512`).** Near the top of `labs/lab8_capstone.py`, `WIN, AWIN, GAP = 36, 24, 32` sets how many tokens before loop onset window A ends (`GAP = 32`). Change `GAP` from `32` to `512` and pass `--try-name gap-512`: |
| ```bash |
| CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --try-name gap-512 |
| ``` |
| Check `results/lab8/try-gap-512/capstone.json` (and reset `GAP` to `32` afterward). Does moving window A 512 tokens before onset raise recipe A's `gap_size` (`0.689` at layer 15 when `GAP = 32`) and bring its cosine with the released $u$ closer to recipe B's `0.782`? |
| 3. **Predict what recipe A would do with 16 random loops instead of the first 16 (`used[:16]`).** We did not run this experiment. In `labs/lab8_capstone.py`, `used = used[:a.max_samples]` takes the first 16 loops in file order (9 GSM8K problems), giving recipe A a cosine of `0.574 / 0.596` with the released $u$ and `17.6%` accuracy on `dev256`. In `results/lab8/capstone.json`, the `answers = 16` entry of `data_needed` shows that across random 16-loop draws from all 225 loops, recipe A averages `0.666` (layer 15) and `0.683` (layer 19) with the released $u$, and `0.919` and `0.914` with the held-out half (`113` answers). Predict before editing: if you replace that line with `used = [used[i] for i in rng.choice(len(used), a.max_samples, replace=False)]` (`rng` is defined earlier) and run `--max-samples 16 --try-name rand16`, how much closer will `cos_with_released_u` and `overlap_with_released_64` move toward `capA-a6` (`0.722 / 0.724`), and will recipe A still trail recipe B at 16 random loops? |
| 4. **Read windows A and B of one loop yourself with `labs/peek.py --tokens`.** Under `"loops"` in `results/lab8/capstone.json`, each of the 225 loops records `id`, answer index `j`, `loop_starts_at`, `period`, and the token ranges of `A`, `B`, and `loop` from the start of the answer. In the median-onset answer `math_train_algebra_1625` (answer `1`), window B covers tokens `2,893` to `2,929`, window A covers `5,754` to `5,790`, and the 51-token loop starts at `5,822`. Print window B and the run-up to the loop on `sampler`: |
| ```bash |
| CUDA_VISIBLE_DEVICES="" python labs/peek.py build/base --id math_train_algebra_1625 --answer 1 --tokens 2893 2929 |
| CUDA_VISIBLE_DEVICES="" python labs/peek.py build/base --id math_train_algebra_1625 --answer 1 --tokens 5200 5822 |
| ``` |
| Each command prints where the loop starts, one cycle, the last 1,500 characters, and the requested tokens. Window B is ordinary algebra (*"The equation is $b + n = 17.8$. So $b = 17.8 -$"*). Soon after token `5,200` the model checks $b = 8.8$ (the right answer), then circles through near-copies of one check (*"Wait, if $b = 8.99$, then $b \le 9$."* alternating with *"... then $b < 9$."*) for a few hundred tokens before the 51-token cycle locks in. Window A sits inside that circling. What does that show about what recipe A subtracts out when using window A as its baseline? |
|
|
