Tool Orchestration Patterns & Interpretability-Reasoning Lab

Steps

1. Step 1: Load GPT-2 via TransformerLens.

   # pip install transformer-lens dspy-ai matplotlib numpy torch
   from transformer_lens import HookedTransformer
   model = HookedTransformer.from_pretrained("gpt2-small")
   print(f"Layers: {model.cfg.n_layers}, Heads: {model.cfg.n_heads}")

   Code Fragment 27.6a.1: Environment setup and model loading.
2. Step 2: Visualize attention patterns and scan for previous-token heads.

   import numpy as np
   prompt = "When Mary and John went to the store, John gave a drink to"
   tokens = model.to_tokens(prompt)
   logits, cache = model.run_with_cache(tokens)
   for layer in range(model.cfg.n_layers):
       for head in range(model.cfg.n_heads):
           pattern = cache["pattern", layer][0, head].detach().cpu().numpy()
           prev_attn = np.mean([pattern[i, i-1] for i in range(1, pattern.shape[0])])
           if prev_attn > 0.3:
               print(f"  L{layer}H{head}: prev-token attn = {prev_attn:.3f}")

     L0H7: prev-token attn = 0.412
     L1H4: prev-token attn = 0.358
     L4H11: prev-token attn = 0.331

   Code Fragment 27.6a.2: Extract attention patterns and identify heads that strongly attend to the immediately preceding token.
3. Step 3: Activation patching to identify circuits. Replace each head's output with its clean-run value when processing a corrupted prompt, then measure how much the target logit difference recovers. Heads with high recovery are causally responsible for the behavior.

   import torch
   clean, corrupt = ("When Mary and John went to the store, John gave a drink to",
                     "When Mary and Bob went to the store, John gave a drink to")
   clean_t, corr_t = model.to_tokens(clean), model.to_tokens(corrupt)
   _, clean_cache = model.run_with_cache(clean_t)
   mary_id = model.to_single_token(" Mary")
   john_id = model.to_single_token(" John")
   results = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
   for layer in range(model.cfg.n_layers):
       for head in range(model.cfg.n_heads):
           def patch(v, hook, l=layer, h=head):
               v[0, :, h, :] = clean_cache[hook.name][0, :, h, :]
               return v
           pl = model.run_with_hooks(corr_t, fwd_hooks=[(f"blocks.{layer}.attn.hook_result", patch)])
           results[layer, head] = (pl[0, -1, mary_id] - pl[0, -1, john_id]).item()

   Code Fragment 27.6a.3: Activation patching to localize the indirect-object-identification circuit. The heads with the largest logit-diff recovery are the causally important ones.
4. Step 4: DSPy structured reasoning pipeline. Use DSPy to turn the raw patching numbers into a testable hypothesis and a follow-up experiment, demonstrating the Right Tool pattern: TransformerLens for probing, DSPy for orchestrating multi-step LLM reasoning.

   import dspy
   dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))
   
   class AnalyzeCircuit(dspy.Signature):
       """Given activation-patching findings, hypothesize the circuit's function."""
       patching_results: str = dspy.InputField()
       task_description: str = dspy.InputField()
       hypothesis: str = dspy.OutputField()
       next_experiment: str = dspy.OutputField()
   
   analyzer = dspy.ChainOfThought(AnalyzeCircuit)
   out = analyzer(
       patching_results="Top heads: L9H9 +1.85, L9H6 +1.52, L10H0 +1.29 (promote Mary)",
       task_description="Indirect object identification: predict 'Mary' from 'John gave a drink to'",
   )
   print("Hypothesis:", out.hypothesis)
   print("Next experiment:", out.next_experiment)

   Output: Hypothesis: Layer 9-10 heads form an indirect-object-identification name-mover circuit; they copy the "Mary" token from its earlier position to the final prediction position. Next experiment: Test with novel name pairs (Alice/Bob) to verify generalization.

   Code Fragment 27.6a.4: DSPy ChainOfThought turning raw activation-patching numbers into a testable hypothesis and a follow-up experiment.