pyvene-interventions

Watching one part of a model light up does not prove it is responsible for anything; this changes that part on purpose and checks whether the model's behaviour changes with it.

WHEN TO USE

SCOPE

USE FOR

Locate which layer stores a specific factual association using causal tracing.
Test which attention layers are necessary for a behavior by patching activations between runs.
Train a small intervention that discovers a causal subspace inside a model.
Share a reproducible intervention experiment with others through HuggingFace.

NOT FOR

Use TransformerLens instead for exploratory activation analysis without interventions.
Use SAELens instead for training or analyzing sparse autoencoders.
Use nnsight instead for remote execution or lower-level control on very large models.

This is a saved copy of the skill's SKILL.md instructions. It may differ from the upstream file or the version you install.

In this document (15)

pyvene: Causal Interventions for Neural Networks

pyvene 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.

GitHub: stanfordnlp/pyvene (840+ stars) Paper: pyvene: A Library for Understanding and Improving PyTorch Models via Interventions (NAACL 2024)

When to Use pyvene

Use pyvene when you need to:

Perform causal tracing (ROME-style localization)
Run activation patching experiments
Conduct interchange intervention training (IIT)
Test causal hypotheses about model components
Share/reproduce intervention experiments via HuggingFace
Work with any PyTorch architecture (not just transformers)

Consider alternatives when:

You need exploratory activation analysis → Use TransformerLens
You want to train/analyze SAEs → Use SAELens
You need remote execution on massive models → Use nnsight
You want lower-level control → Use nnsight

Installation

pip install pyvene

Standard import:

import pyvene as pv

Core Concepts

IntervenableModel

The main class that wraps any PyTorch model with intervention capabilities:

import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load base model
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Define intervention configuration
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,
            component="block_output",
            intervention_type=pv.VanillaIntervention,
        )
    ]
)

# Create intervenable model
intervenable = pv.IntervenableModel(config, model)

Intervention Types

Type	Description	Use Case
`VanillaIntervention`	Swap activations between runs	Activation patching
`AdditionIntervention`	Add activations to base run	Steering, ablation
`SubtractionIntervention`	Subtract activations	Ablation
`ZeroIntervention`	Zero out activations	Component knockout
`RotatedSpaceIntervention`	DAS trainable intervention	Causal discovery
`CollectIntervention`	Collect activations	Probing, analysis

Component Targets

# Available components to intervene on
components = [
    "block_input",      # Input to transformer block
    "block_output",     # Output of transformer block
    "mlp_input",        # Input to MLP
    "mlp_output",       # Output of MLP
    "mlp_activation",   # MLP hidden activations
    "attention_input",  # Input to attention
    "attention_output", # Output of attention
    "attention_value_output",  # Attention value vectors
    "query_output",     # Query vectors
    "key_output",       # Key vectors
    "value_output",     # Value vectors
    "head_attention_value_output",  # Per-head values
]

Workflow 1: Causal Tracing (ROME-style)

Locate where factual associations are stored by corrupting inputs and restoring activations.

Step-by-Step

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 clean and corrupted inputs
clean_prompt = "The Space Needle is in downtown"
corrupted_prompt = "The ##### ###### ## ## ########"  # Noise

clean_tokens = tokenizer(clean_prompt, return_tensors="pt")
corrupted_tokens = tokenizer(corrupted_prompt, return_tensors="pt")

# 2. Get clean activations (source)
with torch.no_grad():
    clean_outputs = model(**clean_tokens, output_hidden_states=True)
    clean_states = clean_outputs.hidden_states

# 3. Define restoration intervention
def run_causal_trace(layer, position):
    """Restore clean activation at specific layer and position."""
    config = pv.IntervenableConfig(
        representations=[
            pv.RepresentationConfig(
                layer=layer,
                component="block_output",
                intervention_type=pv.VanillaIntervention,
                unit="pos",
                max_number_of_units=1,
            )
        ]
    )

    intervenable = pv.IntervenableModel(config, model)

    # Run with intervention
    _, patched_outputs = intervenable(
        base=corrupted_tokens,
        sources=[clean_tokens],
        unit_locations={"sources->base": ([[[position]]], [[[position]]])},
        output_original_output=True,
    )

    # Return probability of correct token
    probs = torch.softmax(patched_outputs.logits[0, -1], dim=-1)
    seattle_token = tokenizer.encode(" Seattle")[0]
    return probs[seattle_token].item()

# 4. Sweep over layers and positions
n_layers = model.config.n_layer
seq_len = clean_tokens["input_ids"].shape[1]

