Lab 7: What else did the edit change

Big thing: You will score the base model and 14 edited versions on 65,536 tokens of Wikipedia articles (WikiText-2) and on the agent's 8 short math texts (4 productive, 4 looping). For each model, you measure how well it predicts the text (perplexity) and how far its next-token probabilities move from the base model's (KL divergence). The released edit (p50c) moves predictions 2.5 times more on looping texts than on productive ones (0.0241 vs 0.0097) and raises WikiText-2 perplexity by only 0.79%.

Why it matters: A loop fix is useless if it ruins ordinary language, since the model still has to read questions and write normal sentences. The 14 versions cover weaker and stronger edits ($\alpha$ from -4 to 12), single-layer edits, edits across all 3,584 neurons, and two random controls with the same weight-change size as the released edit. Comparing them shows what makes the edit act on loops and leave other text alone: the direction $u$, the 64 chosen neurons, or both.

Time and hardware: About 1.5 minutes on one A100 GPU (78 seconds in our run), plus a few seconds on the first run to download WikiText-2.

Theory guide: Section 13.4 (measure more than accuracy).

Before you run: predict#

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

1. The released edit (p50c, $\alpha = 6$) flips the loop push of 64 neurons in layers 15 and 19 and makes it 5 times larger. On plain Wikipedia text (no chat template, no thinking), will perplexity get much worse (say +15% or more), change by less than 1%, or improve?
2. Two random controls test whether any change of the same size works. randdir edits the same 64 neurons along a random direction at right angles to $u$. randcols edits 64 random neurons along $u$. Both match p50c's weight-change size (matched before bf16 rounding, so they agree closely but not exactly; the build script raised their $\alpha$ to between about 20 and 31). On the agent's 8 texts, will they move predictions more on looping texts than on productive ones, or move both about equally?
3. Compare how far WikiText-2 predictions move when editing only layer 15 (l15only) versus both layers 15 and 19 (p50c). Must the two-layer edit move WikiText-2 predictions more than the single-layer edit?

Run it#

Connect to sampler (or trainer). Then use nvidia-smi to pick a free GPU index and pass it in CUDA_VISIBLE_DEVICES.

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

The script loads the base model from /nvme/wsp-lab/models/Qwen3.5-0.8B, reads edit_note.json from each /nvme/wsp-lab/models/qwen35-* folder (built by scripts/build_models.py), and loads the down_proj matrices at the union of edited layers (layers 15 and 19 for the 14 standard edits, plus any extra edited layers from models you built in Labs 3 and 8). Pass --models to score a subset. It evaluates each model on Salesforce/wikitext (wikitext-2-raw-v1, test split) and labs/data/astra_pairs.json, then writes:

• results/lab7/side_effects.json (perplexity, percent perplexity change, and KL divergences for base and each edited model)
• results/figures/lab7_side_effects.png (two-panel bar chart of WikiText-2 perplexity change and productive-versus-looping KL divergence)

What you should see#

The script lists the 15 models in the default run (base plus 14 standard edits) and the swapped down_proj layers, then prints one row per model and the output path and runtime (yours will differ a little):

models: ['base', 'a-4', 'a1', 'a12', 'a2', 'a4', 'a8', 'all1', 'all2', 'l15only', 'l19only', 'p50c', 'randcols', 'randdir', 'rebuilt'] | down_proj swapped at layers [15, 19]
model           ppl  ppl +%   KL wiki   KL prod   KL loop
base         20.009    0.00   0.00000   0.00000   0.00000
a-4          21.088    5.39   0.02058   0.00558   0.01466
a1           19.912   -0.49   0.00160   0.00095   0.00138
a12          22.878   14.34   0.16430   0.03661   0.10268
a2           19.871   -0.69   0.00445   0.00178   0.00346
a4           19.916   -0.47   0.01516   0.00485   0.01120
a8           20.675    3.33   0.06168   0.01645   0.04416
all1         19.925   -0.42   0.00248   0.00227   0.00289
all2         19.937   -0.36   0.00782   0.00634   0.00919
l15only      20.020    0.05   0.03476   0.00550   0.01938
l19only      20.340    1.65   0.01025   0.00464   0.00755
p50c         20.167    0.79   0.03327   0.00967   0.02405
randcols     19.989   -0.10   0.00346   0.00446   0.00383
randdir      20.070    0.30   0.00462   0.00479   0.00387
rebuilt      20.167    0.79   0.03326   0.00961   0.02384
wrote /nvme/wsp-lab/code/results/lab7 in 78 s

