Cómo aprende

Alineación a bajo nivel. Episodio dos: cómo aprende. En el episodio anterior vimos que un modelo es una función con dos mil seiscientos millones de números adentro. Hoy vemos cómo se cambian esos números. Y vas a ver que alinear, en su forma más básica, es elegir qué pérdida bajar y con qué datos.

Entrenar es bajar una pérdida

• Una pérdida: un número que dice qué tan mal predijo
• Gradiente: hacia dónde mover cada peso para que baje
• Optimizador: cuánto moverlo en cada paso
• Datos + máscara: qué tokens cuentan para la pérdida

Todo entrenamiento de una red tiene cuatro piezas. Una pérdida: un único número que mide qué tan mal predijo el modelo. El gradiente: para cada uno de los parámetros, hacia dónde moverlo para que ese número baje. Un optimizador, que decide cuánto moverlo en cada paso. Y los datos, con una máscara que dice qué tokens cuentan. Preentrenar, ajustar con instrucciones y, en los próximos episodios, RLHF y DPO, cambian sobre todo la primera y la última pieza. Las del medio son siempre las mismas.

Parte 1 · La pérdida

Empecemos por el número que queremos bajar: la entropía cruzada del próximo token.

Entropía cruzada del próximo token

promedio de la sorpresa en cada posición

log softmax de los logits

$y$: vector one-hot del token correcto

La pérdida es una sola fórmula. En cada posición, el modelo le asigna una probabilidad al token que de verdad vino después. Tomamos el logaritmo, le cambiamos el signo y promediamos: en el dibujo, una barra de sorpresa por posición, y la línea de abajo es el promedio. Si el modelo estaba seguro y acertó, el término es casi cero. Ese logaritmo de la probabilidad se calcula directo desde los logits: el logit del token correcto menos el logaritmo de la suma de los exponenciales. Es el softmax del episodio uno, en escala logarítmica. Y lo lindo: la derivada respecto de los logits es la probabilidad predicha menos el one-hot del correcto. Sube el logit del token que vino y baja los demás, en proporción a cuánto se los creía.

Equivocarse con confianza cuesta caro

Así se ve la sorpresa en función de la probabilidad del token correcto. Con probabilidad uno, la pérdida es cero. Con un medio, cero coma seis nueve. Y cuando la probabilidad se acerca a cero, la curva se dispara. Equivocarse con confianza cuesta muchísimo más que dudar. Esa asimetría empuja al modelo a repartir probabilidad entre todo lo que podría venir, y por eso un modelo base es tan bueno imitando la variedad de internet.

Si el modelo da probabilidad 1 al token correcto, la pérdida es 0.

La fórmula, contra PyTorch

Verifiquémoslo con números de juguete: cuatro posiciones, un vocabulario de seis. El log softmax, escrito a mano con log sum exp. Indexamos el logaritmo del token correcto en cada fila, promediamos y cambiamos el signo. Después le pedimos a PyTorch que haga backprop sobre su propia entropía cruzada. Y armamos lo que dice la ecuación: probabilidades menos one-hot, dividido por cuatro. Las dos pérdidas dan dos coma cero cuatro dos uno. En la fila cero del gradiente, el único negativo es el del token correcto, el dos: subir ese logit baja la pérdida. Y la comparación da True. Backprop no es magia: es esta resta, propagada hacia atrás capa por capa.

Cada línea, su término:

• logsumexp → $\log p = \ell - \log\textstyle\sum_j e^{\ell_j}$ — log softmax a mano
• a_mano = → $\mathcal{L} = -\frac{1}{T}\sum_t \log p(x_t)$ — tomar el correcto y promediar
• torch_ce.backward() → $\nabla_\ell \mathcal{L}$ — backprop calcula la derivada
• esperado = → $\frac{1}{T}(p - y)$ — la derivada que predijimos

logp = logits - logits.logsumexp(-1, keepdim=True)   # log softmax
a_mano = -logp[torch.arange(T), objetivo].mean()     # −log p del correcto
torch_ce = F.cross_entropy(logits, objetivo)
print(f"a mano:   {a_mano.item():.4f}")
print(f"F.cross_entropy: {torch_ce.item():.4f}")

torch_ce.backward()                                  # backprop
p = logits.softmax(-1).detach()
esperado = (p - F.one_hot(objetivo, V)) / T          # (p − y) / T
print("grad fila 0:", logits.grad[0].numpy().round(3))
print("iguales:", torch.allclose(logits.grad, esperado))