results = torch.zeros(n_layers, seq_len)
for layer in range(n_layers):
    for pos in range(seq_len):
        results[layer, pos] = run_causal_trace(layer, pos)

# 5. Visualize (layer x position heatmap)
# High values indicate causal importance

Checklist

Prepare clean prompt with target factual association
Create corrupted version (noise or counterfactual)
Define intervention config for each (layer, position)
Run patching sweep
Identify causal hotspots in heatmap

Workflow 2: Activation Patching for Circuit Analysis

Test which components are necessary for a specific behavior.

Step-by-Step

import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# IOI task setup
clean_prompt = "When John and Mary went to the store, Mary gave a bottle to"
corrupted_prompt = "When John and Mary went to the store, John gave a bottle to"

clean_tokens = tokenizer(clean_prompt, return_tensors="pt")
corrupted_tokens = tokenizer(corrupted_prompt, return_tensors="pt")

john_token = tokenizer.encode(" John")[0]
mary_token = tokenizer.encode(" Mary")[0]

def logit_diff(logits):
    """IO - S logit difference."""
    return logits[0, -1, john_token] - logits[0, -1, mary_token]

# Patch attention output at each layer
def patch_attention(layer):
    config = pv.IntervenableConfig(
        representations=[
            pv.RepresentationConfig(
                layer=layer,
                component="attention_output",
                intervention_type=pv.VanillaIntervention,
            )
        ]
    )

    intervenable = pv.IntervenableModel(config, model)

    _, patched_outputs = intervenable(
        base=corrupted_tokens,
        sources=[clean_tokens],
    )

    return logit_diff(patched_outputs.logits).item()

# Find which layers matter
results = []
for layer in range(model.config.n_layer):
    diff = patch_attention(layer)
    results.append(diff)
    print(f"Layer {layer}: logit diff = {diff:.3f}")

Workflow 3: Interchange Intervention Training (IIT)

Train interventions to discover causal structure.

Step-by-Step

import pyvene as pv
from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")

# 1. Define trainable intervention
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=6,
            component="block_output",
            intervention_type=pv.RotatedSpaceIntervention,  # Trainable
            low_rank_dimension=64,  # Learn 64-dim subspace
        )
    ]
)

intervenable = pv.IntervenableModel(config, model)

# 2. Set up training
optimizer = torch.optim.Adam(
    intervenable.get_trainable_parameters(),
    lr=1e-4
)

# 3. Training loop (simplified)
for base_input, source_input, target_output in dataloader:
    optimizer.zero_grad()

    _, outputs = intervenable(
        base=base_input,
        sources=[source_input],
    )

    loss = criterion(outputs.logits, target_output)
    loss.backward()
    optimizer.step()

# 4. Analyze learned intervention
# The rotation matrix reveals causal subspace
rotation = intervenable.interventions["layer.6.block_output"][0].rotate_layer

DAS (Distributed Alignment Search)

# Low-rank rotation finds interpretable subspaces
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,
            component="block_output",
            intervention_type=pv.LowRankRotatedSpaceIntervention,
            low_rank_dimension=1,  # Find 1D causal direction
        )
    ]
)

Workflow 4: Model Steering (Honest LLaMA)

Steer model behavior during generation.

import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# Load pre-trained steering intervention
intervenable = pv.IntervenableModel.load(
    "zhengxuanzenwu/intervenable_honest_llama2_chat_7B",
    model=model,
)

# Generate with steering
prompt = "Is the earth flat?"
inputs = tokenizer(prompt, return_tensors="pt")

# Intervention applied during generation
outputs = intervenable.generate(
    inputs,
    max_new_tokens=100,
    do_sample=False,
)

print(tokenizer.decode(outputs[0]))

Saving and Sharing Interventions

# Save locally
intervenable.save("./my_intervention")

# Load from local
intervenable = pv.IntervenableModel.load(
    "./my_intervention",
    model=model,
)

# Share on HuggingFace
intervenable.save_intervention("username/my-intervention")

# Load from HuggingFace
intervenable = pv.IntervenableModel.load(
    "username/my-intervention",
    model=model,
)

Common Issues & Solutions

Issue: Wrong intervention location

# WRONG: Incorrect component name
config = pv.RepresentationConfig(
    component="mlp",  # Not valid!
)

# RIGHT: Use exact component name
config = pv.RepresentationConfig(
    component="mlp_output",  # Valid
)

Issue: Dimension mismatch

# Ensure source and base have compatible shapes
# For position-specific interventions:
config = pv.RepresentationConfig(
    unit="pos",
    max_number_of_units=1,  # Intervene on single position
)

