Mechanistic Interpretability Pyvene – Tutorials

# pyvene Tutorials ## Tutorial 1: Basic Activation Patching ### Goal Swap activations between two prompts to test causal relationships. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer import torch # 1. Load model model = AutoModelForCausalLM.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2") # 2. Prepare inputs base_prompt = "The Colosseum is in the city of" source_prompt = "The Eiffel Tower is in the city of" base_inputs = tokenizer(base_prompt, return_tensors="pt") source_inputs = tokenizer(source_prompt, return_tensors="pt") # 3. Define intervention (patch layer 8) config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=8, component="block_output", intervention_type=pv.VanillaIntervention, ) ] ) intervenable = pv.IntervenableModel(config, model) # 4. Run intervention _, patched_outputs = intervenable( base=base_inputs, sources=[source_inputs], ) # 5. Check predictions patched_logits = patched_outputs.logits probs = torch.softmax(patched_logits[0, -1], dim=-1) rome_token = tokenizer.encode(" Rome")[0] paris_token = tokenizer.encode(" Paris")[0] print(f"P(Rome): {probs[rome_token].item():.4f}") print(f"P(Paris): {probs[paris_token].item():.4f}") ``` --- ## Tutorial 2: Causal Tracing (ROME-style) ### Goal Locate where factual associations are stored by corrupting inputs and restoring activations. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer import torch model = AutoModelForCausalLM.from_pretrained("gpt2-xl") tokenizer = AutoTokenizer.from_pretrained("gpt2-xl") # 1. Define prompts clean_prompt = "The Space Needle is in downtown" # We'll corrupt by adding noise to embeddings clean_inputs = tokenizer(clean_prompt, return_tensors="pt") seattle_token = tokenizer.encode(" Seattle")[0] # 2. Get clean baseline with torch.no_grad(): clean_outputs = model(**clean_inputs) clean_prob = torch.softmax(clean_outputs.logits[0, -1], dim=-1)[seattle_token].item() print(f"Clean P(Seattle): {clean_prob:.4f}") # 3. Sweep over layers - corrupt input, restore at each layer results = [] for restore_layer in range(model.config.n_layer): # Config: add noise at input, restore at target layer config = pv.IntervenableConfig( representations=[ # Noise intervention at embedding pv.RepresentationConfig( layer=0, component="block_input", intervention_type=pv.NoiseIntervention, ), # Restore clean at target layer pv.RepresentationConfig( layer=restore_layer, component="block_output", intervention_type=pv.VanillaIntervention, ), ] ) intervenable = pv.IntervenableModel(config, model) # Source is clean (for restoration), base gets noise _, outputs = intervenable( base=clean_inputs, sources=[clean_inputs], # Restore from clean ) prob = torch.softmax(outputs.logits[0, -1], dim=-1)[seattle_token].item() results.append(prob) print(f"Restore at layer {restore_layer}: P(Seattle) = {prob:.4f}") # 4. Find critical layers (where restoration helps most) import numpy as np results = np.array(results) critical_layers = np.argsort(results)[-5:] print(f"nMost critical layers: {critical_layers}") ``` --- ## Tutorial 3: Trainable Interventions (DAS) ### Goal Learn a low-rank intervention that achieves a target counterfactual behavior. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer import torch model = AutoModelForCausalLM.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2") # 1. Define trainable intervention config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=6, component="block_output", intervention_type=pv.LowRankRotatedSpaceIntervention, low_rank_dimension=64, # Learn 64-dim subspace ) ] ) intervenable = pv.IntervenableModel(config, model) # 2. Setup optimizer optimizer = torch.optim.Adam( intervenable.get_trainable_parameters(), lr=1e-3 ) # 3. Training data (simplified example) # Goal: Make model predict "Paris" instead of "Rome" base_prompt = "The capital of Italy is" target_token = tokenizer.encode(" Paris")[0] base_inputs = tokenizer(base_prompt, return_tensors="pt") # 4. Training loop for step in range(100): optimizer.zero_grad() _, outputs = intervenable( base=base_inputs, sources=[base_inputs], # Self-intervention ) # Loss: maximize probability of target token logits = outputs.logits[0, -1] loss = -torch.log_softmax(logits, dim=-1)[target_token] loss.backward() optimizer.step() if step % 20 == 0: prob = torch.softmax(logits.detach(), dim=-1)[target_token].item() print(f"Step {step}: loss={loss.item():.4f}, P(Paris)={prob:.4f}") # 5. Analyze learned rotation rotation = intervenable.interventions["layer.6.comp.block_output.unit.pos.nunit.1#0"][0] print(f"Learned rotation shape: {rotation.rotate_layer.weight.shape}") ``` --- ## Tutorial 4: Position-Specific Intervention ### Goal Intervene at specific token positions only. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2") # 1. Setup base_prompt = "John and Mary went to the store" source_prompt = "Alice and Bob went to the store" base_inputs = tokenizer(base_prompt, return_tensors="pt") source_inputs = tokenizer(source_prompt, return_tensors="pt") # 2. Position-specific config config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=5, component="block_output", intervention_type=pv.VanillaIntervention, unit="pos", max_number_of_units=1, # Single position ) ] ) intervenable = pv.IntervenableModel(config, model) # 3. Intervene at position 0 only (first name) _, outputs = intervenable( base=base_inputs, sources=[source_inputs], unit_locations={"sources->base": ([[[0]]], [[[0]]])}, ) # 4. Intervene at multiple positions _, outputs = intervenable( base=base_inputs, sources=[source_inputs], unit_locations={"sources->base": ([[[0, 2]]], [[[0, 2]]])}, ) ``` --- ## Tutorial 5: Collecting Activations ### Goal Extract activations without modifying them. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2") # 1. Config with CollectIntervention config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=5, component="block_output", intervention_type=pv.CollectIntervention, ), pv.RepresentationConfig( layer=10, component="attention_output", intervention_type=pv.CollectIntervention, ), ] ) intervenable = pv.IntervenableModel(config, model) # 2. Run and collect inputs = tokenizer("Hello world", return_tensors="pt") _, collected = intervenable(base=inputs) # 3. Access collected activations layer5_output = collected[0] layer10_attn = collected[1] print(f"Layer 5 block output shape: {layer5_output.shape}") print(f"Layer 10 attention output shape: {layer10_attn.shape}") ``` --- ## Tutorial 6: Generation with Interventions ### Goal Apply interventions during text generation. ### Step-by-Step ```python import pyvene as pv from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer.pad_token = tokenizer.eos_token # 1. Get steering direction (happy vs sad) happy_inputs = tokenizer("I am very happy and", return_tensors="pt") sad_inputs = tokenizer("I am very sad and", return_tensors="pt") # Collect activations config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=6, component="mlp_output", intervention_type=pv.CollectIntervention, ) ] ) collector = pv.IntervenableModel(config, model) _, happy_acts = collector(base=happy_inputs) _, sad_acts = collector(base=sad_inputs) steering_direction = happy_acts[0].mean(dim=1) - sad_acts[0].mean(dim=1) # 2. Config for steering during generation config = pv.IntervenableConfig( representations=[ pv.RepresentationConfig( layer=6, component="mlp_output", intervention_type=pv.AdditionIntervention, ) ] ) intervenable = pv.IntervenableModel(config, model) # 3. Generate with steering prompt = "Today I feel" inputs = tokenizer(prompt, return_tensors="pt") # Create source with steering direction # (This is simplified - actual implementation varies) output = intervenable.generate( inputs, max_new_tokens=20, do_sample=True, temperature=0.7, ) print(tokenizer.decode(output[0])) ``` --- ## External Resources ### Official Tutorials - [pyvene 101](https://stanfordnlp.github.io/pyvene/tutorials/pyvene_101.html) - [Causal Tracing](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/Causal_Tracing.html) - [DAS Introduction](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/DAS_Main_Introduction.html) - [IOI Replication](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/IOI_Replication.html) ### Papers - [pyvene Paper](https://arxiv.org/abs/2403.07809) - NAACL 2024 - [ROME](https://arxiv.org/abs/2202.05262) - Meng et al. (2022) - [Inference-Time Intervention](https://arxiv.org/abs/2306.03341) - Li et al. (2023)