Interpretabilidad mecanicista (Checkpoint C3b)

B.1 La premisa

Tenemos modelos que hablan a nivel humano y no sabemos cómo funcionan por dentro. La interpretabilidad mecanicista (mech interp) intenta hacer ingeniería inversa de los algoritmos que el modelo aprendió, a partir de sus pesos y activaciones. No es filosofía: es experimentación con bucles de feedback cortísimos, y conecta con seguridad (entender ⇒ poder controlar).

Tres niveles de análisis que recorrerás:

1. Circuitos: subgrafos de atención+MLP que implementan una función (p.ej. "copiar el nombre correcto").
2. Features: direcciones en el espacio de activaciones que representan conceptos (vía SAEs).
3. Intervención causal: cambiar una activación/feature y verificar el efecto → prueba de que eso causaba aquello.

B.2 Herramientas

• TransformerLens (pip install transformer-lens): carga modelos GPT-style con hooks en cada activación → puedes leerlas y editarlas. La navaja suiza de mech interp.
• SAELens (pip install sae-lens, repo decoderesearch/SAELens, v6): entrenar y analizar Sparse Autoencoders; carga SAEs preentrenados (Gemma Scope) con SAE.from_pretrained.
• nnsight: para modelos grandes que no caben en TransformerLens.
• circuitsvis: visualización de patrones de atención.

B.3 Concepto 1 — Induction heads (el circuito fundacional)

Los induction heads son el circuito más estudiado: cabezas de atención que implementan "si viste [A][B] antes y ahora ves [A], predice [B]" — copia de patrones. Emergen durante el pretraining y explican gran parte del in-context learning.

Laboratorio B.1 — localizar induction heads:

1# lab_n3b_induction.py — detecta induction heads en un modelo pequeño con TransformerLens
2import torch
3from transformer_lens import HookedTransformer
4
5model = HookedTransformer.from_pretrained("gpt2")  # cabe de sobra en la 5090
6# Secuencia repetida: el modelo debería "inducir" la segunda mitad de la primera
7torch.manual_seed(0)
8seq_len = 50
9tokens = torch.randint(1000, 10000, (1, seq_len), device="cuda")
10rep = torch.cat([tokens, tokens], dim=1)   # [random][mismo random] -> patrón inducible
11
12logits, cache = model.run_with_cache(rep)   # cache guarda TODAS las activaciones
13
14# induction score por cabeza: cuánto atiende la posición t al token que siguió a la
15# aparición previa del token actual (offset = seq_len - 1)
16def induction_score(cache, layer, head):
17    pattern = cache["pattern", layer][0, head]   # [dst, src] pesos de atención
18    # diagonal desplazada: posición i (en la 2ª mitad) atiende a i - seq_len + 1
19    idx = torch.arange(seq_len, 2*seq_len, device=pattern.device)
20    src = idx - seq_len + 1
21    return pattern[idx, src].mean().item()
22
23for L in range(model.cfg.n_layers):
24    for H in range(model.cfg.n_heads):
25        s = induction_score(cache, L, H)
26        if s > 0.3:
27            print(f"L{L}H{H}: induction score = {s:.2f}")  # estas son tus induction heads

Líneas no triviales:

• run_with_cache: ejecuta el modelo guardando todas las activaciones intermedias (patrones de atención, residual stream, etc.) para inspeccionarlas. Es el corazón de TransformerLens.
• cache["pattern", layer]: los pesos de atención de esa capa, shape [head, dst, src].
• La "diagonal desplazada": en la 2ª mitad repetida, una induction head atiende a la posición que siguió al token actual la primera vez → alto peso en pattern[i, i-seq_len+1].

B.4 Concepto 2 — El circuito IOI (reproducción canónica)

IOI (Indirect Object Identification): "When Mary and John went to the store, John gave a drink to ___" → el modelo predice "Mary". Reproducir este circuito en GPT-2 small es el ejercicio de iniciación de mech interp, usando dos técnicas centrales:

• Activation patching: sustituyes una activación por la de otro prompt y ves cuánto cambia el output → localiza dónde vive la información.
• Direct logit attribution (DLA): descompones el logit final en la contribución de cada componente → localiza quién empuja la respuesta.

Laboratorio B.2: sigue el tutorial IOI de ARENA (learn.arena.education) con TransformerLens; reproduce la identificación de las "name mover heads" y haz activation patching para confirmar causalidad. (El código completo y los diagramas están en ARENA; tu trabajo es reconstruirlo en el segundo pase, no copiarlo.)

B.5 Concepto 3 — SAEs: de neuronas polisemánticas a features

Problema: una sola neurona se activa con conceptos no relacionados (polisemántica) — "caras humanas, frontales de coches y espaldas de gatos" a la vez. Causa: superposición (el modelo empaqueta más features que dimensiones). Solución: un Sparse Autoencoder proyecta las activaciones a un espacio mucho más ancho y disperso donde cada dimensión (feature) tiende a ser monosemántica:

f = Act(W_enc · a + b_enc)        # código disperso (la mayoría ceros), dim >> d_model
â = W_dec · f + b_dec             # reconstrucción
Loss = ||a − â||²  +  λ·||f||₁    # reconstrucción + penalización de dispersión (o TopK)

