pyvene-interventions

pyvene：神经网络的因果干预

pyvene 是 Stanford NLP 用于对 PyTorch 模型执行因果干预的库。它提供基于声明式字典的框架，用于激活补丁、因果追踪和交互干预训练，让干预实验可复现、可分享。

GitHub：stanfordnlp/pyvene 论文：pyvene：通过干预理解和改进 PyTorch 模型的库（NAACL 2024）

何时使用 pyvene

在以下情况下使用 pyvene：

• 执行因果追踪（ROME 风格的定位）
• 运行激活补丁实验
• 进行交互干预训练（IIT）
• 测试有关模型组件的因果假设
• 通过 HuggingFace 分享／复现实验
• 使用任意 PyTorch 架构（而不只是 transformer）

以下情况可考虑替代方案：

• 需要探索性激活分析 → 使用 TransformerLens
• 想要训练／分析 SAE → 使用 SAELens
• 需要在大型模型上远程执行 → 使用 nnsight
• 想要更底层的控制 → 使用 nnsight

安装

pip install pyvene

标准导入：

import pyvene as pv

核心概念

IntervenableModel

用干预功能封装任意 PyTorch 模型的主类：

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)

干预类型

类型	描述	用例
VanillaIntervention	在不同运行之间交换激活	激活补丁
AdditionIntervention	将激活加到基础运行中	引导、消融
SubtractionIntervention	减去激活	消融
ZeroIntervention	将激活置零	组件敲除
RotatedSpaceIntervention	DAS 可训练干预	因果发现
CollectIntervention	收集激活	探测、分析

组件目标

# 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
]

工作流 1：因果追踪（ROME 风格）

通过破坏输入并恢复激活，定位事实关联的存储位置。

分步说明

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

检查清单

• 准备包含目标事实关联的干净提示词
• 创建损坏版本（噪声或反事实）
• 为每个（层、位置）定义干预配置
• 运行补丁扫描
• 在热图中找出因果热点

工作流 2：用于电路分析的激活补丁

测试哪些组件对于特定行为是必要的。

分步说明

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}")

工作流 3：交互干预训练（IIT）

训练干预以发现因果结构。

分步说明

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（分布式对齐搜索）

# 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
        )
    ]
)

工作流 4：模型引导（Honest LLaMA）

在生成过程中引导模型行为。

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]))

保存和分享干预

详情、代码示例和参数表：references/saving-and-sharing-interventions.md。适用时请阅读。

常见问题与解决方案

问题：干预位置错误

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

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

问题：维度不匹配

# 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
)

问题：大型模型的内存问题

# 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",
        )
    ]
)

问题：LoRA 集成

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

关键类参考

类	用途
IntervenableModel	干预的主封装类
IntervenableConfig	配置容器
RepresentationConfig	单个干预规范
VanillaIntervention	激活交换
RotatedSpaceIntervention	可训练的 DAS 干预
CollectIntervention	激活收集

支持的模型

pyvene 适用于任意 PyTorch 模型。已测试的模型：

• GPT-2（所有尺寸）
• LLaMA / LLaMA-2
• Pythia
• Mistral / Mixtral
• OPT
• BLIP（视觉语言）
• ESM（蛋白质模型）
• Mamba（状态空间模型）

参考文档

如需详细的 API 文档、教程和高级用法，请参阅 references/ 文件夹：

文件	内容
references/README.md	概览和快速入门指南
references/api.md	IntervenableModel、干预类型和配置的完整 API 参考
references/tutorials.md	关于因果追踪、激活修补和 DAS 的分步教程

外部资源

教程

• pyvene 101
• 因果追踪教程
• IOI 电路复现
• DAS 简介

论文

• 定位并编辑 GPT 中的事实关联 - Meng 等（2022）
• 推理时干预 - Li 等（2023）
• 真实环境中的可解释性 - Wang 等（2022）

官方文档

• 官方文档
• API 参考

与其他工具的比较

特性	pyvene	TransformerLens	nnsight
声明式配置	是	否	否
HuggingFace 共享	是	否	否
可训练干预	是	有限	是
任意 PyTorch 模型	是	仅限 Transformers	是
远程执行	否	否	是（NDIF）

Agent 操作流程

1. 检查环境。 确认模型、检查点、基准测试版本和评估框架版本。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。对于会影响结果的值，应先询问，不要猜测。
3. 先运行小规模版本。 在小型子集上进行评估，以验证提示词、解析和评分。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 使用基准测试的官方划分和指标；检查是否存在数据污染；报告不同随机种子或提示词下的方差。
6. 报告结果。 说明运行内容（版本、命令、参数）、检查内容以及尚不确定的事项。

如果出现这种情况	请执行以下操作
分数与已发表的数值不同	在下结论之前，先检查提示词格式、少样本设置和评估框架版本。
当前安装的版本中缺少这些说明所述的函数、标志或端点	查阅已安装版本自身的文档（help()、--help、官方文档），进行调整并告知用户。切勿编造 API。
所需输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿伪造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请如实说明。
• 完整报告评估设置；切勿挑选提示词、随机种子或子集。
• 将此处与特定版本相关的细节视为可能已过时：请对照已安装版本的官方文档进行确认。
• 对于需要付费、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• transformer-lens-interpretability：提供使用 TransformerLens 开展机制可解释性研究的指导，用于通过 HookP… 检查和操控 Transformer 内部结构
• hypothesis-generation：构建受证据约束的科学问题、候选假设、竞争性解释、因果或相关性主张以及具有区分力的 p…
• langsmith-observability：用于追踪、评估和监控的 LLM 可观测性平台。