---
name: transformer-lens-interpretability
source: https://app.decimal.ai/s/transformer-lens-interpretability@1/SKILL.md
source_sha256: 8bda233e866c
---

# TransformerLens API grammar

## Contract

Enforces the current TransformerLens public API when the task is to load a model, cache
activations, read a specific internal activation, or intervene on the forward pass for
mechanistic-interpretability work. The library exposes every activation through a named HookPoint,
and the names follow one deterministic grammar. Apply whenever the task inspects or edits transformer
internals with TransformerLens.

## Rules

1. **Load with the class method.** `from transformer_lens import HookedTransformer`, then
   `model = HookedTransformer.from_pretrained("gpt2-small")`. Never the bare constructor
   `HookedTransformer("gpt2")`, never `HookedTransformer.load(...)`, and never HuggingFace's
   `AutoModelForCausalLM.from_pretrained` — you lose the hooks.

2. **Cache with `run_with_cache`, which returns a tuple.** One call captures every intermediate
   activation: `logits, cache = model.run_with_cache(tokens)`. It returns exactly two values —
   the logits and the `ActivationCache`. Do NOT use HuggingFace `output_hidden_states=True` /
   `output_attentions=True`, and do NOT hand-register PyTorch hooks to collect activations.

3. **Hook names follow ONE grammar.** A full hook name is
   `blocks.{L}.[<submodule>.]hook_<name>` where `{L}` is the layer index. The `hook_` prefix is
   mandatory. Whether there is a submodule depends on WHERE the activation lives:
   - **Block level (no submodule):** the residual stream and the two block outputs —
     `blocks.{L}.hook_resid_pre`, `blocks.{L}.hook_resid_mid`, `blocks.{L}.hook_resid_post`,
     `blocks.{L}.hook_attn_out`, `blocks.{L}.hook_mlp_out`.
   - **Attention internals live under `attn.`:** `blocks.{L}.attn.hook_q`, `hook_k`, `hook_v`,
     `hook_z`, `hook_pattern` (post-softmax weights), `hook_attn_scores` (pre-softmax),
     `hook_result`.
   - **MLP internals live under `mlp.`:** `blocks.{L}.mlp.hook_pre`, `blocks.{L}.mlp.hook_post`.
   - **LayerNorm lives under `ln1`/`ln2`/`ln_final`:** `blocks.{L}.ln1.hook_normalized`,
     `blocks.{L}.ln1.hook_scale`, and the final one `ln_final.hook_normalized`.
   - **Embeddings are top-level:** `hook_embed`, `hook_pos_embed`.

   Two easy-to-miss rules: `mlp_out` and `attn_out` sit on the BLOCK (`blocks.{L}.hook_mlp_out`),
   NOT under `mlp.`/`attn.`; and `ln_final.hook_normalized` has no `blocks.{i}.` prefix.

4. **Read the cache by full name OR the tuple shorthand.** `cache["blocks.5.hook_resid_post"]`
   equals `cache["resid_post", 5]`. The shorthand is `(short_name, layer)`. Useful shorthands:
   `cache["resid_post", L]`, `cache["pattern", L]`, `cache["z", L]`, `cache["q"/"k"/"v", L]`,
   `cache["mlp_out", L]`, `cache["attn_scores", L]`.

5. **Restrict what you cache with `names_filter`.** To cache a subset, pass a predicate over the
   name: `model.run_with_cache(tokens, names_filter=lambda name: "resid_post" in name)`. To offload
   the captured tensors, pass `device="cpu"`.

6. **Intervene with `run_with_hooks`.** To edit activations for one run:
   `model.run_with_hooks(tokens, fwd_hooks=[(name, hook_fn)])`. `fwd_hooks` is a LIST of
   `(hook_name, function)` tuples. Never PyTorch's `module.register_forward_hook`.