Laboratorio B.3 — entrenar un SAE (código completo):

1# lab_n3b_sae.py — entrena un TopK-SAE sobre el residual stream de un modelo pequeño
2from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner
3
4cfg = LanguageModelSAERunnerConfig(
5    model_name="tiny-stories-1L-21M",   # modelo minúsculo: SAE entrenable en minutos en la 5090
6    hook_name="blocks.0.hook_resid_pre",# DÓNDE extraemos activaciones (residual stream, capa 0)
7    d_in=1024,                          # dim del residual stream del modelo
8    expansion_factor=16,                # d_sae = 16 × d_in -> diccionario ancho
9    activation_fn="topk", k=32,         # TopK-SAE: solo 32 features activas por token
10    lr=4e-4, l1_coefficient=0.0,        # TopK no necesita L1 (la dispersión la fija k)
11    train_batch_size_tokens=4096,
12    training_tokens=20_000_000,
13    device="cuda", wandb_project="mi-sae",
14)
15sae = SAETrainingRunner(cfg).run()
16sae.save_model("saes/tinystories-resid-topk")

Para un modelo de verdad (Qwen3-1.7B), cambia model_name/hook_name/d_in y sube training_tokens; cabe en la 5090 entrenando sobre activaciones cacheadas. También puedes cargar un SAE preentrenado para analizar sin entrenar:

1from sae_lens import SAE
2sae = SAE.from_pretrained("gemma-scope-2b-pt-res-canonical", "layer_12/width_16k/canonical", device="cuda")

Líneas no triviales:

• hook_name: el punto exacto del modelo de donde se toman las activaciones. Distintos puntos (resid_pre, resid_post, mlp_out) dan features distintas; el residual stream es el estándar.
• expansion_factor=16 + k=32: el diccionario es 16× más ancho que d_model pero solo 32 features se activan por token → fuerza monosemanticidad.
• TopK fija la dispersión por construcción (mejor que ajustar λ a mano).

B.6 Concepto 4 — Intervención causal (el cierre de C3b)

Una feature solo es interesante si causa comportamiento. Feature steering: identificas una feature monosemántica, la amplificas (o suprimes) añadiendo su dirección al residual stream durante el forward, y verificas que el output cambia de forma predecible.

1# lab_n3b_steering.py — steering: amplifica una feature y verifica el efecto causal
2import torch
3from transformer_lens import HookedTransformer
4# (carga modelo + tu SAE entrenado; identifica una feature 'f_id' que se activa con un concepto)
5
6def steer_hook(resid, hook, sae, f_id, strength):
7    # añade la dirección de la feature f_id (columna f_id del decoder) al residual stream
8    direction = sae.W_dec[f_id]                # dirección de la feature en el espacio del modelo
9    return resid + strength * direction
10
11model.run_with_hooks(
12    tokens,
13    fwd_hooks=[("blocks.0.hook_resid_pre",
14                lambda r, h: steer_hook(r, h, sae, f_id=1234, strength=8.0))],
15)
16# Compara la generación con strength=0 vs strength=8: el concepto de la feature 1234
17# debería aparecer/intensificarse en el texto -> evidencia CAUSAL, no correlacional.

Si al amplificar la feature "1234" el texto se llena del concepto que esa feature representa, has demostrado causalidad. Eso es el checkpoint.

B.7 CHECKPOINT C3b — criterio de aprobado

• Localizas induction heads en un modelo pequeño (B.1) y reproduces el circuito IOI con activation patching (B.2).
• Entrenas un SAE sobre un modelo (TinyStories o Qwen3-1.7B) e identificas al menos una feature monosemántica (sabes qué concepto representa).
• Haces una intervención causal verificable (steering) que cambia el output de forma predecible.
• Sabes explicar superposición y por qué un SAE produce features más interpretables que las neuronas.

Rúbrica: Nivel 3 si reproduces lo anterior; Nivel 4 si abordas un problema de las 200 Concrete Open Problems de Neel Nanda y aportas un hallazgo.

B.8 Ejercicios

E1. Entrena dos SAEs en dos capas distintas del mismo modelo. ¿Las features de capas tempranas vs tardías representan conceptos más concretos o más abstractos?

E2. Encuentra una feature que se active con un concepto claro (p.ej. "código Python") y mide su efecto de steering a varias strength. ¿Hay un punto donde rompe la fluidez?

E3. Usa activation patching para localizar en qué capa "decide" GPT-2 el sujeto correcto en una frase IOI propia.

B.9 Trampas comunes

• Confundir correlación (una feature se activa con X) con causalidad (steering provoca X). El steering es la prueba.
• Entrenar el SAE con muy pocos tokens → features ruidosas/no interpretables.
• Elegir mal el hook_name.

B.10 Referencias

• Neel Nanda, "Concrete Steps to Get Started in Mech Interp" + 200 Concrete Open Problems (neelnanda.io). ARENA (learn.arena.education). TransformerLens / SAELens docs. Papers: A Mathematical Framework for Transformer Circuits (Elhage 2021), IOI (Wang et al. 2022), Towards/Scaling Monosemanticity (Bricken 2023 / Templeton 2024), Gemma Scope.