Demo Scripts
scripts/intervention_experiments.py
#!/usr/bin/env python3
"""
Intervention Experiments for Factual Association Analysis
This script demonstrates intervention techniques including hidden state patching,
sublayer knockout, and causal analysis of factual recall in language models.
Requires: pip install torch transformers numpy matplotlib
"""
import torch
import torch.nn.functional as F
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from typing import List, Dict, Tuple, Optional, Callable
import matplotlib.pyplot as plt
from dataclasses import dataclass
@dataclass
class InterventionResult:
"""Store results from intervention experiments."""
original_prediction: str
original_prob: float
intervened_prediction: str
intervened_prob: float
layer: int
intervention_type: str
class InterventionAnalyzer:
"""
Performs intervention experiments on language models to analyze factual recall.
"""
def __init__(self, model_name: str = "gpt2"):
"""Initialize the intervention analyzer."""
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = GPT2LMHeadModel.from_pretrained(model_name).to(self.device)
self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
self.model.eval()
# Get model configuration
self.n_layers = self.model.config.n_layer
self.n_heads = self.model.config.n_head
self.d_model = self.model.config.n_embd
print(f"Loaded {model_name}: {self.n_layers} layers, {self.n_heads} heads, {self.d_model} dims")
def get_predictions(self, text: str, top_k: int = 1) -> Tuple[List[str], List[float]]:
"""
Get model predictions for the next token.
Args:
text: Input text
top_k: Number of top predictions to return
Returns:
Tuple of (predicted tokens, probabilities)
"""
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits[0, -1]
probs = F.softmax(logits, dim=-1)
top_probs, top_indices = torch.topk(probs, k=top_k)
tokens = [self.tokenizer.decode([idx]) for idx in top_indices]
return tokens, top_probs.cpu().numpy()
def patch_hidden_states(self,
source_text: str,
target_text: str,
layer_idx: int,
position: int = -1) -> InterventionResult:
"""
Patch hidden states from source to target at a specific layer.
Args:
source_text: Text to extract hidden state from
target_text: Text to inject hidden state into
layer_idx: Layer at which to perform patching
position: Token position to patch (-1 for last)
Returns:
InterventionResult with analysis
"""
# Get original predictions
orig_tokens, orig_probs = self.get_predictions(target_text)
# Get source hidden states
source_inputs = self.tokenizer(source_text, return_tensors="pt").to(self.device)
with torch.no_grad():
source_outputs = self.model(**source_inputs, output_hidden_states=True)
source_hidden = source_outputs.hidden_states[layer_idx][0, position].unsqueeze(0).unsqueeze(0)
# Prepare target with intervention
target_inputs = self.tokenizer(target_text, return_tensors="pt").to(self.device)
# Define intervention hook
def patch_hook(module, input, output):
if isinstance(output, tuple):
hidden_states = output[0]
else:
hidden_states = output
# Patch the hidden state at specified position
hidden_states[0, position] = source_hidden[0, 0]
if isinstance(output, tuple):
return (hidden_states,) + output[1:]
return hidden_states
# Register hook and run with intervention
if layer_idx == 0:
handle = self.model.transformer.wte.register_forward_hook(patch_hook)
else:
handle = self.model.transformer.h[layer_idx-1].register_forward_hook(patch_hook)
with torch.no_grad():
intervened_outputs = self.model(**target_inputs)
intervened_logits = intervened_outputs.logits[0, -1]
intervened_probs = F.softmax(intervened_logits, dim=-1)
top_prob, top_idx = torch.topk(intervened_probs, k=1)
handle.remove()
intervened_token = self.tokenizer.decode([top_idx[0]])
return InterventionResult(
original_prediction=orig_tokens[0],
original_prob=orig_probs[0],
intervened_prediction=intervened_token,
intervened_prob=top_prob[0].cpu().item(),
layer=layer_idx,
intervention_type="hidden_state_patch"
)
def knockout_sublayer(self,
text: str,
layer_idx: int,
sublayer: str = "attn") -> InterventionResult:
"""
Knockout (zero out) a specific sublayer's contribution.
Args:
text: Input text
layer_idx: Layer index
sublayer: "attn" or "mlp"
Returns:
InterventionResult with analysis
"""
# Get original predictions
orig_tokens, orig_probs = self.get_predictions(text)
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
# Define knockout hook
def knockout_hook(module, input, output):
if isinstance(output, tuple):
# Return zeros for the sublayer output
zeros = torch.zeros_like(output[0])
return (zeros,) + output[1:] if len(output) > 1 else zeros
else:
return torch.zeros_like(output)
# Register hook on appropriate sublayer
if sublayer == "attn":
target_module = self.model.transformer.h[layer_idx].attn
elif sublayer == "mlp":
target_module = self.model.transformer.h[layer_idx].mlp
else:
raise ValueError(f"Unknown sublayer: {sublayer}")
handle = target_module.register_forward_hook(knockout_hook)
with torch.no_grad():
knockout_outputs = self.model(**inputs)
knockout_logits = knockout_outputs.logits[0, -1]
knockout_probs = F.softmax(knockout_logits, dim=-1)
top_prob, top_idx = torch.topk(knockout_probs, k=1)
handle.remove()
knockout_token = self.tokenizer.decode([top_idx[0]])
return InterventionResult(
original_prediction=orig_tokens[0],
original_prob=orig_probs[0],
intervened_prediction=knockout_token,
intervened_prob=top_prob[0].cpu().item(),
layer=layer_idx,
intervention_type=f"{sublayer}_knockout"
)
def analyze_factual_flow(self,
prompt: str,
expected_answer: str) -> Dict:
"""
Analyze how factual information flows through the network.
Args:
prompt: Question or prompt requiring factual knowledge
expected_answer: Expected factual answer
Returns:
Dictionary with layer-wise analysis
"""
results = {
"prompt": prompt,
"expected": expected_answer,
"layer_analysis": []
}
# Get base prediction
base_tokens, base_probs = self.get_predictions(prompt)
results["base_prediction"] = base_tokens[0]
results["base_confidence"] = base_probs[0]
# Analyze each layer's contribution
for layer_idx in range(self.n_layers):
layer_info = {"layer": layer_idx}
# Test attention knockout
attn_result = self.knockout_sublayer(prompt, layer_idx, "attn")
layer_info["attn_knockout_changes_prediction"] = (
attn_result.intervened_prediction != base_tokens[0]
)
# Test MLP knockout
mlp_result = self.knockout_sublayer(prompt, layer_idx, "mlp")
layer_info["mlp_knockout_changes_prediction"] = (
mlp_result.intervened_prediction != base_tokens[0]
)
# Calculate importance scores
layer_info["attn_importance"] = abs(
base_probs[0] - attn_result.intervened_prob
)
layer_info["mlp_importance"] = abs(
base_probs[0] - mlp_result.intervened_prob
)
results["layer_analysis"].append(layer_info)
return results
def visualize_layer_importance(self, analysis_results: Dict):
"""
Visualize the importance of different layers for factual prediction.
Args:
analysis_results: Results from analyze_factual_flow
"""
layers = [item["layer"] for item in analysis_results["layer_analysis"]]
attn_importance = [item["attn_importance"] for item in analysis_results["layer_analysis"]]
mlp_importance = [item["mlp_importance"] for item in analysis_results["layer_analysis"]]
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
# Plot attention importance
ax1.bar(layers, attn_importance, color='blue', alpha=0.7)
ax1.set_xlabel('Layer')
ax1.set_ylabel('Importance Score')
ax1.set_title('Attention Sublayer Importance for Factual Prediction')
ax1.grid(True, alpha=0.3)
# Plot MLP importance
ax2.bar(layers, mlp_importance, color='green', alpha=0.7)
ax2.set_xlabel('Layer')
ax2.set_ylabel('Importance Score')
ax2.set_title('MLP Sublayer Importance for Factual Prediction')
ax2.grid(True, alpha=0.3)
plt.suptitle(f'Layer Analysis: "{analysis_results["prompt"]}"')
plt.tight_layout()
return fig
def run_intervention_experiments():
"""
Run a series of intervention experiments to analyze factual recall.
"""
print("Initializing Intervention Analyzer...")
analyzer = InterventionAnalyzer("gpt2")
# Test cases for factual associations
test_cases = [
("The capital of France is", "Paris"),
("Water freezes at", "0"),
("The largest planet is", "Jupiter"),
]
for prompt, expected in test_cases:
print(f"\n{'='*60}")
print(f"Analyzing: '{prompt}'")
print(f"Expected: '{expected}'")
print("-" * 60)
# Get baseline prediction
predictions, probs = analyzer.get_predictions(prompt, top_k=3)
print(f"\nTop 3 predictions: {list(zip(predictions, probs))}")
# Test hidden state patching
print("\n--- Hidden State Patching Experiment ---")
# Create a source that strongly predicts the expected answer
source_text = f"The answer is definitely {expected}"
for layer in [0, 6, 11]: # Early, middle, late layers
patch_result = analyzer.patch_hidden_states(
source_text, prompt, layer
)
print(f"Layer {layer}: {patch_result.original_prediction} -> {patch_result.intervened_prediction}")
# Test sublayer knockout
print("\n--- Sublayer Knockout Experiment ---")
critical_layers = []
for layer in range(analyzer.n_layers):
attn_ko = analyzer.knockout_sublayer(prompt, layer, "attn")
mlp_ko = analyzer.knockout_sublayer(prompt, layer, "mlp")
if attn_ko.intervened_prediction != predictions[0]:
critical_layers.append(f"L{layer}-attn")
if mlp_ko.intervened_prediction != predictions[0]:
critical_layers.append(f"L{layer}-mlp")
print(f"Critical sublayers (change prediction): {critical_layers[:5]}...")
# Comprehensive factual flow analysis
print("\n--- Factual Flow Analysis ---")
flow_analysis = analyzer.analyze_factual_flow(prompt, expected)
# Find most important layers
layer_data = flow_analysis["layer_analysis"]
most_important_attn = max(layer_data, key=lambda x: x["attn_importance"])
most_important_mlp = max(layer_data, key=lambda x: x["mlp_importance"])
print(f"Most important attention layer: {most_important_attn['layer']} "
f"(importance: {most_important_attn['attn_importance']:.4f})")
print(f"Most important MLP layer: {most_important_mlp['layer']} "
f"(importance: {most_important_mlp['mlp_importance']:.4f})")
# Visualize if matplotlib is available
try:
fig = analyzer.visualize_layer_importance(flow_analysis)
# Save figure
fig.savefig(f'layer_importance_{prompt[:20].replace(" ", "_")}.png')
print(f"Visualization saved to layer_importance_{prompt[:20].replace(' ', '_')}.png")
plt.close(fig)
except:
print("Visualization skipped (matplotlib issue)")
if __name__ == "__main__":
print("Starting Intervention Experiments for Factual Association Analysis")
print("=" * 60)
# Check GPU availability
if torch.cuda.is_available():
print(f"GPU available: {torch.cuda.get_device_name(0)}")
else:
print("Running on CPU (GPU recommended)")
# Run experiments
run_intervention_experiments()
print("\n" + "=" * 60)
print("Experiments complete!")
print("\nNote: For more detailed analysis with GPT2-xl or GPT-J,")
print("use appropriate GPU resources (V100/A100) and adjust model name.")
scripts/model_analysis.py
#!/usr/bin/env python3
"""
Model Analysis Script for Dissecting Factual Predictions in Language Models
This script demonstrates how to load GPT-2 models, apply hooks for intermediate
information extraction, and perform basic analysis of factual associations.
Requires: pip install torch transformers numpy
"""
import torch
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from typing import List, Dict, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')
class ModelAnalyzer:
"""
Analyzer for dissecting factual predictions in GPT-2 style models.
"""
def __init__(self, model_name: str = "gpt2"):
"""
Initialize the model analyzer with a specific GPT-2 model.
Args:
model_name: Name of the GPT-2 model (e.g., "gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl")
"""
print(f"Loading model: {model_name}")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = GPT2LMHeadModel.from_pretrained(model_name).to(self.device)
self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
self.model.eval()
# Storage for intermediate activations
self.stored_activations = {}
self.hooks = []
def register_hook(self, module_name: str, hook_type: str = "forward"):
"""
Register a hook on a specific module to extract intermediate information.
Args:
module_name: Name of the module to hook (e.g., "transformer.h.0.attn")
hook_type: Type of hook ("forward" or "backward")
"""
def get_activation(name):
def hook(module, input, output):
self.stored_activations[name] = output.detach()
return hook
# Navigate to the module
module = self.model
for name in module_name.split('.'):
if name.isdigit():
module = module[int(name)]
else:
module = getattr(module, name)
# Register the hook
if hook_type == "forward":
handle = module.register_forward_hook(get_activation(module_name))
else:
handle = module.register_backward_hook(get_activation(module_name))
self.hooks.append(handle)
print(f"Hook registered on {module_name}")
def clear_hooks(self):
"""Remove all registered hooks."""
for hook in self.hooks:
hook.remove()
self.hooks = []
self.stored_activations = {}
def analyze_attention_flow(self, text: str, layer_idx: int = 0) -> Dict:
"""
Analyze attention flow for a given text at a specific layer.
Args:
text: Input text to analyze
layer_idx: Layer index to analyze (0 to n_layers-1)
Returns:
Dictionary containing attention weights and analysis
"""
# Tokenize input
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
# Register hook on attention layer
attn_module = f"transformer.h.{layer_idx}.attn"
self.register_hook(attn_module)
# Forward pass
with torch.no_grad():
outputs = self.model(**inputs)
# Extract attention patterns
if attn_module in self.stored_activations:
attn_output = self.stored_activations[attn_module]
# For GPT-2, attention output is a tuple
if isinstance(attn_output, tuple):
attn_weights = attn_output[1] if len(attn_output) > 1 else None
attn_output = attn_output[0]
self.clear_hooks()
return {
"layer": layer_idx,
"attention_output_shape": attn_output.shape if 'attn_output' in locals() else None,
"input_tokens": self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0].cpu())
}
def project_to_vocabulary(self, hidden_state: torch.Tensor) -> torch.Tensor:
"""
Project hidden states to vocabulary space using the language model head.
Args:
hidden_state: Hidden state tensor to project
Returns:
Vocabulary projections (logits)
"""
with torch.no_grad():
# Apply final layer norm if needed (GPT-2 specific)
if hasattr(self.model, 'transformer') and hasattr(self.model.transformer, 'ln_f'):
hidden_state = self.model.transformer.ln_f(hidden_state)
# Project through language model head
vocab_projection = self.model.lm_head(hidden_state)
return vocab_projection
def extract_mlp_contributions(self, text: str, layer_idx: int = 0) -> Dict:
"""
Extract and analyze MLP sublayer contributions.
Args:
text: Input text to analyze
layer_idx: Layer index to analyze
Returns:
Dictionary with MLP analysis results
"""
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
# Register hooks on MLP components
mlp_in = f"transformer.h.{layer_idx}.mlp.c_fc" # First MLP matrix
mlp_out = f"transformer.h.{layer_idx}.mlp.c_proj" # Second MLP matrix
self.register_hook(mlp_in)
self.register_hook(mlp_out)
# Forward pass
with torch.no_grad():
outputs = self.model(**inputs)
results = {
"layer": layer_idx,
"mlp_input_shape": None,
"mlp_output_shape": None
}
if mlp_in in self.stored_activations:
results["mlp_input_shape"] = self.stored_activations[mlp_in].shape
if mlp_out in self.stored_activations:
results["mlp_output_shape"] = self.stored_activations[mlp_out].shape
# Project MLP output to vocabulary
mlp_output = self.stored_activations[mlp_out]
vocab_proj = self.project_to_vocabulary(mlp_output)
# Get top predicted tokens
top_tokens = torch.topk(vocab_proj[0, -1], k=5)
results["top_mlp_predictions"] = [
self.tokenizer.decode([idx]) for idx in top_tokens.indices
]
self.clear_hooks()
return results
def perform_attention_knockout(self, text: str, layer_idx: int, head_idx: int) -> Dict:
"""
Perform attention knockout experiment by zeroing specific attention heads.
Args:
text: Input text
layer_idx: Layer index
head_idx: Head index to knockout
Returns:
Analysis results with and without knockout
"""
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
# Normal forward pass
with torch.no_grad():
normal_outputs = self.model(**inputs)
normal_logits = normal_outputs.logits
normal_probs = torch.softmax(normal_logits[0, -1], dim=-1)
normal_top = torch.topk(normal_probs, k=5)
# Forward pass with attention knockout
def knockout_hook(module, input, output):
# Modify attention output to zero out specific head
if isinstance(output, tuple):
attn_output, attn_weights = output[0], output[1] if len(output) > 1 else None
# Simple knockout: zero contribution from specific head
# Note: This is a simplified version; actual implementation may vary
modified_output = attn_output.clone()
return (modified_output, attn_weights) if attn_weights is not None else modified_output
return output
# Register knockout hook
attn_module = self.model.transformer.h[layer_idx].attn
handle = attn_module.register_forward_hook(knockout_hook)
with torch.no_grad():
knockout_outputs = self.model(**inputs)
knockout_logits = knockout_outputs.logits
knockout_probs = torch.softmax(knockout_logits[0, -1], dim=-1)
knockout_top = torch.topk(knockout_probs, k=5)
handle.remove()
return {
"layer": layer_idx,
"head": head_idx,
"normal_predictions": [self.tokenizer.decode([idx]) for idx in normal_top.indices],
"knockout_predictions": [self.tokenizer.decode([idx]) for idx in knockout_top.indices],
"prediction_change": not torch.equal(normal_top.indices, knockout_top.indices)
}
def demonstrate_analysis():
"""
Demonstrate the model analysis capabilities.
"""
# Initialize analyzer with base GPT-2 model
analyzer = ModelAnalyzer("gpt2")
# Example text for analysis
test_text = "The capital of France is"
print("\n=== Analyzing Factual Association ===")
print(f"Input text: '{test_text}'")
# 1. Analyze attention flow
print("\n--- Attention Flow Analysis ---")
attn_analysis = analyzer.analyze_attention_flow(test_text, layer_idx=0)
print(f"Layer {attn_analysis['layer']} attention analysis:")
print(f" Input tokens: {attn_analysis['input_tokens']}")
print(f" Attention output shape: {attn_analysis['attention_output_shape']}")
# 2. Extract MLP contributions
print("\n--- MLP Contribution Analysis ---")
mlp_analysis = analyzer.extract_mlp_contributions(test_text, layer_idx=5)
print(f"Layer {mlp_analysis['layer']} MLP analysis:")
print(f" MLP input shape: {mlp_analysis['mlp_input_shape']}")
print(f" MLP output shape: {mlp_analysis['mlp_output_shape']}")
if 'top_mlp_predictions' in mlp_analysis:
print(f" Top MLP predictions: {mlp_analysis['top_mlp_predictions']}")
# 3. Perform attention knockout
print("\n--- Attention Knockout Experiment ---")
knockout_results = analyzer.perform_attention_knockout(test_text, layer_idx=5, head_idx=0)
print(f"Layer {knockout_results['layer']}, Head {knockout_results['head']} knockout:")
print(f" Normal predictions: {knockout_results['normal_predictions']}")
print(f" Knockout predictions: {knockout_results['knockout_predictions']}")
print(f" Prediction changed: {knockout_results['prediction_change']}")
# 4. Hidden state projection
print("\n--- Hidden State Vocabulary Projection ---")
inputs = analyzer.tokenizer(test_text, return_tensors="pt").to(analyzer.device)
with torch.no_grad():
outputs = analyzer.model(**inputs, output_hidden_states=True)
last_hidden = outputs.hidden_states[-1] # Last layer hidden states
vocab_proj = analyzer.project_to_vocabulary(last_hidden)
top_tokens = torch.topk(vocab_proj[0, -1], k=5)
predictions = [analyzer.tokenizer.decode([idx]) for idx in top_tokens.indices]
print(f"Top vocabulary projections from last hidden layer: {predictions}")
if __name__ == "__main__":
print("Starting Factual Associations Dissection Analysis")
print("=" * 50)
# Check for GPU
if torch.cuda.is_available():
print(f"GPU available: {torch.cuda.get_device_name(0)}")
else:
print("Running on CPU (GPU recommended for larger models)")
# Run demonstration
demonstrate_analysis()
print("\n" + "=" * 50)
print("Analysis complete!")
print("\nNote: For full experiments with GPT2-xl, use a V100 GPU.")
print("For GPT-J experiments, use an A100 GPU.")
zjunlp