Lab 3: Rebuild the edit from the agent's 8 texts

Big thing: You will feed the agent's 8 short math scratchpads (4 productive, 4 looping) through the unedited model, record every layer's activations (its internal numbers) with forward hooks (functions called as layers run), and rebuild the paper's edit from scratch. You will pick layers 15 and 19, recover u (the loop direction) to a cosine of 0.99998 (nearly the same line as the released direction), and match all 64 of 64 neurons in both layers.

Why it matters: Lab 1 showed what changed in the weights after the fact. Lab 3 shows where those numbers came from, using 8 forward passes (running inputs through the network) and basic vector arithmetic, with no training.

Time and hardware: About 1 minute on 1 NVIDIA A100 GPU on sampler or trainer.

Theory guide: Section 9, steps 1 to 3.

Before you run: predict#

Write down a short guess for each question before you run the script.

1. As a looping scratchpad and a productive scratchpad for the same math problem pass through the 24 layers of Qwen3.5-0.8B, does the distance between their residual streams (the running lists of numbers carried from layer to layer) stay flat, shrink, or grow from layer 0 to layer 22?
2. Each block has two parts that add to the residual stream: a token mixer (attention or linear DeltaNet, which mixes information across word pieces called tokens) and an MLP (multi-layer perceptron, which combines neuron activations into an output for each token). Which two blocks have the MLPs that add the most to the looping-minus-productive gap?
3. When we rank all 3,584 neurons in layers 15 and 19 by (how much neuron j writes along u) x (how much more neuron j fires in loops) on the 8 texts, how many of our top 64 neurons match the 64 columns from Lab 1?

Run it#

Connect to sampler (or trainer) and use nvidia-smi to pick a free GPU.

ssh sampler
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
nvidia-smi
CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py

The script reads labs/data/astra_pairs.json, runs 8 teacher-forced forward passes (feeding fixed text rather than generating new tokens) on the unedited model, writes results/lab3/rebuild.json and configs/rebuilt.json, and saves two figures: results/figures/lab3_where_the_loop_grows.png and results/figures/lab3_neuron_scores.png.

To build a vLLM model folder (/nvme/wsp-lab/models/qwen35-rebuilt) from configs/rebuilt.json and test it on the 256 unseen practice problems (dev256, 2 answers per problem, about 21 minutes on one GPU), run the two commands that labs/lab3_rebuild.py prints at the end:

CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants configs/rebuilt.json
python scripts/run_jobs.py plans/extra.json --gpus 7 --only dev/rebuilt

What you should see#

token counts match the agent's: [124, 123, 131, 136, 107, 146, 98, 135]; hooks vs hidden_states max diff 0.00e+00
layers whose MLP adds most to the gap: [(15, 0.263), (19, 0.249), (21, 0.164), (14, 0.149)]
layer 15: cos(rebuilt u, released u) = +0.99998; neurons shared with the released 64: 64/64; top-8/32/64 share of positive score: [0.193, 0.389, 0.498]
layer 19: cos(rebuilt u, released u) = -0.99998; neurons shared with the released 64: 64/64; top-8/32/64 share of positive score: [0.397, 0.498, 0.555]
cos(u15, u19) with both pointing toward looping: +0.849
wrote /nvme/wsp-lab/code/results/lab3 and configs/rebuilt.json and two figures
next, build the edit and test it on the 256 unseen practice problems (about 20 minutes on one GPU):
  CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants configs/rebuilt.json
  python scripts/run_jobs.py plans/extra.json --gpus 7 --only dev/rebuilt

Block-growth measurements (for the top 5 MLP blocks plus blocks 0 and 22) and the layer 15 and 19 reconstruction checks from results/lab3/rebuild.json:

When you build qwen35-rebuilt and test it on the 256 unseen practice problems (python scripts/run_jobs.py plans/extra.json --gpus 7 --only dev/rebuilt, recorded in results/summary.md), it matches the released edit p50c:

In lab3_where_the_loop_grows.png, the left panel plots the size (Euclidean norm) of the looping-minus-productive gap across the stream. It climbs from 0.16 after block 0 to 2.62 after block 22, growing at every block except block 7 (0.56 to 0.55). The panel omits the final position because the final normalization layer rescales the stream there, and the gap jumps to about 31, which would flatten the rest of the curve; a corner note gives that value. The right panel splits each block's addition along the gap direction into red bars (MLP) and grey bars (attention or DeltaNet). The two tallest red bars are at block 15 (0.263) and block 19 (0.249), while the mixer in block 19 adds 0.281, the tallest bar overall. lab3_neuron_scores.png plots the sorted neuron scores for layers 15 and 19 with the top 64 in red. At layer 19, neuron 11 spikes to 0.134, nearly five times the second neuron (0.028), before a long, smooth tail.

