Repository navigation
docs(demos): add SwiGLU gate-sign demo for double-negative echoes - #1936
Conversation
Adds demos/SwiGLU_Gate_Sign_Demo.ipynb (proposal TransformerLensOrg#1933). Using only the gated-MLP hook aliases (mlp.hook_pre, hook_pre_linear, hook_post), it tags each neuron's top-k activations by sign regime, finds neurons whose top-k is made of gate-closed echoes, compares echo and gate-open peaks, and removes the echo with a forward hook against two controls. It runs SmolLM2-135M on CPU float32 with a small token budget and lists full-size results for SmolLM2-135M, Qwen3-0.6B and Qwen3-1.7B. Registers the notebook in the notebook CI matrix, the makefile, the docs build and its unit test, and links it from the hook system page.
0 * inf is NaN, so masking ineligible neurons by adding (~mask) * inf turned every eligible score into NaN and argmin returned the first all-echo neuron. Mask with torch.where instead, and rewrite the example note for the neuron that is now selected.
|
Good timing on this! Reviewing now while CI completes so I can squeeze it into 4.3.0 before I put out the release this afternoon |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Just a couple detail comments, let me know if you have any questions
| "for layer in range(N_LAYERS):\n", | ||
| " pool = (~echo_dominated[layer]).nonzero().squeeze(-1)\n", | ||
| " n = int(echo_dominated[layer].sum())\n", | ||
| " random_neurons[layer, pool[torch.randperm(len(pool), generator=generator)[:n]]] = True\n", |
There was a problem hiding this comment.
The random set and the selected set have the same neuron count and these neurons carry much less echo. So the gap to the selected set mixes which neurons were edited with how much was zeroed. Can the control be matched on removed echo magnitude, so it tests the selection on its own?
There was a problem hiding this comment.
Good catch. The random control now draws neurons per layer until they carry as much echo on the test rows as the selected set (1.01× in the demo, within 2% for the full-size runs). The selected neurons still cost 10–100× more: +0.334 vs +0.033 live, and in the full-size table +0.21/+1.31/+0.53 vs +0.018/+0.013/+0.030. I also added a note that a count-matched control overstates the selection effect.
| "| Δ loss, echo removed, echo-dominated neurons | +0.21 | +1.31 | +0.53 |\n", | ||
| "| Δ loss, echo removed, random neurons | +0.002 | +0.049 | +0.039 |\n", | ||
| "\n", | ||
| "In the Qwen3 models, layer 1 holds many echo-dominated neurons, and some neurons never open their gate. Splitting the echo-dominated set shows the removal cost does not come from either group. Each cell gives Δ loss, then the same-size random control in parentheses, then the number of neurons. The first two groups overlap (layer 1 contains some never-open neurons), and \"Neither\" excludes both:\n", |
There was a problem hiding this comment.
The Qwen3-0.6B column doesn't support this sentence: the 108 never-open neurons cost +0.366, more than the 1,899 neurons in "Neither". Could this be reworded to match what the split shows for each model?
There was a problem hiding this comment.
Agreed, that sentence overclaimed. I replaced it with a per-model description. In 0.6B the never-open neurons carry the largest cost, because their whole positive side is echo. In SmolLM2 and 1.7B the cost sits mostly in neither group. The split also uses the matched controls now.
The random control matched the echo-dominated set on neuron count, but those neurons carry much less echo, so the gap mixed which neurons were edited with how much was zeroed. Draw random neurons per layer until they carry as much echo on the test rows as the selected set (1.01x here), and use the same matching for the full-size tables. The split paragraph claimed the removal cost came from neither layer 1 nor never-open neurons, which Qwen3-0.6B contradicts (its 108 never-open neurons cost more than the 1,899 in neither group). Describe the split per model instead.
|
Looks great! Merged and including in the next release |
Description
Adds
demos/SwiGLU_Gate_Sign_Demo.ipynb, the demo approved in #1933.A gated MLP neuron outputs
act(gate) * value. When the gate is closed,act(gate)is a small negative number, so a negative value produces a positive "echo" (Goodfire's double-negative regime). Using only the existing gated-MLP hook aliases (mlp.hook_pre,mlp.hook_pre_linear,mlp.hook_post), the notebook:It runs SmolLM2-135M on CPU in float32 with a small token budget (about 2 minutes on 4 threads). The last section has a static table of full-size results for SmolLM2-135M, Qwen3-0.6B and Qwen3-1.7B, including the layer-1 / never-open-gate split promised in #1933. No helper function is added, per the approval comment.
Registration, one line each:
checks.yml) andmakefilenotebook-test;docs/make_docs.pycopy_demos, the docs toctree (index.md), andtests/unit/test_make_docs.py;hook_system.md.Fixes #1933
Type of change
Checklist:
On the unchecked box: no library code changes, so I ran only the affected tests rather than the whole unit suite.
pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/SwiGLU_Gate_Sign_Demo.ipynbpassed 11/11, repeated across thread counts (4 and 32) and two torch builds.tests/unit/test_make_docs.pypassed.The notebook was only validated on Windows, so the Linux notebook job in CI is the first check there.