File size: 21,187 Bytes
bb102c0 97b4e95 bb102c0 97b4e95 bb102c0 76cc783 bb102c0 76cc783 bb102c0 76cc783 bb102c0 76cc783 bb102c0 6e132b3 bb102c0 76cc783 bb102c0 76cc783 bb102c0 97b4e95 bb102c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | # Weight-space programming labs
**Big thing:** These nine hands-on labs reproduce and take apart Experiment 8 of the MIT CSAIL blog post "Sparks of Weight-Space Programming" on two 8-GPU A100 machines. In that experiment, an AI agent fixed a small language model (Qwen3.5-0.8B) that got stuck repeating itself on math problems by editing 131,072 of its 873,438,784 weights with arithmetic, without training.
**Why it matters:** Most model fixes change the prompt or sampling settings, or run gradient descent (training that adjusts weights to reduce errors). Here you do four things:
- Inspect the `131,072` edited weights (`130,188` change after bf16 rounding, which stores numbers at reduced precision).
- Rebuild the edit from 8 short texts in one minute.
- Prove that the weight edit equals an input-driven steering vector (a change to the model's internal numbers whose strength depends on the input).
- Build your own edit from the model's failures.
## The story
Weight-space programming changes what a trained neural network does by reading and editing its weight matrices with arithmetic instead of gradient training.
In Experiment 8 (2 Nov 2026), an AI agent (GPT-6 Astra) got a CPU-only Python session with Qwen3.5-0.8B, an open model with about 0.87 billion weights across 24 repeated layers (blocks). It had practice math problems with answers and a budget of 50 forward passes. A forward pass runs inputs through the network to read internal activations (numbers computed as the input passes through each layer) or generate answers. One pass could answer up to 512 problems or read activations on up to 1,024 texts, and all training was forbidden.
With thinking on, the model writes a `<think>...</think>` scratchpad before its final answer. In this mode, the unedited model kept falling into verbatim loops, repeating the same block of tokens (word pieces) over and over until it hit the 32,768-token cap: 85 of its first 96 practice answers hit that limit.
The agent wrote 4 productive and 4 looping scratchpads and compared the model's activations on them. That contrast gave a unit direction vector $u$: a list of 1,024 numbers whose overall length is 1, pointing from productive text toward looping text. It then applied a rank-one projection edit (one column vector times one row vector) in blocks 15 and 19. Each block's MLP (multi-layer perceptron, the block's second half) has an output matrix `down_proj` of shape 1,024 by 3,584. Each column is the 1,024-number vector one neuron adds to the residual stream (the running list of numbers each token carries from block to block) per unit of activation. The agent picked the 64 neurons that push hardest along $u$ during loops and updated their columns $W_S$ via $W' = W - \alpha u (u^T W_S)$ with strength $\alpha = 6$. That edit touches $2 \times 64 \times 1,024 = 131,072$ weights (`p50c`, from the 50-pass run folder `runs/p50c`).
On three unseen benchmarks, the paper reported these results for `p50c`:
- **MATH-500:** Accuracy rose from 7.9% to 48.1%. Answers cut at the cap fell from 93.0% to 33.05%, and mean tokens fell from 30,868 to 22,037.
- **GSM8K:** Accuracy rose from 13.84% to 67.1%. Answers cut at the cap fell from 87.5% to 7.24%, and mean tokens fell from 29,289 to 11,993.
- **IFEval:** Instruction-following accuracy rose from 21.0% to 62.3%. Answers cut at the cap fell from 74.9% to 2.2%, and mean tokens fell from 25,296 to 4,241.
## What we found
- **The benchmark gains reproduce on our hardware** ([`labs/lab2_reproduce.md`](labs/lab2_reproduce.md), [`labs/lab1_diff.md`](labs/lab1_diff.md)). Diffing all 488 weight tensors (arrays of weights) confirms that only `layers.15.mlp.down_proj.weight` and `layers.19.mlp.down_proj.weight` changed, in 64 columns each, with fitted $\alpha = 6.0003$ and $6.0002$. On our GPUs, `p50c` raises MATH-500 from $7.3 \pm 0.9\%$ to $51.8 \pm 1.9\%$, GSM8K from $14.2 \pm 0.8\%$ to $66.9 \pm 1.2\%$, and IFEval from $21.2 \pm 1.4\%$ to $63.4 \pm 1.8\%$.
- **The 5-minute quick run shows the same effect on 16 problems** ([`QUICKSTART.md`](QUICKSTART.md), [`labs/lab2_reproduce.md`](labs/lab2_reproduce.md)). On the first 16 GSM8K problems of our 256-problem practice test set (`dev256`), the unedited model scores $18.8 \pm 7.7\%$ and `p50c` scores $65.6 \pm 11.8\%$ ($+46.9 \pm 11.6$ percentage points). Answers cut at the 32,768-token cap fall from 27 of 32 (all ending in a verbatim loop) to 3 of 32 (0 in a loop).
- **Rebuilding the edit from the agent's 8 texts recovers all 64 neurons in both layers** ([`labs/lab3_rebuild.md`](labs/lab3_rebuild.md)). Teacher-forcing the agent's 4 productive and 4 looping texts (feeding fixed text through the model and recording activations with hooks, small functions attached to layers) shows that the MLPs of blocks 15 and 19 widen the looping-minus-productive gap the most. The rebuilt edit recovers directions matching the released $u$ to cosine $+0.99998$ and $-0.99998$ (nearly the same line, with opposite orientation at the second layer), picks the same 64 neurons in both layers, and scores $61.5 \pm 2.8\%$ on `dev256` (matching `p50c`'s $61.5 \pm 2.8\%$). Using only the first 2 text pairs still reaches $62.3 \pm 2.7\%$.
- **Accuracy rises with $\alpha$ up to a plateau from 4 to 8, while random controls move accuracy by about +2 points** ([`labs/lab4_dose.md`](labs/lab4_dose.md), [`labs/lab7_side_effects.md`](labs/lab7_side_effects.md)). On `dev256`, accuracy goes from $12.1\%$ ($\alpha = 0$, unedited) to $22.7\%$ ($\alpha = 1$, erase), $35.2\%$ ($\alpha = 2$, mirror), $59.2\%$ ($\alpha = 4$), $61.5\%$ ($\alpha = 6$), and $60.7\%$ ($\alpha = 8$). Accuracy then drops to $51.8\%$ at $\alpha = 12$, and WikiText-2 perplexity (prediction uncertainty on ordinary text, lower is better) rises by $+14.34\%$, versus $+0.79\%$ at $\alpha = 6$. Matched-norm controls with a random orthogonal direction (`randdir`) or 64 random neurons (`randcols`) score $14.5\%$ ($+2.3 \pm 1.6$ points) and $14.6\%$ ($+2.5 \pm 1.6$ points).
- **Setting $\alpha = -4$ gains +7.6 points for the wrong reason** ([`labs/lab4_dose.md`](labs/lab4_dose.md)). Amplifying the 64 neurons' push toward the loop direction 5 times raises `dev256` accuracy from $12.1\%$ to $19.7 \pm 1.9\%$ ($+7.6 \pm 2.2$ points) because 197 of 512 answers abruptly emit `</think>` out of a loop. Yet it loses 24 problems the unedited model solved, drops accuracy among finished answers from $100.0\%$ to $55.4\%$, and raises WikiText-2 perplexity by $+5.39\%$.
- **Turning thinking off reaches 52.1% using 13 times fewer tokens, and the weight edit adds nothing there** ([`labs/lab5_decoding.md`](labs/lab5_decoding.md)). With thinking on, setting temperature to 1.0 ($27.9\%$) or repetition penalty to 1.1 ($34.4\%$) falls well short of the weight edit ($61.5\%$), and dropping the 1.5 presence penalty cuts accuracy to $4.5\%$, with $99.0\%$ of answers cut. Without thinking, the base model reaches $52.1 \pm 2.7\%$ on `dev256` with 2,270 tokens per answer on average, versus 29,683 with thinking on (16,104 with the edit). Adding `p50c` with thinking off scores $53.5 \pm 2.8\%$ ($+1.4 \pm 1.9$ points), because answers rarely loop without thinking.
- **The static weight edit is algebraically equal to token-by-token activation steering** ([`labs/lab6_steering.md`](labs/lab6_steering.md), [`labs/lab0_toy.md`](labs/lab0_toy.md)). Multiplying the edited matrix by neuron activations $a(t)$ gives $W' a(t) = W a(t) - \alpha s(t) u$, where $s(t) = u \cdot (W_S a_S(t))$ is one number per token measuring how hard the 64 chosen neurons push along $u$. In 64-bit floats on one layer, the matrix edit and a runtime steering hook match to $3.33 \times 10^{-16}$. On 48 real looping answers from the base model, $s(t)$ stays low early in the answer (0.096 at layer 15, 0.116 at layer 19) and jumps inside the loop (0.424 and 0.758), so the edit pushes back hardest inside loops.
- **An edit built from the model's own loops reaches 58.2% without hand-written texts** ([`labs/lab8_capstone.md`](labs/lab8_capstone.md)). On 225 looping answers the unedited model wrote for 128 separate practice problems (`build128`), comparing loop tokens with tokens well before the loop (recipe B) matches $u$ to cosine 0.782 and 0.773 and scores $58.2 \pm 2.9\%$ on `dev256` ($+46.1 \pm 2.6$ points). Comparing with the 32 tokens just before the loop starts (recipe A) reaches only $43.0 \pm 2.6\%$ because pre-loop text has already drifted along $u$.
## Start here
Start with [**`QUICKSTART.md`**](QUICKSTART.md) for a one-page, 2-hour path through the labs with exact commands and expected numbers.
The Quickstart also shows how to download the code from this site. Lab 0 runs on any laptop. The other labs need NVIDIA GPUs; we used two machines with 8 A100 80 GB GPUs each, and most steps use 1 or 2 of them.
## The nine labs
```mermaid
flowchart LR
L0["Lab 0: Toy 4x6 matrix"] --> L1["Lab 1: Diff weight files"]
L1 --> L2["Lab 2: Reproduce benchmarks"]
L1 --> L3["Lab 3: Rebuild from 8 texts"]
L3 --> L4["Lab 4: Dose curve and controls"]
L2 --> L5["Lab 5: Decoding settings"]
L3 --> L6["Lab 6: Edit equals steering"]
L4 --> L7["Lab 7: Side effects on text"]
L6 --> L8["Lab 8: Capstone from own loops"]
L7 --> L8
```
| Lab | Title and sheet | Question it answers | What you run | Time and hardware |
|---|---|---|---|---|
| 0 | [The toy edit on a 4 x 6 matrix](labs/lab0_toy.md) | What does $W' = W - \alpha u u^T W_S$ do on a matrix small enough to print? | `uv run --with numpy --with matplotlib python labs/lab0_toy.py` | 5 seconds on any laptop CPU |
| 1 | [Find the edit by diffing the two weight files](labs/lab1_diff.md) | What did the agent change in the 488 weight tensors? | `python labs/lab1_diff.py` | Under 1 minute on VM CPU |
| 2 | [Reproduce the paper's benchmark numbers](labs/lab2_reproduce.md) | Do the looping failure and the `p50c` fix reproduce on MATH-500, GSM8K, and IFEval? | `python scripts/run_jobs.py plans/trainer.json --gpus 0-7` (or `plans/quick.json --gpus 0,1`) | 5 minutes on 2 A100s (quick run) or 1.5 hours on 8 A100s (`trainer`) |
| 3 | [Rebuild the edit from the agent's 8 texts](labs/lab3_rebuild.md) | Can we recover layers 15 and 19, $u$, and the 64 neurons from the agent's 8 texts with hooks? | `CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py` | About 1 minute on 1 A100 GPU (plus 21 minutes on 1 A100 to test on `dev256`) |
| 4 | [Turn the alpha knob and check controls](labs/lab4_dose.md) | How does accuracy change across $\alpha \in \{-4, 1, 2, 4, 6, 8, 12\}$, single layers, all neurons, and random controls? | `python scripts/run_jobs.py plans/sampler.json --gpus 0-6` and `python labs/lab4_endings.py` | 17 to 45 minutes per run on 1 A100 (`sampler`), seconds on CPU for endings |
| 5 | [Can decoding settings do the same job?](labs/lab5_decoding.md) | Can temperature, presence or repetition penalties, a 4,096-token cap, or turning thinking off replace the weight edit? | `python scripts/run_jobs.py plans/sampler.json --gpus 0-6` (decoding runs) | About 3 to 49 minutes per run on 1 A100 (`sampler`) |
| 6 | [The weight edit is steering in disguise](labs/lab6_steering.md) | Is the weight edit identical to subtracting $\alpha s(t) u$, and when does $s(t)$ fire on real loops? | `CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py` (and `--parts 3,4`) | About 1 minute on 1 A100 for parts 1 and 2; about 3 minutes for parts 3 and 4 (`sampler`) |
| 7 | [What else did the edit change?](labs/lab7_side_effects.md) | How much does each edit move WikiText-2 perplexity and KL divergence on productive versus looping text? | `CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py` | About 1.5 minutes (78 seconds) on 1 A100 GPU |
| 8 | [Capstone: build your own edit from the model's own loops](labs/lab8_capstone.md) | Can we extract $u$ and 64 neurons from the base model's own looping answers and fix unseen problems? | `CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py` | About 5 minutes on 1 A100 to build, 30 minutes on 2 A100s to test (`sampler`) |
## Project layout
- `configs/`: JSON definitions of edited models (`variants.json` for Lab 4 controls and $\alpha$ sweep, `rebuilt.json` from Lab 3, `capstone.json` from Lab 8).
- `data/`: Prepared JSONL datasets (`dev256.jsonl` with 256 unseen practice problems, `build128.jsonl` with 128 practice problems for Lab 8, plus `math500.jsonl`, `gsm8k.jsonl`, and `ifeval.jsonl`).
- `external/`: Local checkout of the paper's log repository (`lm_understand-logs`, ignored by `.gitignore`; `exp8_math/` holds the agent's scripts, edit manifest, and work notes).
- `labs/`: Lab scripts (`lab0_toy.py` through `lab8_capstone.py`), answer inspection tools (`peek.py`, `lab4_endings.py`), lab sheets (`lab0_toy.md` through `lab8_capstone.md`), and the agent's 8 contrastive texts (`labs/data/astra_pairs.json`).
- `plans/`: Generation job plans for `scripts/run_jobs.py` (`quick.json`, `trainer.json`, `sampler.json`, `extra.json`, and `lab3-tries.json`).
- `reference/`: The first-principles companion guide [`reference/theory-guide.md`](reference/theory-guide.md).
- `results/`: Saved JSON outputs for Labs 1 through 8 (including committed `results/*/try-*` folders to compare your runs with ours), per-run graded records in `results/runs/`, the master table [`results/summary.md`](results/summary.md) (`results/summary.json`), and all plots in `results/figures/*.png`.
- `scripts/`: VM setup and sync scripts (`setup_vm.sh`, `sync_vm.sh`, `pull_runs.sh`, `env.sh`), dataset and model builders (`prepare_data.py`, `build_models.py`), the multi-GPU vLLM runner (`run_jobs.py`, `gen_worker.py`), and reporting scripts (`summarize.py`, `plot_runs.py`).
- `wsp/`: Shared Python library (`hf.py` for model loading and PyTorch hooks, `edits.py` for rank-one edits and SVD recovery, `loops.py` for verbatim loop detection, `protocol.py` for presets and sampler probabilities, `grading.py`, and `summary.py`).
## Running things
```mermaid
flowchart LR
Trainer["trainer (8x A100)<br/>/nvme/wsp-lab/runs/<br/>bench/*, quick/*, lab3 tries"] <-->|"sync_vm.sh / pull_runs.sh"| Laptop["This machine (CPU)<br/>Project root and results/<br/>Lab 0 and plotting"]
Laptop <-->|"sync_vm.sh / pull_runs.sh"| Sampler["sampler (8x A100)<br/>/nvme/wsp-lab/runs/<br/>build/base, dev/*"]
```
**Cloning the paper's logs.** `.gitignore` ignores `external/`, so the paper's log repository is not committed here. If `external/lm_understand-logs` is missing, clone it from the project root before syncing or building models:
```bash
git clone https://github.com/RikiyaT/lm_understand-logs external/lm_understand-logs
git -C external/lm_understand-logs log -1
```
We used commit `4d09ba6d` (29 October 2026). `scripts/prepare_data.py` and `scripts/build_models.py` read from `external/lm_understand-logs/exp8_math`, and `scripts/sync_vm.sh` copies the parts the GPU machines need (`code`, `practice`, and `runs/p50c`).
**The two GPU VMs and environment setup.** GPU work runs on `trainer` and `sampler`, each with 8 NVIDIA A100 80 GB GPUs. All VM files live under `/nvme/wsp-lab` (`venv/`, `models/`, `data/`, `runs/`, `logs/`, and `code/`). Start every VM command by sourcing `env.sh`, which activates Python 3.12 (vLLM 0.33.0, PyTorch 2.14.0, transformers 5.20.2, flash-linear-attention 0.6.0) and keeps caches inside `/nvme/wsp-lab`:
```bash
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
```
**Syncing code and pulling results.** Files move between each VM and this machine. From the project root on this machine:
```bash
bash scripts/sync_vm.sh trainer sampler # push wsp, scripts, configs, plans, labs, and exp8_math files
bash scripts/pull_runs.sh trainer sampler # pull merged.jsonl, summary.json, and *.done.json into results/runs/
python3 scripts/summarize.py results/runs # rebuild results/summary.md and results/summary.json
uv run --with matplotlib python scripts/plot_runs.py
```
`sync_vm.sh` mirrors the five code folders onto the VM with `rsync --delete`, so any file you create or edit only on the VM is deleted or overwritten on the next sync. Edit on this machine first, then sync. The one exception is `configs/capstone.json`, which Lab 8 writes on the VM: syncing never deletes it, but a local copy will overwrite the VM's copy, so pull the VM version first. `pull_runs.sh` copies only run records. Lab scripts write their outputs to `results/labN/` and `results/figures/` on the VM where they ran; copy those back by name:
```bash
rsync -az sampler:/nvme/wsp-lab/code/results/lab6/ results/lab6/ </dev/null
rsync -az sampler:/nvme/wsp-lab/code/results/figures/lab6_real_loop.png results/figures/ </dev/null
```
Copy figures individually instead of syncing the whole folder, since each VM holds only the figures it generated and an older copy would overwrite a newer one.
**Setting up a new machine.** To provision a fresh GPU machine (about 10 minutes and 15 GB under `/nvme/wsp-lab`, no root needed), run these commands from the project root (see `scripts/setup_vm.sh`):
```bash
scp scripts/setup_vm.sh <host>:/tmp/ && ssh <host> bash /tmp/setup_vm.sh
bash scripts/sync_vm.sh <host>
ssh <host> 'cd /nvme/wsp-lab/code && . scripts/env.sh &&
python scripts/prepare_data.py --exp8 external/lm_understand-logs/exp8_math --out /nvme/wsp-lab/data &&
python scripts/build_models.py && python scripts/build_models.py --variants configs/rebuilt.json'
```
**Where full answer texts live.** Full answer texts (`runs/<name>/shard*.texts.jsonl.gz`) are ignored by `.gitignore` and skipped by `scripts/pull_runs.sh`, so they stay on the VM that generated them (while `results/*/try-*` folders are committed to git). Because `labs/peek.py` and `labs/lab4_endings.py` read `.texts.jsonl.gz` files, run them on the VM that produced the run:
- On `trainer`: `bench/*`, `quick/*`, `dev/rebuilt-pairs2`, `dev/rebuilt-windows36-36`, and `dev/rebuilt-layers15-19-21`.
- On `sampler`: `build/base` and every other `dev/*` run (`dev/base`, `dev/p50c`, `dev/rebuilt`, the Lab 4 dose and control runs, the Lab 5 decoding runs, and the Lab 8 `dev/cap*` runs).
**Rerunning a finished job.** `scripts/run_jobs.py` skips any shard (a batch of problems) that has a `shard*.done.json` file, so rerunning a finished plan reports `0 shards to run`. To run a finished job again, move its output folder aside on the VM first:
```bash
mv /nvme/wsp-lab/runs/quick /nvme/wsp-lab/runs/quick-ours
```
Then run the plan; `python labs/peek.py quick/base` reads the new run and `python labs/peek.py quick-ours/base` reads the earlier one.
**How to read `results/summary.md`.** Every generation run has one row in [`results/summary.md`](results/summary.md) with these nine columns:
- `run`: Run identifier (`bench/<model>/<dataset>`, `dev/<variant>`, `build/base`, or `quick/<model>`).
- `problems x samples`: Distinct problems times sampled answers per problem (for example `256x2` = 512 answers on `dev256`).
- `accuracy %`: Percentage of correct answers, plus or minus one standard error (SE, typical sampling uncertainty; a gap of 2 SE or more is unlikely to be chance).
- `cut at the cap %`: Percentage of answers that hit the token limit (32,768 by default) without finishing.
- `mean tokens`: Average generated tokens per answer.
- `correct when finished %`: Accuracy among answers that finished before the token limit (97.5% to 100% for the unedited model on math benchmarks, showing its main failure is not finishing).
- `cut answers ending in repetition %`: Percentage of cap-cut answers ending in repetition. An answer counts if it ends in a word-for-word loop (a block of words repeated at least 3 times in a row) or in looser circling. Looser circling means at least half of the 4-word phrase positions in the last 200 words hold a phrase repeated 3 or more times there (see `results/summary.md`).
- `gain vs baseline (pp)`: Accuracy change in percentage points versus the paired baseline on the same problems, plus or minus paired SE (`dev/*` against `dev/base`, `dev/p50c+nothink` against `dev/base+nothink`, `bench/p50c/X` against `bench/base/X`, and `quick/*` against `quick/base`).
- `rescued / lost problems`: `rescued` counts problems where the run got at least one answer right and the baseline got none; `lost` counts the reverse.
## Sources
- **Blog post:** Rikiya Takahashi, ["Sparks of Weight-Space Programming"](https://rikiyat.github.io/blog/weight-space-programming/) (MIT CSAIL, 2 Nov 2026).
- **Experiment logs and code:** [`https://github.com/RikiyaT/lm_understand-logs`](https://github.com/RikiyaT/lm_understand-logs) (local copy in `external/lm_understand-logs/`, with Experiment 8 in `external/lm_understand-logs/exp8_math/`).
- **Released edited model:** [`https://huggingface.co/RikiyaT/Qwen3.5-0.8B-weight-space-programming`](https://huggingface.co/RikiyaT/Qwen3.5-0.8B-weight-space-programming).
- **Interactive demo Space:** [`https://huggingface.co/spaces/ML-Intern-lab/weight-space-programming-demo`](https://huggingface.co/spaces/ML-Intern-lab/weight-space-programming-demo).
- **First-principles theory guide:** [`reference/theory-guide.md`](reference/theory-guide.md).
|