El modelo por dentro

Alineación a bajo nivel. Episodio uno: el modelo por dentro. En esta serie vamos a ver cómo se alinea un modelo de lenguaje mirando el código y la matemática de verdad, sin magia. Hoy, la base: qué es exactamente el objeto que queremos alinear.

¿Qué es alinear?

• Que el modelo haga lo que queremos, no solo lo que pedimos literalmente
• Alineación externa: ¿el objetivo que le damos es el correcto?
• Alineación interna: ¿lo que aprendió es ese objetivo, o algo parecido?
• Verificación: ¿cómo sabemos cuál de las dos cosas pasó?

Alinear es lograr que un modelo haga lo que de verdad queremos, no solo lo que le pedimos al pie de la letra. El problema se parte en tres. Primero, la alineación externa: si el objetivo con el que lo entrenamos es el correcto. Segundo, la alineación interna: si lo que el modelo aprendió es ese objetivo, o algo que se le parece en los datos de entrenamiento pero se separa afuera. Y tercero, la verificación: cómo nos damos cuenta de cuál de las dos cosas pasó. Toda la serie gira alrededor de estas tres preguntas.

El mapa de la serie

Este es el recorrido. Los dos primeros episodios son la base: qué es el modelo y cómo aprende. Después, las técnicas para alinearlo: modelos de recompensa, RLHF con PPO, y DPO. Luego, qué sale mal cuando optimizás un objetivo imperfecto. Y la segunda mitad es abrir la caja: interceptar activaciones con hooks de PyTorch, encontrar direcciones, y separar conceptos con sparse autoencoders. Cerramos con cómo se supervisa un modelo que sabe más que vos.

Modelo por dentro → Cómo aprende → Recompensa → RLHF · PPO → DPO · constitución → Qué sale mal → Hooks → Direcciones → SAEs · circuitos → Supervisar

Parte 1 · Un modelo es una función

Arranquemos por lo concreto. Un modelo de lenguaje es una función: entra una lista de números y sale una distribución de probabilidad sobre la palabra siguiente. Nada más, y nada menos.

Texto → números

El primer paso es el tokenizador. Cargamos el de Gemma dos, de dos mil millones de parámetros, que es el modelo que vamos a usar en toda la serie. Le pasamos una frase y nos devuelve una lista de enteros: los input ids. Si los traducimos de vuelta, vemos los pedazos: un token especial de inicio, y después palabras, con ese guión bajo raro que marca el espacio. Fijate en la salida: cinco tokens, y un vocabulario de doscientos cincuenta y seis mil. Para el modelo, tu frase es esta lista de enteros.

tok = AutoTokenizer.from_pretrained("google/gemma-2-2b")
ids = tok("El paciente tiene fiebre")["input_ids"]
print(ids)
print(tok.convert_ids_to_tokens(ids))
print("vocabulario:", tok.vocab_size)

[2, 3275, 48147, 8819, 180081]
['<bos>', 'El', '▁paciente', '▁tiene', '▁fiebre']
vocabulario: 256000

Cada token se vuelve un vector

una fila de la tabla de embeddings

una fila por token del vocabulario

la frase entera: T filas de 2304 números

Cada entero es un índice en una tabla gigante. El vector inicial de cada token es simplemente una fila de la matriz de embeddings, W sub E. En el dibujo, el id cuarenta y ocho mil ciento cuarenta y siete, paciente, elige su fila y la copia: no hay ninguna cuenta, es buscar. En Gemma dos esa matriz tiene doscientas cincuenta y seis mil filas, una por token, y dos mil trescientas cuatro columnas. Así que una frase de T tokens se convierte en una matriz de T por dos mil trescientos cuatro: una fila por token, apiladas en orden. Esa matriz es lo que viaja por dentro del modelo.

Parte 2 · El residual stream

Ahora viene la idea más importante del episodio, y la que más vamos a usar después: el residual stream, la corriente por la que viaja todo.

Cada capa lee y escribe en el mismo lugar