Why it works#

1. The 8 contrastive texts and teacher forcing#

In request fwd_4 (the 3rd forward pass that ran, after fwd_3 failed; see exp8_math/runs/p50c/requests/fwd_4.done.json and tool_calls.jsonl), the agent wrote 4 pairs of short math responses (labs/data/astra_pairs.json). Each pair shares the same user prompt and opening calculation inside <think>, then diverges:

• The productive text (even indices 0, 2, 4, 6) checks its answer once and stops.
• The looping text (odd indices 1, 3, 5, 7) repeats one or two short sentences three or four times, such as "But is 20 the total? No, it is the remaining time."

The script uses teacher forcing instead of generating new tokens: it feeds each fixed text into the unedited model in a single forward pass and records the internal activations at every token.

To read inside the model without changing its code, wsp/hf.py attaches two PyTorch hooks (small functions that PyTorch calls automatically as each layer runs):

• record_outputs(m, range(n_layers)) attaches a forward hook to every block l and copies the 1,024-number residual stream leaving that block (hidden[l+1]). The script verifies that the hook outputs match Hugging Face's output_hidden_states with a maximum difference of 0.00e+00.
• record_mlp_acts(m, range(n_layers)) attaches a pre-hook to every down_proj matrix and copies the 3,584 neuron activations a entering down_proj at each token.

2. Where the loop signal grows across layers#

Because the two texts in a pair share their prompt and opening, they differ only near the end, where one loops and the other checks its answer. Following the agent's script, lab3_rebuild.py averages the residual stream over the last H_WIN = 36 tokens of each text and computes the mean looping-minus-productive difference across the 4 pairs:

Here $\Delta h_l$ is the 1,024-number vector pointing from productive to looping text entering block l, and $\Delta h_{l+1}$ is the gap vector leaving block l.

We can measure how much block l pushes the stream along the new gap direction $d = \Delta h_{l+1} / \|\Delta h_{l+1}\|$. Block l adds its mixer output plus its MLP output to the stream, so the total growth along $d$ splits cleanly into two pieces:

Sorting all blocks by the MLP share gives [(15, 0.263), (19, 0.249), (21, 0.164), (14, 0.149)]. The MLPs in blocks 15 and 19 push hardest toward the loop, so the agent chose those two layers from this table. Both happen to be full-attention blocks (Qwen3.5-0.8B uses full attention in every fourth block: 3, 7, 11, 15, 19, and 23, and DeltaNet in the rest), though the choice came from the MLP numbers rather than the block type.

3. Computing u and scoring the 3,584 neurons#

At each chosen layer $L \in \{15, 19\}$, the loop direction u is the gap vector leaving block $L$, scaled to length 1:

By construction, $+u_L$ points from productive toward looping text, resolving the sign ambiguity from Lab 1's SVD (singular value decomposition, which splits a matrix change into independent pieces):

• At layer 15, $\cos(u_{15}^{\text{rebuilt}}, u_{15}^{\text{released}}) = +0.99998$.
• At layer 19, $\cos(u_{19}^{\text{rebuilt}}, u_{19}^{\text{released}}) = -0.99998$ (the SVD convention in results/edit_info.json happened to pick the productive sign for layer 19).
• With both vectors pointing toward looping, $\cos(u_{15}, u_{19}) = +0.849$, so the two loop directions are similar, about 32 degrees apart, and not opposite.

Next, we find which neurons in layer $L$'s MLP push the stream along $u_L$ during a loop. A neuron adds to the loop gap along $u_L$ when two things hold at once: its column $W[:, j]$ writes along $u_L$, and its activation differs between looping and productive text.

Over the agent's window of the last A_WIN = 24 tokens, neuron j fires $\Delta a_{24}[j]$ more on average in loops than in productive text, and it writes $(u_L \cdot W[:, j])$ along $u_L$ per unit of activation. Their product is the amount neuron j adds to the loop direction:

A neuron with a negative write[j] and a negative gapA24[j] (writing toward productive text while firing less in loops) also gets a positive score. Taking the top 64 of the 3,584 neurons by score reproduces 64 of 64 columns from the released edit in layer 15 and 64 of 64 in layer 19.

4. Concentration of scores and behavior of the rebuilt edit#