Grouping the 14 edited variants by what they test (rounded to 4 decimal places for KL, with base WikiText-2 perplexity 20.009):

The script also saves results/figures/lab7_side_effects.png:

Why it works#

1. Testing 15 models (the base model plus 14 edits) by swapping matrices in memory#

A language model processes text one token (word piece) at a time through 24 blocks (or layers), each reading from and adding to the 1,024-number residual stream. Inside each block, the MLP (multi-layer perceptron) expands the stream to 3,584 neurons and projects them back through down_proj (a $1,024 \times 3,584$ matrix stored in bf16, 16-bit floating point).

The base model has 488 weight tensors. In the default 15-model run (base plus 14 edits in results/lab7/side_effects.json), each edited model differs from the base model in at most two tensors: layers.15.mlp.down_proj.weight and layers.19.mlp.down_proj.weight (single-layer edits change only one). Instead of loading 15 separate 1.7 GB models from disk, labs/lab7_side_effects.py loads the base model onto the GPU once. It reads edit_note.json in each model folder, takes the union of edited layers (LAYERS), and swaps down_proj at those layers with set_weights(w) (printing down_proj swapped at layers ...). Any layer a model leaves untouched gets the base weights, which is how Lab 3's layer-21 experiment (rebuilt-layers15-19-21, layers 15, 19, and 21) is scored alongside 2-layer models.

For each chunk of text, the script uses teacher forcing: it feeds a fixed sequence of $1,024$ tokens in one forward pass and reads next-token predictions at positions $0$ through $1,022$ (predicting tokens $1$ through $1,023$). On WikiText-2 test, 64 chunks of 1,024 tokens ($65,536$ tokens total) yield 65,472 predictions, since the first token of each chunk has no prior context.

2. Perplexity and KL divergence from scratch#

At each token position, the model outputs logits (one raw score per vocabulary token), and softmax converts those scores into probabilities that sum to 1. We compute two metrics from those probabilities.

Imagine a two-token vocabulary, "38" and "But", at a step where "38" is the true next token:

• The base model predicts $p_{\text{base}}(\text{"38"}) = 0.80$ and $p_{\text{base}}(\text{"But"}) = 0.20$.
• An edited model predicts $p_{\text{edit}}(\text{"38"}) = 0.50$ and $p_{\text{edit}}(\text{"But"}) = 0.50$.

Perplexity measures how well a model predicts the true next token in ground-truth text. First take the negative natural log of the true token's probability ("38"), called the negative log-likelihood (NLL):

• Base model: $\text{NLL} = -\ln(0.80) = 0.2231$.
• Edited model: $\text{NLL} = -\ln(0.50) = 0.6931$.

Average the NLL across all tokens and raise $e$ to that average: $\text{perplexity} = e^{\text{mean NLL}}$. On this single token, perplexity is $e^{0.2231} = 1 / 0.80 = 1.25$ for the base model and $e^{0.6931} = 1 / 0.50 = 2.00$ for the edited model. A model as uncertain as a fair $K$-sided die at every token (assigning probability $1/K$ to the true token) has perplexity $K$. Lower perplexity means better prediction. On the 65,536 WikiText-2 tokens, the base model's perplexity is 20.009.

KL divergence, written $\text{KL}(\text{base} \parallel \text{edited})$, measures how far the edited model's full probability distribution moved from the base model's, regardless of which token came next. At one position, it averages $\ln(p_{\text{base}}(v) / p_{\text{edit}}(v))$ weighted by the base probabilities:

In our two-token example:

Identical distributions have KL divergence 0. Even when perplexity barely changes because gains on some tokens offset losses on others, KL divergence is strictly positive whenever predictions move.

3. Reading the four findings in the table#

