Sparse Autoencoder Training

SAELens: Autoencodificadores Esparsos para Interpretabilidade Mecanicista

SAELens é a biblioteca principal para treinar e analisar Autoencodificadores Esparsos (SAEs) — uma técnica para decompor ativações polissemânticas de redes neurais em features esparsas e interpretáveis. Baseado na pesquisa groundbreaking da Anthropic sobre monossemânticidade.

GitHub: jbloomAus/SAELens (1.100+ stars)

O Problema: Polissemânticidade & Superposição

Neurônios individuais em redes neurais são polissemânticos — ativam em múltiplos contextos semanticamente distintos. Isso ocorre porque modelos usam superposição para representar mais features do que neurônios possuem, dificultando a interpretabilidade.

SAEs resolvem isso decompondo ativações densas em features esparsas e monossemânticas — tipicamente apenas um pequeno número de features ativa para qualquer entrada, e cada feature corresponde a um conceito interpretável.

Quando Usar SAELens

Use SAELens quando precisar:

• Descobrir features interpretáveis em ativações de modelos
• Entender quais conceitos um modelo aprendeu
• Estudar superposição e geometria de features
• Realizar direcionamento ou ablação baseados em features
• Analisar features relevantes para segurança (engano, viés, conteúdo prejudicial)

Considere alternativas quando:

• Precisar de análise básica de ativações → Use TransformerLens diretamente
• Quiser experimentos de intervenção causal → Use pyvene ou TransformerLens
• Precisar de direcionamento em produção → Considere engenharia direta de ativações

Instalação

pip install sae-lens

Requisitos: Python 3.10+, transformer-lens>=2.0.0

Conceitos Principais

O Que SAEs Aprendem

SAEs são treinados para reconstruir ativações de modelos através de um gargalo esparso:

Ativação Entrada → Encoder → Features Esparsas → Decoder → Ativação Reconstruída
    (d_model)       ↓        (d_sae >> d_model)    ↓         (d_model)
                 penalidade                     perda de
                 de                          reconstrução
                 esparsidade

Função de Perda: MSE(original, reconstruída) + L1_coefficient × L1(features)

Validação Chave (Pesquisa Anthropic)

Em "Towards Monosemanticity", avaliadores humanos encontraram 70% das features de SAE genuinamente interpretáveis. Features descobertas incluem:

• Sequências de DNA, linguagem jurídica, requisições HTTP
• Texto em hebraico, declarações nutricionais, sintaxe de código
• Sentimento, entidades nomeadas, estruturas gramaticais

Workflow 1: Carregando e Analisando SAEs Pré-treinadas

Passo a Passo

from transformer_lens import HookedTransformer
from sae_lens import SAE

# 1. Carregar modelo e SAE pré-treinada
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, cfg_dict, sparsity = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
    device="cuda"
)

# 2. Obter ativações do modelo
tokens = model.to_tokens("The capital of France is Paris")
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]  # [batch, pos, d_model]

# 3. Codificar para features de SAE
sae_features = sae.encode(activations)  # [batch, pos, d_sae]
print(f"Active features: {(sae_features > 0).sum()}")

# 4. Encontrar features principais para cada posição
for pos in range(tokens.shape[1]):
    top_features = sae_features[0, pos].topk(5)
    token = model.to_str_tokens(tokens[0, pos:pos+1])[0]
    print(f"Token '{token}': features {top_features.indices.tolist()}")

# 5. Reconstruir ativações
reconstructed = sae.decode(sae_features)
reconstruction_error = (activations - reconstructed).norm()

SAEs Pré-treinadas Disponíveis

Checklist

• Carregar modelo com TransformerLens
• Carregar SAE correspondente para camada alvo
• Codificar ativações para features esparsas
• Identificar features principais ativadas por token
• Validar qualidade de reconstrução

Workflow 2: Treinando uma SAE Personalizada

Passo a Passo

from sae_lens import SAE, LanguageModelSAERunnerConfig, SAETrainingRunner

# 1. Configurar treinamento
cfg = LanguageModelSAERunnerConfig(
    # Modelo
    model_name="gpt2-small",
    hook_name="blocks.8.hook_resid_pre",
    hook_layer=8,
    d_in=768,  # Dimensão do modelo

    # Arquitetura SAE
    architecture="standard",  # ou "gated", "topk"
    d_sae=768 * 8,  # Fator de expansão de 8
    activation_fn="relu",

    # Treinamento
    lr=4e-4,
    l1_coefficient=8e-5,  # Penalidade de esparsidade
    l1_warm_up_steps=1000,
    train_batch_size_tokens=4096,
    training_tokens=100_000_000,

    # Dados
    dataset_path="monology/pile-uncopyrighted",
    context_size=128,

    # Logging
    log_to_wandb=True,
    wandb_project="sae-training",

    # Checkpointing
    checkpoint_path="checkpoints",
    n_checkpoints=5,
)

# 2. Treinar
trainer = SAETrainingRunner(cfg)
sae = trainer.run()

