| |
| """Lab 0: weight-space programming on a toy layer small enough to print. |
| |
| uv run --with numpy --with matplotlib python labs/lab0_toy.py # seconds, any laptop |
| |
| An MLP's last matrix, down_proj, is W (here 4 x 6; in Qwen3.5-0.8B it is 1,024 x 3,584). Neuron j fires with |
| strength a[j] and writes its column W[:, j] into the stream, so the layer adds y = W a = sum_j a[j] * W[:, j]. |
| The paper's edit picks a direction u and some neurons S and sets W' = W - alpha * u (u^T W[:, S]) on those |
| columns. Every claim below is checked with an assert, so if this script finishes, they all hold. |
| """ |
| import os |
|
|
| import numpy as np |
|
|
| np.set_printoptions(precision=3, suppress=True) |
| rng = np.random.default_rng(0) |
| d, n = 4, 6 |
| W = rng.normal(size=(d, n)) |
| u = rng.normal(size=d) |
| u /= np.linalg.norm(u) |
| S = [1, 4] |
| others = [j for j in range(n) if j not in S] |
|
|
|
|
| def edit(W, u, alpha, S=None): |
| g = u @ W |
| if S is not None: |
| mask = np.zeros(W.shape[1]) |
| mask[S] = 1 |
| g = g * mask |
| return W - alpha * np.outer(u, g) |
|
|
|
|
| print("1. How much each neuron writes along u (u . W[:, j]):", u @ W) |
|
|
| print("\n2. The edit scales the chosen neurons' push along u by (1 - alpha) and leaves everything else alone") |
| P = np.eye(d) - np.outer(u, u) |
| for alpha, name in [(0, "no edit"), (1, "erase"), (2, "mirror"), (6, "flip and amplify 5x (the paper)"), |
| (-4, "amplify 5x the other way")]: |
| W2 = edit(W, u, alpha, S) |
| before, after = (u @ W)[S], (u @ W2)[S] |
| assert np.allclose(after, (1 - alpha) * before) |
| assert np.allclose(W2[:, others], W[:, others]) |
| assert np.allclose(P @ W2, P @ W) |
| print(f" alpha={alpha:>2} ({name:31s}) push along u of neurons {S}: {before} -> {after}") |
|
|
| print("\n3. The change is rank one: one direction out (u), one pattern in (u^T W on the chosen columns)") |
| D = edit(W, u, 6, S) - W |
| sv = np.linalg.svd(D, compute_uv=False) |
| assert sv[1] < 1e-12 * sv[0] |
| print(" singular values of W' - W:", sv) |
|
|
| print("\n4. The weight edit is steering: W' a = W a - alpha * s * u, with s = u . (W[:, S] a[S])") |
| alpha = 6 |
| for name, a in [("chosen neurons quiet", np.where(np.isin(np.arange(n), S), 0.0, rng.normal(size=n))), |
| ("chosen neurons firing", np.where(np.isin(np.arange(n), S), 3.0, rng.normal(size=n)))]: |
| s = u @ (W[:, S] @ a[S]) |
| y_edit, y_steer = edit(W, u, alpha, S) @ a, W @ a - alpha * s * u |
| assert np.allclose(y_edit, y_steer) |
| print(f" {name:22s} s = {s:+.3f} change to the output |W'a - Wa| = {np.linalg.norm(y_edit - W @ a):.3f}") |
| print(" Same weights, different tokens, different push: the edit only acts when the chosen neurons fire.") |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| fig, ax = plt.subplots(1, 2, figsize=(11, 3.6)) |
| im = ax[0].imshow(edit(W, u, 6, S) - W, cmap="RdBu", vmin=-6, vmax=6) |
| ax[0].set_title("W' - W: only the chosen columns change") |
| ax[0].set_xlabel("neuron (column)") |
| ax[0].set_ylabel("stream dimension (row)") |
| fig.colorbar(im, ax=ax[0]) |
| al = np.linspace(-5, 8, 100) |
| ax[1].plot(al, 1 - al, color="#0969da") |
| for x, lab in [(0, "no edit"), (1, "erase"), (2, "mirror"), (6, "the paper")]: |
| ax[1].plot(x, 1 - x, "o", color="#cf222e") |
| ax[1].annotate(lab, (x, 1 - x), xytext=(6, 4), textcoords="offset points") |
| ax[1].axhline(0, color="black", lw=0.5) |
| ax[1].set_xlabel("alpha") |
| ax[1].set_ylabel("push along u after / before") |
| ax[1].set_title("What alpha does to the chosen neurons") |
| fig.tight_layout() |
| out = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results", "figures") |
| os.makedirs(out, exist_ok=True) |
| fig.savefig(os.path.join(out, "lab0_toy.png"), dpi=130) |
| print("\nall checks passed; figure in results/figures/lab0_toy.png") |
|
|