7. **A hook function is `(activation, hook)` and returns the activation.**
   ```python
   def hook_fn(activation, hook):
       activation[:] = 0.0        # edit in place (or build a new tensor)
       return activation          # must return it
   ```
   The two parameters are the activation TENSOR and the HookPoint object (conventionally named
   `hook`; `hook.name` is its name). This is NOT the PyTorch `(module, input, output)` signature.

8. **Reset hooks between runs.** Hooks added by `add_hook` (or a run that did not clean up) persist.
   Call `model.reset_hooks()` for a clean state before the next run.

9. **Tokenize with the model's helpers.** `model.to_tokens(text)` → token tensor;
   `model.to_single_token(" Paris")` → one id; `model.to_str_tokens(text)` → list of piece strings;
   `model.to_string(tokens)` → text. Control the leading BOS with `prepend_bos=False`. Prefer these
   over a raw HuggingFace tokenizer.

10. **Config is on `model.cfg`; weights are named properties.** Shapes and counts live on the
    `HookedTransformerConfig`: `model.cfg.n_layers`, `model.cfg.n_heads`, `model.cfg.d_model`,
    `model.cfg.d_head` — NOT `model.config.num_hidden_layers` / `hidden_size`. Weights are exposed
    directly: `model.W_U` (unembed), `model.W_E` (token embed), `model.W_pos` (positional),
    `model.W_Q`, `model.W_K`, `model.W_V`, `model.W_O`, `model.W_in`, `model.W_out`. The per-head
    stacks are indexed `[layer, head]`, e.g. `model.W_O[9, 4]`.

11. **Prefer the `patching` helpers for causal tracing.** `from transformer_lens import patching`,
    then `patching.get_act_patch_resid_pre(...)` (and the `get_act_patch_attn_out` /
    `get_act_patch_mlp_out` / `get_act_patch_attn_head_*` family) rather than hand-rolling the
    patch loop.

## Worked examples

BEFORE = the base model's wrong default → AFTER = the conforming TransformerLens grammar.

**Load + cache.**
```python
# BEFORE
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("gpt2")
out = model(tokens, output_hidden_states=True)      # not the TL way
# AFTER
from transformer_lens import HookedTransformer
model = HookedTransformer.from_pretrained("gpt2-small")
logits, cache = model.run_with_cache(tokens)         # (logits, cache) tuple
```

**Read an activation.**
```python
# BEFORE
resid = cache["blocks.5.residual"]           # wrong name
patt  = cache["blocks.3.attention.weights"]  # wrong name
# AFTER
resid = cache["resid_post", 5]               # == cache["blocks.5.hook_resid_post"]
patt  = cache["pattern", 3]                  # == cache["blocks.3.attn.hook_pattern"]
mlp   = cache["mlp_out", 7]                  # == cache["blocks.7.hook_mlp_out"] (block level!)
```

**Selective / offloaded cache.**
```python
# AFTER
logits, cache = model.run_with_cache(
    tokens,
    names_filter=lambda name: "resid_post" in name,
    device="cpu",
)
```

**Intervene (ablate a head output).**
```python
# BEFORE (torch native)
def hook(module, input, output):      # wrong signature
    return output * 0
handle = model.blocks[6].register_forward_hook(hook)
# AFTER
def ablate(activation, hook):
    activation[:] = 0.0
    return activation
logits = model.run_with_hooks(
    tokens, fwd_hooks=[("blocks.6.hook_attn_out", ablate)]
)
model.reset_hooks()
```

**Tokenize + config + weights.**
```python
# AFTER
tokens = model.to_tokens("The Eiffel Tower is in", prepend_bos=True)
paris  = model.to_single_token(" Paris")
for layer in range(model.cfg.n_layers):
    for head in range(model.cfg.n_heads):
        contribution = cache["z", layer][0, -1, head] @ model.W_O[layer, head] @ model.W_U
```

## Edge cases & exceptions

