transformer-lens-interpretability

TransformerLens：トランスフォーマーのメカニズム解釈

TransformerLens は、GPT スタイルの言語モデルのメカニズム解釈研究における事実上の標準ライブラリです。Neel Nanda が作成し、Bryce Meyer がメンテナンスしているこのライブラリは、すべての活性化に設けられた HookPoints を通じて、モデル内部を調査・操作するための使いやすいインターフェースを提供します。

GitHub： TransformerLensOrg/TransformerLens

TransformerLens を使用する場面

次のような場合に TransformerLens を使用してください：

• 学習中に獲得したアルゴリズムをリバースエンジニアリングする
• 活性化パッチング／因果トレース実験を行う
• アテンションパターンと情報の流れを調べる
• 回路を分析する（例：誘導ヘッド、IOI 回路）
• 中間活性化をキャッシュして調査する
• 直接ロジット帰属を適用する

次のような場合は、代替手段を検討してください：

• トランスフォーマー以外のアーキテクチャを扱う必要がある → nnsight または pyvene を使用してください
• 疎オートエンコーダーを学習／分析したい → SAELens を使用してください
• 大規模モデルをリモートで実行する必要がある → NDIF と nnsight を使用してください
• より高水準の因果介入の抽象化が必要である → pyvene を使用してください

インストール

pip install transformer-lens

開発版の場合：

pip install git+https://github.com/TransformerLensOrg/TransformerLens

主な概念

HookedTransformer

すべての活性化に HookPoints を設け、トランスフォーマーモデルをラップする主要クラスです：

from transformer_lens import HookedTransformer

# Load a model
model = HookedTransformer.from_pretrained("gpt2-small")

# For gated models (LLaMA, Mistral)
import os
os.environ["HF_TOKEN"] = "your_token"
model = HookedTransformer.from_pretrained("meta-llama/Llama-2-7b-hf")

対応モデル（50 以上）

ファミリー	モデル
GPT-2	gpt2, gpt2-medium, gpt2-large, gpt2-xl
LLaMA	llama-7b, llama-13b, llama-2-7b, llama-2-13b
EleutherAI	pythia-70m to pythia-12b, gpt-neo, gpt-j-6b
Mistral	mistral-7b, mixtral-8x7b
その他	phi, qwen, opt, gemma

活性化キャッシュ

モデルを実行し、すべての中間活性化をキャッシュします：

# Get all activations
tokens = model.to_tokens("The Eiffel Tower is in")
logits, cache = model.run_with_cache(tokens)

# Access specific activations
residual = cache["resid_post", 5]  # Layer 5 residual stream
attn_pattern = cache["pattern", 3]  # Layer 3 attention pattern
mlp_out = cache["mlp_out", 7]  # Layer 7 MLP output

# Filter which activations to cache (saves memory)
logits, cache = model.run_with_cache(
    tokens,
    names_filter=lambda name: "resid_post" in name
)

ActivationCache のキー

キーパターン	形状	説明
resid_pre, layer	[batch, pos, d_model]	アテンション前の残差
resid_mid, layer	[batch, pos, d_model]	アテンション後の残差
resid_post, layer	[batch, pos, d_model]	MLP 後の残差
attn_out, layer	[batch, pos, d_model]	アテンション出力
mlp_out, layer	[batch, pos, d_model]	MLP 出力
pattern, layer	[batch, head, q_pos, k_pos]	アテンションパターン（ソフトマックス後）
q, layer	[batch, pos, head, d_head]	クエリベクトル
k, layer	[batch, pos, head, d_head]	キーベクトル
v, layer	[batch, pos, head, d_head]	バリューベクトル

ワークフロー 1：活性化パッチング（因果トレース）

破損した実行にクリーンな活性化をパッチングし、どの活性化がモデル出力に因果的な影響を与えるかを特定します。

手順

from transformer_lens import HookedTransformer, patching
import torch

model = HookedTransformer.from_pretrained("gpt2-small")