Imaginá una cinta transportadora por la que viaja el vector de cada token. Cada capa del modelo tiene dos partes: una atención, y un MLP, multilayer perceptron, que es una red densa: multiplicar por una matriz, aplicar una no linealidad, y multiplicar por otra. Ninguna de las dos reemplaza al vector: leen de la cinta, calculan algo, y lo suman de vuelta. Por eso se llama residual stream: la corriente residual. Al final, el vector es la suma de todo lo que escribieron todas las capas. Y acá está la clave para alineación: todo lo que el modelo calcula pasa por esta cinta. Si querés mirar qué piensa, o cambiarlo, es acá donde enganchás.

Un bloque transformer, línea por línea

Así se ve en PyTorch. Es un bloque de juguete, pero la estructura es la misma que en Gemma. Primero normalizamos el vector con layer norm, para que la escala no explote. Después la atención calcula algo y fijate el signo más: x igual a x más la atención. Eso es el residual stream en una línea. Lo mismo con el MLP: x más MLP de x. Y la máscara triangular impide que un token mire a los que vienen después. Entra un tensor de uno por cinco por sesenta y cuatro, y sale exactamente la misma forma. Un bloque nunca cambia la forma del flujo, solo su contenido.

Cada línea, su término:

• h = self.ln1(x) → $h = \mathrm{LN}(x_l)$ — normalizar antes de leer
• x = x + self.attn → $x' = x_l + \mathrm{Attn}(h)$ — la atención suma al flujo
• x = x + self.mlp → $x_{l+1} = x' + \mathrm{MLP}(\mathrm{LN}(x'))$ — el MLP también suma
• mask = torch.triu → $M_{ij} = -\infty \;\text{ si } j > i$ — máscara causal

class Bloque(nn.Module):
    def __init__(self, d=64, cabezas=4):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, cabezas, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):                      # x: [batch, pos, d]
        T = x.shape[1]
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask)[0]   # suma al residual stream
        x = x + self.mlp(self.ln2(x))                    # y vuelve a escribir
        return x

x = torch.randn(1, 5, 64)
print(Bloque()(x).shape)

torch.Size([1, 5, 64])

Parte 3 · Dentro de una capa

Abramos una capa. Tiene dos piezas: la atención, que es la única parte donde los tokens se hablan entre sí, y el MLP.

La fórmula de la atención

tres proyecciones del mismo residual stream

cuánto le interesa cada token a cada otro

cada fila: pesos que suman 1

promedio pesado de los valores

La fórmula tiene cuatro pasos. Del mismo flujo sacamos tres proyecciones: consultas, claves y valores. Q, K y V. En el dibujo: cada fila de X, un token, por las columnas de W sub Q, da su consulta. Multiplicamos consultas por claves: la fila i de Q por la columna j de K transpuesta da un solo número, la celda i j: cuánto le interesa el token i al token j. Dividimos por raíz de la dimensión para que los números no crezcan, y sumamos la máscara, que pone menos infinito en el futuro. Softmax convierte cada fila en pesos positivos que suman uno: arriba de la diagonal no queda nada. Y la salida de cada token es el promedio de los valores, pesado por su fila de A, proyectado de vuelta al flujo con W sub O.

La misma fórmula en seis líneas

Y esto es la fórmula traducida a torch, sin ninguna librería de por medio. Tres multiplicaciones de matrices para Q, K y V. Consultas por claves transpuestas, divididas por raíz de d. La máscara: llenamos con menos infinito todo lo que está arriba de la diagonal. Softmax por filas. Y la matriz de atención por los valores. Mirá la matriz impresa: es triangular, porque cada token solo ve hacia atrás, y la suma de cada fila da exactamente uno. Esa matriz A es algo que después vamos a poder leer, y hasta editar, en un modelo real.

