Lab 0: The toy edit on a 4 x 6 matrix

Big thing: You will run the paper's weight edit W' = W - alpha * u (u^T W_S) on a 4 x 6 matrix that fits on one screen. The edit changes how hard two chosen neurons push along one direction, leaves everything else alone, and the script checks each claim with assert.

Why it matters: In Qwen3.5-0.8B the same matrix is 1,024 x 3,584, too large to read by eye, but the linear algebra is identical.

Time and hardware: 5 seconds on any laptop CPU.

Theory guide: Sections 6 and 8.

Before you run: predict#

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

1. In the toy layer, neuron 1 pushes along u by -0.555 per unit of activation. With edit strength alpha = 1, what does neuron 1's push along u become? What if alpha = 6?
2. The toy matrix has 6 columns (one per neuron), and the edit touches only neurons S = [1, 4]. How many columns of W' - W will be non-zero?
3. On a token where neurons 1 and 4 have activation 0.0 and the other four neurons fire normally, how much does the edited output W' a differ from W a?

Run it#

Run this from the project root (uv creates a temporary environment with numpy and matplotlib):

uv run --with numpy --with matplotlib python labs/lab0_toy.py

The script prints four checks, verifies each with assert, and writes results/figures/lab0_toy.png.

What you should see#

1. How much each neuron writes along u (u . W[:, j]): [ 1.062 -0.555  1.043 -0.495  0.318  0.162]

2. The edit scales the chosen neurons' push along u by (1 - alpha) and leaves everything else alone
   alpha= 0 (no edit                        ) push along u of neurons [1, 4]: [-0.555  0.318] -> [-0.555  0.318]
   alpha= 1 (erase                          ) push along u of neurons [1, 4]: [-0.555  0.318] -> [ 0. -0.]
   alpha= 2 (mirror                         ) push along u of neurons [1, 4]: [-0.555  0.318] -> [ 0.555 -0.318]
   alpha= 6 (flip and amplify 5x (the paper)) push along u of neurons [1, 4]: [-0.555  0.318] -> [ 2.776 -1.592]
   alpha=-4 (amplify 5x the other way       ) push along u of neurons [1, 4]: [-0.555  0.318] -> [-2.776  1.592]

3. The change is rank one: one direction out (u), one pattern in (u^T W on the chosen columns)
   singular values of W' - W: [3.84 0.   0.   0.  ]

4. The weight edit is steering: W' a = W a - alpha * s * u, with s = u . (W[:, S] a[S])
   chosen neurons quiet   s = +0.000   change to the output |W'a - Wa| = 0.000
   chosen neurons firing  s = -0.710   change to the output |W'a - Wa| = 4.261
   Same weights, different tokens, different push: the edit only acts when the chosen neurons fire.

all checks passed; figure in results/figures/lab0_toy.png

On the left of lab0_toy.png, the W' - W heatmap has four white columns at neurons 0, 2, 3, and 5 (exact zeros) and two colored columns at neurons 1 and 4. In every row, column 4 has the opposite color of column 1. Neuron 1 started with a negative push along u (-0.555) and neuron 4 started with a positive push (0.318), so the edit moves them in opposite directions along u. On the right, the blue line plots the multiplier 1 - alpha: 1 at alpha = 0 (no edit), 0 at alpha = 1 (erase), -1 at alpha = 2 (mirror), and -5 at alpha = 6 (the paper's setting).

Why it works#

1. Neurons write columns into the residual stream#

A language model processes text one token (word piece) at a time. Across the 24 blocks (layers) of Qwen3.5-0.8B, the model updates a running list of numbers for each token called the residual stream. That stream holds 1,024 numbers per token in Qwen3.5-0.8B and d = 4 in lab0_toy.py.

Each block ends with an MLP (multi-layer perceptron). Inside it, two input matrices (gate_proj and up_proj) and a SiLU nonlinearity turn the stream into n neuron activations a[0], ..., a[n-1] (n = 3,584 in Qwen3.5-0.8B, n = 6 in the toy). Each activation a[j] is a single number measuring how strongly neuron j fires on the current token. The MLP's last matrix, down_proj, is W of shape d x n. Column W[:, j] is the d-number vector that neuron j adds to the stream when it fires with activation 1. The MLP output is the weighted sum of its columns:

2. Projecting a column onto the direction u#

A direction u is a residual-stream vector scaled to length 1 (u . u = 1). In the paper, u points toward repetitive looping. The dot product of u with column W[:, j] measures how much neuron j pushes the stream along u per unit of activation:

Across all columns, the row vector g = u @ W holds those six numbers: [1.062, -0.555, 1.043, -0.495, 0.318, 0.162].

Every column W[:, j] splits into two perpendicular parts: g_j u along u, and W[:, j] - g_j u at right angles to u.

3. What the edit does to one chosen column#

The edit picks a neuron subset S (here S = [1, 4]) and a strength alpha, then subtracts alpha * g_j u from each column j in S:

Columns outside S stay unchanged (W'[:, j] = W[:, j]). Writing W_S for W with columns outside S zeroed out gives the paper's matrix formula:

To find an edited column's new push along u, take the dot product with u and use u . u = 1:

The part of W[:, j] perpendicular to u does not change, because u has zero dot product with any perpendicular direction.

Compare this with the printed numbers for neuron 1 (g_1 = -0.555) and neuron 4 (g_4 = 0.318):

• At alpha = 1 (erase), the multiplier 1 - 1 = 0 sends -0.555 and 0.318 to 0. Both neurons still write perpendicular to u, while their push along u is zeroed out.
• At alpha = 2 (mirror), the multiplier 1 - 2 = -1 flips -0.555 to +0.555 and 0.318 to -0.318.
• At alpha = 6 (the paper's value), the multiplier is 1 - 6 = -5. For neuron 1, -5 x -0.555 = +2.775, while the script prints 2.776 from the unrounded -0.555157 (-5 x -0.555157 = 2.776). For neuron 4, -5 x 0.318416 = -1.592.

4. Why the change is rank one, and why editing weights equals steering#

The matrix change D = W' - W equals -alpha * outer(u, g_masked). Every column of D is a multiple of u, and a single column vector times a row vector has rank one. Running SVD (singular value decomposition, which splits a matrix into independent rank-one pieces ordered by size) on D gives a first singular value of 3.84 and three zeros (0). Since u has length 1, 3.84 equals |alpha| * ||g_S|| = 6 * sqrt((-0.5552)^2 + (0.3184)^2) = 3.840.

Now multiply W' by an activation vector a on one token:

Here s(t) = u . (W_S a) is how hard the chosen neurons push along u on token t. Since W_S keeps only columns in S, W_S a equals W[:, S] a[S], the activation-weighted sum of those columns. When neurons 1 and 4 are quiet (a[1] = a[4] = 0), s = 0.000 and the output is unchanged (|W'a - Wa| = 0.000). When both fire at 3.0, s = 3.0 x (-0.5552 + 0.3184) = -0.710, shifting the output along u by |alpha * s| = 6 x 0.710 = 4.261. A permanent weight edit therefore acts like adding a steering vector -alpha * s(t) * u whose strength s(t) comes from the token's own activations.

Check yourself#

1. What value of alpha makes the chosen neurons push 3 times as hard in their original direction along u?

2. Why does the SVD of W' - W have only one non-zero singular value even though we changed two columns (S = [1, 4])?

3. If you replace the direction u with -u and keep alpha = 6 and S = [1, 4], how does the edited matrix W' change?

4. On a token where neuron 1 fires with a[1] = 2.0 and neuron 4 fires with a[4] = 0.0, what is s, and how big is the change |W'a - Wa| at alpha = 6?

Try this next#

1. Test that u and -u give the exact same edit. In labs/lab0_toy.py, after D = edit(W, u, 6, S) - W (line 48), add assert np.allclose(edit(W, u, 6, S), edit(W, -u, 6, S)) and re-run. All assertions pass, showing why SVD sign flips in Lab 1 do not change the edit.
2. Touch 4 neurons instead of 2. Change line 21 from S = [1, 4] to S = [0, 1, 2, 4] and re-run. The heatmap in results/figures/lab0_toy.png shows four colored columns instead of two, and step 3 still prints 0. for the second, third, and fourth singular values, because the edit is still rank one.
3. Add alpha = -2 to the table. On lines 38 and 39, add (-2, "amplify 3x original way") and re-run. The push along u of neurons [1, 4] triples from [-0.555, 0.318] to [-1.665, 0.955].
4. Scale up to the real layer dimensions. Change line 17 from d, n = 4, 6 to d, n = 1024, 3584 and line 21 to S = list(range(64)), then re-run. Every assert passes in a few seconds. Numpy shortens arrays over 1,000 numbers with ... in step 1, while step 2 prints 64 numbers per line. Step 4 prints s = -35.717 and a change of 214.301 when the chosen neurons fire. The output moves much more than in the toy because 64 neurons push along u together.