# Specify locations explicitly
intervenable(
    base=base_tokens,
    sources=[source_tokens],
    unit_locations={"sources->base": ([[[5]]], [[[5]]])},  # Position 5
)

Issue: Memory with large models

# Use gradient checkpointing
model.gradient_checkpointing_enable()

# Or intervene on fewer components
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,  # Single layer instead of all
            component="block_output",
        )
    ]
)

Issue: LoRA integration

# pyvene v0.1.8+ supports LoRAs as interventions
config = pv.RepresentationConfig(
    intervention_type=pv.LoRAIntervention,
    low_rank_dimension=16,
)

Key Classes Reference

Class	Purpose
`IntervenableModel`	Main wrapper for interventions
`IntervenableConfig`	Configuration container
`RepresentationConfig`	Single intervention specification
`VanillaIntervention`	Activation swapping
`RotatedSpaceIntervention`	Trainable DAS intervention
`CollectIntervention`	Activation collection

Supported Models

pyvene works with any PyTorch model. Tested on:

GPT-2 (all sizes)
LLaMA / LLaMA-2
Pythia
Mistral / Mixtral
OPT
BLIP (vision-language)
ESM (protein models)
Mamba (state space)

Reference Documentation

For detailed API documentation, tutorials, and advanced usage, see the references/ folder:

File	Contents
references/README.md	Overview and quick start guide
references/api.md	Complete API reference for IntervenableModel, intervention types, configurations
references/tutorials.md	Step-by-step tutorials for causal tracing, activation patching, DAS

External Resources

Tutorials

Papers

Locating and Editing Factual Associations in GPT - Meng et al. (2022)
Inference-Time Intervention - Li et al. (2023)
Interpretability in the Wild - Wang et al. (2022)

Official Documentation

Comparison with Other Tools

Feature	pyvene	TransformerLens	nnsight
Declarative config	Yes	No	No
HuggingFace sharing	Yes	No	No
Trainable interventions	Yes	Limited	Yes
Any PyTorch model	Yes	Transformers only	Yes
Remote execution	No	No	Yes (NDIF)

This skill ships inside the AI Research Skills plugin.

Run this outside your AI tool, using Node.js 22.20.0 or newer. Skills CLI 1.7.2 skips skill selection; prompts cancel installation without an interactive terminal. If choosing Project when asked about installation scope, open your project's folder before running this.

Install checklist

Check selected tools if a picker appears; Claude Code may be unchecked under Additional agents. ↑/↓ moves, Space toggles selections, Enter confirms.
Without a tool picker, check Installing to: for your intended tool. Stop if it names a different tool.
If asked, choose Symlink (Recommended) for shared files or Copy to all agents for separate copies; review the summary before confirming installation.
Install the find-skills skill? defaults to Yes, adding another skill for all projects. Choose No if you don't want that extra installation.

Skills CLI 1.7.2 sends install data to skills.sh by default, including the source, skill names, target tools and file details; GitHub sources are reported only when GitHub confirms the repository is public, while other remote sources can also be reported. Confirmed-public GitHub sources also trigger a security-audit request. Before installing, set the environment variable DO_NOT_TRACK=1 or DISABLE_TELEMETRY=1 to disable both.

npx skills add orchestra-research/ai-research-skills --skill "pyvene-interventions" --full-depth

After installing for Claude Code, open it in the installation project if you chose Project; Global makes the skill available across projects on this computer. Once Claude Code loads the skill, its description helps Claude decide when to use it.

If installation created a top-level skills folder that did not exist when your Claude Code session started, run /reload-skills in that session. Run it again after later changes in that folder.

What's inside

Files included with this skill.

View on GitHub

pyvene: Causal Interventions for Neural Networks

When to Use pyvene

Installation

Core Concepts

IntervenableModel

Intervention Types

Component Targets

Workflow 1: Causal Tracing (ROME-style)

Step-by-Step

Checklist

Workflow 2: Activation Patching for Circuit Analysis

Step-by-Step

Workflow 3: Interchange Intervention Training (IIT)

Step-by-Step

DAS (Distributed Alignment Search)

Workflow 4: Model Steering (Honest LLaMA)

Saving and Sharing Interventions

Common Issues & Solutions

Issue: Wrong intervention location

Issue: Dimension mismatch

Issue: Memory with large models

Issue: LoRA integration

Key Classes Reference

Supported Models

Reference Documentation

External Resources

Tutorials

Papers

Official Documentation

Comparison with Other Tools

Local tool required

What's inside

More skills

academic-plotting

ara-compiler

ara-research-manager

ara-rigor-reviewer

audiocraft-audio-generation

autogpt-agents