Q, K, V = X @ W_Q, X @ W_K, X @ W_V               # [T, d] cada uno
scores = Q @ K.T / math.sqrt(d)                    # [T, T]: quién mira a quién
mask = torch.triu(torch.ones(T, T), 1).bool()      # no mirar el futuro
scores = scores.masked_fill(mask, float("-inf"))
A = scores.softmax(dim=-1)                         # cada fila suma 1
out = A @ V                                        # mezcla de valores
print(A.round(decimals=2))
print(A.sum(dim=-1))

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.6300, 0.3700, 0.0000, 0.0000, 0.0000],
        [0.4600, 0.1200, 0.4200, 0.0000, 0.0000],
        [0.0200, 0.9300, 0.0500, 0.0000, 0.0000],
        [0.2300, 0.4700, 0.2100, 0.0400, 0.0500]])
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

El MLP: donde se guarda lo que sabe

la versión con compuerta que usa Gemma

cuatro veces más ancho por dentro

suma de columnas, cada una pesada por una neurona

La otra mitad de cada capa es el MLP. En Gemma es una versión con compuerta: dos proyecciones hacia arriba, una pasa por una no linealidad y multiplica a la otra elemento a elemento, y una proyección hacia abajo. Por dentro es cuatro veces más ancho que el residual stream: nueve mil doscientas dieciséis neuronas. Cada neurona es una fila de W up multiplicada por x: un solo número. Una forma útil de leerlo: la salida es una suma de columnas de W down, cada una pesada por cuánto se activó su neurona. Las neuronas en cero no aportan nada. Se cree que buena parte del conocimiento factual vive acá, y en el episodio siete vamos a ver cómo se encuentra.

Parte 4 · De vuelta a palabras

Ya sabemos qué pasa adentro. Falta el último paso: cómo un vector de dos mil trescientos números se convierte en la palabra siguiente.

Del último vector a la palabra siguiente

un logit por token del vocabulario

softmax: logits → probabilidades

y x_L es la suma de todo lo escrito

Después de veintiséis capas, tomamos el vector del último token. Lo normalizamos y lo multiplicamos por la matriz de unembedding: sale un número por cada token del vocabulario, los logits: el logit del token k es la fila k de W sub U por el vector. Softmax los convierte en probabilidades: exponencial de cada uno, dividido por la suma, así todo suma uno. Y recordá que ese vector final es la suma del embedding más todo lo que escribieron las capas. Por eso podés preguntarte cuánto empujó cada capa hacia una palabra: es una suma, se puede descomponer.

Gemma 2 2B, de verdad

Ahora con el modelo de verdad. Cargamos Gemma dos, dos B, con transformers. Le pedimos que además de los logits nos devuelva los estados ocultos: el residual stream después de cada capa. Son veintisiete tensores: el embedding más uno por cada una de las veintiséis capas, cada uno de seis tokens por dos mil trescientos cuatro. Sacamos las probabilidades del último token. Y mirá lo que pasa: para La capital de Francia es, lo más probable es una, un, la. París aparece recién en el puesto dieciséis. No está mal: está completando texto de internet, donde después de esa frase suele venir una ciudad hermosa o algo así. Un modelo base no responde preguntas, continúa documentos.

nombre = "google/gemma-2-2b"
tok = AutoTokenizer.from_pretrained(nombre)
model = AutoModelForCausalLM.from_pretrained(nombre, dtype=torch.float32)

ids = tok("La capital de Francia es", return_tensors="pt")
with torch.no_grad():
    out = model(**ids, output_hidden_states=True)

print(len(out.hidden_states), out.hidden_states[0].shape)  # embeddings + 26 capas
print(out.logits.shape)                                     # [1, pos, vocabulario]
p = out.logits[0, -1].softmax(-1)
top = p.topk(5)
for prob, i in zip(top.values, top.indices):
    print(f"{tok.decode(i)!r:12} {prob:.3f}")
paris = tok(" París", add_special_tokens=False)["input_ids"][0]
print("' París'", f"{p[paris]:.3f}", "puesto", int((p > p[paris]).sum()) + 1)

