Coverage for transformer_lens/tools/analysis/direct_path_patching.py: 94%
60 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 w = getattr(ln1, "weight", None)
88 if w is not None and not torch.allclose(w, torch.ones_like(w), atol=1e-3):
89 warnings.warn(
90 "get_act_patch_direct_path is most accurate when LayerNorm parameters "
91 "are folded into the weight matrices: call "
92 "model.process_weights(fold_ln=True) on the TransformerBridge. "
93 "Results may be inaccurate with unfolded LayerNorm.",
94 UserWarning,
95 stacklevel=3,
96 )
97 except (AttributeError, TypeError):
98 pass # non-standard model — cannot inspect LN weights, proceed
101# ---------------------------------------------------------------------------
102# Core hook factory
103# ---------------------------------------------------------------------------
106def _make_direct_path_hook(
107 delta_resid: Float[torch.Tensor, "batch pos d_model"],
108 dst_head: int,
109 W_component: Float[torch.Tensor, "d_model d_head"],
110 ln_scale_name: str,
111 corrupted_cache: ActivationCache,
112 component: Literal["q", "k", "v"],
113) -> Callable:
114 """Return a hook function that adds the linearised delta to one head's Q, K, or V.
116 Parameters
117 ----------
118 delta_resid:
119 (clean_A_result - corrupted_A_result), shape [batch, pos, d_model].
120 dst_head:
121 Index of the destination attention head to patch.
122 W_component:
123 The weight matrix for the component being patched:
124 W_Q[dst_head], W_K[dst_head], or W_V[dst_head].
125 Shape [d_model, d_head].
126 ln_scale_name:
127 Cache key for the layer-norm scale at the destination layer,
128 e.g. "blocks.3.ln1.hook_scale".
129 corrupted_cache:
130 Cache from the corrupted forward pass (used to look up ln1 scale).
131 component:
132 One of "q", "k", "v" — determines which QKV tensor is hooked.
133 """
135 def hook_fn(
136 value: Float[torch.Tensor, "batch pos n_heads d_head"],
137 hook, # HookPoint, unused but required by TransformerLens
138 ) -> Float[torch.Tensor, "batch pos n_heads d_head"]:
139 # ln scale: [batch, pos, 1]
140 ln_scale = corrupted_cache[ln_scale_name] # [batch, pos, 1]
142 # Linearised delta in query/key/value space
143 # delta_resid: [batch, pos, d_model]
144 # W_component: [d_model, d_head]
145 delta = (delta_resid / ln_scale) @ W_component # [batch, pos, d_head]
147 if value.requires_grad: 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 value = value.clone()
149 value[:, :, dst_head, :] = value[:, :, dst_head, :] + delta
150 return value
152 return hook_fn
155# ---------------------------------------------------------------------------
156# Public API
157# ---------------------------------------------------------------------------
160def get_act_patch_direct_path(
161 model: TransformerLensModel,
162 corrupted_tokens: torch.Tensor,
163 clean_cache: ActivationCache,
164 corrupted_cache: ActivationCache,
165 patching_metric: Callable[[torch.Tensor], torch.Tensor],
166 src_layer: int,
167 src_head: int,
168 component: Literal["q", "k", "v"] = "q",
169 verbose: bool = True,
170) -> Float[torch.Tensor, "n_layers n_heads"]:
171 """Sweep direct path patches from one source head to all downstream heads.
173 For every destination head B = (dst_layer, dst_head) where dst_layer > src_layer,
174 patch the contribution of source head A = (src_layer, src_head) into B's query
175 (or key / value) input, and record the patching metric.
177 The patch is a linear approximation:
179 delta_resid = clean_A_result - corrupted_A_result [batch, pos, d_model]
180 delta_B_comp = (delta_resid / ln1_scale) @ W_comp[dst_head]
182 where W_comp is W_Q, W_K, or W_V according to `component`.
184 Parameters
185 ----------
186 model:
187 Any TransformerLens model (e.g. TransformerBridge).
188 corrupted_tokens:
189 Token IDs for the corrupted input, shape [batch, seq_len].
190 clean_cache:
191 Cached activations from the clean (unpatched) run.
192 corrupted_cache:
193 Cached activations from the corrupted run (needed for ln1 scale).
194 patching_metric:
195 A function mapping the model's logits tensor to a scalar.
196 src_layer:
197 Layer index of the source attention head.
198 src_head:
199 Head index of the source attention head.
200 component:
201 Which input to patch at the destination head — "q" (default), "k", or "v".
202 verbose:
203 Whether to show a tqdm progress bar.
205 Returns
206 -------
207 results : Float[Tensor, "n_layers n_heads"]
208 results[dst_layer, dst_head] is the patching metric when the direct path
209 A → B is patched in. Entries for dst_layer <= src_layer are left as 0.0
210 (no causal path from A to those layers).
211 """
212 _check_fold_ln(model)
214 n_layers = model.cfg.n_layers
215 n_heads = model.cfg.n_heads
217 results = torch.zeros(n_layers, n_heads, device=model.cfg.device)
219 # Residual stream delta from source head A.
220 #
221 # hook_result (per-head residual contribution) requires cfg.use_hook_result=True
222 # and is not in the default cache. We compute it instead from hook_z and W_O,
223 # which are always available:
224 # result_h = z[:, :, h, :] @ W_O[h] shape [batch, pos, d_model]
225 src_z_name = f"blocks.{src_layer}.attn.hook_z"
226 # Static-only view of the weights surface (runtime checks use the base protocol:
227 # nn.Module submodules are invisible to getattr_static-based isinstance).
228 weights_model = cast(TransformerLensModelWithWeights, model)
229 W_O = weights_model.blocks[src_layer].attn.W_O # [n_heads, d_head, d_model]
231 def _head_result(cache, h):
232 z = cache[src_z_name][:, :, h, :] # [batch, pos, d_head]
233 return z @ W_O[h] # [batch, pos, d_model]
235 delta_resid = _head_result(clean_cache, src_head) - _head_result(corrupted_cache, src_head)
236 # shape: [batch, pos, d_model]
238 # Weight matrix for the component being patched
239 _comp_map = {
240 "q": lambda attn: attn.W_Q, # [n_heads, d_model, d_head]
241 "k": lambda attn: attn.W_K,
242 "v": lambda attn: attn.W_V,
243 }
244 _hook_name_map = {
245 "q": lambda lb: f"blocks.{lb}.attn.hook_q",
246 "k": lambda lb: f"blocks.{lb}.attn.hook_k",
247 "v": lambda lb: f"blocks.{lb}.attn.hook_v",
248 }
249 W_all = _comp_map[component] # callable: attn → [n_heads, d_model, d_head]
250 hook_name_fn = _hook_name_map[component]
252 dst_pairs = [(lb, hb) for lb in range(src_layer + 1, n_layers) for hb in range(n_heads)]
254 for dst_layer, dst_head in tqdm(
255 dst_pairs,
256 desc=f"Direct path patch ({src_layer},{src_head}) → * [{component}]",
257 disable=not verbose,
258 ):
259 ln_scale_name = f"blocks.{dst_layer}.ln1.hook_scale"
260 W_comp = W_all(weights_model.blocks[dst_layer].attn)[dst_head] # [d_model, d_head]
262 hook_fn = _make_direct_path_hook(
263 delta_resid=delta_resid,
264 dst_head=dst_head,
265 W_component=W_comp,
266 ln_scale_name=ln_scale_name,
267 corrupted_cache=corrupted_cache,
268 component=component,
269 )
271 patched_logits = model.run_with_hooks(
272 corrupted_tokens,
273 fwd_hooks=[(hook_name_fn(dst_layer), hook_fn)],
274 )
276 results[dst_layer, dst_head] = patching_metric(patched_logits).item()
278 return results
281def get_act_patch_direct_path_all_sources(
282 model: TransformerLensModel,
283 corrupted_tokens: torch.Tensor,
284 clean_cache: ActivationCache,
285 corrupted_cache: ActivationCache,
286 patching_metric: Callable[[torch.Tensor], torch.Tensor],
287 component: Literal["q", "k", "v"] = "q",
288 verbose: bool = True,
289) -> Float[torch.Tensor, "n_layers n_heads n_layers n_heads"]:
290 """Full sweep: all (src_layer, src_head) → (dst_layer, dst_head) direct paths.
292 Returns a 4-D tensor of shape [n_layers, n_heads, n_layers, n_heads].
293 result[sl, sh, dl, dh] = patching metric when head (sl,sh)'s output is
294 patched directly into head (dl,dh)'s query/key/value input.
296 Entries where dl <= sl are 0 (no causal path).
298 This runs O(n_layers * n_heads * n_layers * n_heads) forward passes and is
299 intended for small models or targeted sub-sweeps. For large models prefer
300 calling get_act_patch_direct_path per source head.
301 """
302 _check_fold_ln(model)
304 n_layers = model.cfg.n_layers
305 n_heads = model.cfg.n_heads
306 results = torch.zeros(n_layers, n_heads, n_layers, n_heads, device=model.cfg.device)
308 src_pairs = [(sl, sh) for sl in range(n_layers) for sh in range(n_heads)]
309 for src_layer, src_head in tqdm(
310 src_pairs,
311 desc=f"Direct path patch — all sources [{component}]",
312 disable=not verbose,
313 ):
314 results[src_layer, src_head] = get_act_patch_direct_path(
315 model=model,
316 corrupted_tokens=corrupted_tokens,
317 clean_cache=clean_cache,
318 corrupted_cache=corrupted_cache,
319 patching_metric=patching_metric,
320 src_layer=src_layer,
321 src_head=src_head,
322 component=component,
323 verbose=False,
324 )
326 return results