transformer_lens.tools.analysis.direct_path_patching module

Deprecated location of direct path patching, now in transformer_lens.patching.

transformer_lens.tools.analysis.direct_path_patching.get_act_patch_direct_path(model: TransformerLensModel, corrupted_tokens: Tensor, clean_cache: ActivationCache, corrupted_cache: ActivationCache, patching_metric: Callable[[Tensor], Tensor], src_layer: int, src_head: int, component: Literal['q', 'k', 'v'] = 'q', verbose: bool = True) → Float[Tensor, 'n_layers n_heads']

Patch one source head’s output into each later head’s Q, K or V input alone.

Direct path patching (Wang et al. 2022) isolates the edge A -> B: A’s clean-minus-corrupted write (z @ W_O) is added at every position to B’s own pre-ln1 copy of the residual, while every other reader still sees A’s corrupted output. The result is the exact edge intervention, folded LayerNorm or not.

The sweep temporarily enables cfg.use_split_qkv_input (restoring the caller’s flags on exit), which costs about 3 * n_heads times the usual ln1 activation memory per forward.

Parameters:

  • model – A TransformerBridge whose attention exposes the pre-Q/K/V fork.

  • corrupted_tokens – Token ids for the corrupted run, [batch, pos].

  • clean_cache – Cache of the clean run; must hold blocks.{src_layer}.attn.hook_z.

  • corrupted_cache – Cache of the corrupted run; must hold the same hook_z.

  • patching_metric – Maps the patched logits to a scalar.

  • src_layer – Layer of the source head.

  • src_head – Index of the source head.

  • component – Destination input to patch: "q" (default), "k" or "v".

Returns:

results[l, h] is the metric with the path into head (l, h) patched. Rows l <= src_layer are 0.0 placeholders, not measurements.

Return type:

[n_layers, n_heads]

Raises:

  • ValueError – if a block from src_layer on has no attention, if the source block normalises, scales or joins its attention output before the residual add (post-norm, scaled-residual or parallel-mixer blocks), if component is "k" or "v" on a GQA/MQA model, or if any module is in training mode.

  • NotImplementedError – if a destination layer’s attention has no pre-Q/K/V fork.

transformer_lens.tools.analysis.direct_path_patching.get_act_patch_direct_path_all_sources(model: TransformerLensModel, corrupted_tokens: Tensor, clean_cache: ActivationCache, corrupted_cache: ActivationCache, patching_metric: Callable[[Tensor], Tensor], component: Literal['q', 'k', 'v'] = 'q', verbose: bool = True) → Float[Tensor, 'n_layers n_heads n_layers n_heads']

Run get_act_patch_direct_path() from every source head.

results[sl, sh, dl, dh] is the metric with the path from head (sl, sh) into head (dl, dh) patched, and 0.0 where dl <= sl. One forward per (source, later destination) pair, so it suits small models or targeted sub-sweeps.