{"id":"pyvene","name":"pyvene-interventions","summary":"pyveneの宣言的介入フレームワークを用いたPyTorchモデルに対する因果介入の実施に関する指針を提供します。","body":"# pyvene: Causal Interventions for Neural Networks\n\npyvene is Stanford NLP's library for performing causal interventions on PyTorch models. It provides a declarative, dict-based framework for activation patching, causal tracing, and interchange intervention training - making intervention experiments reproducible and shareable.\n\n**GitHub**: [stanfordnlp/pyvene](https://github.com/stanfordnlp/pyvene) (840+ stars)\n**Paper**: [pyvene: A Library for Understanding and Improving PyTorch Models via Interventions](https://aclanthology.org/2024.naacl-demo.16) (NAACL 2024)\n\n## When to Use pyvene\n\n**Use pyvene when you need to:**\n- Perform causal tracing (ROME-style localization)\n- Run activation patching experiments\n- Conduct interchange intervention training (IIT)\n- Test causal hypotheses about model components\n- Share/reproduce intervention experiments via HuggingFace\n- Work with any PyTorch architecture (not just transformers)\n\n**Consider alternatives when:**\n- You need exploratory activation analysis → Use **TransformerLens**\n- You want to train/analyze SAEs → Use **SAELens**\n- You need remote execution on massive models → Use **nnsight**\n- You want lower-level control → Use **nnsight**\n\n## Installation\n\n```bash\npip install pyvene\n```\n\nStandard import:\n```python\nimport pyvene as pv\n```\n\n## Core Concepts\n\n### IntervenableModel\n\nThe main class that wraps any PyTorch model with intervention capabilities:\n\n```python\nimport pyvene as pv\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n# Load base model\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2\")\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# Define intervention configuration\nconfig = pv.IntervenableConfig(\n    representations=[\n        pv.RepresentationConfig(\n            layer=8,\n            component=\"block_output\",\n            intervention_type=pv.VanillaIntervention,\n        )\n    ]\n)\n\n# Create intervenable model\nintervenable = pv.IntervenableModel(config, model)\n```\n\n### Intervention Types\n\n| Type | Description | Use Case |\n|------|-------------|----------|\n| `VanillaIntervention` | Swap activations between runs | Activation patching |\n| `AdditionIntervention` | Add activations to base run | Steering, ablation |\n| `SubtractionIntervention` | Subtract activations | Ablation |\n| `ZeroIntervention` | Zero out activations | Component knockout |\n| `RotatedSpaceIntervention` | DAS trainable intervention | Causal discovery |\n| `CollectIntervention` | Collect activations | Probing, analysis |\n\n### Component Targets\n\n```python\n# Available components to intervene on\ncomponents = [\n    \"block_input\",      # Input to transformer block\n    \"block_output\",     # Output of transformer block\n    \"mlp_input\",        # Input to MLP\n    \"mlp_output\",       # Output of MLP\n    \"mlp_activation\",   # MLP hidden activations\n    \"attention_input\",  # Input to attention\n    \"attention_output\", # Output of attention\n    \"attention_value_output\",  # Attention value vectors\n    \"query_output\",     # Query vectors\n    \"key_output\",       # Key vectors\n    \"value_output\",     # Value vectors\n    \"head_attention_value_output\",  # Per-head values\n]\n```\n\n## Workflow 1: Causal Tracing (ROME-style)\n\nLocate where factual associations are stored by corrupting inputs and restoring activations.\n\n### Step-by-Step\n\n```python\nimport pyvene as pv\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nimport torch\n\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2-xl\")\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2-xl\")\n\n# 1. Define clean and corrupted inputs\nclean_prompt = \"The Space Needle is in downtown\"\ncorrupted_prompt = \"The ##### ###### ## ## ########\"  # Noise\n\nclean_tokens = tokenizer(clean_prompt, return_tensors=\"pt\")\ncorrupted_tokens = tokenizer(corrupted_prompt, return_tensors=\"pt\")\n\n# 2. Get clean activations (source)\nwith torch.no_grad():\n    clean_outputs = model(**clean_tokens, output_hidden_states=True)\n    clean_states = clean_outputs.hidden_states\n\n# 3. Define restoration intervention\ndef run_causal_trace(layer, position):\n    \"\"\"Restore clean activation at specific layer and position.\"\"\"\n    config = pv.IntervenableConfig(\n        representations=[\n            pv.RepresentationConfig(\n                layer=layer,\n                component=\"block_output\",\n                intervention_type=pv.VanillaIntervention,\n                unit=\"pos\",\n                max_number_of_units=1,\n            )\n        ]\n    )\n\n    intervenable = pv.IntervenableModel(config, model)\n\n    # Run with intervention\n    _, patched_outputs = intervenable(\n        base=corrupted_tokens,\n        sources=[clean_tokens],\n        unit_locations={\"sources->base\": ([[[position]]], [[[position]]])},\n        output_original_output=True,\n    )\n\n    # Return probability of correct token\n    probs = torch.softmax(patched_outputs.logits[0, -1], dim=-1)\n    seattle_token = tokenizer.encode(\" Seattle\")[0]\n    return probs[seattle_token].item()\n\n# 4. Sweep over layers and positions\nn_layers = model.config.n_layer\nseq_len = clean_tokens[\"input_ids\"].shape[1]\n\nresults = torch.zeros(n_layers, seq_len)\nfor layer in range(n_layers):\n    for pos in range(seq_len):\n        results[layer, pos] = run_causal_trace(layer, pos)\n\n# 5. Visualize (layer x position heatmap)\n# High values indicate causal importance\n```\n\n### Checklist\n- [ ] Prepare clean prompt with target factual association\n- [ ] Create corrupted version (noise or counterfactual)\n- [ ] Define intervention config for each (layer, position)\n- [ ] Run patching sweep\n- [ ] Identify causal hotspots in heatmap\n\n## Workflow 2: Activation Patching for Circuit Analysis\n\nTest which components are necessary for a specific behavior.\n\n### Step-by-Step\n\n```python\nimport pyvene as pv\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nimport torch\n\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2\")\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# IOI task setup\nclean_prompt = \"When John and Mary went to the store, Mary gave a bottle to\"\ncorrupted_prompt = \"When John and Mary went to the store, John gave a bottle to\"\n\nclean_tokens = tokenizer(clean_prompt, return_tensors=\"pt\")\ncorrupted_tokens = tokenizer(corrupted_prompt, return_tensors=\"pt\")\n\njohn_token = tokenizer.encode(\" John\")[0]\nmary_token = tokenizer.encode(\" Mary\")[0]\n\ndef logit_diff(logits):\n    \"\"\"IO - S logit difference.\"\"\"\n    return logits[0, -1, john_token] - logits[0, -1, mary_token]\n\n# Patch attention output at each layer\ndef patch_attention(layer):\n    config = pv.IntervenableConfig(\n        representations=[\n            pv.RepresentationConfig(\n                layer=layer,\n                component=\"attention_output\",\n                intervention_type=pv.VanillaIntervention,\n            )\n        ]\n    )\n\n    intervenable = pv.IntervenableModel(config, model)\n\n    _, patched_outputs = intervenable(\n        base=corrupted_tokens,\n        sources=[clean_tokens],\n    )\n\n    return logit_diff(patched_outputs.logits).item()\n\n# Find which layers matter\nresults = []\nfor layer in range(model.config.n_layer):\n    diff = patch_attention(layer)\n    results.append(diff)\n    print(f\"Layer {layer}: logit diff = {diff:.3f}\")\n```\n\n## Workflow 3: Interchange Intervention Training (IIT)\n\nTrain interventions to discover causal structure.\n\n### Step-by-Step\n\n```python\nimport pyvene as pv\nfrom transformers import AutoModelForCausalLM\nimport torch\n\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n\n# 1. Define trainable intervention\nconfig = pv.IntervenableConfig(\n    representations=[\n        pv.RepresentationConfig(\n            layer=6,\n            component=\"block_output\",\n            intervention_type=pv.RotatedSpaceIntervention,  # Trainable\n            low_rank_dimension=64,  # Learn 64-dim subspace\n        )\n    ]\n)\n\nintervenable = pv.IntervenableModel(config, model)\n\n# 2. Set up training\noptimizer = torch.optim.Adam(\n    intervenable.get_trainable_parameters(),\n    lr=1e-4\n)\n\n# 3. Training loop (simplified)\nfor base_input, source_input, target_output in dataloader:\n    optimizer.zero_grad()\n\n    _, outputs = intervenable(\n        base=base_input,\n        sources=[source_input],\n    )\n\n    loss = criterion(outputs.logits, target_output)\n    loss.backward()\n    optimizer.step()\n\n# 4. Analyze learned intervention\n# The rotation matrix reveals causal subspace\nrotation = intervenable.interventions[\"layer.6.block_output\"][0].rotate_layer\n```\n\n### DAS (Distributed Alignment Search)\n\n```python\n# Low-rank rotation finds interpretable subspaces\nconfig = pv.IntervenableConfig(\n    representations=[\n        pv.RepresentationConfig(\n            layer=8,\n            component=\"block_output\",\n            intervention_type=pv.LowRankRotatedSpaceIntervention,\n            low_rank_dimension=1,  # Find 1D causal direction\n        )\n    ]\n)\n```\n\n## Workflow 4: Model Steering (Honest LLaMA)\n\nSteer model behavior during generation.\n\n```python\nimport pyvene as pv\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\ntokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n\n# Load pre-trained steering intervention\nintervenable = pv.IntervenableModel.load(\n    \"zhengxuanzenwu/intervenable_honest_llama2_chat_7B\",\n    model=model,\n)\n\n# Generate with steering\nprompt = \"Is the earth flat?\"\ninputs = tokenizer(prompt, return_tensors=\"pt\")\n\n# Intervention applied during generation\noutputs = intervenable.generate(\n    inputs,\n    max_new_tokens=100,\n    do_sample=False,\n)\n\nprint(tokenizer.decode(outputs[0]))\n```\n\n## Saving and Sharing Interventions\n\n```python\n# Save locally\nintervenable.save(\"./my_intervention\")\n\n# Load from local\nintervenable = pv.IntervenableModel.load(\n    \"./my_intervention\",\n    model=model,\n)\n\n# Share on HuggingFace\nintervenable.save_intervention(\"username/my-intervention\")\n\n# Load from HuggingFace\nintervenable = pv.IntervenableModel.load(\n    \"username/my-intervention\",\n    model=model,\n)\n```\n\n## Common Issues & Solutions\n\n### Issue: Wrong intervention location\n```python\n# WRONG: Incorrect component name\nconfig = pv.RepresentationConfig(\n    component=\"mlp\",  # Not valid!\n)\n\n# RIGHT: Use exact component name\nconfig = pv.RepresentationConfig(\n    component=\"mlp_output\",  # Valid\n)\n```\n\n### Issue: Dimension mismatch\n```python\n# Ensure source and base have compatible shapes\n# For position-specific interventions:\nconfig = pv.RepresentationConfig(\n    unit=\"pos\",\n    max_number_of_units=1,  # Intervene on single position\n)\n\n# Specify locations explicitly\nintervenable(\n    base=base_tokens,\n    sources=[source_tokens],\n    unit_locations={\"sources->base\": ([[[5]]], [[[5]]])},  # Position 5\n)\n```\n\n### Issue: Memory with large models\n```python\n# Use gradient checkpointing\nmodel.gradient_checkpointing_enable()\n\n# Or intervene on fewer components\nconfig = pv.IntervenableConfig(\n    representations=[\n        pv.RepresentationConfig(\n            layer=8,  # Single layer instead of all\n            component=\"block_output\",\n        )\n    ]\n)\n```\n\n### Issue: LoRA integration\n```python\n# pyvene v0.1.8+ supports LoRAs as interventions\nconfig = pv.RepresentationConfig(\n    intervention_type=pv.LoRAIntervention,\n    low_rank_dimension=16,\n)\n```\n\n## Key Classes Reference\n\n| Class | Purpose |\n|-------|---------|\n| `IntervenableModel` | Main wrapper for interventions |\n| `IntervenableConfig` | Configuration container |\n| `RepresentationConfig` | Single intervention specification |\n| `VanillaIntervention` | Activation swapping |\n| `RotatedSpaceIntervention` | Trainable DAS intervention |\n| `CollectIntervention` | Activation collection |\n\n## Supported Models\n\npyvene works with any PyTorch model. Tested on:\n- GPT-2 (all sizes)\n- LLaMA / LLaMA-2\n- Pythia\n- Mistral / Mixtral\n- OPT\n- BLIP (vision-language)\n- ESM (protein models)\n- Mamba (state space)\n\n## Reference Documentation\n\nFor detailed API documentation, tutorials, and advanced usage, see the `references/` folder:\n\n| File | Contents |\n|------|----------|\n| [references/README.md](references/README.md) | Overview and quick start guide |\n| [references/api.md](references/api.md) | Complete API reference for IntervenableModel, intervention types, configurations |\n| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for causal tracing, activation patching, DAS |\n\n## External Resources\n\n### Tutorials\n- [pyvene 101](https://stanfordnlp.github.io/pyvene/tutorials/pyvene_101.html)\n- [Causal Tracing Tutorial](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/Causal_Tracing.html)\n- [IOI Circuit Replication](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/IOI_Replication.html)\n- [DAS Introduction](https://stanfordnlp.github.io/pyvene/tutorials/advanced_tutorials/DAS_Main_Introduction.html)\n\n### Papers\n- [Locating and Editing Factual Associations in GPT](https://arxiv.org/abs/2202.05262) - Meng et al. (2022)\n- [Inference-Time Intervention](https://arxiv.org/abs/2306.03341) - Li et al. (2023)\n- [Interpretability in the Wild](https://arxiv.org/abs/2211.00593) - Wang et al. (2022)\n\n### Official Documentation\n- [Official Docs](https://stanfordnlp.github.io/pyvene/)\n- [API Reference](https://stanfordnlp.github.io/pyvene/api/)\n\n## Comparison with Other Tools\n\n| Feature | pyvene | TransformerLens | nnsight |\n|---------|--------|-----------------|---------|\n| Declarative config | Yes | No | No |\n| HuggingFace sharing | Yes | No | No |\n| Trainable interventions | Yes | Limited | Yes |\n| Any PyTorch model | Yes | Transformers only | Yes |\n| Remote execution | No | No | Yes (NDIF) |","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/04-mechanistic-interpretability/pyvene","license":"MIT","category":"coding","lang":"en","tokens":3238,"stars":0,"calls30d":1,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/api.md","size":7871,"sha256":"e10295e644d9855c3954fcc51b80f1d42448b60bb56ef991961a1d11e44f1003"},{"path":"references/README.md","size":2105,"sha256":"4b842fea60747f23a54986d9bbf20b13f38071f62ed24b93dce1227e2e35782b"},{"path":"references/tutorials.md","size":10111,"sha256":"04132b155e440f95acf08f1349651355ade3d2b0ad46d19370c59857bb38a5d4"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["aclanthology.org","arxiv.org","stanfordnlp.github.io"]}}