The cumulative share of positive score (positive_score_share_top8_32_64) shows how spread out the effect is across neurons:

• In layer 15, the top 8 neurons hold 19.3% of the positive score, the top 32 hold 38.9%, and the top 64 hold 49.8% (about 50%). No single neuron dominates; the top one, neuron 294, holds 5.6%.
• In layer 19, the top 8 neurons hold 39.7%, the top 32 hold 49.8%, and the top 64 hold 55.5%. Neuron 11 alone scores 0.134, or 24% of the layer's positive score (0.552), ahead of the second-highest, neuron 1, at 0.028.

The 64th neuron in layer 15 scores 0.00136 and the 65th scores 0.00134, so the cut at 64 goes through a smooth tail rather than a sharp drop to zero. It captures about half of the positive loop contribution while touching under 1.8% of the neurons. In Lab 7, rebuilt matches the released p50c edit: both shift WikiText-2 perplexity (prediction uncertainty on ordinary text, lower is better) by +0.79% and give a WikiText KL divergence (change in next-token probabilities) of 0.0333 (0.0096 vs 0.0097 on productive texts, 0.0238 vs 0.0241 on looping texts). On dev256, both score 61.5%.

Check yourself#

1. Why does the neuron score multiply write = u @ Wd[L] by gapA24[L] instead of selecting neurons by |u @ Wd[L]| alone?

2. In layer 19, cos(rebuilt u, released u) is -0.99998. Since the signs are opposite, why doesn't the sign matter for the edit formula?

3. Look at the curve in lab3_neuron_scores.png. Why did the agent stop at 64 neurons per layer instead of editing only neuron 11, or editing all 3,584 neurons?

Try this next#

You can test recipe changes without editing labs/lab3_rebuild.py. Passing any non-default flag (--pairs N, --windows STREAM NEURONS, or --layers ...) writes to results/lab3/try-<name>/, named after the flags you changed, and leaves results/lab3/rebuild.json, configs/rebuilt.json, and results/figures/ untouched.

Each try folder holds rebuild.json, the two figures (lab3_where_the_loop_grows.png and lab3_neuron_scores.png), variant.json (edit rebuilt-<name>), and plan.json (a one-GPU dev256 run dev/rebuilt-<name>). The script prints the two commands to build /nvme/wsp-lab/models/qwen35-rebuilt-<name> on CPU in a few seconds and test it on the 256 unseen practice problems (about 20 minutes on one GPU):

CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants results/lab3/try-<name>/variant.json
python scripts/run_jobs.py results/lab3/try-<name>/plan.json --gpus 7 --only dev/rebuilt-<name>

We ran all three tries below, built their models, and evaluated them in about 11 minutes across 6 GPUs (2 shards per run) with plans/lab3-tries.json. We then scored their side effects on WikiText-2 and the 8 texts in Lab 7 (results/lab7/try-lab3-tries/) and checked how their 512 answers ended (results/lab3/endings-tries.json):

python scripts/run_jobs.py plans/lab3-tries.json --gpus 0,1,2,3,4,6
CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --models rebuilt rebuilt-pairs2 rebuilt-windows36-36 rebuilt-layers15-19-21 --try-name lab3-tries

(scripts/run_jobs.py skips any shard with a .done.json marker; move the existing run folder under /nvme/wsp-lab/runs/dev/ aside first to generate fresh answers.)

1. Use only the agent's first 2 pairs (--pairs 2, folder results/lab3/try-pairs2/).

   • Predict first: With half the data (the first 4 texts instead of all 8), how far will cos(rebuilt u, released u) and the 64-neuron overlap drop, and will dev256 accuracy suffer?
   • Run it:
     bash

     CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --pairs 2
     CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants results/lab3/try-pairs2/variant.json
     python scripts/run_jobs.py results/lab3/try-pairs2/plan.json --gpus 7 --only dev/rebuilt-pairs2
     

   • What we saw: Blocks 15 and 19 still lead in MLP growth ([(15, 0.201), (19, 0.188), (22, 0.167), (14, 0.163)]), while the direction and neuron list get rougher:
     • Layer 15: cos(rebuilt u, released u) = +0.82157 (0.822), 42/64 neurons match, and the top 8 / 32 / 64 neurons hold [0.144, 0.278, 0.386] (14.4% / 27.8% / 38.6%) of the positive score.
     • Layer 19: cos(rebuilt u, released u) = -0.78384 (-0.784), 31/64 neurons match, and the top 8 / 32 / 64 neurons hold [0.219, 0.348, 0.419] (21.9% / 34.8% / 41.9%).
     • On dev256, dev/rebuilt-pairs2 reaches 62.3% +/- 2.7 accuracy (+50.2 +/- 2.6 points over dev/base, 129 rescued, 0 lost), matching the 4-pair rebuild (61.5%). What changes is how the edit works and how hard it pushes. More answers hit the 32,768-token cap (25.2% vs 17.6%). More cut answers also end in repetition (74.4% vs 32.2%, or 68.2% vs 15.6% for verbatim loops). Across all 512 answers in results/lab3/endings-tries.json, 86 answers from rebuilt-pairs2 are cut while thinking and 74 close out of a loop, vs 55 and 118 for rebuilt in results/lab3/endings.json. In Lab 7, rebuilt-pairs2 disturbs WikiText-2 less (perplexity +0.43% vs +0.79%, KL 0.0194 vs 0.0333) and is less selective on the 8 texts (looping KL over productive KL is 0.0109 / 0.0075 = 1.46, vs 0.0238 / 0.0096 = 2.48 for rebuilt).

