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:
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):
- SSH hosts: Connect with
ssh trainerorssh sampler(8x A100 80 GB each).- Environment line (run first in every VM shell):
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code- Which terminal: Steps 1 and 7 run on this machine. Step 2 runs on
trainerand ends withexit. Steps 3 to 6 share onesamplerterminal, and Step 6 ends withexit. Alwayssshto 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).
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.000when the chosen neurons are quiet vs4.261when firing) and openresults/figures/lab0_toy.png. Singular values measure the sizes of independent pieces of a matrix change; only one is nonzero here. Readlabs/lab0_toy.mdfor 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 forpeek.py). If finished answers forquick/base(the unedited model) andquick/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/quickaside first.
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/basegets 6 of 32 right (18.8%) with 27 of 32 cut in a verbatim loop, whilequick/p50cgets 21 of 32 (65.6%,+46.9points) with 3 cut and 0 loops. Ongsm8k_train_5628(truth240),quick/baserepeats an 83-token block about292times from token 8,568 to 32,768, whilequick/p50cfinishes in5,825tokens after closing</think>and writing\boxed{240}. Seelabs/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.
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.15andlayers.19mlp.down_proj.weight), and that all changes sit in 64 columns per layer. The first singular value holds0.999994and0.999991of the squared change, and the fitted (edit strength) is6.0003and6.0002. These paired numbers refer to the two layers in the order listed. Readlabs/lab1_diff.mdand inspectresults/figures/lab1_diff.png.
Step 4: Rebuild the edit from the agent's 8 texts (20 minutes)
- Where: the same
samplerterminal (1 GPU, about 1 minute). Run this before Step 5 because Lab 6 reads which way (the direction from productive toward looping text in the model's internal state) points fromresults/lab3/rebuild.json.
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.263and0.249), the rebuilt directions match the released to cosine+0.99998and-0.99998, and the top 64 neuron scores match64/64released columns in both layers. Readlabs/lab3_rebuild.mdand checkresults/figures/lab3_where_the_loop_grows.pngandlab3_neuron_scores.png.
Step 5: Verify that the weight edit is activation steering (15 minutes)
- Where: the same
samplerterminal (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 frombuild/base, the unedited model's answers to 128 practice problems, in/nvme/wsp-lab/runs/build/baseonsampler. 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 thatbuild/baserun. The optionalrun_jobs.pyline below makes it if it is missing (about 24 minutes on 1 GPU). If the run is already there, the line prints0 shards to run on GPUs [7]:and thenALL DONE.
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_projvs subtracting in a hook (a function attached to a layer) matches to3.33e-16in float64 (1.4e-8mean 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, (how hard the chosen neurons push toward looping on each token) is near zero on productive text (0.00to0.04at layer 15) and positive on looping text (0.19to0.42). Parts 3 and 4 show the same split on real answers: at layer 19, averages0.76inside loops and about0.1on answers that went well (small, not zero). Readlabs/lab6_steering.md, which coversresults/figures/lab6_real_loop.pngandlab6_finished_vs_loops.png.
Step 6: Measure side effects on ordinary text and selective KL (15 minutes)
- Where: the same
samplerterminal (1 GPU, about 1.5 minutes). Then leave the VM.
CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py
exit # back to this machine
- What to look at:
p50candrebuiltraise WikiText-2 perplexity (prediction uncertainty on ordinary text, where lower is better and the base is20.009) by+0.79%. At the same time, they shift predictions 2.5 times more on looping texts than on productive ones (0.0241vs0.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 controlsranddir/randcolsshift looping text no more than productive text (randdir:0.0039vs0.0048). Readlabs/lab7_side_effects.mdand viewresults/figures/lab7_side_effects.png.
Step 7: Read the finished results of Labs 2, 4, 5, and 8 (30 minutes)
- Where: This machine (no GPU;
results/summary.mdandresults/figures/hold all runs).
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 gains+7.6points by breaking loops).labs/lab5_decoding.md(lab5_decoding.png, where turning thinking off reaches52.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 reaches58.2%ondev256).
To re-run any sheet's GPU jobs, use the move-aside rule in
README.md.