Weight-space programming lab

Quick start: a 2-hour path through the labs

On this page

Big thing: Run the core weight-space programming checks in 2 hours (about 10 minutes of GPU compute), from a 5-second toy matrix on this machine to the full benchmark, dose, decoding, and capstone results.

Get the code first. This site is a public copy of the project. On your own machine, the first line below installs uv, which runs the lab's Python scripts with the packages they need and can fetch Python itself. The second line puts uv on your PATH (the folders the shell searches for commands) in this terminal. The uvx command then runs hf, the download tool from huggingface_hub, without installing it, and fetches the code/ folder:

bash
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env
uvx --from huggingface_hub hf download thegovind/weight-space-programming \
  --repo-type space --include "code/*" --local-dir ~/weight-space-programming
cd ~/weight-space-programming/code

If you already have uv, skip the first two lines. Commands for this machine run from that code folder. Commands on the GPU machines run from /nvme/wsp-lab/code, where the setup copies the code.

Machines. Step 1 runs on any laptop. The other steps ran on two GPU machines, each with 8 NVIDIA A100 80 GB GPUs, and most steps use only 1 or 2 of them. These pages call the two machines trainer and sampler, so use your own host names wherever you see those. To set up a machine, follow "Cloning the paper's logs" and then "Setting up a new machine", both under Running things. The setup puts everything under /nvme/wsp-lab (about 15 GB) and needs no root access.

Data. The problem files in data/ are not included. data/README.md explains why and shows how to rebuild them.

Before you start (VM access and environment):

  1. SSH hosts: Connect with ssh trainer or ssh sampler (8x A100 80 GB each).
  2. Environment line (run first in every VM shell): . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
  3. Which terminal: Steps 1 and 7 run on this machine. Step 2 runs on trainer and ends with exit. Steps 3 to 6 share one sampler terminal, and Step 6 ends with exit. Always ssh to each VM from this machine.

Step 1: Run the toy edit on a 4 x 6 matrix (15 minutes)

  • Where: This machine (CPU, 5 seconds to run, 15 minutes with the sheet).
