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.

1. 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.

1. 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.

1. 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].

1. 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".

1. 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.

1. 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.

1. 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.

1. 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.

1. 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].

1. 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.