Hooks: mirar y tocar activaciones

Alineación a bajo nivel. Episodio siete: hooks. Hasta acá cambiamos al modelo entrenándolo. Desde hoy lo abrimos: vamos a interceptar el forward de PyTorch para leer lo que el modelo calcula por dentro, y para cambiarlo mientras corre.

El plan de hoy

• Cómo funciona un hook de PyTorch por dentro
• Guardar y modificar activaciones de Gemma 2
• TransformerLens y nnsight: lo mismo, con nombres
• Logit lens y activation patching

Cuatro partes. Primero, el mecanismo: qué es un hook en PyTorch, qué recibe y qué pasa con lo que devuelve. Después lo usamos en Gemma dos de verdad, para guardar una activación y para modificarla. Tercero, las dos librerías que usa la comunidad, TransformerLens y nnsight, que por debajo hacen lo mismo. Y al final, dos técnicas que salen directo de esto: el logit lens y el activation patching.

Parte 1 · El hook de PyTorch

Todo módulo de PyTorch tiene un lugar donde podés colgar funciones propias, antes y después de su forward.

Qué pasa cuando llamás a un módulo

Cuando escribís capa de x, no se llama directo al forward. Se llama al método call de nn.Module, que es el que orquesta. Primero corre los pre-hooks, que ven los argumentos y pueden reemplazarlos. Después, el forward que escribiste vos. Después, los forward hooks, que ven la entrada y la salida, y también pueden reemplazar la salida. Y recién ahí te devuelve el resultado. Los hooks son eso: funciones que call ejecuta alrededor del forward.

capa(x) → call → pre-hooks → forward → hooks → salida

Anatomía de un hook

Probémoslo en una capa lineal de juguete. Un forward hook recibe tres cosas: el módulo, la entrada, y la salida. Si no devuelve nada, solo mira. Si devuelve algo, eso reemplaza la salida, y es lo que recibe la capa siguiente. El pre-hook corre antes y recibe solo el módulo y los argumentos; lo que devuelve reemplaza los argumentos. Y register devuelve un handle: si no llamás a remove, el hook queda colgado para siempre. En la salida: la entrada llega como tupla de un elemento, no como tensor. Mirar se imprime dos veces porque en la segunda llamada siguen los dos hooks, en orden de registro. La suma da cero: ganó anular. Y con el pre-hook, la salida es la de la entrada duplicada.

def mirar(module, input, output):        # se llama DESPUÉS del forward
    print(type(module).__name__, type(input).__name__, len(input))
    # devolver None = no tocar nada

def anular(module, input, output):
    return output * 0                    # lo que devuelve REEMPLAZA la salida

def duplicar_entrada(module, args):      # pre-hook: ANTES del forward
    return (args[0] * 2,)                # reemplaza los argumentos

h1 = capa.register_forward_hook(mirar)
print(capa(x).abs().sum().item() > 0)
h2 = capa.register_forward_hook(anular)
print(capa(x).sum().item())
h1.remove(); h2.remove()                 # sin remove, el hook queda para siempre
h3 = capa.register_forward_pre_hook(duplicar_entrada)
print(torch.allclose(capa(x), capa.weight @ (2 * x[0]) + capa.bias))

Linear tuple 1
True
Linear tuple 1
0.0
True

La semántica, en tres líneas

pre-hook

el forward de siempre

forward hook: y' sigue viaje

Escrito como matemática es muy simple. El pre-hook es una función g de los argumentos: si devuelve algo, el forward usa eso. El forward calcula y igual a f sub theta de x, sin enterarse de nada. Y el hook es una función h de la entrada y la salida: si devuelve algo, ese y prima es lo que sigue viaje por el resto de la red. Mirar es devolver None. Tocar es devolver otro tensor. No hay más.

Parte 2 · En Gemma 2, de verdad

Ahora colgamos hooks en una capa de Gemma dos. Primero hay que saber qué devuelve esa capa.

Guardar la salida de la capa 12

