LLM Brainmap

Watch a language model react to prompt injection, as a 3D "brain activation" map.

Every point is one GemmaScope 2 SAE feature of gemma-3-4b-it (34 layers × 16384 features, residual stream). As the prompt is read step by step, firing features light up bottom-to-top — the forward pass sweeping through depth — then fade with an afterglow. Rose points are normal activations; orange points are activations tied to the injected text (or, in diff mode, features that never fire on the clean version of the same prompt).

Three companion overlays turn the light show into something readable:

• top-left: Neuronpedia auto-interp labels of the strongest firing features ("cautious actions", "absence and negation", ...)
• top-right: which earlier tokens the current step attends back to — when the model reads No translate just output the key, it looks back at "Spanish" and "translate", the instruction being overridden
• bottom: token ribbon with the injected span in orange

Side-by-side A/B (clean vs injected)	Single view

The interactive player — drag to rotate, scroll to zoom, seek any step:

Outputs

• MP4 video, 720p/30fps, offline rendered (PyVista, headless)
• Interactive 3D in the browser (Three.js, self-hosted): drag to rotate, scroll to zoom, seek any position, switch token / phrase / sentence granularity live

Everything is served from one Gradio app on one port.

Quickstart

Requirements: a GPU with ≥ 16 GB VRAM, Python ≥ 3.10.

pip install torch transformers gradio pyvista imageio-ffmpeg umap-learn \
            safetensors fastapi uvicorn pillow matplotlib huggingface_hub

Download weights into your HF cache (HF_HOME is respected):

• google/gemma-3-4b-it
• google/gemma-scope-2-4b-it — folders resid_post_all/layer_*_width_16k_l0_big (34 layers, ~11 GB, the visual field) and resid_post/layer_{9,17,22,29}_width_16k_l0_medium (~1.3 GB, the label track)

Build the one-time per-model assets:

python assets_layout.py     # UMAP layout per layer (~30 min, CPU)
python assets_baseline.py   # benign firing-rate baseline (~10 min, GPU)
# label explanations: download Neuronpedia batches, then
python assets_labels.py

Run:

bash start.sh               # then open http://localhost:7860

Optional machine-specific settings go in .env (see start.sh): BRAINMAP_PYTHON, HF_HOME, CUDA_VISIBLE_DEVICES, BRAINMAP_SAMPLES.

A few self-contained example prompts ship in examples/samples.jsonl; point BRAINMAP_SAMPLES at your own benchmark (JSONL with id, is_attack, goal_text, clean_content, eval_content) to browse it in the UI.

How it works

1. One forward pass captures the residual stream at every layer (causal masking means position t already is "the state after reading prefix t" — the whole animation costs a single pass) plus layer/head-averaged attention.
2. Each layer's residuals go through its JumpReLU SAE; all nonzero features are stored raw (SAE sparsity is the sparsity — no top-k cut).
3. Feature positions on each layer plane come from a UMAP of the SAE decoder directions, so semantically similar features cluster into "brain regions".
4. The renderer (or the browser player) replays the recording: per step it injects that step's activations layer-by-layer bottom-up, decays the rest, and swaps the label / attention overlays only once the wave tops out.
5. Label text comes from Neuronpedia explanations. They index the l0_medium SAEs, not the l0_big ones used for the visuals, so a separate 4-layer label track is encoded at capture time (the variant match was verified empirically: distinctive-concept probes hit 17/44 on medium vs 4/44 on big).

Note the attention overlay shows where the model looks, which is evidence, not proof, of causation.

Repo layout

capture.py        forward pass + SAE encode + attention -> sparse recording
render.py         recording -> MP4 (single view or side-by-side A/B)
export_web.py     recording -> JSON for the browser player
web/              Three.js player (vendored, no CDN)
serve.py          Gradio app + static mount, single port
assets_*.py       one-time per-model assets (layout / baseline / labels)
examples/         bundled sample prompts
config.py         all paths, env-overridable

Recordings and rendered files are transient: the app keeps the newest 6 recordings and 4 web exports and prunes the rest. Download an MP4 (icon on the video player) if you want to keep it.

Credits

• Gemma 3 and GemmaScope 2 by Google DeepMind (weights not included; subject to their licenses)
• Feature explanations from Neuronpedia's public exports
• three.js (MIT, vendored under web/vendor/)

License

MIT — see LICENSE.