- **`resid_pre` vs `resid_mid` vs `resid_post`.** `resid_pre` is before attention, `resid_mid` is
  after attention / before MLP, `resid_post` is after MLP (the block's output). Pick by which one
  the analysis needs; all three are block-level (`blocks.{L}.hook_resid_post`).
- **`pattern` vs `attn_scores`.** `hook_pattern` is post-softmax (a probability distribution);
  `hook_attn_scores` is the pre-softmax logits. They are different hooks under `attn.`.
- **`z` vs `attn_out`.** `hook_z` (under `attn.`) is per-head, shape `[batch, pos, head, d_head]`,
  before the output projection; `hook_attn_out` (block level) is the summed, projected attention
  output, shape `[batch, pos, d_model]`.
- **Shorthand aliases.** `cache["attn", L]` resolves to `hook_pattern`; `cache["key"/"query"/"value", L]`
  resolve to `hook_k`/`hook_q`/`hook_v`. Prefer the canonical short names.
- **A returned tensor is required.** A hook that edits in place still must `return activation`; a
  hook that returns `None` leaves the activation unchanged.

## Do / Don't

- DON'T load with `AutoModelForCausalLM.from_pretrained`. ALWAYS `HookedTransformer.from_pretrained`.
- DON'T collect activations with `output_hidden_states=True`. ALWAYS `run_with_cache` (a tuple).
- DON'T drop the `hook_` prefix or guess `blocks.L.residual`. ALWAYS `blocks.{L}.hook_resid_post`.
- DON'T nest `mlp_out`/`attn_out` under `mlp.`/`attn.`. They sit at block level.
- DON'T prefix `ln_final` with `blocks.{i}.`. It is top-level: `ln_final.hook_normalized`.
- DON'T use PyTorch `register_forward_hook` or the `(module, input, output)` signature. ALWAYS
  `run_with_hooks(fwd_hooks=[(name, fn)])` with `def fn(activation, hook): … return activation`.
- DON'T read `model.config.num_hidden_layers` / `hidden_size`. ALWAYS `model.cfg.n_layers` /
  `model.cfg.d_model`.
- DON'T reach for `model.lm_head.weight`. ALWAYS `model.W_U` (and `W_E`/`W_pos`/`W_O`/…).

## Common mistakes (the base model's wrong defaults)

- Loading via HuggingFace `AutoModelForCausalLM` and losing the hooks.
- Treating `run_with_cache` as if it returned only the cache (it returns `(logits, cache)`).
- Hallucinated hook names: `blocks.5.residual`, `blocks.3.attention.weights`,
  `layers[5].attn_pattern`, `blocks.7.mlp.hook_out` for the MLP output.
- The torch `(module, input, output)` hook signature instead of `(activation, hook)`, or forgetting
  to return the activation.
- Forgetting `model.reset_hooks()` so edits leak into later runs.
- Reading `model.config.num_hidden_layers` / `hidden_size` instead of `model.cfg.n_layers` /
  `d_model`.
- Using `model.lm_head.weight` instead of `model.W_U`.

## Quick checklist

- [ ] `HookedTransformer.from_pretrained("gpt2-small")` (not `AutoModelForCausalLM`).
- [ ] `logits, cache = model.run_with_cache(tokens)` (a `(logits, cache)` tuple).
- [ ] Hook names `blocks.{L}.hook_resid_post` / `blocks.{L}.attn.hook_pattern` / `…attn.hook_z` /
      `blocks.{L}.hook_mlp_out` / `blocks.{L}.mlp.hook_pre` / `ln_final.hook_normalized`.
- [ ] Cache shorthand `cache["resid_post", L]`, `cache["pattern", L]`, `cache["z", L]`.
- [ ] `names_filter=` to subset, `device="cpu"` to offload.
- [ ] `run_with_hooks(fwd_hooks=[(name, fn)])`; `def fn(activation, hook): … return activation`.
- [ ] `model.reset_hooks()` between runs.
- [ ] `to_tokens` / `to_single_token` / `to_str_tokens`, `prepend_bos=False` when needed.
- [ ] `model.cfg.n_layers` / `n_heads` / `d_model`; `model.W_U` / `W_O` / `W_E` / `W_pos`.