Este hook tiene una firma más larga, porque lo registramos con with kwargs igual a true. El motivo: la capa recibe el flujo como argumento posicional, pero la máscara y la posición le llegan por nombre, y un hook común no las ve. Corremos, y lo sacamos. Primera lección: en transformers cinco, la capa decoder de Gemma dos devuelve un tensor, no una tupla. En versiones viejas era una tupla, y un hook que modifica tiene que devolver lo mismo que recibe. Después, los kwargs: la máscara, el caché, las posiciones. Y lo que guardamos es idéntico al estado oculto trece. Mirá las normas: el primer token mide casi dos mil, entre diez y veinte veces más que los demás. Ese token de inicio va a molestar en todos los episodios que siguen.

cache = {}
def guardar(module, args, kwargs, output):
    cache["out"] = output.detach()
    cache["kwargs"] = sorted(kwargs)

capa = model.model.layers[12]                 # Gemma2DecoderLayer
h = capa.register_forward_hook(guardar, with_kwargs=True)
with torch.no_grad():
    out = model(**ids, output_hidden_states=True)
h.remove()

print(type(cache["out"]).__name__, tuple(cache["out"].shape))
print(cache["kwargs"][:4])
print(torch.equal(cache["out"], out.hidden_states[13]))
print(f"norma por posición: {cache['out'][0].norm(dim=-1).round().tolist()}")

Tensor (1, 9, 2304)
['attention_mask', 'output_hidden_states', 'past_key_values', 'position_embeddings']
True
norma por posición: [1946.0, 689.0, 146.0, 187.0, 136.0, 110.0, 103.0, 133.0, 119.0]

Un hook que cambia la respuesta

Ahora uno que toca. La idea: que ciertas capas no escriban nada en el último token. Como la capa suma al flujo, alcanza con devolver la entrada en esa posición: la contribución de la capa desaparece. Lo colgamos en las capas ocho a diecinueve. Y al final los sacamos. Normal, el modelo dice París con noventa y tres por ciento. Sin esas doce capas en la última posición, París cae a doce por ciento y gana France. Sacamos los hooks, y vuelve exactamente a noventa y tres. Ojo con la lectura: eso no prueba que el dato esté en esas capas; prueba que el último token las necesita para terminar la respuesta.

def apagar_ultimo(module, input, output):
    nuevo = output.clone()
    nuevo[:, -1, :] = input[0][:, -1, :]   # la capa no escribe en el último token
    return nuevo                            # tensor, no tupla (transformers 5.x)

with torch.no_grad():
    print("normal  ", top3(model(**ids).logits))
    hs = [model.model.layers[l].register_forward_hook(apagar_ultimo)
          for l in range(8, 20)]
    print("sin 8-19", top3(model(**ids).logits))
    for h in hs:
        h.remove()
    print("de nuevo", top3(model(**ids).logits))

normal   [(' Paris', 0.93), (' ', 0.02), (' the', 0.01)]
sin 8-19 [(' France', 0.23), (' Paris', 0.12), (' X', 0.1)]
de nuevo [(' Paris', 0.93), (' ', 0.02), (' the', 0.01)]

Parte 3 · Las librerías

Escribir hooks a mano escala mal. Las dos librerías que se usan resuelven lo mismo de dos maneras distintas.

HookPoint: una identidad con nombre

TransformerLens se apoya en una idea muy simple. Un HookPoint es un módulo que devuelve su entrada tal cual. No calcula nada: existe para que haya un lugar con nombre donde colgar un hook de PyTorch. Setup recorre el modelo y le pone a cada uno su nombre. Run with cache cuelga un hook que guarda en cada punto. Y run with hooks recibe pares de nombre y función; la firma es activación y hook, y lo que devolvés reemplaza la activación. El HookPoint es un nn.Module llamado hook mid, el caché tiene esa clave, y con la función que multiplica por cero la salida queda en cero.

class Mini(HookedRootModule):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(3, 3)
        self.hook_mid = HookPoint()          # identidad con nombre
        self.setup()                         # registra los nombres

    def forward(self, x):
        return self.hook_mid(self.lin(x)) * 10

m = Mini()
x = torch.randn(1, 3)
print(isinstance(m.hook_mid, nn.Module), m.hook_mid.name)
_, cache = m.run_with_cache(x)
print(list(cache.keys()))
cero = lambda act, hook: act * 0             # firma TL: (act, hook)
print(m.run_with_hooks(x, fwd_hooks=[("hook_mid", cero)]))

True hook_mid
['hook_mid']
tensor([[0., -0., 0.]], grad_fn=<MulBackward0>)

TransformerLens sobre Gemma 2

