Direcciones: probes y steering

Alineación a bajo nivel. Episodio ocho: direcciones. En el episodio anterior aprendimos a colgar hooks. Hoy los usamos para algo concreto: encontrar un vector en el residual stream que corresponde a un concepto, y después leerlo, sumarlo o borrarlo.

La hipótesis lineal

• Un concepto ≈ una dirección r en el residual stream
• Leer: proyectar x sobre r (probe)
• Escribir: sumar α·r (steering)
• Borrar: quitar la componente sobre r (ablación)

Todo el episodio se apoya en una hipótesis. Que muchos conceptos viven como direcciones: un vector r de dos mil trescientos cuatro números. Es una creencia con bastante evidencia, no un teorema. Si es cierta, hay tres operaciones. Leer: proyectar el residual stream sobre r, que es lo que hace una probe. Escribir: sumarle alfa por r, que es el steering. Y borrar: sacarle al flujo su componente sobre r. Las tres son una línea de álgebra lineal dentro de un hook.

Parte 1 · Leer: probes

Una probe es un clasificador chiquito que mira activaciones. Si un clasificador lineal acierta, la información está ahí de forma lineal.

Dos nubes y la flecha que las separa

Un esquema en dos dimensiones, con puntos inventados. Cada punto es la activación de una frase: verdes las verdaderas, rojas las falsas. La forma más simple de encontrar la dirección es restar los promedios: la media de las verdes menos la de las rojas. Esa flecha es r. Y la frontera de decisión es perpendicular a esa flecha. Una regresión logística busca la mejor frontera; la diferencia de medias solo resta dos promedios. En la práctica, las dos se parecen mucho.

Una probe de verdad sobre Gemma

El dataset son sesenta frases sintéticas del tipo la capital de Francia es París, la mitad verdaderas y la mitad con la capital cambiada. Con un hook guardamos el residual stream de la capa doce en el último token, el punto final. Entrenamos con veinte países y probamos con otros diez que la probe nunca vio. Primero, una logística de scikit-learn. Después, la diferencia de medias, con un umbral. Y comparamos las dos direcciones. Las dos aciertan las veinte frases de prueba, y el coseno entre ellas es cero coma setenta y cinco: apuntan parecido sin ser iguales. Ojo: veinte frases es un test chico y fácil. Esto muestra cómo se hace, no cuánto generaliza.

def activacion(frase, capa=12):               # residual stream, último token
    cache = {}
    h = model.model.layers[capa].register_forward_hook(
        lambda m, i, o: cache.update(x=o[0, -1]))
    model(**tok(frase, return_tensors="pt")); h.remove()
    return cache["x"]

X = torch.stack([activacion(f) for f in frases])        # [60, 2304]
tr, te = slice(0, 40), slice(40, 60)                    # 20 países / 10 nuevos

probe = LogisticRegression(max_iter=2000).fit(X[tr], y[tr])
print("logística, acc test:", probe.score(X[te], y[te]))

r = X[tr][y[tr] == 1].mean(0) - X[tr][y[tr] == 0].mean(0)   # μ+ − μ−
umbral = (X[tr] @ r).median()
acc = ((X[te] @ r > umbral) == y[te]).float().mean()
print("dif. de medias, acc test:", acc.item())
w = torch.tensor(probe.coef_[0], dtype=torch.float32)
print(f"coseno(w, r) = {torch.cosine_similarity(w, r, dim=0):.2f}")

logística, acc test: 1.0
dif. de medias, acc test: 1.0
coseno(w, r) = 0.75

Dos probes, una pregunta

regresión logística: w aprendido

diferencia de medias: sin entrenar

leer = proyectar y comparar con un umbral

Las dos probes en matemática. La logística aprende w y b para que sigma de w por x más b sea la probabilidad de verdadera. En el dibujo, la sigmoide aplasta cualquier número entre cero y uno: bien negativo es falsa, bien positivo es verdadera. La diferencia de medias no aprende nada: promedia las activaciones de las verdaderas, promedia las de las falsas, y resta. En el dibujo, las verdaderas suman con un tercio, las falsas restan con un tercio, y sale r. Para clasificar, en los dos casos, se proyecta x sobre el vector y se compara con un umbral tau. En el dibujo, el producto escalar: multiplicar coordenada por coordenada y sumar. Una probe lineal es eso.

Las decisiones que esconde una probe

• Qué capa: en las del medio suele haber más señal
• Qué token: el último, o un promedio sin el <bos>
• Qué control: etiquetas al azar deberían dar 50 %
• Leer ≠ usar: que esté ahí no prueba que el modelo lo use

