pyvene-interventions

pyvene: ニューラルネットワークのための因果的介入

pyvene は Stanford NLP による、PyTorch モデルに対して因果的介入を行うためのライブラリです。活性化パッチング、因果トレーシング、交換介入（interchange intervention）トレーニングのための宣言的な dict ベースのフレームワークを提供し、介入実験を再現可能かつ共有可能にします。

GitHub: stanfordnlp/pyvene（840+ スター） 論文: pyvene: A Library for Understanding and Improving PyTorch Models via Interventions（NAACL 2024）

pyvene を使うべき場合

次のことが必要な場合は pyvene を使用してください:

• 因果トレーシング（ROME 方式の局在化）を行う
• 活性化パッチング実験を実行する
• 交換介入トレーニング（IIT）を行う
• モデルのコンポーネントに関する因果仮説を検証する
• HuggingFace を介して介入実験を共有・再現する
• 任意の PyTorch アーキテクチャで作業する（transformers に限らない）

次の場合は代替手段を検討してください:

• 探索的な活性化分析が必要 → 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

チェックリスト

• 対象となる事実的関連を含むクリーンなプロンプトを用意する
• 破壊されたバージョン（ノイズまたは反事実）を作成する
• 各 (layer, position) に対する介入設定を定義する
• パッチングのスイープを実行する
• ヒートマップで因果的ホットスポットを特定する

ワークフロー 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（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
        )
    ]
)

ワークフロー 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]))

介入の保存と共有

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

よくある問題と解決策

問題: 介入位置が誤っている

# 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
• Causal Tracing Tutorial
• IOI Circuit Replication
• DAS Introduction

論文

• 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 Docs
• API Reference

他のツールとの比較

機能	pyvene	TransformerLens	nnsight
宣言的な設定	Yes	No	No
HuggingFace での共有	Yes	No	No
学習可能な介入	Yes	Limited	Yes
任意の PyTorch モデル	Yes	Transformers only	Yes
リモート実行	No	No	Yes (NDIF)