ab490/activation-patching-playground

Activation Patching Playground

Live demo: [huggingface.co/spaces/ab490/activation-patching-playground](https://huggingface.co/spaces/ab490/activation-patching-playground)

A browser-based mechanistic interpretability tool.

Type two prompts, pick a target token, and watch a live heatmap reveal which layers and token positions causally drive the model's prediction.

What is activation patching?

Activation patching is a causal intervention technique from mechanistic interpretability. Given two prompts that produce different outputs (e.g. "Eiffel Tower -> Paris" vs "Colosseum -> Rome"), it finds exactly where inside the model the correct information is stored.

For every (layer, token position) pair, it asks: "If I surgically swap in the clean run's internal activation at this spot, does the model recover the correct answer?" The result is a recovery score per cell, visualised as a heatmap.

• —Score 1.0 - this spot causally encodes the answer; injecting the clean activation here fully restores the correct prediction
• —Score 0.0 - this spot has no effect on the output

Bright purple cells are where the model "knows" the answer. That's mechanistic interpretability in one heatmap.

What you can do

• —Factual recall - run the built-in "Eiffel Tower / Colosseum" example and see how factual knowledge localises to specific mid-to-late layers at the subject token
• —Indirect object identification (IOI) - load the "Mary and John" example to see the classic IOI circuit signal at the name tokens
• —Bring your own prompts - type any clean + corrupted pair and any target token to explore arbitrary model behaviours
• —Explore predictions - click any cell after the run to see the top-5 tokens the model predicts at that exact (layer, position) patch
• —Compare models - switch between GPT-2 variants, Pythia, and Phi-2 to see how circuit structure changes with scale
• —Export - download the heatmap as a PDF, or export the raw scores as CSV or JSON

Features

• —Live streaming heatmap - rows appear layer by layer as patching runs, no waiting for the full result
• —Click any cell: bar chart of top-5 predicted tokens for that patch (fetched lazily on demand)
• —Hover tooltips on every cell showing layer, token, and score
• —Score threshold slider - filter out low-signal cells to focus on what matters
• —Zoom controls on the heatmap
• —Run history - previous runs are saved in the browser and can be restored with one click
• —Real-time target token validation - warns if your target encodes to more than one token
• —Cancel mid-run
• —PDF, CSV, and JSON export

Models

All models run on CPU. First use downloads weights automatically; subsequent runs load from cache.

Performance

Uses batched patching - all token positions for a layer are patched in a single forward pass, reducing total forward passes from n_layers × seq_len down to n_layers + 2. For GPT-2 small with a 10-token prompt that's ~14 passes instead of ~120.

Expected time on the HF Spaces free CPU tier: ~25–50s for GPT-2 small. Larger models take proportionally longer.

Running locally

Requirements: Python 3.11+, Node 18+, uv

# Install Python deps
uv sync

# Terminal 1 - backend
uv run uvicorn backend.main:app --reload --port 8000

# Terminal 2 - frontend
cd frontend && npm install && npm run dev

Frontend at http://localhost:5173, backend at http://localhost:8000. Model weights download automatically on first use to ~/.cache/huggingface/hub/.

How it works

Clean prompt: run_with_cache() - cache residual stream at every layer
Corrupted prompt: baseline logits

For each layer:
    batch = [corrupted_prompt] × seq_len
    patch position N in batch item N with clean activation
    one forward pass: scores for all positions in this layer

score = (patched_prob − corrupted_prob) / (clean_prob − corrupted_prob)

Hook point: blocks.{layer}.hook_resid_post (residual stream after attention + MLP).

Results stream to the browser via SSE, one layer at a time. Top-5 token distributions are fetched lazily only when a cell is clicked.

Stack

• —Backend: Python, FastAPI, TransformerLens
• —Frontend: React, TypeScript, Vite, D3.js, Tailwind CSS v4, jsPDF