Detrás de esas tres líneas hay decisiones que cambian el resultado. La capa: la comunidad suele encontrar más señal en las capas del medio que en las primeras o las últimas, pero hay que barrerlas. El token: el último, donde el modelo ya leyó toda la frase, o un promedio, siempre sin el token de inicio, que con su norma gigante domina cualquier media. El control: si entrenás la misma probe con etiquetas mezcladas al azar y también acierta, tu probe memoriza, no lee. Y la más importante: que la información se pueda leer no prueba que el modelo la use. Para eso hay que intervenir, y eso es la parte que sigue.

Marks & Tegmark, 2023 — The Geometry of Truth · arXiv:2310.06824

μ₊ − μ₋ · «Simple difference-in-mean probes generalize as well as other probing techniques», y las direcciones que encuentran son más causales.

Esto no es un truco de juguete. En The Geometry of Truth, Marks y Tegmark mostraron que las probes de diferencia de medias generalizan tan bien como otras técnicas, y que las direcciones que encuentran están más ligadas causalmente a lo que el modelo responde. Si viste Caja blanca, ya probaste esto con experimentos propios; acá nos quedamos con el mecanismo.

Parte 2 · Escribir: steering

Si una dirección representa un concepto, sumarla al flujo debería empujar al modelo hacia ese concepto.

Steering con un hook

Ahora con un concepto que se ve en el texto: positivo contra negativo. La dirección es la media de seis frases positivas menos la de seis negativas, promediando todos los tokens menos el de inicio, que tiene la norma gigante. El hook devuelve la salida más alfa por r, en todas las posiciones. Lo colgamos después de la capa doce. Y generamos de forma greedy, sin azar, para comparar justo. La norma de r es treinta y dos. Sin steering, el restaurante es un gran lugar para comer. Con alfa menos cuatro: el peor restaurante del mundo. Con más cuatro sigue siendo positivo, porque ya lo era. Fijate que alfa por la norma de r da ciento veintiocho: alfa no sale de una fórmula, se elige probando y mirando lo que genera.

Cada línea, su término:

• r = media(pos) - media(neg) → $r = \bar{x}_{\text{pos}} - \bar{x}_{\text{neg}}$ — seis frases de cada lado
• o + alfa * r → $x'_t = x_t + \alpha\, r \quad \forall t$ — en todas las posiciones
• register_forward_hook(sumar(alfa)) → $x_{13} \leftarrow x_{13} + \alpha r$ — después de la capa 12
• do_sample=False → $y = \arg\max_k p(k \mid x')$ — decodificación greedy

r = media(pos) - media(neg)                 # «positivo − negativo»
print(f"|r| = {r.norm():.1f}")

def sumar(alfa):
    return lambda m, i, o: o + alfa * r     # x + α r, toda posición

ids = tok("I think that this restaurant is", return_tensors="pt")
for alfa in (-4, 0, 4):
    h = model.model.layers[CAPA].register_forward_hook(sumar(alfa))
    g = model.generate(**ids, max_new_tokens=9, do_sample=False)
    h.remove()
    print(f"{alfa:+d}", repr(tok.decode(g[0, ids.input_ids.shape[1]:])))

|r| = 32.0
-4 ' the worst restaurant in the world. I don'
+0 ' a great place to eat. The food is'
+4 ' a great place to enjoy the food and the'

Qué hace la suma

En el plano, esquemático. El flujo del último token es un vector x. Sumarle alfa por r, con alfa negativo, lo corre en la dirección negativa del concepto. Y el resultado, x prima, es lo que leen todas las capas siguientes. No cambiamos ningún peso: cambiamos el estado por el que pasa el cálculo. Si alfa es muy grande, el texto se degrada; por eso alfa se elige mirando generaciones, no con una fórmula.

Rimsky et al., 2023 — Steering Llama 2 via Contrastive Activation Addition · arXiv:2312.06681

CAA · «CAA computes "steering vectors" by averaging the difference in residual stream activations»; lo evalúan sobre Llama 2 Chat.

Esta receta tiene nombre. Contrastive Activation Addition: promediar la diferencia de activaciones entre pares de ejemplos de un comportamiento y su opuesto, y sumarla durante la inferencia. Lo evaluaron sobre Llama dos Chat, con preguntas de opción múltiple sobre distintos comportamientos.

Parte 3 · Borrar: ablación direccional

La operación inversa: dejar el flujo intacto en todas las direcciones menos una.

Proyectar afuera una dirección

primero, vector unitario

restar la sombra de x sobre r̂

porque r̂·r̂ = 1

un proyector: aplicarlo dos veces no cambia nada

La matemática es corta. Normalizamos r para que mida uno. x punto r sombrero es cuánto de x apunta en esa dirección; multiplicado por r sombrero, es la sombra de x sobre la recta. Se la restamos. En el dibujo: la flecha punteada es la sombra, y lo que queda es x prima. Escrito como matriz: identidad menos r sombrero por r sombrero transpuesto, por x. La verificación: el producto de x prima con r sombrero da cero, porque r sombrero con sí mismo da uno. En el dibujo, x prima queda perpendicular a la recta. Esa matriz P es un proyector: aplicarla dos veces es lo mismo que una.

Ablación en todas las capas

Misma dirección positivo menos negativo, ahora unitaria. El hook resta la proyección, tal cual la fórmula. Y lo colgamos en las veintiséis capas, para que ninguna pueda volver a escribirla. Un segundo hook, registrado después, lee el resultado. Primera línea, la clave: proyectadas sobre r sombrero, las frases positivas dan menos tres coma nueve en promedio y las negativas menos treinta y cinco coma nueve. Después: la comida era horrible y el mozo era grosero. Con la ablación, el mozo fue muy amable. Y la proyección en la capa doce da exactamente cero.

r = A_pos.mean(0) - A_neg.mean(0)
r_hat = r / r.norm()                                  # vector unitario
print(f"x·r̂ medio: positivas {(A_pos @ r_hat).mean():.1f}, "
      f"negativas {(A_neg @ r_hat).mean():.1f}")

def ablar(m, i, x):
    return x - (x @ r_hat)[..., None] * r_hat         # x − (x·r̂) r̂

ids = tok("The food was awful and the waiter was", return_tensors="pt")
print("normal :", continuar(ids))
hs = [c.register_forward_hook(ablar) for c in model.model.layers]
print("ablada :", continuar(ids))
leer = model.model.layers[CAPA].register_forward_hook(   # corre después
    lambda m, i, x: print(f"x·r̂ en capa {CAPA}: {(x[0, -1] @ r_hat):.1e}"))
model(**ids)
for h in hs + [leer]: h.remove()

x·r̂ medio: positivas -3.9, negativas -35.9
normal : ' rude. We ordered a pizza and a'
ablada : ' very friendly and helpful. The atmosphere was'
x·r̂ en capa 12: 0.0e+00

Cero no es neutral

¿Por qué borrar lo negativo lo volvió positivo? Las negativas están a treinta y seis unidades del cero sobre esa recta. Las positivas, a menos de cuatro. La ablación manda todo al cero, y el cero queda pegado a las positivas. Proyectar afuera una dirección no la neutraliza: la fija en un valor, y ese valor tiene significado. Por eso algunas variantes restan hasta la media en lugar de hasta cero.

• negativas: −35,9
• positivas: −3,9
• con ablación: 0

Parte 4 · Ortogonalizar los pesos

Un hook vive mientras corre tu script. Si querés un modelo que directamente no pueda escribir esa dirección, la cambiás en los pesos.

Del hook a la matriz

El truco: si una matriz W escribe en el flujo, cualquier salida W por a es un vector del flujo. Entonces podemos aplicar el proyector directamente a W, una sola vez. Por asociatividad, W prima por a es lo mismo que calcular W por a y después ablar. Con una matriz de juguete de ocho por ocho: la diferencia entre las dos cuentas es seis por diez a la menos siete, ruido de punto flotante. La salida sobre r sombrero, también cero. Y el proyector tiene rango siete de ocho: borra una dimensión y no toca las demás. Si la capa tiene sesgo, se proyecta igual, porque también se suma al flujo.

Cada línea, su término:

• P = torch.eye(d) - torch.outer(r_hat, r_hat) → $P = I - \hat r \hat r^\top$ — el mismo proyector
• W_orto = P @ W → $W' = (I - \hat r \hat r^\top)\,W$ — se aplica una vez, a los pesos
• ablada = salida - (salida @ r_hat) → $W'a = Wa - (Wa\cdot\hat r)\,\hat r$ — igual a la ablación en el forward
• matrix_rank(P) → $\mathrm{rango}(P) = d - 1$ — se pierde exactamente una dimensión

W = torch.randn(d, d)                 # ESCRIBE en el residual stream
r_hat = torch.randn(d); r_hat /= r_hat.norm()
P = torch.eye(d) - torch.outer(r_hat, r_hat)          # I − r̂ r̂ᵀ
W_orto = P @ W                                        # W' = (I − r̂ r̂ᵀ) W

a = torch.randn(5, d)                 # entradas cualesquiera
salida = a @ W.T
ablada = salida - (salida @ r_hat)[:, None] * r_hat   # ablación en el forward
print("máx |W'a − ablación| =", f"{(a @ W_orto.T - ablada).abs().max():.1e}")
print("máx |(W'a)·r̂|       =", f"{(a @ W_orto.T @ r_hat).abs().max():.1e}")
print("rango de P:", torch.linalg.matrix_rank(P).item(), "de", d)

máx |W'a − ablación| = 6.0e-07
máx |(W'a)·r̂|       = 4.8e-07
rango de P: 7 de 8

Quién escribe en el residual stream

En un transformer clásico, todo termina en el residual stream. ¿Quién escribe ahí? Solo tres tipos de matriz: los embeddings, la salida de la atención y la proyección de bajada del MLP, el multilayer perceptron. Ortogonalizás esas, y el modelo ya no puede escribir la dirección r en ningún lado. Cuidado con los embeddings en Gemma: están atados a la matriz de salida, así que tocarlos también cambia los logits. Y Gemma dos tiene otro detalle: después de cada atención y cada MLP hay una normalización con pesos propios, antes de sumar al flujo.

embed_tokens → o_proj (×26) → down_proj (×26) → post-norms Gemma 2 → residual stream

La normalización que rompe la receta

Lo verificamos en la capa doce, con una dirección al azar. Ortogonalizamos down proj. Leemos la salida del MLP y la salida de la normalización que viene después, que es lo que realmente se suma al flujo. El MLP sobre r sombrero da uno coma cuatro por diez a la menos siete: cero. Pero después de la normalización, cero coma dos uno. No es cero. La norma RMS, root mean square norm, divide por la raíz cuadrática media y multiplica cada coordenada por su propio peso, y eso tuerce la dirección. En Gemma dos, ortogonalizar down proj no alcanza: hay que proyectar después de la norma, con un hook, o rehacer la cuenta incluyendo esos pesos.

capa = model.model.layers[12]
W = capa.mlp.down_proj.weight                         # [2304, 9216]
W.copy_((torch.eye(2304) - torch.outer(r_hat, r_hat)) @ W)   # ortogonalizar

proy = {}
capa.mlp.register_forward_hook(
    lambda m, i, o: proy.update(mlp=(o[0, -1] @ r_hat).item()))
capa.post_feedforward_layernorm.register_forward_hook(
    lambda m, i, o: proy.update(post=(o[0, -1] @ r_hat).item()))
model(**ids)
print(f"MLP·r̂            = {proy['mlp']:.1e}")       # cero: W' no escribe r̂
print(f"post-norm(MLP)·r̂ = {proy['post']:.1e}")      # va al residual stream
print(type(capa.post_feedforward_layernorm).__name__)

MLP·r̂            = 1.4e-07
post-norm(MLP)·r̂ = 2.1e-01
Gemma2RMSNorm

Arditi et al., 2024 — Refusal in Language Models Is Mediated by a Single Direction · arXiv:2406.11717

13 modelos · hasta 72B · El rechazo está mediado por una dirección «across 13 popular open-source chat models up to 72B parameters in size».

Esta es la receta del paper de la dirección de rechazo. En trece modelos de chat abiertos, de hasta setenta y dos mil millones de parámetros, encontraron una dirección por diferencia de medias entre instrucciones dañinas e inofensivas. Borrarla del flujo hace que el modelo deje de negarse; sumarla hace que se niegue a pedidos inofensivos. Y la ortogonalización de pesos es su forma de dejar ese cambio fijo en el modelo.

Qué significa para alinear

• Monitorear: una probe barata sobre el residual stream
• Corregir: steering sin reentrenar
• Riesgo: si el rechazo es una dirección, se puede borrar
• Siempre: verificar en tu arquitectura, no suponer

¿Por qué importa esto para alinear? Para monitorear: una probe es un producto escalar, se puede correr en cada token de producción. Para corregir: el steering cambia el comportamiento sin tocar un peso. Pero también es un riesgo: si el entrenamiento de seguridad termina siendo una dirección, cualquiera con los pesos la puede borrar; el paper lo dice como fragilidad. Y la lección de hoy: las recetas dependen de la arquitectura. Verificá con un hook antes de confiar.

Lo que te llevás

• Probe = proyección + umbral; dif. de medias ≈ logística
• Steering: hook que devuelve x + α·r
• Ablación: x − (x·r̂)r̂, y el cero no es neutral
• Ortogonalizar: W′ = (I − r̂r̂ᵀ)W, igual a ablar
• Gemma 2: las post-norms tuercen la dirección
• Próximo: superposición, SAEs y circuitos

Papers del episodio

• Marks & Tegmark, 2023 — The Geometry of Truth
• Rimsky et al., 2023 — Steering Llama 2 via Contrastive Activation Addition
• Arditi et al., 2024 — Refusal in Language Models Is Mediated by a Single Direction