Una advertencia: en TransformerLens cuatro ya no existe HookedTransformer. Ahora se usa TransformerBridge, que envuelve el modelo de transformers y le agrega HookPoints. Run with cache corre una vez y te deja leer cosas como blocks punto doce punto hook resid post: el residual stream después del bloque doce. Mil cuatrocientas diecisiete activaciones guardadas en una sola pasada. Después borramos la salida del MLP, el multilayer perceptron de cada capa, en el último token, en tres capas. Sin hook, París está en cero coma nueve dos seis. Borrar el MLP de la capa veintidós lo baja a cero coma ocho tres dos; en la doce, hasta sube un poco. Un MLP solo casi nunca decide.

model = TransformerBridge.boot_transformers("google/gemma-2-2b",
                                            dtype=torch.float32, device="cpu")
texto = "The Eiffel Tower is in the city of"
logits, cache = model.run_with_cache(texto)
print(len(cache), "activaciones guardadas")
print(cache["blocks.12.hook_resid_post"].shape)
print([k[10:] for k in cache if k.startswith("blocks.12.hook_")][:5])

paris = model.to_single_token(" Paris")
def borrar_mlp(act, hook):                 # act: [batch, pos, d]
    act[:, -1] = 0
    return act
for capa in (2, 12, 22):
    l = model.run_with_hooks(texto, fwd_hooks=[
        (f"blocks.{capa}.hook_mlp_out", borrar_mlp)])
    print(capa, f"{l[0, -1].softmax(-1)[paris]:.3f}  (sin hook: "
          f"{logits[0, -1].softmax(-1)[paris]:.3f})")

1417 activaciones guardadas
torch.Size([1, 9, 2304])
['hook_in', 'hook_resid_pre', 'hook_attn_out', 'hook_resid_mid', 'hook_mlp_out']
2 0.910  (sin hook: 0.926)
12 0.942  (sin hook: 0.926)
22 0.832  (sin hook: 0.926)

nnsight: intervenir asignando

nnsight lo resuelve distinto. Adentro de model trace no se ejecuta nada todavía: cada cosa que escribís arma un grafo de intervenciones, que corre cuando cerrás el bloque. Save marca qué valores querés que sobrevivan afuera del with. Y para intervenir, asignás: acá multiplicamos por un medio el flujo del último token en la capa doce. Comparamos la probabilidad de París. Mismas formas que antes, París en cero coma nueve dos seis, el mismo número que TransformerLens, y con la intervención cae a cero coma seis ocho nueve. Una trampa que pisé: si en vez de la salida del modelo guardás la de lm head, te salteás el softcap final de Gemma y París da uno coma cero.

model = LanguageModel("google/gemma-2-2b", dtype=torch.float32,
                      device_map="cpu", dispatch=True)
texto = "The Eiffel Tower is in the city of"
with model.trace(texto):                       # arma un grafo de intervención
    h12 = model.model.layers[12].output.save() # .save(): sobrevive al with
    logits = model.output.logits.save()     # salida final del modelo
print(h12.shape, logits.shape)

with model.trace(texto):
    model.model.layers[12].output[:, -1, :] *= 0.5   # intervenir = asignar
    logits2 = model.output.logits.save()
paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0]
for l in (logits, logits2):
    print(f"p(' Paris') = {l[0, -1].softmax(-1)[paris]:.3f}")
with model.trace(texto):                       # trampa: lm_head no tiene el softcap
    crudos = model.lm_head.output.save()
print(f"lm_head: p(' Paris') = {crudos[0, -1].softmax(-1)[paris]:.3f}")

torch.Size([1, 9, 2304]) torch.Size([1, 9, 256000])
p(' Paris') = 0.926
p(' Paris') = 0.689
lm_head: p(' Paris') = 1.000

Dos estilos, el mismo mecanismo

Resumiendo. TransformerLens te da nombres estables para todo, iguales entre arquitecturas, y funciones que reciben la activación. Es cómodo para recorrer capa por capa. nnsight respeta los nombres del modelo de transformers y te deja escribir la intervención como si editaras el tensor. Debajo de los dos hay lo mismo que viste en la parte uno: un register forward hook.

Parte 4 · El logit lens

Primera técnica: preguntarle a cada capa qué palabra diría si el modelo terminara ahí.