# 1. Define clean and corrupted prompts
clean_prompt = "The Eiffel Tower is in the city of"
corrupted_prompt = "The Colosseum is in the city of"

clean_tokens = model.to_tokens(clean_prompt)
corrupted_tokens = model.to_tokens(corrupted_prompt)

# 2. Get clean activations
_, clean_cache = model.run_with_cache(clean_tokens)

# 3. Define metric (e.g., logit difference)
paris_token = model.to_single_token(" Paris")
rome_token = model.to_single_token(" Rome")

def metric(logits):
    return logits[0, -1, paris_token] - logits[0, -1, rome_token]

# 4. Patch each position and layer
results = torch.zeros(model.cfg.n_layers, clean_tokens.shape[1])

for layer in range(model.cfg.n_layers):
    for pos in range(clean_tokens.shape[1]):
        def patch_hook(activation, hook):
            activation[0, pos] = clean_cache[hook.name][0, pos]
            return activation

        patched_logits = model.run_with_hooks(
            corrupted_tokens,
            fwd_hooks=[(f"blocks.{layer}.hook_resid_post", patch_hook)]
        )
        results[layer, pos] = metric(patched_logits)

# 5. Visualize results (layer x position heatmap)

チェックリスト

• 最小限の差異があるクリーンな入力と破損した入力を定義する
• 振る舞いの違いを捉える指標を選ぶ
• クリーンな活性化をキャッシュする
• 各（レイヤー、位置）の組み合わせを体系的にパッチングする
• 結果をヒートマップとして可視化する
• 因果的なホットスポットを特定する

ワークフロー 2：回路分析（間接目的語の特定）

「Interpretability in the Wild」で紹介された IOI 回路の発見を再現します。

手順

from transformer_lens import HookedTransformer
import torch

model = HookedTransformer.from_pretrained("gpt2-small")

# IOI task: "When John and Mary went to the store, Mary gave a bottle to"
# Model should predict "John" (indirect object)

prompt = "When John and Mary went to the store, Mary gave a bottle to"
tokens = model.to_tokens(prompt)

# 1. Get baseline logits
logits, cache = model.run_with_cache(tokens)

john_token = model.to_single_token(" John")
mary_token = model.to_single_token(" Mary")

# 2. Compute logit difference (IO - S)
logit_diff = logits[0, -1, john_token] - logits[0, -1, mary_token]
print(f"Logit difference: {logit_diff.item():.3f}")

# 3. Direct logit attribution by head
def get_head_contribution(layer, head):
    # Project head output to logits
    head_out = cache["z", layer][0, :, head, :]  # [pos, d_head]
    W_O = model.W_O[layer, head]  # [d_head, d_model]
    W_U = model.W_U  # [d_model, vocab]

    # Head contribution to logits at final position
    contribution = head_out[-1] @ W_O @ W_U
    return contribution[john_token] - contribution[mary_token]

# 4. Map all heads
head_contributions = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)
for layer in range(model.cfg.n_layers):
    for head in range(model.cfg.n_heads):
        head_contributions[layer, head] = get_head_contribution(layer, head)

# 5. Identify top contributing heads (name movers, backup name movers)

チェックリスト

• IO／S トークンが明確なタスクを設定する
• ベースラインのロジット差を計算する
• アテンションヘッドの寄与ごとに分解する
• 主要な回路構成要素（名前移動ヘッド、S 抑制、誘導）を特定する
• アブレーション実験で検証する

ワークフロー 3：誘導ヘッドの検出

[A][B]...[A] → [B] パターンを実装する誘導ヘッドを見つけます。

from transformer_lens import HookedTransformer
import torch

model = HookedTransformer.from_pretrained("gpt2-small")

# Create repeated sequence: [A][B][A] should predict [B]
repeated_tokens = torch.tensor([[1000, 2000, 1000]])  # Arbitrary tokens

_, cache = model.run_with_cache(repeated_tokens)

# Induction heads attend from final [A] back to first [B]
# Check attention from position 2 to position 1
induction_scores = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)

