Weight-space programming lab

Lab 8 (capstone): build your own edit from the model's own loops

On this page

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 uu 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 explains the four groups), and labs/lab7_side_effects.py measures side effects on WikiText-2 and the agent's 8 texts (Lab 7). 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 uu Released neurons picked (of 64) Split-half cosine (stability) Gap size (∣Δh‾∣|\overline{\Delta h}|) Agent's 4 pairs: looping minus productive along uu 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 uu 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 NN random answers and a direction from a held-out half (113 answers) climbs fast:
Answers used (NN) 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

Cosine against the number of answers used to build u, at layers 15 and 19, for recipes A (red) and B (purple). Solid lines compare with the other half of the answers and dashed lines with the released u.

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 uu, layer 15 1.14 1.39 1.58 1.95 2.20
Stream along released uu, layer 19 2.27 2.66 2.95 3.50 3.96
Trigger s(t)s(t) (released 64 neurons), layer 15 0.135 0.173 0.206 0.288 0.366
Trigger s(t)s(t) (released 64 neurons), layer 19 0.183 0.237 0.277 0.376 0.512

Both the stream projection along uu and the 64 neurons' push s(t)=u⋅(WSaS(t))s(t) = u \cdot (W_S a_S(t)) (the Lab 6 trigger measuring how hard those 64 neurons push along uu at token tt) rise hundreds of tokens before repetition starts, then jump at token 0.

Median and middle half of 225 loops lined up at the token where each loop starts. The stream along the released u and the 64 neurons' push s(t) rise slowly at layers 15 and 19, then jump once the loop starts.

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).

One answer that loops from token 5,822: the stream along the released u and the push s(t) at layers 15 and 19, with shaded strips for window B, window A and the loop window.

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 (ee to the mean negative log-probability of the true next token; base model 20.009 on WikiText-2, lower is better) and KL divergence (KL(base∥edited)\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). 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.

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∈{15,19}L \in \{15, 19\}, the remaining steps follow Lab 3:

  1. Direction uu: For each answer ii and recipe k∈{A,B}k \in \{A, B\}, subtract the mean residual stream leaving block LL over window kk from the mean over the loop window: Δhi,L(k)=hˉi,Lloop−hˉi,L(k)\Delta h_{i, L}^{(k)} = \bar{h}_{i, L}^{\text{loop}} - \bar{h}_{i, L}^{(k)}. Average Δhi,L(k)\Delta h_{i, L}^{(k)} across all 225 answers and scale to unit length: uL(k)=Δh‾L(k)∥Δh‾L(k)∥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 Δa‾L(k)[j]\overline{\Delta a}_L^{(k)}[j] (loop window minus comparison window) for each of the 3,584 MLP neurons. Multiply by how much column jj of down_proj writes along uL(k)u_L^{(k)}: scorej=(uL(k)⋅Wdown,L[:,j])×Δa‾L(k)[j]\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 scorej\text{score}_j descending, keep the top 64 indices SS, and apply the rank-one update (one column vector uu times one row vector uTWSu^T W_S) W′=W−α u (uTWS)W' = W - \alpha \, u \, (u^T W_S) with edit strength α=6\alpha = 6. This scales those 64 neurons' push along uu by 1−α=−51 - \alpha = -5.

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 uu 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 s15(t)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 ∥Δh‾∥\|\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 uAu_A, lowering its cosine with the released uu (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 Δa‾24\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 (H1H_1 and H2H_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∈{1,2,4,8,16,32,64,112}N \in \{1, 2, 4, 8, 16, 32, 64, 112\} answers from H1H_1 and measures the cosine of their direction with the full held-out half H2H_2:

At N=1N = 1 (one looping answer), recipe B already has a 0.68 cosine with the held-out half. By N=16N = 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 uu) 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=1N = 1. From N=16N = 16 on, held-out agreement shifts by under 0.006 and released-uu agreement by under 0.009. When the half that the NN 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 uu 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 pp (the fraction of sample answers solved by model 1 minus the fraction solved by model 2 on problem pp) 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 uu and 64 neurons and raises edit strength from α=6\alpha = 6 to α=8\alpha = 8. The relative weight change (Frobenius norm (the overall size) of W′−WW' - W divided by Frobenius norm of WW, 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 uu (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 α=6\alpha = 6 to +1.61% and +2.76% at α=8\alpha = 8. Even the released edit gains nothing at α=8\alpha = 8 (-0.8 +/- 1.6 points, perplexity +3.33%). Raising α\alpha scales the whole update, including any part of uu 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 uu. 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 uu. Its WikiText KL is 0.0081 (vs 0.0204 for capB-max-samples-16 and 0.0178 for capA-a6), but Lab 4 and Lab 7 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 α=2\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 uu 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?
Answer

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 uu 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.

  1. 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?
Answer

The residual stream along uu and the 64 neurons' push s(t)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 ∥Δh‾∥\|\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.

  1. 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?
Answer

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-uu 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-uu 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-uu cosines of 0.666 / 0.683 (recipe A) and 0.757 / 0.746 (recipe B).

  1. 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%)?
Answer

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.

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 uu 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 uu 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 uu 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 uu, 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.8b + n = 17.8. So b=17.8−b = 17.8 -"). Soon after token 5,200 the model checks b=8.8b = 8.8 (the right answer), then circles through near-copies of one check ("Wait, if b=8.99b = 8.99, then b≤9b \le 9." alternating with "... then b<9b < 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?