a mano:   2.0421
F.cross_entropy: 2.0421
grad fila 0: [ 0.013  0.012 -0.22   0.025  0.091  0.078]
iguales: True

Parte 2 · Gradientes y AdamW

Ya sabemos hacia dónde mover cada peso. Falta decidir cuánto, y repetirlo miles de veces.

Del gradiente al paso

descenso por gradiente: $\eta$ es la tasa de aprendizaje

promedios móviles del gradiente y de su cuadrado

AdamW: paso normalizado + decaimiento de pesos

La regla más simple es el descenso por gradiente. Restale al peso el gradiente multiplicado por un número chico, eta, la tasa de aprendizaje. En la práctica casi todos los modelos de lenguaje se entrenan con AdamW, que guarda dos promedios móviles por parámetro: el del gradiente, m, y el de su cuadrado, v. El paso es m dividido por la raíz de v: cada parámetro se mueve en una escala parecida, tenga gradientes grandes o chicos. Y la W es el término lambda por theta, que achica los pesos un poquito en cada paso. Ojo con la memoria: AdamW guarda dos números extra por cada parámetro.

El loop de entrenamiento

Todo entrenamiento es este ciclo. Tomás un lote de secuencias. Hacés el forward y obtenés logits. Calculás la pérdida. Backward llena el atributo grad de cada parámetro. El optimizador da un paso. Y borrás los gradientes, porque PyTorch los acumula si no lo hacés. Preentrenar un modelo grande es este ciclo repetido cientos de miles de veces sobre billones de tokens, es decir, millones de millones, repartido en miles de GPUs. La lógica no cambia.

El loop, con un modelo de juguete

Ahora el ciclo en código. Un modelo de la familia Llama, de dos capas y sesenta y cuatro dimensiones, con los pesos al azar. El optimizador AdamW. Los datos son sintéticos: secuencias que cuentan de a uno, de a dos o de a tres, módulo treinta y dos. Le pasamos los mismos ids como labels: Hugging Face corre los labels una posición y calcula la entropía cruzada por nosotros. Backward. Paso. Arranca en tres coma cinco, que es el logaritmo de treinta y dos: adivinar al azar. En doscientos pasos baja a cero coma cero ocho. No llega a cero porque el segundo número de cada secuencia no se puede adivinar: todavía no sabés de a cuánto cuenta.

model = LlamaForCausalLM(cfg)                        # pesos al azar
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)

for paso in range(301):
    ids = lote()                                     # [32, 16] tokens
    loss = model(input_ids=ids, labels=ids).loss     # entropía cruzada
    loss.backward()                                  # ∇L en cada parámetro
    opt.step()                                       # θ ← θ − η·(paso de Adam)
    opt.zero_grad()
    if paso % 100 == 0:
        print(f"paso {paso:3d}  loss {loss.item():.3f}")
print(f"azar puro: {torch.log(torch.tensor(32.)):.3f}")

paso   0  loss 3.506
paso 100  loss 0.132
paso 200  loss 0.084
paso 300  loss 0.090
azar puro: 3.466

¿Qué aprendió?

Veamos qué aprendió ese mismo modelo después del loop. Generamos en modo greedy: en cada paso, el token más probable. Le damos tres, seis. Y veintisiete, veintinueve, para ver si entiende la vuelta en treinta y dos. Por último, miramos la distribución después de un solo número. Tres, seis, nueve, doce: cuenta de a tres. Veintisiete, veintinueve, treinta y uno, uno, tres: da la vuelta. Y después de un cinco solo, reparte un tercio entre seis, siete y ocho. Aprendió exactamente la incertidumbre que tenían los datos. Eso es lo que hace la entropía cruzada: no premia una respuesta, premia la distribución del corpus.

def continuar(inicio, n=6):
    x = torch.tensor([inicio])
    out = model.generate(x, max_new_tokens=n, do_sample=False,
                         pad_token_id=0)   # greedy: el más probable
    return out[0].tolist()

print(continuar([3, 6]))                             # ¿de a tres?
print(continuar([27, 29]))                           # ¿y da la vuelta en 32?
with torch.no_grad():
    p = model(torch.tensor([[5]])).logits[0, -1].softmax(-1)
top = p.topk(3)
print("después de [5]:", [(int(i), round(v.item(), 2)) for v, i in zip(*top)])

[3, 6, 9, 12, 15, 18, 21, 24]
[27, 29, 31, 1, 3, 5, 7, 9]
después de [5]: [(6, 0.35), (7, 0.34), (8, 0.29)]

