pyvene-interventions

pyvene：ニューラルネットワークへの因果的介入

pyveneは、PyTorchモデルに因果的介入を行うためのStanford NLPのライブラリです。宣言的な辞書ベースのフレームワークにより、活性化パッチング、因果トレーシング、インターチェンジ介入学習を実行でき、介入実験の再現性と共有を実現します。

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

pyveneを使う場面

次のことが必要な場合は、pyveneをお使いください。

• 因果トレーシング（ROME形式の局所化）を実行する
• 活性化パッチング実験を行う
• インターチェンジ介入学習（IIT）を実施する
• モデルコンポーネントに関する因果仮説を検証する
• HuggingFaceを介して介入実験を共有・再現する
• PyTorchのあらゆるアーキテクチャ（トランスフォーマーに限りません）を扱う

次の場合は、別のツールをご検討ください。

• 探索的な活性化分析が必要 → 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 et al. (2022)
• 推論時の介入 - Li et al. (2023)
• 実環境における解釈可能性 - Wang et al. (2022)

公式ドキュメント

• 公式ドキュメント
• API リファレンス

他のツールとの比較

機能	pyvene	TransformerLens	nnsight
宣言型の設定	はい	いいえ	いいえ
HuggingFace との共有	はい	いいえ	いいえ
学習可能な介入	はい	制限あり	はい
任意の PyTorch モデル	はい	Transformers のみ	はい
リモート実行	いいえ	いいえ	はい（NDIF）

エージェントの実行手順

1. 環境を確認します。 モデル、チェックポイント、ベンチマークのバージョン、評価ハーネスのバージョンを確認します。
2. 入力を確定します。 データまたはユーザーから、形式、識別子、パラメーターを確認します。結果を左右する値を推測せず、確認してください。
3. まず小規模に実行します。 小さなサブセットで評価し、プロンプト、解析、スコアリングを検証します。
4. 上記の手順とリファレンスを使用して、タスク全体を実行します。
5. 結果を検証します。 ベンチマークの公式分割と指標を使用し、データの混入を確認して、シードまたはプロンプトごとのばらつきを報告します。
6. 報告します。 実行内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の事項を記載します。

このような場合	対応
スコアが公開済みの数値と異なる	結論を出す前に、プロンプト形式、少数例設定、ハーネスのバージョンを確認します。
この手順に記載された関数、フラグ、エンドポイントがインストール済みバージョンに存在しない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、適宜調整してユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに確認するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• 評価設定をすべて報告し、プロンプト、シード、サブセットを恣意的に選んではいけません。
• ここに記載されたバージョン固有の詳細は古くなっている可能性があります。インストール済みバージョンの公式ドキュメントと照合してください。
• 費用が発生する、共有 GPU やクラウドのクォータを消費する、個人データまたは患者データを扱う、あるいは取り消せない操作を行う場合は、事前に確認してください。

関連スキル

• transformer-lens-interpretability：HookP…を介して Transformer の内部を調査・操作する TransformerLens を用いたメカニスティック解釈可能性研究のガイダンスを提供します。
• hypothesis-generation：証拠に基づく科学的な問い、候補となる仮説、対立する説明、因果的または関連性に関する主張、識別力のある p…を作成します。
• langsmith-observability：トレーシング、評価、モニタリングのための LLM オブザーバビリティプラットフォームです。