Leer el residual stream como si fuera el final

El código. Guardamos la salida de cada capa con un hook. Ojo: no uso hidden states para la última, porque transformers la devuelve ya normalizada. A cada x sub l le aplicamos la norma final, la RMSNorm, root mean square norm, que divide por la raíz cuadrática media, y la matriz de unembedding, como si fuera la última capa. Gemma dos además aplica un softcap: treinta por tangente hiperbólica de logits sobre treinta. Y softmax. Capas cero y cinco: tokens sin sentido. En la diez y la quince, un token en árabe, el ruido típico del logit lens en capas intermedias. En la veinte ya aparece París con cero coma nueve uno cinco, y en la veinticinco coincide con la salida real del modelo.

Cada línea, su término:

• register_forward_hook(lambda → $x_l = \text{salida de la capa } l$ — guardada con un hook
• model.lm_head(model.model.norm(h)) → $\ell_l = W_U\,\mathrm{RMSNorm}(x_l)$ — el mismo final del modelo
• torch.tanh(l / cap) * cap → $\tilde\ell = 30\,\tanh(\ell/30)$ — softcap de Gemma 2
• softmax(-1) → $p_l = \mathrm{softmax}(\tilde\ell_l)$ — qué diría la capa l

hs = []                              # salida de cada capa, con hooks
for capa in model.model.layers:
    capa.register_forward_hook(lambda m, i, o: hs.append(o.detach()))
with torch.no_grad():
    final = model(**ids).logits[0, -1].softmax(-1)

def logit_lens(h):                  # el mismo final que usa el modelo
    l = model.lm_head(model.model.norm(h))
    return torch.tanh(l / cap) * cap             # softcap de Gemma 2

curva = []
for capa, h in enumerate(hs):
    p = logit_lens(h[0, -1]).softmax(-1)
    curva.append(p[paris].item())
    if capa % 5 == 0:
        print(f"capa {capa:2}: {tok.decode(p.argmax())!r:9} p(Paris)={p[paris]:.3f}")
print(f"modelo  : {tok.decode(final.argmax())!r:9} p(Paris)={final[paris]:.3f}")

capa  0: ' of'     p(Paris)=0.000
capa  5: '<bos>'   p(Paris)=0.000
capa 10: 'بوابة'   p(Paris)=0.000
capa 15: 'بوابة'   p(Paris)=0.000
capa 20: ' Paris'  p(Paris)=0.915
capa 25: ' Paris'  p(Paris)=0.926
modelo  : ' Paris'  p(Paris)=0.926

p(París) según el logit lens, capa por capa

Toda la curva. Hasta la capa dieciocho, la probabilidad de París es prácticamente cero. Entre la diecinueve y la veinte salta a más de cero coma nueve, y se queda ahí salvo una caída en la veintitrés. La lectura prudente: la respuesta se vuelve legible en el vocabulario recién en el último cuarto del modelo. Eso no dice que antes no esté: puede estar en una forma que W sub U no sabe leer.

Parte 5 · Activation patching

El logit lens es observación. Para hablar de causas hay que intervenir: correr el modelo con un pedazo de otra corrida.

La métrica: logit diff

en el último token

Eiffel vs. Colosseum

0 = nada, 1 = todo restaurado

Necesitamos dos prompts del mismo largo, uno limpio con la Torre Eiffel y uno corrupto con el Coliseo. La métrica es la diferencia de logits entre París y Roma. En el dibujo: la fila de París de W sub U menos la de Roma da una sola dirección, así que la diferencia es lineal en el residual stream y cancela lo que no importa. En el limpio vale diez coma ocho seis; en el corrupto, menos doce coma tres dos. En el dibujo, ese es el recorrido entero: veintitrés coma dieciocho. Corremos el corrupto, pegamos la activación limpia de una capa y una posición, y normalizamos: cero es que no cambió nada, uno es que volvió todo el comportamiento limpio.

Parchear capa por capa

El código son tres pasos. Primero verificamos que los dos prompts tengan el mismo largo: nueve tokens. Corremos el limpio guardando la salida de cada capa con hooks. Después, para cada capa y posición, un hook que pega en el corrupto la activación limpia en esa posición, y normalizamos. Cincuenta y dos pasadas, unos segundos. En la posición de Tower, el último token del sujeto, parchear la capa cinco o la diez restaura casi el noventa por ciento. Desde la veinte ya no sirve. En la última posición pasa lo contrario: nada hasta la diez, y uno coma cero en la veinticinco.

limpio = tok("The Eiffel Tower is in the city of", return_tensors="pt")
corrupto = tok("The Colosseum is in the city of", return_tensors="pt")
print(limpio.input_ids.shape, corrupto.input_ids.shape)   # mismo largo
ld = lambda logits: (logits[0, -1, PARIS] - logits[0, -1, ROMA]).item()

h_limpio = []                                   # 1) correr limpio y guardar
hs = [c.register_forward_hook(lambda m, i, o: h_limpio.append(o.clone()))
      for c in capas]
ld_limpio = ld(model(**limpio).logits); [h.remove() for h in hs]
ld_corr = ld(model(**corrupto).logits)          # 2) correr corrupto

def parchear(capa, pos):                        # 3) corrupto + un pedazo limpio
    def pegar(m, i, o):
        o[:, pos] = h_limpio[capa][:, pos]; return o
    h = capas[capa].register_forward_hook(pegar)
    efecto = ld(model(**corrupto).logits); h.remove()
    return (efecto - ld_corr) / (ld_limpio - ld_corr)

tabla = {pos: [parchear(c, pos) for c in range(26)] for pos in (3, 8)}
print(f"logit diff limpio {ld_limpio:.2f}, corrupto {ld_corr:.2f}")
for c in (0, 5, 10, 15, 20, 25):
    print(f"capa {c:2}  pos 'Tower' {tabla[3][c]:5.2f}   pos 'of' {tabla[8][c]:5.2f}")

torch.Size([1, 9]) torch.Size([1, 9])
logit diff limpio 10.86, corrupto -12.32
capa  0  pos 'Tower'  0.56   pos 'of' -0.00
capa  5  pos 'Tower'  0.86   pos 'of' -0.00
capa 10  pos 'Tower'  0.89   pos 'of'  0.02
capa 15  pos 'Tower'  0.68   pos 'of'  0.17
capa 20  pos 'Tower'  0.03   pos 'of'  0.71
capa 25  pos 'Tower'  0.00   pos 'of'  1.00

Dónde vive la información, según el patching

Las dos curvas juntas cuentan una historia. En las primeras quince capas, lo que distingue a Eiffel de Coliseo está en la posición del sujeto. Entre la quince y la veinte se cruzan: la información se mueve al último token, que es donde se lee la respuesta. Es una sola pareja de prompts, así que es una ilustración, no un resultado; para afirmarlo harían falta muchos pares.

Meng et al., 2022 — Locating and Editing Factual Associations in GPT · arXiv:2202.05262

MLP de capas medias · Una intervención causal encuentra «a distinct set of steps in middle-layer feed-forward modules that mediate factual predictions while processing subject tokens».

Esta técnica tiene nombre. En ROME, rank-one model editing, la llamaron causal tracing: corromper el sujeto, restaurar una activación limpia y medir. Con eso encontraron que las predicciones factuales pasan por los MLP de capas medias mientras procesan los tokens del sujeto, y usaron ese hallazgo para editar un dato cambiando una sola matriz.

Wang et al., 2022 — Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small · arXiv:2211.00593

26 heads · 7 clases · Su explicación de la tarea IOI en GPT-2 small «encompasses 26 attention heads grouped into 7 main classes».

Y el patching llevado al extremo. En el paper de IOI, indirect object identification, la tarea de completar a quién se le da algo en la frase, parchearon attention heads uno por uno en GPT dos small, con la misma métrica de logit diff, y armaron un circuito de veintiséis heads agrupados en siete clases. El mismo hook que viste hoy, repetido miles de veces.

Lo que te llevás

• Hook = función que call corre alrededor de forward
• Devolver algo reemplaza la salida; remove() o queda colgado
• Gemma 2 en transformers 5: la capa devuelve un tensor
• TransformerLens (Bridge) y nnsight: nombres sobre los mismos hooks
• Logit lens: W_U·norm(x_l); patching: efecto normalizado del LD
• Próximo: direcciones, probes y steering

Papers del episodio

• Meng et al., 2022 — Locating and Editing Factual Associations in GPT
• Wang et al., 2022 — Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small