27 torch.Size([1, 6, 2304])
torch.Size([1, 6, 256000])
' una'       0.217
' un'        0.118
' la'        0.116
' uno'       0.081
' el'        0.077
' París' 0.006 puesto 16

¿Dónde están los 2.600 millones de números?

Una pregunta que conviene hacerse antes de querer cambiar un modelo: dónde están sus parámetros. Recorremos todos los tensores con named parameters y los agrupamos por nombre. La tabla de embeddings, que en Gemma se comparte con la de unembedding: la misma matriz sirve para entrar y para salir. Las proyecciones de la atención. Y las del MLP. El resultado: dos mil seiscientos millones de números, y casi dos tercios están en los MLP. La atención, que es lo más famoso, es apenas el catorce por ciento. Cuando en el próximo episodio entrenemos con LoRA, low-rank adaptation, vamos a elegir a cuáles de estas matrices les agregamos parámetros.

grupos = Counter()
for nombre, p in model.named_parameters():
    if "embed_tokens" in nombre:  grupos["embeddings (W_E = W_U)"] += p.numel()
    elif "self_attn" in nombre:   grupos["atención (Q, K, V, O)"] += p.numel()
    elif "mlp" in nombre:         grupos["MLP (gate, up, down)"] += p.numel()
    else:                         grupos["normalizaciones"] += p.numel()

total = sum(grupos.values())
for g, n in grupos.most_common():
    print(f"{g:24} {n / 1e6:8.1f} M  {n / total:6.1%}")
print(f"{'total':24} {total / 1e6:8.1f} M")

MLP (gate, up, down)       1656.2 M   63.4%
embeddings (W_E = W_U)      589.8 M   22.6%
atención (Q, K, V, O)       368.1 M   14.1%
normalizaciones               0.2 M    0.0%
total                      2614.3 M

Parte 5 · Qué tiene que ver con alinear

Volvamos a la pregunta del principio, ahora que sabemos qué hay adentro.

Por qué esto importa para alinear

• El modelo base imita su corpus: no tiene la intención de ayudar
• Alinear = cambiar los pesos W para cambiar qué distribución sale
• Todo lo que calcula pasa por el residual stream: ahí se mira y se interviene
• Las dos palancas de la serie: entrenar (ep. 2–5) e inspeccionar (ep. 7–10)

Ese resultado es el punto de partida de la serie. Un modelo base imita su corpus. No tiene ninguna intención de ayudarte: tiene una distribución. Alinear es cambiar los pesos para que la distribución que sale sea la que queremos: que conteste, que no invente, que se niegue cuando corresponde. Y como todo lo que calcula pasa por el residual stream, ese es el lugar para mirar si el cambio funcionó por las razones correctas. Esas son las dos palancas: entrenar, en los episodios dos a cinco, e inspeccionar, del siete al diez.

Ouyang et al., 2022 — Training language models to follow instructions with human feedback · arXiv:2203.02155

1,3B > 175B · Las respuestas del InstructGPT de 1.300 millones de parámetros fueron preferidas por humanos a las de GPT-3 de 175.000 millones.

El ejemplo clásico de por qué esto importa. En el paper de InstructGPT, de OpenAI, los evaluadores humanos prefirieron las respuestas de un modelo alineado de mil trescientos millones de parámetros a las del GPT tres original, que era más de cien veces más grande. Alinear no es solo seguridad: es lo que convierte un predictor de texto en algo útil.

Lo que te llevás

• Texto → ids → filas de W_E → matriz T × 2304
• Residual stream: cada capa lee y suma, nunca reemplaza
• Atención = softmax(QKᵀ/√d + M)·V: los tokens se mezclan
• MLP = suma de columnas pesadas por neuronas
• Logits = W_U·LN(x_L) → softmax → próximo token
• Próximo: cómo se cambian esos pesos (entropía cruzada, SFT, LoRA)

Papers del episodio

• Ouyang et al., 2022 — Training language models to follow instructions with human feedback