2. Average neuron activations over 36 tokens instead of 24 (--windows 36 36, folder results/lab3/try-windows36-36/).

   • Predict first: The agent used 36 tokens for the stream (H_WIN = 36) and 24 tokens for the neuron scores (A_WIN = 24). If you set A_WIN to 36, will u change, and will the extra 12 tokens shift the 64 chosen neurons enough to affect dev256?
   • Run it:
     bash

     CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --windows 36 36
     CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants results/lab3/try-windows36-36/variant.json
     python scripts/run_jobs.py results/lab3/try-windows36-36/plan.json --gpus 7 --only dev/rebuilt-windows36-36
     

   • What we saw: Because H_WIN remains 36, the block growth numbers and u are unchanged (+0.99998 at layer 15, -0.99998 at layer 19). The 36-token neuron window swaps a few tail neurons: 58/64 match at layer 15 (top 8 / 32 / 64 score shares [0.187, 0.375, 0.481]) and 51/64 match at layer 19 (shares [0.378, 0.474, 0.533]). On dev256, dev/rebuilt-windows36-36 scores 61.9% +/- 2.8 (+49.8 +/- 2.7 points, 21.3% cut, 122 rescued, 0 lost), with a Lab 7 WikiText-2 perplexity change of +0.71% (KL 0.0326). The 24-token window is not fragile.

3. Also edit layer 21, the third-strongest MLP (--layers 15 19 21, folder results/lab3/try-layers15-19-21/).

   • Predict first: Block 21 had the third-largest MLP addition to the loop gap (0.164, behind 0.263 at block 15 and 0.249 at block 19). Does applying the 64-neuron, alpha = 6 edit to layer 21 cut more loops and raise accuracy, or only add side effects?
   • Run it:
     bash

     CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --layers 15 19 21
     CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants results/lab3/try-layers15-19-21/variant.json
     python scripts/run_jobs.py results/lab3/try-layers15-19-21/plan.json --gpus 7 --only dev/rebuilt-layers15-19-21
     

   • What we saw: Layers 15 and 19 match the main rebuild (64/64 neurons in both). At layer 21, scores are flatter. The top 8 / 32 / 64 neurons hold [0.173, 0.292, 0.371] (17.3% / 29.2% / 37.1%) of the positive score (top 12: [2, 1, 11, 99, 14, 85, 9, 109, 168, 335, 212, 169]). The edit's magnitude is 5.8% of layer 21's down_proj matrix magnitude (the square root of the sum of squared entries, not the fraction of entries changed). On dev256, dev/rebuilt-layers15-19-21 scores 61.3% +/- 2.8 (+49.2 +/- 2.6 points, 21.1% cut, 122 rescued, 0 lost, with 3 answers stopping mid-thinking before </think>), showing no gain over the 2-layer rebuild (61.5%, difference -0.2 +/- 1.9 points). In Lab 7, WikiText-2 perplexity rises by +1.49% (nearly double the 2-layer rebuild's +0.79%) and KL climbs from 0.0333 to 0.0394.

The main rebuild and the three tries line up as follows across Lab 3, dev256, and Lab 7:

Lesson: The recipe is forgiving. Using half the pairs (--pairs 2) or widening the neuron window (--windows 36 36) still reaches about 62% accuracy on unseen problems. Those variants mainly shift side effects on ordinary text and the balance between preventing loops early and breaking them after they start. Editing a third layer with flatter neuron scores (--layers 15 19 21) increases ordinary-text damage without improving math accuracy.