bash
cd ~/weight-space-programming/code
uv run --with numpy --with matplotlib python labs/lab0_toy.py
  • What to look at: Check the 4 numbered outputs (singular values of W' - W: [3.84 0. 0. 0.] and |W'a - Wa| = 0.000 when the chosen neurons are quiet vs 4.261 when firing) and open results/figures/lab0_toy.png. Singular values measure the sizes of independent pieces of a matrix change; only one is nonzero here. Read labs/lab0_toy.md for why the edit is rank one (one column vector times one row vector) and acts like steering (adding a vector to the model's internal state).

Step 2: Run the 5-minute quick test and inspect a looping answer (15 minutes)

  • Where: trainer (2 GPUs for 5 minutes, then CPU for peek.py). If finished answers for quick/base (the unedited model) and quick/p50c (the released edit) are already on the machine, the run skips them. To have vLLM (the program that generates the answers) write fresh ones, move /nvme/wsp-lab/runs/quick aside first.
bash
ssh trainer
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
nvidia-smi   # check which GPUs are free
mv /nvme/wsp-lab/runs/quick /nvme/wsp-lab/runs/quick-ours   # optional: only if generating fresh answers
python scripts/run_jobs.py plans/quick.json --gpus 0,1
CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/base
CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/p50c
CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/base --id gsm8k_train_5628
CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/p50c --id gsm8k_train_5628
exit   # back to this machine
  • What to look at: On our run, quick/base gets 6 of 32 right (18.8%) with 27 of 32 cut in a verbatim loop, while quick/p50c gets 21 of 32 (65.6%, +46.9 points) with 3 cut and 0 loops. On gsm8k_train_5628 (truth 240), quick/base repeats an 83-token block about 292 times from token 8,568 to 32,768, while quick/p50c finishes in 5,825 tokens after closing </think> and writing \boxed{240}. See labs/lab2_reproduce.md.

Step 3: Find the edit in the weight files on CPU (15 minutes)

  • Where: sampler, from this machine (CPU only, under 1 minute). Keep this terminal open for Steps 4 to 6.
bash
ssh sampler
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
python labs/lab1_diff.py
  • What to look at: Confirm that only 2 of 488 tensors (arrays of weights) differ (layers.15 and layers.19 mlp.down_proj.weight), and that all changes sit in 64 columns per layer. The first singular value holds 0.999994 and 0.999991 of the squared change, and the fitted α\alpha (edit strength) is 6.0003 and 6.0002. These paired numbers refer to the two layers in the order listed. Read labs/lab1_diff.md and inspect results/figures/lab1_diff.png.

Step 4: Rebuild the edit from the agent's 8 texts (20 minutes)

  • Where: the same sampler terminal (1 GPU, about 1 minute). Run this before Step 5 because Lab 6 reads which way uu (the direction from productive toward looping text in the model's internal state) points from results/lab3/rebuild.json.
bash
CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py
  • What to look at: Verify that blocks 15 and 19 have the largest MLP gap growth (0.263 and 0.249), the rebuilt directions match the released uu to cosine +0.99998 and -0.99998, and the top 64 neuron scores match 64/64 released columns in both layers. Read labs/lab3_rebuild.md and check results/figures/lab3_where_the_loop_grows.png and lab3_neuron_scores.png.

Step 5: Verify that the weight edit is activation steering (15 minutes)

  • Where: the same sampler terminal (1 GPU, about 1 minute for parts 1 and 2). Run --parts 3,4 (about 3 minutes) to trace real loops and correct finished answers. A loop is a stretch of an answer where the model repeats itself. Both come from build/base, the unedited model's answers to 128 practice problems, in /nvme/wsp-lab/runs/build/base on sampler. In our run, part 3 traced 48 loops and part 4 traced 26 correct finished answers, and a new run may give different counts. Parts 3 and 4 need that build/base run. The optional run_jobs.py line below makes it if it is missing (about 24 minutes on 1 GPU). If the run is already there, the line prints 0 shards to run on GPUs [7]: and then ALL DONE.
bash
CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py
python scripts/run_jobs.py plans/sampler.json --only build/base --gpus 7   # optional: makes build/base if missing
CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 3,4   # optional
  • What to look at: Editing down_proj vs subtracting αs(t)u\alpha s(t) u in a hook (a function attached to a layer) matches to 3.33e-16 in float64 (1.4e-8 mean KL in float32). Float64 and float32 are number formats with different precision; KL measures the change in next-token probabilities. Over the last 36 tokens of the agent's 8 texts, s(t)s(t) (how hard the chosen neurons push toward looping on each token) is near zero on productive text (0.00 to 0.04 at layer 15) and positive on looping text (0.19 to 0.42). Parts 3 and 4 show the same split on real answers: at layer 19, s(t)s(t) averages 0.76 inside loops and about 0.1 on answers that went well (small, not zero). Read labs/lab6_steering.md, which covers results/figures/lab6_real_loop.png and lab6_finished_vs_loops.png.

Step 6: Measure side effects on ordinary text and selective KL (15 minutes)

  • Where: the same sampler terminal (1 GPU, about 1.5 minutes). Then leave the VM.
bash
CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py
exit   # back to this machine
  • What to look at: p50c and rebuilt raise WikiText-2 perplexity (prediction uncertainty on ordinary text, where lower is better and the base is 20.009) by +0.79%. At the same time, they shift predictions 2.5 times more on looping texts than on productive ones (0.0241 vs 0.0097, measured by KL divergence, the change in next-token probabilities). By contrast, a12 (a stronger edit) costs +14.34%. The random-direction and random-neuron controls randdir/randcols shift looping text no more than productive text (randdir: 0.0039 vs 0.0048). Read labs/lab7_side_effects.md and view results/figures/lab7_side_effects.png.

Step 7: Read the finished results of Labs 2, 4, 5, and 8 (30 minutes)

bash
cd ~/weight-space-programming/code
ls results/figures/
  • What to look at: Read these four lab pages and their figures:

    • labs/lab2_reproduce.md (lab2_benchmarks.png, lab2_lengths.png).
    • labs/lab4_dose.md (lab4_dose.png, lab4_controls.png, including why α=−4\alpha = -4 gains +7.6 points by breaking loops).
    • labs/lab5_decoding.md (lab5_decoding.png, where turning thinking off reaches 52.1% at 2,270 tokens).
    • labs/lab8_capstone.md (lab8_loop_onset.png, lab8_data_needed.png, lab8_one_answer.png, where recipe B built from 225 of the model's own loops reaches 58.2% on dev256).

    To re-run any sheet's GPU jobs, use the move-aside rule in README.md.