Superposición, SAEs y circuitos

Alineación a bajo nivel. Episodio nueve: superposición, sparse autoencoders y circuitos. En el episodio anterior buscamos una dirección por concepto. Hoy el problema de fondo: hay muchos más conceptos que dimensiones, y entonces las direcciones se pisan.

El problema y las herramientas

• Una neurona responde a cosas que no tienen nada que ver
• Superposición: más features que dimensiones
• SAE: un diccionario grande y disperso que las separa
• Atribución: gradientes para saber cuáles importan

El punto de partida es una observación incómoda. Si mirás una neurona sola, muchas veces responde a cosas que no tienen relación entre sí. La explicación más aceptada es la superposición: el modelo representa más conceptos que dimensiones tiene, y los guarda en direcciones casi perpendiculares. La herramienta para separarlos es un sparse autoencoder, un SAE. Y para saber cuáles de esos conceptos importan en una respuesta, usamos gradientes. Hoy vemos las tres cosas en código.

Parte 1 · Superposición

Empecemos por el modelo de juguete donde se estudió esto: cinco conceptos que tienen que pasar por un cuello de botella de dos dimensiones.

Cinco features en dos dimensiones

Esta es la versión del paper Toy Models of Superposition, achicada. Cinco features que valen un número al azar, pero la mayoría de las veces valen cero: esa probabilidad de cero es la dispersión. El modelo las comprime con una matriz W de dos por cinco y las reconstruye con la transpuesta y una ReLU, rectified linear unit: el máximo entre cero y el valor. La pérdida pesa más a las primeras features, que son más importantes. Con datos densos, solo dos columnas de W tienen norma uno: el modelo guarda las dos features más importantes y tira el resto. Con noventa y cinco por ciento de ceros, guarda las cinco, separadas por entre sesenta y ocho y setenta y cinco grados: cerca de los setenta y dos de un pentágono.

Cada línea, su término:

• torch.rand(1024, n) > p_cero → $x_i \sim U(0,1)\ \text{con prob. } 1-S,\ \text{si no } 0$ — S: cuán dispersas son
• torch.relu(x @ W.T @ W + b) → $\hat x = \mathrm{ReLU}(W^\top W x + b),\ W \in \mathbb{R}^{2\times 5}$ — comprimir a 2 y descomprimir
• (importancia * (x - x_hat) ** 2) → $L = \textstyle\sum_i I_i\,(x_i - \hat x_i)^2$ — algunas features importan más