Parte 3 · Preentrenar vs. ajustar

Misma pérdida, mismo loop. Lo que cambia entre un modelo base y uno que contesta son los datos y qué tokens cuentan.

Dos usos del mismo loop

La diferencia está en los datos. En el preentrenamiento entra texto crudo, billones de tokens, y cada token cuenta para la pérdida. El resultado es lo que vimos en el episodio uno: un modelo que completa documentos. En el SFT, supervised fine-tuning, el ajuste fino supervisado, entran pares de instrucción y respuesta escritos por personas. Y acá viene el detalle técnico clave: solo la respuesta cuenta para la pérdida. No queremos que el modelo aprenda a escribir preguntas de usuario; queremos que aprenda a contestarlas.

SFT: la máscara de la pérdida

Así se implementa. Primero, una plantilla de chat con los marcadores de turno de Gemma. El modelo instruido trae su propia plantilla; el base no, así que acá la escribimos nosotros. Tokenizamos el prompt solo, para saber cuántos tokens ocupa. Y a esas posiciones les ponemos menos cien en los labels. Menos cien es el valor que la entropía cruzada de PyTorch ignora. Para estar seguros, recalculamos a mano con ignore index. De diecinueve tokens, entrenan tres: la respuesta y el cierre de turno. Las dos pérdidas coinciden. Que valgan doce coma cinco es lógico: el modelo de juguete está al azar, y el logaritmo de doscientos cincuenta y seis mil es doce coma cuatro.

plantilla = "<start_of_turn>user\n{p}<end_of_turn>\n<start_of_turn>model\n"
prompt = plantilla.format(p="¿Cuánto es dos más dos?")
respuesta = "Cuatro.<end_of_turn>"

ids_p = tok(prompt)["input_ids"]
ids = tok(prompt + respuesta, return_tensors="pt")["input_ids"]
labels = ids.clone()
labels[0, :len(ids_p)] = -100                        # el prompt no cuenta
print("tokens:", ids.shape[1], "| entrenan:", int((labels != -100).sum()))

loss_hf = model(input_ids=ids, labels=labels).loss
lg = model(input_ids=ids).logits[0, :-1]             # posición t predice t+1
loss_mano = F.cross_entropy(lg, labels[0, 1:], ignore_index=-100)
print(f"HF {loss_hf.item():.4f} | a mano {loss_mano.item():.4f}")

tokens: 19 | entrenan: 3
HF 12.5783 | a mano 12.5783

La pérdida de SFT

$x$: el prompt; $y$: la respuesta

sin gradiente desde el prompt

en esperanza: imitar las demostraciones

En matemática. La pérdida de SFT es la misma entropía cruzada, pero la suma corre solo sobre los tokens de la respuesta, condicionados al prompt. En el dibujo, las barras del prompt quedan tachadas: solo suman las dos últimas. Un label en menos cien borra ese término. El prompt sigue entrando al forward, el modelo lo lee, pero no genera gradiente. Y una lectura útil: minimizar esto es bajar la divergencia KL, de Kullback-Leibler, entre la distribución de quienes escribieron las demostraciones y la del modelo. SFT es imitación: el objetivo apunta a parecerse a los ejemplos, no a superarlos. Ese techo es lo que el modelo de recompensa del episodio tres va a intentar romper.

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

~13k prompts · El conjunto de SFT de InstructGPT tiene unos 13.000 prompts de entrenamiento (de la API y escritos por los etiquetadores).

Para darte una escala. En InstructGPT, el conjunto de ajuste supervisado tenía unos trece mil prompts de entrenamiento. Comparalo con los cientos de miles de millones de tokens del preentrenamiento. El SFT no le enseña conocimiento nuevo al modelo: le enseña un formato y un rol, sobre lo que ya sabía.

Parte 4 · LoRA

Ajustar los dos mil seiscientos millones de parámetros es caro. LoRA, low-rank adaptation, congela el modelo y entrena un parche chico.

Un parche de rango bajo

$W$ congelada; solo se entrenan $A$ y $B$

el producto $BA$ tiene rango $\le r$

una matriz del tamaño de Gemma, con $r = 8$

