transformer-lens-interpretability

TransformerLens：Transformer 的机制可解释性

TransformerLens 是 GPT 风格语言模型机制可解释性研究中事实上的标准库。它由 Neel Nanda 创建、由 Bryce Meyer 维护，可通过每个激活上的 HookPoints 提供简洁的接口，用于检查和操纵模型内部状态。

GitHub：TransformerLensOrg/TransformerLens

何时使用 TransformerLens

在以下情况下使用 TransformerLens：

• 逆向工程模型在训练期间学到的算法
• 执行激活修补／因果追踪实验
• 研究注意力模式和信息流
• 分析电路（例如，归纳头、IOI 电路）
• 缓存并检查中间激活
• 应用直接 logit 归因

在以下情况下考虑其他方案：

• 需要使用非 Transformer 架构 → 使用 nnsight 或 pyvene
• 想要训练／分析稀疏自编码器 → 使用 SAELens
• 需要在超大型模型上远程执行 → 使用带 NDIF 的 nnsight
• 想要更高层次的因果干预抽象 → 使用 pyvene

安装

pip install transformer-lens

开发版本：

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

核心概念

HookedTransformer

在每个激活上都通过 HookPoints 封装 Transformer 模型的主要类：

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 到 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]	注意力模式（softmax 后）
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 标记的任务
• 计算基准 logit 差值
• 按注意力头的贡献进行分解
• 找出关键电路组件（名称移动头、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)

常见问题与解决方案

问题：调试后 Hook 仍然存在

# 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	激活修补、电路分析、logit lens 的分步教程

外部资源

教程

• 主演示 Notebook
• 激活修补演示
• ARENA 机械可解释性课程 - 200 多小时的教程

论文

• Transformer 电路的数学框架
• 上下文学习与归纳头
• 实地可解释性（IOI）

官方文档

• 官方文档
• 模型属性表
• Neel Nanda 的术语表

版本说明

• v2.0：移除了 HookedSAE（已移至 SAELens）
• v3.0 (alpha)：用于加载任意 nn.Module 的 TransformerBridge

Agent 操作流程

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

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

诚信规则

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

相关技能

• pyvene-interventions：提供使用 pyvene 的声明式干预框架对 PyTorch 模型执行因果干预的指导。
• nnsight-remote-interpretability：提供使用 nnsight 解释和操作神经网络内部机制的指导，并支持可选的 NDIF 远程执行。
• sparse-autoencoder-training：提供使用 SAELens 训练和分析稀疏自编码器（SAE），将神经网络激活分解为可解释……的指导。