for layer in range(model.cfg.n_layers):
    pattern = cache["pattern", layer][0]  # [head, q_pos, k_pos]
    # Attention from pos 2 to pos 1
    induction_scores[layer] = pattern[:, 2, 1]

# Heads with high scores are induction heads
top_heads = torch.topk(induction_scores.flatten(), k=5)

よくある問題と解決策

問題：デバッグ後もフックが残る

# WRONG: Old hooks remain active
model.run_with_hooks(tokens, fwd_hooks=[...])  # Debug, add new hooks
model.run_with_hooks(tokens, fwd_hooks=[...])  # Old hooks still there!

# RIGHT: Always reset hooks
model.reset_hooks()
model.run_with_hooks(tokens, fwd_hooks=[...])

問題：トークン化に関する注意点

# WRONG: Assuming consistent tokenization
model.to_tokens("Tim")  # Single token
model.to_tokens("Neel")  # Becomes "Ne" + "el" (two tokens!)

# RIGHT: Check tokenization explicitly
tokens = model.to_tokens("Neel", prepend_bos=False)
print(model.to_str_tokens(tokens))  # ['Ne', 'el']

問題：分析で LayerNorm が無視される

# WRONG: Ignoring LayerNorm
pre_activation = residual @ model.W_in[layer]

# RIGHT: Include LayerNorm
ln_scale = model.blocks[layer].ln2.w
ln_out = model.blocks[layer].ln2(residual)
pre_activation = ln_out @ model.W_in[layer]

問題：大規模モデルでメモリが急増する

# Use selective caching
logits, cache = model.run_with_cache(
    tokens,
    names_filter=lambda n: "resid_post" in n or "pattern" in n,
    device="cpu"  # Cache on CPU
)

主要クラスのリファレンス

クラス	目的
HookedTransformer	フックを備えたメインのモデルラッパー
ActivationCache	活性化の辞書形式のキャッシュ
HookedTransformerConfig	モデル設定
FactoredMatrix	効率的な行列分解演算

SAELens との統合

TransformerLens は、スパースオートエンコーダーの分析に SAELens と統合されています。

from transformer_lens import HookedTransformer
from sae_lens import SAE

model = HookedTransformer.from_pretrained("gpt2-small")
sae = SAE.from_pretrained("gpt2-small-res-jb", "blocks.8.hook_resid_pre")

# Run with SAE
tokens = model.to_tokens("Hello world")
_, cache = model.run_with_cache(tokens)
sae_acts = sae.encode(cache["resid_pre", 8])

リファレンスドキュメント

詳しい API ドキュメント、チュートリアル、高度な使い方については、references/ フォルダーを参照してください。

ファイル	内容
references/README.md	概要とクイックスタートガイド
references/api.md	HookedTransformer、ActivationCache、HookPoints の完全な API リファレンス
references/tutorials.md	活性化パッチング、回路分析、ロジットレンズのステップごとのチュートリアル

外部リソース

チュートリアル

• メインデモノートブック
• 活性化パッチングのデモ
• ARENA 機械論的解釈可能性コース - 200 時間以上のチュートリアル

論文

• Transformer 回路の数学的フレームワーク
• インコンテキスト学習と誘導ヘッド
• 実世界における解釈可能性（IOI）

公式ドキュメント

• 公式ドキュメント
• モデルプロパティの表
• Neel Nanda の用語集

バージョンに関する注意

• v2.0：HookedSAE を削除（SAELens に移動）
• v3.0 (alpha)：任意の nn.Module を読み込むための TransformerBridge

エージェントの運用手順

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

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

誠実性に関するルール

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

関連スキル

• pyvene-interventions：pyvene の宣言的介入フレームワークを使用して、PyTorch モデルに因果介入を行う方法を説明します。
• nnsight-remote-interpretability：必要に応じて NDIF のリモート実行を利用し、nnsight でニューラルネットワークの内部を解釈および操作する方法を説明します。
• sparse-autoencoder-training：SAELens を使用して Sparse Autoencoders（SAE）をトレーニングおよび分析し、ニューラルネットワークの活性化を分解する方法を説明します。