First, compare selectivity on looping versus productive text. Each variant in configs/variants.json modifies down_proj along a direction $u$ (a 1,024-number unit vector) with edit strength $\alpha$ (scaling the chosen neurons' push along $u$ by $1 - \alpha$). In the right panel of lab7_side_effects.png, compare the red bars (kl_looping_texts) with the green bars (kl_productive_texts). The released edit p50c ($\alpha = 6$) moves next-token predictions about 2.5 times more on looping texts (0.0241) than on productive ones (0.0097). Lab 3's rebuild (rebuilt) matches p50c closely (0.0238 vs 0.0096, +0.79% perplexity).

Now compare p50c with the random controls randdir and randcols. Both set norm_match: true, rescaling $\alpha$ so the Frobenius size of the weight change $\|W' - W\|_F$ matches p50c's change size before bf16 rounding (10.6% of the matrix at layer 15 and 8.1% at layer 19, from Lab 1). Despite equal change size, randdir (0.0039 looping vs 0.0048 productive) and randcols (0.0038 looping vs 0.0045 productive) move both text types by nearly the same small amount, shifting productive texts slightly more. Editing all 3,584 neurons (all1, all2) is also less selective than editing only the 64 chosen ones. At $\alpha = 2$, a2 (64 neurons) moves productive text by 0.0018 and looping text by 0.0035 (ratio about 1.9), whereas all2 (3,584 neurons) moves them by 0.0063 and 0.0092 (ratio about 1.5), shifting productive text 3.5 times more than a2.

KL has no sign: it measures how far predictions move, not which way. a-4 pushes the 64 neurons toward loops instead of away, yet it is selective too (looping texts move 2.6 times more than productive ones). It is selective because it acts in the same place as p50c: where those 64 neurons push hard along $u$, which happens on looping text. Selectivity shows where an edit acts; Lab 4's accuracy runs show whether it helps or hurts.

Second, trace the WikiText-2 perplexity dose curve across $\alpha$ (a-4, a1, a2, a4, p50c at $\alpha = 6$, a8, a12) in the left panel of lab7_side_effects.png:

• Setting $\alpha$ to -4 (a-4, pushing 5 times harder in the loop direction) raises perplexity by +5.39% (21.088).
• Small positive edits (a1 -0.49%, a2 -0.69%, a4 -0.47%, all1 -0.42%, all2 -0.36%) lower WikiText-2 perplexity slightly below the base model's 20.009, reaching 19.871 at a2.
• At the released strength $\alpha = 6$ (p50c), flipping and amplifying 5x costs +0.79% perplexity (20.167).
• Above $\alpha = 6$, perplexity climbs fast: a8 ($\alpha = 8$) adds +3.33% (20.675), and a12 ($\alpha = 12$) adds +14.34% (22.878) with WikiText KL 0.1643.

Third, the single-layer rows l15only and l19only show that KL does not add across edits. On WikiText-2, editing layer 15 alone moves predictions slightly more (wikitext_kl = 0.0348) than editing both layers 15 and 19 (p50c, 0.0333). Yet l15only barely changes perplexity (+0.05%), while l19only moves predictions less (KL 0.0102) and raises perplexity more (+1.65%). This table cannot say why. One plausible story: layer 19 sits four blocks after layer 15, so the layer-15 edit changes the input reaching layer 19's 64 neurons. Because the two loop directions point along similar lines (absolute cosine similarity 0.849 in Lab 1, where 1 means the same direction and 0 means at right angles), the layer-19 edit may undo part of the layer-15 shift on ordinary text. l15only is also the most selective edit in the table, moving looping texts 3.5 times more than productive ones, and in Lab 4 it fixes almost as many problems as the full edit (60.4% correct on the dev set vs 61.5% for p50c).

Fourth, WikiText-2 is plain encyclopedia text scored with no chat template and no thinking. This table measures damage to ordinary language modeling, not to math reasoning or following instructions. Lab 4's dose runs measure the math side. Here, $\alpha = 12$ hurts ordinary text a lot and $\alpha = 6$ barely touches it.

Check yourself#

1. Suppose a model assigns probability 0.25 to the true next token at every position of a test passage. What is its average negative log-likelihood (NLL), and what is its perplexity?

2. How can a2 ($\alpha = 2$) have a positive KL divergence on WikiText-2 (0.0044) while its WikiText-2 perplexity change is negative (-0.69%)?

3. What specific alternative explanation does randdir rule out, and what specific alternative explanation does randcols rule out?

4. Why can l15only have a larger wikitext_kl (0.0348) than p50c (0.0333), even though p50c includes the exact same layer-15 edit plus the layer-19 edit?

Try this next#

Each experiment writes its JSON file and figure to results/lab7/try-<name>/, leaving results/lab7/side_effects.json untouched. Our try folders are committed so you can compare your run with ours.

1. Rank the edits by how selective they are. For each model, divide its KL on the looping texts by its KL on the productive texts, and sort:

   bash

   python -c 'import json; m=json.load(open("results/lab7/side_effects.json"))["models"]; rows=[(v["kl_looping_texts"]/v["kl_productive_texts"], k, v["ppl_change_pct"]) for k,v in m.items() if k != "base"]; print("\n".join(f"{k:10s} ratio={r:.2f}x  ppl_change={p:+.2f}%" for r,k,p in sorted(rows, reverse=True)))'
   

   Our output:

   text

   l15only    ratio=3.52x  ppl_change=+0.05%
   a12        ratio=2.81x  ppl_change=+14.34%
   a8         ratio=2.68x  ppl_change=+3.33%
   a-4        ratio=2.63x  ppl_change=+5.39%
   p50c       ratio=2.49x  ppl_change=+0.79%
   rebuilt    ratio=2.48x  ppl_change=+0.79%
   a4         ratio=2.31x  ppl_change=-0.47%
   a2         ratio=1.95x  ppl_change=-0.69%
   l19only    ratio=1.63x  ppl_change=+1.65%
   all2       ratio=1.45x  ppl_change=-0.36%
   a1         ratio=1.44x  ppl_change=-0.49%
   all1       ratio=1.28x  ppl_change=-0.42%
   randcols   ratio=0.86x  ppl_change=-0.10%
   randdir    ratio=0.81x  ppl_change=+0.30%
   

   The layer-15-only edit ranks first, ahead of the released edit. The a-4 edit, which pushes toward loops, also ranks above p50c because KL measures how far predictions move, not which way. Both random controls rank last, below 1, moving productive text slightly more than looping text.

2. Score only the last 36 tokens of each text. In each pair, the productive and looping texts share their prompt and opening lines before the looping text starts repeating. Those shared tokens get the same KL, pulling the two KL columns together. --last 36 scores only the last 36 tokens of each of the 8 texts (the shortest has 98 tokens):

   bash

   CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --last 36
   

   This writes results/lab7/try-last-36/, with WikiText-2 columns unchanged. Point the ranking one-liner at results/lab7/try-last-36/side_effects.json to compare:

   Model	Ratio over all tokens	Ratio over the last 36 tokens
   l15only	3.52	12.96
   p50c	2.49	6.31
   rebuilt	2.48	6.18
   all1	1.28	2.07
   randcols	0.86	0.49
   randdir	0.81	0.34

   Every edit along $u$ gets more selective, while the random controls get less selective. For p50c, productive KL barely moves (0.0097 to 0.0093), while looping KL more than doubles (0.0241 to 0.0584). This fits Lab 6, where the trigger $s(t)$ rises in the last 512 tokens before a loop and peaks inside the repetition ($s_{15}$ goes 0.096 -> 0.178 -> 0.424 from early in a looping answer, to just before the loop, to inside it).

3. Score less WikiText-2. --chunks 8 scores the first 8 chunks (8,192 tokens) instead of all 64:

   bash

   CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --chunks 8
   

   This writes results/lab7/try-chunks-8/. Base perplexity is 22.804 on these 8,192 tokens versus 20.009 on all 65,536, which is fine because every model is compared with the base model on the same tokens. The WikiText KL ranking matches the 64-chunk run except that p50c and rebuilt (tied to 4 decimal places) swap places. Perplexity changes are less stable: the four largest keep their order (a12 +12.21%, a-4 +5.16%, a8 +2.43%, l19only +1.47%), while shifts under 1% shuffle (p50c drops from +0.79% to +0.20%, and l15only flips from +0.05% to -0.54%). Because KL averages non-negative per-position values, positions cannot cancel out, which helps explain why its ranking held up. A perplexity change subtracts two nearly equal averages, so when it is under 1%, even its sign depends on which text you score. The run took 70 seconds versus 78 for 64 chunks: scoring the text is the fast part.

4. Score Lab 3's rebuild experiments (try-lab3-tries). After building the three Lab 3 experiment models (rebuilt-pairs2, rebuilt-windows36-36, and rebuilt-layers15-19-21), compare them against the main 4-pair rebuild:

   bash

   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
   

   Because rebuilt-layers15-19-21 also edits layer 21, the script swaps down_proj at layers [15, 19, 21] and writes results/lab7/try-lab3-tries/side_effects.json (base perplexity 20.009):

   Model	Perplexity	Change	KL WikiText	KL productive	KL looping
   rebuilt	20.167	+0.79%	0.0333	0.0096	0.0238
   rebuilt-pairs2	20.096	+0.43%	0.0194	0.0075	0.0109
   rebuilt-windows36-36	20.151	+0.71%	0.0326	0.0096	0.0241
   rebuilt-layers15-19-21	20.307	+1.49%	0.0394	0.0114	0.0280

   All four models score between 61.3% and 62.3% on dev256, yet their side effects separate cleanly. Using only 2 pairs (rebuilt-pairs2) produces a gentler, less selective edit (+0.43% perplexity, looping-to-productive KL ratio 1.46 vs 2.48), while widening the neuron window to 36 tokens (rebuilt-windows36-36) changes almost nothing. Adding layer 21 (rebuilt-layers15-19-21) nearly doubles the WikiText-2 perplexity cost (+1.49% vs +0.79%) without helping accuracy (61.3% vs 61.5%).

5. Score your main capstone edits (after Lab 8). Lab 8 adds two models built from the unedited model's 225 loops at $\alpha = 6$: capA-a6 (loop minus just before) and capB-a6 (loop minus well before). Compare them with the released edit and the most selective edit:

   bash

   CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py --models p50c l15only capA-a6 capB-a6
   

   This writes results/lab7/try-models-p50c-l15only-capA-a6-capB-a6/. Our output:

   text

   models: ['base', 'p50c', 'l15only', 'capA-a6', 'capB-a6']
   model           ppl  ppl +%   KL wiki   KL prod   KL loop
   base         20.009    0.00   0.00000   0.00000   0.00000
   p50c         20.167    0.79   0.03327   0.00967   0.02405
   l15only      20.020    0.05   0.03476   0.00550   0.01938
   capA-a6      20.111    0.51   0.01778   0.01025   0.01624
   capB-a6      20.194    0.92   0.02711   0.01166   0.02078
   wrote /nvme/wsp-lab/code/results/lab7/try-models-p50c-l15only-capA-a6-capB-a6 in 34 s
   

   Both capstone edits move productive texts more than the released edit does, and looping texts less. Their selectivity ratios are 1.58 and 1.78 versus 2.49 for p50c. In Lab 8, capA-a6 also helps less: it scores 43.0% on the dev set with 52.5% of answers cut off, versus 61.5% accuracy and 19.7% cut for p50c (capB-a6 reaches 58.2% accuracy with 30.7% cut).

6. Score Lab 8's capstone experiments (try-lab8-tries). Lab 8 tests building the capstone edits from only the first 16 loops (capA-max-samples-16, capB-max-samples-16) or raising the strength from $\alpha = 6$ to $\alpha = 8$ on all 225 loops (capA-alpha-8, capB-alpha-8):

   bash

   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
   

   This writes results/lab7/try-lab8-tries/side_effects.json in about 33 seconds. Here are those four models alongside the main $\alpha = 6$ capstone edits (capA-a6 and capB-a6, base perplexity 20.009):

   Model	Loops used	$\alpha$	dev256 accuracy %	Perplexity	Change	KL WikiText	KL productive	KL looping
   capA-a6 (main)	225	6	43.0 +/- 2.6	20.111	+0.51%	0.0178	0.0103	0.0162
   capA-max-samples-16	16	6	17.6 +/- 1.9	20.029	+0.10%	0.0081	0.0063	0.0086
   capA-alpha-8	225	8	42.8 +/- 2.8	20.331	+1.61%	0.0298	0.0183	0.0265
   capB-a6 (main)	225	6	58.2 +/- 2.9	20.194	+0.92%	0.0271	0.0117	0.0208
   capB-max-samples-16	16	6	53.7 +/- 2.8	20.021	+0.06%	0.0204	0.0102	0.0173
   capB-alpha-8	225	8	54.5 +/- 2.9	20.562	+2.76%	0.0468	0.0206	0.0340

   Raising $\alpha$ from 6 to 8 roughly triples the WikiText-2 perplexity cost (+0.51% to +1.61% for recipe A, +0.92% to +2.76% for recipe B) and nearly doubles the KL shifts (0.0178 to 0.0298 and 0.0271 to 0.0468 on WikiText; 0.0103 to 0.0183 and 0.0117 to 0.0206 on productive texts). Yet dev256 accuracy does not improve: recipe A scores 42.8% vs 43.0% (paired change -0.2 +/- 2.3 points), and recipe B scores 54.5% vs 58.2% (paired drop -3.7 +/- 1.6 points). Pushing harder along an imperfect direction scales the off-target part of $u$ along with the useful part. Using only 16 loops (capB-max-samples-16) keeps ordinary-text side effects tiny (+0.06% perplexity, 0.0204 WikiText KL) while reaching 53.7% on dev256 (+41.6 points over the unedited model).

   capA-max-samples-16 shows why you should not read this table on its own. It has some of the smallest side effects here (+0.10% perplexity, 0.0081 WikiText KL), yet it reaches only 17.6% on dev256 (5.5 +/- 1.9 points above the unedited model) with 82.6% of answers still cut off. By comparison, all2 has a similar WikiText KL (0.0078) and reaches 49.4%. Small side effects are good news only when the edit also does its job. Read side effects next to the accuracy column. Lab 8 explains why this edit fails.