# 3. Avaliar
print(f"L0 (features ativas em média): {trainer.metrics['l0']}")
print(f"CE Loss Recuperada: {trainer.metrics['ce_loss_score']}")

Hiperparâmetros Principais

Métricas de Avaliação

Checklist

• Escolher camada alvo e hook point
• Definir fator de expansão (d_sae = 4-16× d_model)
• Ajustar coeficiente L1 para esparsidade desejada
• Habilitar warm-up de L1 para prevenir dead features
• Monitorar métricas durante treinamento (W&B)
• Validar L0 e recuperação de CE loss
• Verificar razão de dead features

Workflow 3: Análise de Features e Direcionamento

Analisando Features Individuais

from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch

model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
    device="cuda"
)

# Encontrar o que ativa uma feature específica
feature_idx = 1234
test_texts = [
    "The scientist conducted an experiment",
    "I love chocolate cake",
    "The code compiles successfully",
    "Paris is beautiful in spring",
]

for text in test_texts:
    tokens = model.to_tokens(text)
    _, cache = model.run_with_cache(tokens)
    features = sae.encode(cache["resid_pre", 8])
    activation = features[0, :, feature_idx].max().item()
    print(f"{activation:.3f}: {text}")

Direcionamento de Features

def steer_with_feature(model, sae, prompt, feature_idx, strength=5.0):
    """Adicionar direção de feature de SAE ao fluxo residual."""
    tokens = model.to_tokens(prompt)

    # Obter direção de feature do decoder
    feature_direction = sae.W_dec[feature_idx]  # [d_model]

    def steering_hook(activation, hook):
        # Adicionar direção de feature escalada em todas as posições
        activation += strength * feature_direction
        return activation

    # Gerar com direcionamento
    output = model.generate(
        tokens,
        max_new_tokens=50,
        fwd_hooks=[("blocks.8.hook_resid_pre", steering_hook)]
    )
    return model.to_string(output[0])

Atribuição de Features

# Quais features mais afetam um output específico?
tokens = model.to_tokens("The capital of France is")
_, cache = model.run_with_cache(tokens)

# Obter features na posição final
features = sae.encode(cache["resid_pre", 8])[0, -1]  # [d_sae]

# Obter atribuição de logit por feature
# Contribuição = ativação_feature × peso_decoder × unembedding
W_dec = sae.W_dec  # [d_sae, d_model]
W_U = model.W_U    # [d_model, vocab]

# Contribuição para logit "Paris"
paris_token = model.to_single_token(" Paris")
feature_contributions = features * (W_dec @ W_U[:, paris_token])

top_features = feature_contributions.topk(10)
print("Features principais para previsão de 'Paris':")
for idx, val in zip(top_features.indices, top_features.values):
    print(f"  Feature {idx.item()}: {val.item():.3f}")

Problemas Comuns & Soluções

Problema: Alta razão de dead features

# ERRADO: Sem warm-up, features morrem cedo
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,
    l1_warm_up_steps=0,  # Ruim!
)

# CORRETO: Warm-up de penalidade L1
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=8e-5,
    l1_warm_up_steps=1000,  # Aumentar gradualmente
    use_ghost_grads=True,   # Reviver dead features
)

Problema: Reconstrução ruim (baixa recuperação de CE)

# Reduzir penalidade de esparsidade
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=5e-5,  # Menor = melhor reconstrução
    d_sae=768 * 16,       # Mais capacidade
)

Problema: Features não interpretáveis

# Aumentar esparsidade (L1 maior)
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,  # Maior = mais esparso, mais interpretável
)
# Ou usar arquitetura TopK
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn_kwargs={"k": 50},  # Exatamente 50 features ativas
)

Problema: Erros de memória durante treinamento

cfg = LanguageModelSAERunnerConfig(
    train_batch_size_tokens=2048,  # Reduzir tamanho de batch
    store_batch_size_prompts=4,    # Menos prompts no buffer
    n_batches_in_buffer=8,         # Buffer de ativações menor
)

Integração com Neuronpedia

Navegar por features de SAE pré-treinadas em neuronpedia.org:

# Features são indexadas por SAE ID
# Exemplo: gpt2-small layer 8 feature 1234
# → neuronpedia.org/gpt2-small/8-res-jb/1234

Referência de Classes Principais

Documentação de Referência

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

Recursos Externos

Tutoriais

• Basic Loading & Analysis
• Training a Sparse Autoencoder
• ARENA SAE Curriculum

Artigos

• Towards Monosemanticity - Anthropic (2023)
• Scaling Monosemanticity - Anthropic (2024)
• Sparse Autoencoders Find Highly Interpretable Features - Cunningham et al. (ICLR 2024)

Documentação Oficial

• SAELens Docs
• Neuronpedia - Navegador de features

Arquiteturas de SAE

# SAE TopK (exatamente 50 features ativas)
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn="topk",
    activation_fn_kwargs={"k": 50},
)