La idea es esta. En vez de modificar una matriz de pesos W, la dejamos congelada y le sumamos un producto B por A, escalado por alfa sobre r. B es alta y flaca, A es baja y ancha, con r chiquito en el medio: ocho, dieciséis. El producto tiene el tamaño de W, pero su rango es como mucho r. En el dibujo: cada celda de B por A es la fila de B por la columna de A, apenas r productos. Para una matriz de dos mil trescientos cuatro por dos mil trescientos cuatro, son cinco millones de parámetros contra treinta y siete mil. La apuesta es que el cambio que necesita el ajuste vive en pocas direcciones.

LoRA con peft, verificado a mano

Con la librería peft. Elegimos rango cuatro, alfa ocho, y las proyecciones de query y value de la atención, las mismas matrices que contamos en el episodio uno. get peft model congela todo lo demás. Un detalle de diseño: B arranca en cero, así que al principio el modelo es idéntico al original y el ajuste parte de ahí. Para verificar la fórmula, le ponemos valores al azar a B y comparamos la capa de peft con la cuenta a mano. Se entrena el dos coma tres por ciento de los parámetros. Las formas son las esperadas, B arranca en cero, y la salida coincide.

Cada línea, su término:

• LoraConfig(r=4 → $r = 4,\;\; \alpha = 8$ — en qué matrices y con qué rango
• get_peft_model → $\nabla_W = 0$ — W se congela
• B arranca en cero → $B_0 = 0 \;\Rightarrow\; W' = W$ — al inicio no cambia nada
• a_mano = → $x\,(W + \tfrac{\alpha}{r} BA)^\top$ — la fórmula, contra la capa real

conf = LoraConfig(r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"])
model = get_peft_model(base, conf)                   # congela W, agrega A y B
model.print_trainable_parameters()

capa = model.base_model.model.model.layers[0].self_attn.q_proj
W = capa.base_layer.weight                           # congelada
A = capa.lora_A["default"].weight
B = capa.lora_B["default"].weight
print("W", tuple(W.shape), "A", tuple(A.shape), "B", tuple(B.shape))
print("B arranca en cero:", bool((B == 0).all()))

with torch.no_grad():
    B.normal_()                                      # simulamos que ya entrenó
    x = torch.randn(3, 64)
    a_mano = x @ (W + (8 / 4) * B @ A).T             # W' = W + (α/r)·BA
    print("igual a peft:", torch.allclose(capa(x), a_mano, atol=1e-5))

trainable params: 2,048 || all params: 88,384 || trainable%: 2.3172
W (64, 64) A (4, 64) B (64, 4)
B arranca en cero: True
igual a peft: True

Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models · arXiv:2106.09685

10.000× menos · Comparado con GPT-3 175B ajustado con Adam, LoRA puede reducir la cantidad de parámetros entrenables 10.000 veces y la memoria de GPU 3 veces.

El paper original lo midió sobre GPT tres. Frente a ajustar los ciento setenta y cinco mil millones de parámetros con Adam, LoRA reduce diez mil veces los parámetros entrenables y tres veces la memoria de GPU. Por eso casi todo el ajuste fino que vas a ver fuera de los grandes laboratorios es con LoRA.

Parte 5 · Qué tiene que ver con alinear

Con estas herramientas ya podés hacer la primera etapa de cualquier alineación.

Lo que SFT puede y no puede

• Puede: dar formato, rol y estilo con pocos miles de ejemplos
• No apunta a superar a sus demostraciones
• Solo ve respuestas buenas: nunca aprende qué es peor
• Por eso sigue: preferencias → recompensa → RL

Entonces, qué te da el SFT. Con pocos miles de ejemplos convierte un modelo que completa documentos en uno que contesta, con el formato y el tono que quieras. Pero es imitación: su objetivo es parecerse a quienes escribieron las demostraciones, no superarlos. Y tiene un punto ciego: solo ve respuestas buenas. Nunca le mostrás una mala para que aprenda a evitarla. Para eso hace falta comparar respuestas entre sí. Y es más fácil para una persona decir cuál de dos respuestas es mejor que escribir la ideal. De ahí sale todo lo que viene: un modelo de recompensa y reinforcement learning, aprendizaje por refuerzo.

Lo que te llevás

• Entropía cruzada: $-\log p$ del token correcto, promediado
• Su gradiente sobre los logits es $(p - y)/T$
• Loop: forward → loss → backward → step → zero_grad
• SFT = misma pérdida, labels = −100 en el prompt
• LoRA: $W + (\alpha/r)BA$, con $B$ en cero al inicio
• Próximo: modelos de recompensa (Bradley-Terry)

Papers del episodio

• Ouyang et al., 2022 — Training language models to follow instructions with human feedback
• Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models