TorchLens for mechanistic interpretability¶

Mechanistic interpretability needs three things from a model's activations: an exact address for every site (a head's query, a block's residual stream, one MLP neuron), the ability to edit that site and measure the effect on the output, and enough speed to sweep the edit over many sites and prompts. TorchLens records every operation of a forward pass with a stable label, lets you patch, ablate or steer any of them with the same call whether you run once or through generation, and hands the result to the SAE, probing and steering libraries you already use.

The lists below are in the order most people want them: see one site, edit it, then scale up. Every link is to a file in this repository.

Start here¶

• Tutorial notebooks: eleven short notebooks on laptop-sized models. Notebook 05 is the first intervention, 07 compares runs in a bundle, 08 is the semantic shortcuts ("head 3's query") and 10 has a mechanistic-interpretability branch.
• The decoder running example: the small decoder on the IOI prompt that every Field guide page starts from.
• See the graph and inspect any operation: a picture of every operation the forward ran, and the output, module and arguments of any one of them.
• Intervene on a site and measure the effect: ablation, the hook check, and the same edit bound live for sweeps and generate().
• Hugging Face tutorial notebook: capture a transformer from transformers, browse its structure, read activations from any point, draw the graph.

Patching, ablation and steering¶

• Intervention explainers: the three ways TorchLens applies one edit (replay, live hook, rerun) and which question each answers.
• Intervention API reference: selectors, edit helpers, hooks, replay and bundles.
• Stochastic and population edits: mean ablation, resampling ablation and causal scrubbing from one population noun and one edit.
• Editing generation and the KV cache: what "intervene during generation" can mean and how each is spelled.
• Performance: choosing a fast path: full capture, replay and capture-free live hooks differ in cost by two orders of magnitude; pick before a sweep.
• Paired-prompt activation patching notebook: patches each of GPT-2 small's 144 attention heads from a France prompt into an Italy prompt and maps which heads carry " Paris".
• Causal trace notebook: a full clean-versus-corrupt causal trace with a per-site heatmap.
• Ablation grid notebook: an ablation sweep over discovered sites.
• Steering notebook: build a steering vector and apply it at replay time.
• Intervention worked examples: eighteen short scripts, from the first five minutes to per-position steering, bundles of three or more prompts, SAE attachment and publishing a portable .tlspec.

Circuits, attribution and lenses¶

• The mech-interp kit: residual decomposition, direct logit contributions, per-head results, the component edge graph and patching grids.
• ACDC circuit discovery as a driver loop: edge pruning on IOI written as a loop over TorchLens calls.
• A tuned lens in TorchLens's logit lens and the train a tuned lens notebook.
• A Jacobian lens in a few lines and the Jacobian lens notebook: what a block's residual vector is disposed to make the model say, fitted on GPT-2 small and compared with the logit lens.
• Per-step attribution for generated text: attribute each generated token to the tokens before it.
• Attribution graphs with transcoders notebook: GPT-2 small with public transcoders spliced in.
• Function vectors notebook: score every head's causal indirect effect and sum the top heads into a task vector.
• DAS causal basis notebook: distributed alignment search through a bound hook, tested against a random basis.
• Measure faithfulness notebook: ranked removal, output drop and locality checks on an attribution.
• EAP-IG circuit notebook: edge attribution patching with integrated gradients on GPT-2 small's IOI task, with the circuit's faithfulness replayed by patch_edges.
• Logit lens notebook: every GPT-2 small block's next-token prediction, with TorchLens and with plain PyTorch hooks side by side.
• Vision logit lens notebook: per-block class votes of a ViT classifier and CLIP's open-vocabulary readout of every vision block.
• LEACE concept erasure notebook: scrub grammatical number from every GPT-2 small block with bound tl.when rules and watch subject-verb agreement collapse.
• Layer-wise relevance propagation for transformers and the epsilon-LRP litmus test.
• Conditional steering with a per-token audit: steer only when a condition direction crosses a threshold.
• Transformer pictures: token-labelled attention heatmaps and grids, the attention atlas, coloured tokens.

Sparse autoencoders¶

• SAE splicing tutorial notebook: test whether an SAE's features are causally used, not only whether they reconstruct.
• Pretrained SAEs: SAELens loads, TorchLens runs: Gemma Scope and every other published SAE in the SAELens registry.
• Train an SAE on TorchLens activations and SAE trainers: dictionary_learning and sparsify: extract, train elsewhere, splice back and measure.
• SAE latents and top examples for SAEDashboard and delphi.
• Model diffing with a crosscoder: one dictionary over two models' activations.
• CLIP SAEs with Prisma.

Probes and representations¶

• Linear probes: leakage-safe splits, an obfuscation check and a causal check.
• Probe accuracy by layer notebook: accuracy against depth with an untrained-network control.
• Contrastive direction notebook: a positive-minus-negative direction with replay-time steering.
• Streaming CKA notebook: compare representations on streamed data in constant memory.
• Streaming statistics: dataset-level statistics of activations or gradients, one batch at a time.
• Training external machinery against TorchLens surfaces: TorchLens ships no probe fitters or SAE loops; this is how your own trainer plugs in.

Replications on GPT-2¶

• Paper replications: scripts that state one published claim, measure it and print a verdict against a rule fixed before the first run: induction heads (Olsson et al. 2022), the IOI name movers (Wang et al. 2022) and causal tracing (Meng et al. 2022).
• The same three lines on any architecture: the trace, draw and edit calls on six architecture families.

Experiments at scale¶

• Experiments: lineage, vary, site sweeps and grids: the bundle as the record of an experiment.
• Episode capture for multi-step generation loops.
• Facets and the facets tutorial notebook: one access surface for heads, gates and other structural views of an activation, including TransformerLens aliases.
• Fast activation recording: the sparse recording path for when you know what to keep before the forward runs.

Connectors and migration¶

• Connectors to other tools: SAELens, steering-vectors, repeng, dialz, nnsight, inseq, LIT, Captum, SHAP, treescope and the rest, one entry each.
• Coming from TransformerLens (and the migration notebook), nnsight, pyvene, SAELens, baukit, AutoCircuit, RepE or plain forward hooks.
• Which tool for this job? and where to go instead: the jobs another tool does better.
• Glossary: the definition of every public name used above.