Pyvene Interventions

pyvene: Intervenções Causais para Redes Neurais

pyvene é a biblioteca do Stanford NLP para realizar intervenções causais em modelos PyTorch. Ela fornece um framework declarativo baseado em dicionário para activation patching, rastreamento causal e treinamento de intervenção de intercâmbio — tornando experimentos de intervenção reproduzíveis e compartilháveis.

GitHub: stanfordnlp/pyvene (840+ stars) Paper: pyvene: A Library for Understanding and Improving PyTorch Models via Interventions (NAACL 2024)

Quando Usar pyvene

Use pyvene quando você precisar:

• Realizar rastreamento causal (localização estilo ROME)
• Executar experimentos de activation patching
• Conduzir treinamento de intervenção de intercâmbio (IIT)
• Testar hipóteses causais sobre componentes do modelo
• Compartilhar/reproduzir experimentos de intervenção via HuggingFace
• Trabalhar com qualquer arquitetura PyTorch (não apenas transformers)

Considere alternativas quando:

• Você precisa de análise exploratória de ativações → Use TransformerLens
• Você quer treinar/analisar SAEs → Use SAELens
• Você precisa executar em larga escala em modelos massivos → Use nnsight
• Você quer controle de nível inferior → Use nnsight

Instalação

pip install pyvene

Import padrão:

import pyvene as pv

Conceitos Principais

IntervenableModel

A classe principal que encapsula qualquer modelo PyTorch com capacidades de intervenção:

import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer

# Carrega o modelo base
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Define configuração de intervenção
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,
            component="block_output",
            intervention_type=pv.VanillaIntervention,
        )
    ]
)

# Cria modelo intervenável
intervenable = pv.IntervenableModel(config, model)

Tipos de Intervenção

Alvos de Componentes

# Componentes disponíveis para intervir
components = [
    "block_input",      # Entrada para bloco transformer
    "block_output",     # Saída do bloco transformer
    "mlp_input",        # Entrada para MLP
    "mlp_output",       # Saída de MLP
    "mlp_activation",   # Ativações ocultas do MLP
    "attention_input",  # Entrada para attention
    "attention_output", # Saída de attention
    "attention_value_output",  # Vetores de valor de attention
    "query_output",     # Vetores de query
    "key_output",       # Vetores de chave
    "value_output",     # Vetores de valor
    "head_attention_value_output",  # Valores por cabeça
]

Workflow 1: Rastreamento Causal (estilo ROME)

Localize onde associações factuais são armazenadas corrompendo entradas e restaurando ativações.

Passo a Passo

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 entradas limpas e corrompidas
clean_prompt = "The Space Needle is in downtown"
corrupted_prompt = "The ##### ###### ## ## ########"  # Ruído

clean_tokens = tokenizer(clean_prompt, return_tensors="pt")
corrupted_tokens = tokenizer(corrupted_prompt, return_tensors="pt")

# 2. Obtém ativações limpas (fonte)
with torch.no_grad():
    clean_outputs = model(**clean_tokens, output_hidden_states=True)
    clean_states = clean_outputs.hidden_states

# 3. Define intervenção de restauração
def run_causal_trace(layer, position):
    """Restaura ativação limpa em camada e posição específicas."""
    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)

    # Executa com intervenção
    _, patched_outputs = intervenable(
        base=corrupted_tokens,
        sources=[clean_tokens],
        unit_locations={"sources->base": ([[[position]]], [[[position]]])},
        output_original_output=True,
    )

    # Retorna probabilidade do token correto
    probs = torch.softmax(patched_outputs.logits[0, -1], dim=-1)
    seattle_token = tokenizer.encode(" Seattle")[0]
    return probs[seattle_token].item()

# 4. Varre sobre camadas e posições
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. Visualiza (mapa de calor camada x posição)
# Valores altos indicam importância causal

Checklist

• Prepare prompt limpo com associação factual alvo
• Crie versão corrompida (ruído ou contrafactual)
• Define config de intervenção para cada (camada, posição)
• Executa varredura de patching
• Identifica hotspots causais no mapa de calor

Workflow 2: Activation Patching para Análise de Circuitos

Teste quais componentes são necessários para um comportamento específico.

Passo a Passo

import pyvene as pv
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Setup tarefa IOI
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):
    """Diferença de logits IO - S."""
    return logits[0, -1, john_token] - logits[0, -1, mary_token]

# Patch saída de attention em cada camada
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()

