Weight-Space Programming

This cell loads PyTorch and the lab helpers, fixes the random seed, and prints the GPU this session runs on. The helpers handle model loading, sampling, grading and hooks so the TODO cells can focus on the edit itself.

We load the real Qwen3.5-0.8B in bfloat16. The model has 24 decoder layers, and each MLP ends in a down_proj matrix of shape 1024 by 3584. Every column of that matrix is the 1024-number vector one MLP neuron adds to the residual stream, which is why a column is the natural unit to edit. A vLLM engine holds a second copy of the model for fast sampling, and getting it ready takes about three minutes. The model starts loading as soon as your session starts, so it is usually ready by the time you reach this cell. If it is still loading, the cell waits for it and prints how long loading has taken so far. Every edit you make later is copied into the engine in place, so it never has to load again.

The problems come from GSM8K. The 64 held-out problems are drawn from the test split and are only used to measure the model before and after the edit. The practice set comes from the training split, so you can read extra problems without touching the held-out ones.

The model now answers every held-out problem with thinking turned on and the paper's sampling settings, which are temperature 0.6, top-p 0.95, top-k 20 and a presence penalty of 1.5. Each answer may use up to 8,192 tokens, and the 64 answers take a little over a minute. In our test run only 7 of the 64 answers were correct and 58 ran into the cap. While the batch runs, the output follows one question, problem 12 about fabric for curtains, token by token. It comes back when that answer starts repeating itself, so you can watch the loop form. The edited model answers the same question later in the lab, so you can compare the two. Short progress lines report how many answers have finished, and the summary prints once all 64 are done. Pass live=False to generate_answers if you only want the summary. Read the tail that gets printed at the end and you will usually see the model repeating the same sentence instead of finishing.

Before changing anything we need one number that says how often the model gets stuck. An answer counts as stuck when it runs into the token cap or when the helper flags its ending as a loop. Each result is a dictionary with the keys hit_cap and looping.

A looping answer keeps writing the same phrases. You can catch that by splitting the last few hundred words into overlapping groups of n words and checking how many of those groups appear more than once. A healthy answer has almost no repeats, while a stuck answer is mostly repeats.

To find what drives a loop we compare two versions of the same reasoning. Each pair shares a problem and an opening, then one version loops and the other keeps working. These four pairs are the ones the authors used, taken from their published logs. The helper feeds each text through the model and records the residual stream after layers 15 and 19, averaged over the last 36 tokens, and the input to each down_proj, averaged over the last 24 tokens.

If looping has a direction in the residual stream, the average gap between the repetitive and productive texts should point along it. You will turn that gap into a unit vector u for each layer. The cell then checks that u separates the two sides and compares the two layers.

Not every neuron matters. A column is worth editing when it writes a lot along u and its neuron fires more on the repetitive texts. The helper multiplies those two numbers for each of the 3584 columns and keeps the top 64 per layer, then reports how many match the 64 columns the authors chose. With the authors' four pairs you should share all 64 columns in layer 15 and 63 of 64 in layer 19.

Now we write the edit by hand. For each chosen column we take its component along u and subtract six times that amount, which flips the column so it pushes away from the loop instead of toward it. The math is done in float32 on the original weights and only then copied into the bfloat16 model and the vLLM engine. Columns that were not chosen must stay exactly the same.

The edited model answers the same held-out problems with the same sampling settings, the same seed and the same 8,192-token cap. This run is a little faster because fewer answers loop to the cap. The live view follows the same question as before the edit, problem 12 about fabric for curtains. In our test run it stopped thinking after about 1,800 tokens instead of looping, and the output shows it writing its final answer token by token. The cell prints the first problem that looped before the edit and is now correct, so you can read how its ending changed.

The last step puts the two runs side by side. Accuracy is the share of correct answers, the loop rate is your TODO 1 function, and mean tokens is the average answer length.

The authors published the 64 columns they picked in each layer. This cell keeps your directions and your subtract_direction function but swaps in their columns, runs the held-out problems once more with the same live view and prints the three runs side by side. If your columns match theirs, the last two rows should be close. In our test run their columns reached 0.38 accuracy against 0.33 for your edit, which is within the noise of 64 problems.

The edit lives only in this session. This cell copies the original weights back and confirms that no weight is still changed, so you can try a different alpha or a different number of columns.