def entrenar(p_cero, pasos=4000):
    torch.manual_seed(0)
    W = torch.nn.Parameter(torch.randn(m, n) * 0.1)
    b = torch.nn.Parameter(torch.zeros(n))
    opt = torch.optim.Adam([W, b], lr=1e-2)
    for _ in range(pasos):
        x = torch.rand(1024, n) * (torch.rand(1024, n) > p_cero)  # dispersas
        x_hat = torch.relu(x @ W.T @ W + b)                    # comprimir y volver
        loss = (importancia * (x - x_hat) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return W.detach()

for p_cero in (0.0, 0.95):
    W = entrenar(p_cero)
    print(f"P(x_i=0)={p_cero}: |W_i| =", [round(v, 2) for v in W.norm(dim=0).tolist()])
ang = torch.atan2(W[1], W[0]).rad2deg().sort().values
print("ángulos entre vecinas:", torch.diff(ang).round().tolist())

P(x_i=0)=0.0: |W_i| = [1.0, 1.0, 0.01, 0.01, 0.01]
P(x_i=0)=0.95: |W_i| = [1.14, 1.12, 1.12, 1.12, 1.1]
ángulos entre vecinas: [68.0, 73.0, 75.0, 73.0]

Las cinco columnas de W, en el plano

Así quedan las cinco columnas, tal cual salieron del entrenamiento. Un pentágono: cinco direcciones en un plano que solo tiene lugar para dos perpendiculares. ¿Cómo funciona? Como las features casi nunca están activas al mismo tiempo, la interferencia entre ellas casi nunca aparece, y la ReLU con el sesgo recorta la que aparece. Ojo: con otra semilla, el mismo código puede caer en otra configuración. Lo que vale es la idea: si los conceptos son raros, entran más que las dimensiones.

Parte 2 · El sparse autoencoder

Si el modelo comprimió muchas features en pocas dimensiones, la idea es aprender la descompresión: un diccionario mucho más ancho que el residual stream.

La forma de un SAE

La forma es simple. Entra el residual stream de una capa, dos mil trescientos cuatro números. Un encoder lo proyecta a un espacio mucho más grande. En el SAE de Gemma Scope que vamos a cargar, dieciséis mil trescientas ochenta y cuatro features. Pero entrenado para que casi todas valgan cero. Un decoder las vuelve a combinar, y la reconstrucción tiene que parecerse al original. Cada columna del decoder es una dirección candidata a concepto.

x: 2304 → encoder W_e → f: 16384 → casi todo cero → decoder W_d → x̂ ≈ x

Un SAE en veinte líneas

El código entero. Codificar: restamos el sesgo del decoder, multiplicamos por W sub e, sumamos un sesgo, y ReLU, para que las features sean positivas o cero. Decodificar: W sub d por f más el sesgo. Leído de otra forma: x sombrero es una suma de columnas del diccionario, cada una pesada por su feature. Y la pérdida tiene dos términos: el error de reconstrucción al cuadrado, más lambda por la norma uno de f, que castiga tener muchas features prendidas. Las formas: dos mil cuarenta y ocho vectores de dieciséis entran, salen sesenta y cuatro features por vector, y la reconstrucción vuelve a dieciséis.

Cada línea, su término:

• f = torch.relu((x - self.b_d) @ self.W_e.T + self.b_e) → $f = \mathrm{ReLU}\big(W_e (x - b_d) + b_e\big)$ — codificar: features ≥ 0
• return f, f @ self.W_d.T + self.b_d → $\hat x = W_d f + b_d = b_d + \textstyle\sum_j f_j\, d_j$ — suma de columnas del diccionario
• lam * f.abs().sum(-1).mean() → $L = \lVert x - \hat x \rVert^2 + \lambda \lVert f \rVert_1$ — reconstruir bien, con pocas features

class SAE(nn.Module):
    def __init__(self, d, d_sae):
        super().__init__()
        self.W_e = nn.Parameter(torch.randn(d_sae, d) * 0.1)
        self.W_d = nn.Parameter(torch.randn(d, d_sae) * 0.1)
        self.b_e = nn.Parameter(torch.zeros(d_sae))
        self.b_d = nn.Parameter(torch.zeros(d))

    def forward(self, x):
        f = torch.relu((x - self.b_d) @ self.W_e.T + self.b_e)   # features
        return f, f @ self.W_d.T + self.b_d                       # x̂

def perdida(x, f, x_hat, lam):
    return ((x - x_hat) ** 2).sum(-1).mean() + lam * f.abs().sum(-1).mean()

x = datos()
f, x_hat = SAE(d, d_sae)(x)
print("x", tuple(x.shape), "→ f", tuple(f.shape), "→ x̂", tuple(x_hat.shape))
print(f"pérdida al azar: {perdida(x, f, x_hat, 0.3).item():.2f}")

x (2048, 16) → f (2048, 64) → x̂ (2048, 16)
pérdida al azar: 1.15

El dial λ

Lo entrenamos sobre datos donde conocemos la verdad: cuarenta y ocho features escondidas en dieciséis dimensiones, unas dos activas por vez. Después de cada paso normalizamos las columnas del decoder; ya vamos a ver por qué. Probamos dos valores de lambda. Y medimos cuántas de las cuarenta y ocho direcciones verdaderas encontró, con coseno mayor a cero coma nueve. Con lambda chico: reconstrucción perfecta, veintiséis features prendidas por vector, y cero recuperadas. Con lambda cero coma tres: explica el ochenta y uno por ciento, usa tres coma tres features, y encuentra cuarenta y una de las cuarenta y ocho. Reconstruir bien no alcanza: lo que separa los conceptos es la dispersión.

def entrenar(lam, pasos=3000):
    torch.manual_seed(0)
    sae = SAE(d, d_sae)
    opt = torch.optim.Adam(sae.parameters(), lr=3e-3)
    for _ in range(pasos):
        x = datos(); f, x_hat = sae(x)
        opt.zero_grad(); perdida(x, f, x_hat, lam).backward(); opt.step()
        with torch.no_grad():                      # columnas de W_d con norma 1
            sae.W_d /= sae.W_d.norm(dim=0, keepdim=True)
    return sae

for lam in (0.003, 0.3):
    sae = entrenar(lam)
    x = datos(); f, x_hat = sae(x)
    var = 1 - ((x - x_hat) ** 2).sum() / ((x - x.mean(0)) ** 2).sum()
    rec = ((D @ sae.W_d).max(dim=1).values > 0.9).sum()   # ¿encontró las 48?
    print(f"λ={lam}: var. expl. {var:.2f}  L0 {(f > 0).sum(-1).float().mean():.1f}"
          f"  recuperadas {rec}/48")

λ=0.003: var. expl. 1.00  L0 26.2  recuperadas 0/48
λ=0.3: var. expl. 0.81  L0 3.3  recuperadas 41/48

Por qué L1, y por qué normalizar

lo que queremos: pocas features

lo que optimizamos: derivable

la trampa si no fijás la norma

por eso: columnas de norma uno

Dos detalles de la pérdida. Lo que queremos es la norma cero: cuántas features están prendidas. En el dibujo, cuatro de seis. Pero contar no tiene gradiente, así que usamos la norma uno, la suma de valores absolutos, que sí lo tiene. Esa sustitución abre una trampa. En el dibujo, el decoder: x sombrero es la suma de las columnas d sub j pesadas por cada f, y las que están en cero no aportan. si achicás f por un factor c y agrandás la columna del decoder por el mismo factor, la reconstrucción es idéntica y la norma uno baja. El optimizador lo encontraría. Por eso se fija la norma de cada columna en uno.

ReLU, TopK y JumpReLU

La norma uno tiene un efecto secundario: además de apagar features, achica las que quedan prendidas. Por eso aparecieron alternativas. Las tres sobre el mismo vector de seis pre-activaciones. ReLU deja pasar todo lo positivo, aunque sea chiquito. TopK se queda con las k más grandes, acá dos, y apaga el resto: la dispersión se fija a mano, sin lambda. JumpReLU tiene un umbral aprendido por feature: si pasás el umbral, vale tu valor completo; si no, cero. ReLU deja cuatro prendidas, incluido un cero coma cero cinco. TopK, exactamente dos. JumpReLU, tres: el cero coma cero cinco pasa porque esa feature tiene umbral cero coma cero uno, y el cero coma tres no, porque el suyo es cero coma cinco.

relu = torch.relu(pre)

k = 2                                                      # TopK: k fijo
val, idx = pre.topk(k)
topk = torch.zeros_like(pre).scatter(0, idx, torch.relu(val))

jump = torch.where(pre > theta, pre, 0.0)                    # JumpReLU: z·H(z − θ)

for nombre, f in [("ReLU", relu), ("TopK", topk), ("JumpReLU", jump)]:
    print(f"{nombre:9}", [round(v, 2) for v in f.tolist()], "L0 =", int((f > 0).sum()))

ReLU      [2.5, 0.0, 0.3, 1.2, 0.05, 0.0] L0 = 4
TopK      [2.5, 0.0, 0.0, 1.2, 0.0, 0.0] L0 = 2
JumpReLU  [2.5, 0.0, 0.0, 1.2, 0.05, 0.0] L0 = 3

Las tres activaciones

clásica: necesita λ

L0 = k, siempre

JumpReLU: H es el escalón

En fórmulas, y en el dibujo con las mismas seis pre-activaciones. La clásica: ReLU de z, donde z es el encoder aplicado a x. Las negativas quedan en cero y pasan cuatro. TopK: quedarse con las k más grandes, así la norma cero vale k por construcción: pasan dos. Y JumpReLU: z por el escalón de Heaviside de z menos theta. Pasan tres: el cero coma tres no llega a su umbral. Como el escalón no tiene derivada útil, para entrenar theta se usan estimadores especiales; para usar un SAE ya entrenado, alcanza con esta línea.

Gao et al., 2024 — Scaling and evaluating sparse autoencoders · arXiv:2406.04093

16M latentes · 40B tokens · Con TopK, «we train a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens».

TopK no es un juguete. En el paper de OpenAI sobre escalar SAEs, usaron autoencoders TopK y entrenaron uno de dieciséis millones de features sobre activaciones de GPT cuatro, durante cuarenta mil millones de tokens. La razón de tanto tamaño: un modelo aprende muchísimos conceptos, y el diccionario tiene que tener lugar para todos.

Parte 3 · Gemma Scope, sin librerías

Ahora un SAE de verdad, entrenado por Google DeepMind sobre Gemma dos. Lo cargamos con numpy, sin ninguna librería de SAEs.

Cargar el npz y codificar una frase

El SAE es un archivo npz con cinco matrices. Lo pasamos a tensores. Con un hook tomamos la salida de la capa doce, que es justo donde se entrenó, y descartamos el token de inicio. Codificar es una multiplicación, un sesgo y el umbral de JumpReLU. Fijate que acá no se resta b dec antes de codificar: cada SAE trae su convención. Decodificamos. Las formas: dos mil trescientos cuatro por dieciséis mil trescientos ochenta y cuatro. Por token se prenden entre diez y ciento cuarenta features, setenta y dos en promedio, cerca del ochenta y dos que dice el nombre de la carpeta. La reconstrucción explica el noventa y tres coma ocho por ciento de la varianza. Y las tres features más fuertes en Tower y en París.

p = {k: torch.from_numpy(v) for k, v in np.load(ruta).items()}
for k in ("W_enc", "W_dec", "threshold"):
    print(k, tuple(p[k].shape))

cache = {}
capa = model.model.layers[12]
capa.register_forward_hook(lambda m, i, o: cache.update(x=o[0]))
ids = tok("The Eiffel Tower is in the city of Paris", return_tensors="pt")
model(**ids)
x = cache["x"][1:]                         # sin <bos>: [9, 2304]

pre = x @ p["W_enc"] + p["b_enc"]          # [9, 16384]
f = pre * (pre > p["threshold"])           # JumpReLU
x_hat = f @ p["W_dec"] + p["b_dec"]
var = 1 - ((x - x_hat) ** 2).sum() / ((x - x.mean(0)) ** 2).sum()
L0 = (f > 0).sum(-1)
print(f"L0 por token: {L0.tolist()}, media {L0.float().mean():.0f}")
print(f"varianza explicada: {var:.3f}")
for t in (2, 8):                           # ' Tower' y ' Paris'
    v, i = f[t].topk(3)
    nombre = tok.decode(ids.input_ids[0, t + 1])
    print(repr(nombre), list(zip(i.tolist(), v.round().tolist())))

W_enc (2304, 16384)
W_dec (16384, 2304)
threshold (16384,)
L0 por token: [10, 140, 112, 45, 46, 44, 79, 65, 108], media 72
varianza explicada: 0.938
' Tower' [(1041, 32.0), (2620, 30.0), (2291, 23.0)]
' Paris' [(2620, 26.0), (3999, 24.0), (2291, 23.0)]

Cómo se lee una feature

• Un índice no es un concepto: hay que interpretarlo
• Ejemplos que más la activan, en muchos textos
• Steering con su columna d_j: ¿qué cambia?
• Las 2620 y 2291 aparecen en Tower y en Paris

¿Qué es la feature dos mil seiscientos veinte? Con lo que corrimos, no sabemos: un índice no es un concepto. Para interpretarla se corre el SAE sobre muchísimo texto y se miran los fragmentos donde más se activa; visores públicos como Neuronpedia hacen eso. Y para ver si es causal, se usa su columna del decoder como vector de steering, igual que en el episodio anterior. Una pista sí tenemos: la dos mil seiscientos veinte y la dos mil doscientos noventa y uno se prenden en los dos tokens. Algo compartido entre la torre y la ciudad, quizás. Sin mirar ejemplos, es solo una hipótesis.

Lieberum et al., 2024 — Gemma Scope · arXiv:2408.05147

2B y 9B, todas las capas · «An open suite of JumpReLU SAEs trained on all layers and sub-layers of Gemma 2 2B and 9B», más capas selectas de 27B.

De dónde sale ese archivo. Gemma Scope es un conjunto abierto de SAEs JumpReLU entrenados sobre todas las capas y subcapas de Gemma dos de dos mil y nueve mil millones de parámetros. El que cargamos es uno solo: capa doce, residual stream, dieciséis mil features.

Parte 4 · Atribución con gradientes

En el episodio siete parcheamos capa por capa: una pasada por cada lugar. Con miles de features, eso no escala. Los gradientes dan una aproximación casi gratis.

Parchear, linealizado

Taylor de primer orden

el efecto de pegar lo limpio

todas las capas y posiciones a la vez

La idea es una aproximación de Taylor. Si cambiás la activación a en un delta chico, la métrica cambia aproximadamente delta por el gradiente de m respecto de a. Y el delta de parchear es justo la activación limpia menos la corrupta. Así que el efecto del parche es un producto escalar. En el dibujo: cada coordenada del delta por su gradiente, y la suma estima cuánto se mueve la métrica. Lo bueno: una pasada limpia, una corrupta y una hacia atrás dan el gradiente en todas las capas y posiciones a la vez, en lugar de una pasada por lugar. La misma cuenta sirve para las features de un SAE.

Attribution patching en Gemma

Mismos prompts del episodio siete, Eiffel contra Coliseo, misma logit diff. Antes de empezar, congelamos los pesos: solo queremos gradientes de activaciones, no diez gigas de gradientes de parámetros. Corremos el corrupto guardando cada capa con retain grad. Una sola pasada hacia atrás. La atribución es delta por gradiente, normalizada como antes. Y para comparar, el parche real en tres capas. En la última posición, la estimación acierta: capa veinte, cero coma siete dos contra cero coma siete uno. Pero en Tower, capa cinco, estima cero coma dos tres y el parche real da cero coma ocho seis. La aproximación lineal puede fallar por mucho.

a_limpio, a_corr = [], []
ld_limpio = correr(limpio, guardar=a_limpio).item()
m = correr(corrupto, guardar=a_corr)               # una pasada hacia adelante
m.backward()                                       # y UNA hacia atrás
ld_corr, escala = m.item(), ld_limpio - m.item()

def atribucion(capa, pos):                         # (a_limpio − a_corr) · ∂m/∂a
    delta = a_limpio[capa][0, pos] - a_corr[capa][0, pos]
    return (delta @ a_corr[capa].grad[0, pos]).item() / escala

with torch.no_grad():
    for capa in (5, 15, 20):
        for pos, nombre in ((3, "Tower"), (8, "of")):
            real = (correr(corrupto, pegar=(capa, pos, a_limpio[capa][0, pos]))
                    - ld_corr) / escala
            est = atribucion(capa, pos)
            print(f"capa {capa:2} {nombre:5}  estimado {est:5.2f}   real {real:5.2f}")

capa  5 Tower  estimado  0.23   real  0.86
capa  5 of     estimado -0.01   real -0.00
capa 15 Tower  estimado  0.26   real  0.68
capa 15 of     estimado  0.07   real  0.17
capa 20 Tower  estimado  0.04   real  0.03
capa 20 of     estimado  0.72   real  0.71

Estimado vs. real

La comparación fila por fila. En la posición del sujeto, al principio, la estimación se queda corta por casi cuatro veces. En la capa quince, por más del doble. En la última posición, capa quince, también se queda corta. Y en la veinte, coincide. La lectura modesta: cuando el delta es grande y el modelo está lejos de comportarse linealmente, Taylor falla. Por eso attribution patching se usa para filtrar candidatos rápido, y después se verifican los mejores con parches de verdad.

De features a circuitos

• Nodos: features de SAEs en cada capa
• Aristas: atribución entre features
• Podar lo que no mueve la métrica
• Verificar el circuito con intervenciones reales

Juntando todo lo de hoy se llega a los circuitos. En lugar de capas o neuronas, los nodos son features de SAEs, que se interpretan mejor. Las aristas se calculan con la misma cuenta de delta por gradiente, entre una feature y otra. Se poda lo que casi no mueve la métrica, y lo que queda se verifica con parches y steering de verdad. El resultado se suele llamar grafo de atribución: un mapa de qué conceptos internos llevaron a una respuesta.

Lo que te llevás

• Superposición: features raras entran más que dimensiones
• SAE: f = ReLU(W_e(x − b_d) + b_e), x̂ = W_d f + b_d
• L1 + columnas de norma 1; o TopK; o JumpReLU
• Gemma Scope: npz, hook en la capa 12, L0 ≈ 72
• Δm ≈ (a_limpio − a_corr)·∂m/∂a: rápido, a verificar
• Próximo: supervisar lo que no podés verificar

Papers del episodio

• Gao et al., 2024 — Scaling and evaluating sparse autoencoders
• Lieberum et al., 2024 — Gemma Scope