# Encontra quais camadas importam
results = []
for layer in range(model.config.n_layer):
    diff = patch_attention(layer)
    results.append(diff)
    print(f"Layer {layer}: logit diff = {diff:.3f}")

Workflow 3: Treinamento de Intervenção de Intercâmbio (IIT)

Treine intervenções para descobrir estrutura causal.

Passo a Passo

import pyvene as pv
from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")

# 1. Define intervenção treinável
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=6,
            component="block_output",
            intervention_type=pv.RotatedSpaceIntervention,  # Treinável
            low_rank_dimension=64,  # Aprende subespaço 64-dim
        )
    ]
)

intervenable = pv.IntervenableModel(config, model)

# 2. Configura treinamento
optimizer = torch.optim.Adam(
    intervenable.get_trainable_parameters(),
    lr=1e-4
)

# 3. Loop de treinamento (simplificado)
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. Analisa intervenção aprendida
# A matriz de rotação revela subespaço causal
rotation = intervenable.interventions["layer.6.block_output"][0].rotate_layer

DAS (Distributed Alignment Search)

# Rotação de baixo rank encontra subespaços interpretáveis
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,
            component="block_output",
            intervention_type=pv.LowRankRotatedSpaceIntervention,
            low_rank_dimension=1,  # Encontra direção causal 1D
        )
    ]
)

Workflow 4: Steering de Modelo (Honest LLaMA)

Conduza o comportamento do modelo durante a geração.

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

# Carrega intervenção de steering pré-treinada
intervenable = pv.IntervenableModel.load(
    "zhengxuanzenwu/intervenable_honest_llama2_chat_7B",
    model=model,
)

# Gera com steering
prompt = "Is the earth flat?"
inputs = tokenizer(prompt, return_tensors="pt")

# Intervenção aplicada durante geração
outputs = intervenable.generate(
    inputs,
    max_new_tokens=100,
    do_sample=False,
)

print(tokenizer.decode(outputs[0]))

Salvando e Compartilhando Intervenções

# Salva localmente
intervenable.save("./my_intervention")

# Carrega de local
intervenable = pv.IntervenableModel.load(
    "./my_intervention",
    model=model,
)

# Compartilha no HuggingFace
intervenable.save_intervention("username/my-intervention")

# Carrega do HuggingFace
intervenable = pv.IntervenableModel.load(
    "username/my-intervention",
    model=model,
)

Problemas Comuns & Soluções

Problema: Local de intervenção incorreto

# ERRADO: Nome de componente incorreto
config = pv.RepresentationConfig(
    component="mlp",  # Inválido!
)

# CORRETO: Use nome exato do componente
config = pv.RepresentationConfig(
    component="mlp_output",  # Válido
)

Problema: Incompatibilidade de dimensão

# Garanta que fonte e base tenham formas compatíveis
# Para intervenções específicas de posição:
config = pv.RepresentationConfig(
    unit="pos",
    max_number_of_units=1,  # Intervém em posição única
)

# Especifique localizações explicitamente
intervenable(
    base=base_tokens,
    sources=[source_tokens],
    unit_locations={"sources->base": ([[[5]]], [[[5]]])},  # Posição 5
)

Problema: Memória com modelos grandes

# Use gradient checkpointing
model.gradient_checkpointing_enable()

# Ou intervenha em menos componentes
config = pv.IntervenableConfig(
    representations=[
        pv.RepresentationConfig(
            layer=8,  # Camada única em vez de todas
            component="block_output",
        )
    ]
)

Problema: Integração LoRA

# pyvene v0.1.8+ suporta LoRAs como intervenções
config = pv.RepresentationConfig(
    intervention_type=pv.LoRAIntervention,
    low_rank_dimension=16,
)

Referência de Classes Principais

Modelos Suportados

pyvene funciona com qualquer modelo PyTorch. Testado em:

• GPT-2 (todos os tamanhos)
• LLaMA / LLaMA-2
• Pythia
• Mistral / Mixtral
• OPT
• BLIP (vision-language)
• ESM (modelos de proteína)
• Mamba (state space)

Documentação de Referência

Para documentação detalhada da API, tutoriais e uso avançado, veja a pasta references/:

Recursos Externos

Tutoriais

• pyvene 101
• Tutorial de Rastreamento Causal
• Replicação de Circuito IOI
• Introdução a DAS

Papers

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

Documentação Oficial

• Docs Oficiais
• Referência da API

Comparação com Outras Ferramentas