| # 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`): |
|
|
| ```bash |
| 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 |
|
|
| ```text |
| 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: |
|
|
| $$y = W a = \sum_{j=0}^{n-1} a_j W[:, j]$$ |
|
|
| ```mermaid |
| flowchart LR |
| A["Neuron activations<br/>a[0] ... a[5]"] --> B["Column j of W is W[:, j]<br/>(what neuron j writes)"] |
| B --> C["Weighted sum<br/>y = sum of a[j] * W[:, j]"] |
| C --> D["Added to the<br/>residual stream"] |
| ``` |
|
|
| ### 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: |
|
|
| $$g_j = u \cdot W[:, j]$$ |
| |
| 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`: |
|
|
| $$W'[:, j] = W[:, j] - \alpha (u \cdot W[:, j]) u$$ |
|
|
| 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: |
|
|
| $$W' = W - \alpha u (u^T W_S)$$ |
| |
| To find an edited column's new push along `u`, take the dot product with `u` and use `u . u = 1`: |
| |
| $$u \cdot W'[:, j] = u \cdot W[:, j] - \alpha (u \cdot W[:, j])(u \cdot u) = (1 - \alpha)(u \cdot W[:, j])$$ |
| |
| The part of `W[:, j]` perpendicular to `u` does not change, because `u` has zero dot product with any perpendicular direction. |
| |
| ```mermaid |
| flowchart TD |
| Col["Chosen column W[:, j]<br/>(neuron 1: g_1 = -0.555)"] --> Split["Split into two parts at right angles"] |
| Split --> Par["Part along u:<br/>g_j * u"] |
| Split --> Perp["Part at right angles to u:<br/>W[:, j] - g_j * u"] |
| Par --> Scale["Multiply by (1 - alpha)<br/>At alpha = 6: -5 * g_j * u"] |
| Perp --> Keep["Keep unchanged"] |
| Scale --> Out["Edited column W'[:, j]<br/>(neuron 1 push along u: 2.776)"] |
| Keep --> Out |
| ``` |
| |
| 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: |
| |
| $$W' a = W a - \alpha u (u^T W_S a) = W a - \alpha s(t) u$$ |
|
|
| 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`? |
| <details><summary>Answer</summary> |
|
|
| Set `1 - alpha = 3`, which gives `alpha = -2`. Any negative `alpha` amplifies the original push along `u` instead of opposing it. |
| </details> |
|
|
| 2. Why does the SVD of `W' - W` have only one non-zero singular value even though we changed two columns (`S = [1, 4]`)? |
| <details><summary>Answer</summary> |
|
|
| Both changed columns were shifted along the same direction `u` (column 1 moved by `+3.331 u` and column 4 moved by `-1.910 u`). Because every column of `W' - W` is a multiple of one vector `u`, the entire matrix change has rank one. |
| </details> |
|
|
| 3. If you replace the direction `u` with `-u` and keep `alpha = 6` and `S = [1, 4]`, how does the edited matrix `W'` change? |
| <details><summary>Answer</summary> |
|
|
| It does not change. In `W' = W - alpha * u (u^T W_S)`, `u` appears twice in `u u^T`, and `(-u)(-u)^T = u u^T`. Only the line through `u` matters, not which end gets the plus sign. |
| </details> |
|
|
| 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`? |
| <details><summary>Answer</summary> |
|
|
| `s = a[1] * g_1 = 2.0 * (-0.5552) = -1.110`. Since `u` has length 1, the output changes by `|alpha * s| = 6 * 1.110 = 6.662` along `u`. |
| </details> |
|
|
| ## 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. |
|
|
