Coverage for transformer_lens/tools/analysis/direct_path_patching.py: 97%
62 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Direct Path Patching.
3Implements direct path patching — a finer-grained variant of activation patching
4introduced for circuit analysis.
6Background
7----------
8Standard activation patching (see patching.py) replaces an activation at a given
9layer/position with its value from a clean run, and measures how much the model's
10output shifts. But patching the *residual stream* affects ALL downstream components,
11making it hard to isolate the direct information flow between two specific heads.
13Direct path patching isolates the path A → B: it patches *only* the contribution of
14source head A (at layer src_layer) into the input of destination head B (at layer
15dst_layer > src_layer), leaving every other component's view of A's output unchanged.
17The linear approximation used here (following Neel Nanda's description in issue #111)
18is:
20 delta_resid = clean_A_result - corrupted_A_result # [batch, pos, d_model]
21 delta_q = (delta_resid / ln1_scale) @ W_Q[hb] # [batch, pos, d_head]
22 patched_q = corrupted_q + delta_q
24This is an approximation, not exact even with LayerNorm folded: it freezes the
25corrupted run's ln1 scale, whereas the true scale shifts with the patched
26residual. The error grows with how much ``delta_resid`` changes the residual
27norm at the destination layer.
29Usage
30-----
31 # 1. Cache clean and corrupted activations
32 _, clean_cache = model.run_with_cache(clean_tokens)
33 _, corrupted_cache = model.run_with_cache(corrupted_tokens)
35 # 2. Define your metric (same as activation patching)
36 def metric(logits):
37 return logit_diff(logits, ...)
39 # 3. Sweep all (dst_layer, dst_head) pairs for a fixed source head
40 results = get_act_patch_direct_path(
41 model, corrupted_tokens, clean_cache, corrupted_cache,
42 metric, src_layer=9, src_head=9,
43 component="q", # patch into Q; also supports "k", "v"
44 )
45 # results.shape == (n_layers, n_heads)
46 # results[dst_layer, dst_head] = metric when A→B path is patched
48References
49----------
50- Neel Nanda, TransformerLens issue #111 (2022)
51- Wang et al., "Interpretability in the Wild: a Circuit for Indirect Object
52 Identification in GPT-2 small" (2022)
53"""
55from __future__ import annotations
57import warnings
58from typing import Any, Callable, Literal, cast
60import torch
61from jaxtyping import Float
62from tqdm.auto import tqdm
64from transformer_lens.ActivationCache import ActivationCache
65from transformer_lens.model_protocol import (
66 TransformerLensModel,
67 TransformerLensModelWithWeights,
68)
70# ---------------------------------------------------------------------------
71# Internal helpers
72# ---------------------------------------------------------------------------
75def _check_fold_ln(model: Any) -> None:
76 """Warn if the model's LayerNorm weights have not been folded in.
78 TransformerBridge wraps the original HuggingFace module, which stores the
79 learned scale as ``.weight``; legacy TransformerLens components store it as
80 ``.w``. We check both so the guard works for either layout.
81 """
82 try:
83 ln1 = model.blocks[0].ln1 # type: ignore[index]
84 # .weight → TransformerBridge (wraps HF module); .w → legacy TL components
85 w = getattr(ln1, "w", None)
86 if w is None:
87 w = getattr(ln1, "weight", None)
88 if w is not None and getattr(getattr(model, "cfg", None), "rmsnorm_uses_offset", False):
89 # Offset-RMS models (Gemma family) store w as an offset from 1:
90 # the effective scale is (1 + w) and the folded identity is 0.
91 w = 1.0 + w
92 if w is not None and not torch.allclose(w, torch.ones_like(w), atol=1e-3):
93 warnings.warn(
94 "get_act_patch_direct_path is most accurate when LayerNorm parameters "
95 "are folded into the weight matrices: call "
96 "model.process_weights(fold_ln=True) on the TransformerBridge. "
97 "Results may be inaccurate with unfolded LayerNorm.",
98 UserWarning,
99 stacklevel=3,
100 )
101 except (AttributeError, TypeError):
102 pass # non-standard model — cannot inspect LN weights, proceed
105# ---------------------------------------------------------------------------
106# Core hook factory
107# ---------------------------------------------------------------------------
110def _make_direct_path_hook(
111 delta_resid: Float[torch.Tensor, "batch pos d_model"],
112 dst_head: int,
113 W_component: Float[torch.Tensor, "d_model d_head"],
114 ln_scale_name: str,
115 corrupted_cache: ActivationCache,
116 component: Literal["q", "k", "v"],
117) -> Callable:
118 """Return a hook function that adds the linearised delta to one head's Q, K, or V.
120 Parameters
121 ----------
122 delta_resid:
123 (clean_A_result - corrupted_A_result), shape [batch, pos, d_model].
124 dst_head:
125 Index of the destination attention head to patch.
126 W_component:
127 The weight matrix for the component being patched:
128 W_Q[dst_head], W_K[dst_head], or W_V[dst_head].
129 Shape [d_model, d_head].
130 ln_scale_name:
131 Cache key for the layer-norm scale at the destination layer,
132 e.g. "blocks.3.ln1.hook_scale".
133 corrupted_cache:
134 Cache from the corrupted forward pass (used to look up ln1 scale).
135 component:
136 One of "q", "k", "v" — determines which QKV tensor is hooked.
137 """
139 def hook_fn(
140 value: Float[torch.Tensor, "batch pos n_heads d_head"],
141 hook, # HookPoint, unused but required by TransformerLens
142 ) -> Float[torch.Tensor, "batch pos n_heads d_head"]:
143 # ln scale: [batch, pos, 1]
144 ln_scale = corrupted_cache[ln_scale_name] # [batch, pos, 1]
146 # Linearised delta in query/key/value space
147 # delta_resid: [batch, pos, d_model]
148 # W_component: [d_model, d_head]
149 delta = (delta_resid / ln_scale) @ W_component # [batch, pos, d_head]
151 if value.requires_grad: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 value = value.clone()
153 value[:, :, dst_head, :] = value[:, :, dst_head, :] + delta
154 return value
156 return hook_fn
159# ---------------------------------------------------------------------------
160# Public API
161# ---------------------------------------------------------------------------
164def get_act_patch_direct_path(
165 model: TransformerLensModel,
166 corrupted_tokens: torch.Tensor,
167 clean_cache: ActivationCache,
168 corrupted_cache: ActivationCache,
169 patching_metric: Callable[[torch.Tensor], torch.Tensor],
170 src_layer: int,
171 src_head: int,
172 component: Literal["q", "k", "v"] = "q",
173 verbose: bool = True,
174) -> Float[torch.Tensor, "n_layers n_heads"]:
175 """Sweep direct path patches from one source head to all downstream heads.
177 For every destination head B = (dst_layer, dst_head) where dst_layer > src_layer,
178 patch the contribution of source head A = (src_layer, src_head) into B's query
179 (or key / value) input, and record the patching metric.
181 The patch is a linear approximation:
183 delta_resid = clean_A_result - corrupted_A_result [batch, pos, d_model]
184 delta_B_comp = (delta_resid / ln1_scale) @ W_comp[dst_head]
186 where W_comp is W_Q, W_K, or W_V according to `component`.
188 Parameters
189 ----------
190 model:
191 Any TransformerLens model (e.g. TransformerBridge).
192 corrupted_tokens:
193 Token IDs for the corrupted input, shape [batch, seq_len].
194 clean_cache:
195 Cached activations from the clean (unpatched) run.
196 corrupted_cache:
197 Cached activations from the corrupted run (needed for ln1 scale).
198 patching_metric:
199 A function mapping the model's logits tensor to a scalar.
200 src_layer:
201 Layer index of the source attention head.
202 src_head:
203 Head index of the source attention head.
204 component:
205 Which input to patch at the destination head — "q" (default), "k", or "v".
206 verbose:
207 Whether to show a tqdm progress bar.
209 Returns
210 -------
211 results : Float[Tensor, "n_layers n_heads"]
212 results[dst_layer, dst_head] is the patching metric when the direct path
213 A → B is patched in. Entries for dst_layer <= src_layer are left as 0.0
214 (no causal path from A to those layers).
215 """
216 _check_fold_ln(model)
218 n_layers = model.cfg.n_layers
219 n_heads = model.cfg.n_heads
221 results = torch.zeros(n_layers, n_heads, device=model.cfg.device)
223 # Residual stream delta from source head A.
224 #
225 # hook_result (per-head residual contribution) requires cfg.use_hook_result=True
226 # and is not in the default cache. We compute it instead from hook_z and W_O,
227 # which are always available:
228 # result_h = z[:, :, h, :] @ W_O[h] shape [batch, pos, d_model]
229 src_z_name = f"blocks.{src_layer}.attn.hook_z"
230 # Static-only view of the weights surface (runtime checks use the base protocol:
231 # nn.Module submodules are invisible to getattr_static-based isinstance).
232 weights_model = cast(TransformerLensModelWithWeights, model)
233 W_O = weights_model.blocks[src_layer].attn.W_O # [n_heads, d_head, d_model]
235 def _head_result(cache, h):
236 z = cache[src_z_name][:, :, h, :] # [batch, pos, d_head]
237 return z @ W_O[h] # [batch, pos, d_model]
239 delta_resid = _head_result(clean_cache, src_head) - _head_result(corrupted_cache, src_head)
240 # shape: [batch, pos, d_model]
242 # Weight matrix for the component being patched
243 _comp_map = {
244 "q": lambda attn: attn.W_Q, # [n_heads, d_model, d_head]
245 "k": lambda attn: attn.W_K,
246 "v": lambda attn: attn.W_V,
247 }
248 _hook_name_map = {
249 "q": lambda lb: f"blocks.{lb}.attn.hook_q",
250 "k": lambda lb: f"blocks.{lb}.attn.hook_k",
251 "v": lambda lb: f"blocks.{lb}.attn.hook_v",
252 }
253 W_all = _comp_map[component] # callable: attn → [n_heads, d_model, d_head]
254 hook_name_fn = _hook_name_map[component]
256 dst_pairs = [(lb, hb) for lb in range(src_layer + 1, n_layers) for hb in range(n_heads)]
258 for dst_layer, dst_head in tqdm(
259 dst_pairs,
260 desc=f"Direct path patch ({src_layer},{src_head}) → * [{component}]",
261 disable=not verbose,
262 ):
263 ln_scale_name = f"blocks.{dst_layer}.ln1.hook_scale"
264 W_comp = W_all(weights_model.blocks[dst_layer].attn)[dst_head] # [d_model, d_head]
266 hook_fn = _make_direct_path_hook(
267 delta_resid=delta_resid,
268 dst_head=dst_head,
269 W_component=W_comp,
270 ln_scale_name=ln_scale_name,
271 corrupted_cache=corrupted_cache,
272 component=component,
273 )
275 patched_logits = model.run_with_hooks(
276 corrupted_tokens,
277 fwd_hooks=[(hook_name_fn(dst_layer), hook_fn)],
278 )
280 results[dst_layer, dst_head] = patching_metric(patched_logits).item()
282 return results
285def get_act_patch_direct_path_all_sources(
286 model: TransformerLensModel,
287 corrupted_tokens: torch.Tensor,
288 clean_cache: ActivationCache,
289 corrupted_cache: ActivationCache,
290 patching_metric: Callable[[torch.Tensor], torch.Tensor],
291 component: Literal["q", "k", "v"] = "q",
292 verbose: bool = True,
293) -> Float[torch.Tensor, "n_layers n_heads n_layers n_heads"]:
294 """Full sweep: all (src_layer, src_head) → (dst_layer, dst_head) direct paths.
296 Returns a 4-D tensor of shape [n_layers, n_heads, n_layers, n_heads].
297 result[sl, sh, dl, dh] = patching metric when head (sl,sh)'s output is
298 patched directly into head (dl,dh)'s query/key/value input.
300 Entries where dl <= sl are 0 (no causal path).
302 This runs O(n_layers * n_heads * n_layers * n_heads) forward passes and is
303 intended for small models or targeted sub-sweeps. For large models prefer
304 calling get_act_patch_direct_path per source head.
305 """
306 _check_fold_ln(model)
308 n_layers = model.cfg.n_layers
309 n_heads = model.cfg.n_heads
310 results = torch.zeros(n_layers, n_heads, n_layers, n_heads, device=model.cfg.device)
312 src_pairs = [(sl, sh) for sl in range(n_layers) for sh in range(n_heads)]
313 for src_layer, src_head in tqdm(
314 src_pairs,
315 desc=f"Direct path patch — all sources [{component}]",
316 disable=not verbose,
317 ):
318 results[src_layer, src_head] = get_act_patch_direct_path(
319 model=model,
320 corrupted_tokens=corrupted_tokens,
321 clean_cache=clean_cache,
322 corrupted_cache=corrupted_cache,
323 patching_metric=patching_metric,
324 src_layer=src_layer,
325 src_head=src_head,
326 component=component,
327 verbose=False,
328 )
330 return results