Paso 1: SFT

InstructGPT. Episodio tres: paso uno, el SFT, supervised fine-tuning. La mecánica ya la viste en Alineación, episodio dos: entropía cruzada solo sobre la respuesta. Hoy nos quedamos con lo que decidió el paper: cuántas epochs, qué learning rate, qué dropout, y una decisión rara: entrenar de más a propósito y elegir el checkpoint con otra regla.

Las decisiones del paso 1

• 16 epochs sobre unas 13 mil demostraciones
• Learning rate con coseno hasta el 10 %, sin warmup
• Residual dropout de 0,2
• El checkpoint se elige por RM score, no por val loss

Cuatro decisiones, todas en la sección tres punto cinco y en el apéndice C. Dieciséis epochs: dieciséis pasadas completas por los datos de SFT, unos trece mil prompts. Un learning rate que baja con forma de coseno hasta el diez por ciento del inicial, sin warmup, sin rampa de subida al principio. Residual dropout de cero coma dos. Y el checkpoint final, la copia de los pesos que se queda, no se elige por la pérdida de validación, sino por el puntaje que le da el reward model, el RM, que vemos en el episodio cuatro.

Parte 1 · Qué se entrena, en chico

Para ver cada decisión con números, armamos un SFT de juguete que corre en segundos.

La tarea de juguete

El prompt: seis símbolos al azar, ids del siete al dieciséis. La demostración arranca con un comentario libre de tres símbolos, sacados del prompt. Es nuestra versión de que cada etiquetador escribe a su manera: hay muchas respuestas válidas. Después viene lo que importa: los seis símbolos ordenados, y un token de fin. Y la máscara: las posiciones del prompt llevan menos cien, que PyTorch ignora en la pérdida. Diecisiete posiciones; las del prompt, en menos cien. Entrenan diez: el comentario, los seis ordenados y el fin.

def ejemplos(n):                     # tarea de juguete: ordenar 6 símbolos
    x = torch.randint(S0, S0 + 10, (n, L))                # el prompt
    coment = x.gather(1, torch.randint(0, L, (n, F)))     # estilo libre:
    resp = torch.cat([coment, x.sort(1).values,           # muchas válidas
                      torch.full((n, 1), EOS)], 1)
    ids = torch.cat([x, torch.full((n, 1), SEP), resp], 1)
    labels = ids.clone()
    labels[:, :L + 1] = -100          # el prompt no entra en la pérdida
    return ids, labels

ids, labels = ejemplos(1)
nombre = lambda i: {1: "SEP", 2: "EOS"}.get(i, str(i))
print("prompt:", *map(nombre, ids[0, :L + 1].tolist()))
print("resp:  ", *map(nombre, ids[0, L + 1:].tolist()))
print("labels del prompt:", *labels[0, :3].tolist(), "...")
print("entrenan", int((labels != -100).sum()), "de", ids.shape[1])

prompt: 11 16 10 7 10 16 SEP
resp:   16 16 16 7 10 10 11 16 16 EOS
labels del prompt: -100 -100 -100 ...
entrenan 10 de 17

La pérdida, con máscara

La pérdida, a mano y contra Hugging Face. Los logits de cada posición: la posición t predice el token t más uno. La entropía cruzada por posición, sin promediar: menos log de la probabilidad del token correcto. Promediamos solo sobre R, el conjunto de posiciones de la respuesta; las barras verticales alrededor de R quieren decir cuántas son. Diez de dieciséis cuentan, y a mano da lo mismo que Hugging Face. Un aviso: el paper no dice si enmascara el prompt. Pero mirá por qué importa: en la tabla once, la mediana de un prompt de etiquetador es trescientos veinticuatro tokens y la de una demostración, dieciocho. Sin máscara, el noventa y cinco por ciento de la pérdida sería copiar prompts.

Cada línea, su término:

• logits = model(input_ids=ids).logits[0, :-1] → $\ell_t \;\to\; \pi(\cdot \mid x, y_{<t})$ — la posición t predice la t+1
• reduction="none" → $-\log \pi(y_t \mid x, y_{<t})$ — un número por posición
• mano = nll[cuentan].sum() → $\mathcal{L}_{\text{SFT}} = \frac{1}{|R|} \sum_{t \in R} -\log \pi(y_t \mid \cdot)$ — $R$: posiciones de la respuesta
• prompt_med, demo_med = 324, 18 → $\frac{324}{324 + 18} \approx 95\,\%$ — medianas de la tabla 11

logits = model(input_ids=ids).logits[0, :-1]       # posición t predice t+1
objetivo = labels[0, 1:]
nll = F.cross_entropy(logits, objetivo, reduction="none",
                      ignore_index=-100)            # 0 donde no cuenta
cuentan = objetivo != -100
mano = nll[cuentan].sum() / cuentan.sum()           # promedio sobre R
hf = model(input_ids=ids, labels=labels).loss
print(f"tokens que cuentan: {int(cuentan.sum())} de {len(objetivo)}")
print(f"a mano {mano.item():.4f} | HF {hf.item():.4f}")

prompt_med, demo_med = 324, 18                      # medianas, tabla 11
print(f"sin máscara, el prompt sería {prompt_med/(prompt_med+demo_med):.0%}"
      " de los tokens")

tokens que cuentan: 10 de 16
a mano 2.8560 | HF 2.8560
sin máscara, el prompt sería 95% de los tokens

Qué cuenta y qué no

el ejemplo del snippet, con el modelo sin entrenar

adivinar entre 10 símbolos al azar

El dibujo muestra cada posición del ejemplo. Las seis primeras, el prompt, quedan grises: no cuentan. Las diez de la respuesta tienen su menos log pi, y el promedio da dos coma ochenta y seis. Para tener una referencia: si el modelo eligiera al azar entre los diez símbolos, la pérdida sería logaritmo de diez, dos coma treinta. Un modelo recién inicializado arranca un poco peor que el azar, porque reparte probabilidad entre los diecisiete ids del vocabulario. Ahora, a bajarla.

Parte 2 · Las perillas

Learning rate, su forma en el tiempo, y dropout. Cada una con su número del paper.

Apéndice C.1: los números del SFT

El apéndice C da los números. Para el modelo de mil trescientos millones: learning rate nueve coma sesenta y cinco por diez a la menos seis, y batch de treinta y dos. Lo mismo para el de seis mil millones. Y para el de ciento setenta y cinco mil millones, un learning rate más chico, cinco coma cero tres por diez a la menos seis, con batch de ocho. Los eligieron con una búsqueda geométrica: valores espaciados por un factor constante. Todos los modelos usan Adam, adaptive moment estimation, con beta uno cero coma nueve y beta dos cero coma noventa y cinco.

El learning rate con forma de coseno

El schedule que lo hace bajar, en PyTorch, con cuatrocientos pasos de ejemplo. Lo ves entero en Optimizadores, episodio ocho. Eta es el learning rate; eta máximo, el del paper. Eta mínimo, el piso: el diez por ciento. CosineAnnealingLR implementa esta fórmula. t es el paso y T el total. El coseno de pi t sobre T va de uno a menos uno, así que el paréntesis va de dos a cero y eta baja suave del máximo al mínimo. Imprimimos cinco momentos. Cien por ciento al empezar, ochenta y siete al cuarto, cincuenta y cinco a la mitad, veintitrés a los tres cuartos y diez al final. A la mitad, el coseno vale cero: queda el piso más la mitad del rango.

Cada línea, su término:

• lr, pasos = 9.65e-6, 400 → $\eta_{\max} = 9{,}65 \cdot 10^{-6}$ — $\eta$: learning rate
• eta_min=0.1 * lr → $\eta_{\min} = 0{,}1\,\eta_{\max}$ — el piso: 10 %
• CosineAnnealingLR( → $\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max} - \eta_{\min})\big(1 + \cos\tfrac{\pi t}{T}\big)$ — $t$: paso · $T$: pasos totales
• for t in [0, 100, 200, 300, 400] → $t = \tfrac{T}{2}:\; \cos\tfrac{\pi}{2} = 0 \Rightarrow 55\,\%$ — a mitad de camino

w = torch.nn.Parameter(torch.zeros(1))
lr, pasos = 9.65e-6, 400                       # el lr del paper (1,3B y 6B)
opt = torch.optim.Adam([w], lr=lr, betas=(0.9, 0.95))
sched = torch.optim.lr_scheduler.CosineAnnealingLR(
    opt, T_max=pasos, eta_min=0.1 * lr)        # baja hasta el 10 %, sin warmup

lrs = []
for paso in range(pasos + 1):
    lrs.append(sched.get_last_lr()[0])
    opt.step(); sched.step()
for t in [0, 100, 200, 300, 400]:
    print(f"paso {t:3d}  lr = {lrs[t]:.3e}  ({lrs[t] / lr:.0%})")

paso   0  lr = 9.650e-06  (100%)
paso 100  lr = 8.378e-06  (87%)
paso 200  lr = 5.308e-06  (55%)
paso 300  lr = 2.237e-06  (23%)
paso 400  lr = 9.650e-07  (10%)

El learning rate en el tiempo

La misma curva, con los números del snippet. Arranca en el máximo sin ninguna rampa, baja despacio, se acelera en el medio, y se aplana llegando al diez por ciento. Pasos grandes cuando el modelo está lejos, pasos chicos al final, para asentarse.

Residual dropout

$p = 0{,}2$: cada número se apaga con probabilidad 0,2

se aplica a lo que la subcapa suma, no a $x_l$

Dropout lo viste en Deep learning a secas, episodio ocho. Durante el entrenamiento, cada número de un vector h se multiplica por m sub i, que vale uno con probabilidad uno menos p y cero con probabilidad p: eso es una Bernoulli. El círculo con punto es multiplicar elemento a elemento. Lo que sobrevive se divide por uno menos p, para que el promedio no cambie. En el dibujo, h dos se apaga y el resto ya está escalado. Residual quiere decir dónde: en la salida de cada subcapa, la atención o el MLP, multilayer perceptron, justo antes de sumarla al residual stream. LN es layer norm. El residual stream pasa entero; lo nuevo llega con agujeros.

Dónde está el residual dropout en GPT-2

Usamos GPT dos de Hugging Face, de la misma familia que GPT tres, donde la perilla tiene ese nombre. resid pdrop igual a cero coma dos, como el paper. Apagamos los otros dos dropouts para aislarlo. Lo encontramos en dos lugares de cada bloque: a la salida de la atención y a la salida del MLP. En modo train, dos forwards con la misma entrada dan distinto, porque cada uno sortea su máscara. En modo eval, el dropout se apaga. Dos Dropout de cero coma dos; en train, distintos; en eval, iguales.

cfg = GPT2Config(vocab_size=17, n_positions=32, n_embd=64, n_layer=2,
                 n_head=4, resid_pdrop=0.2,          # residual dropout
                 embd_pdrop=0.0, attn_pdrop=0.0)
model = GPT2LMHeadModel(cfg)
print(model.transformer.h[0].attn.resid_dropout)     # antes de sumar
print(model.transformer.h[0].mlp.dropout)            # al residual stream

ids = torch.tensor([[12, 8, 15, 8, 10, 13, 1]])
model.train()
a, b = model(input_ids=ids).logits, model(input_ids=ids).logits
print("train: dos forwards iguales?", torch.allclose(a, b))
model.eval()
a, b = model(input_ids=ids).logits, model(input_ids=ids).logits
print("eval:  dos forwards iguales?", torch.allclose(a, b))

Dropout(p=0.2, inplace=False)
Dropout(p=0.2, inplace=False)
train: dos forwards iguales? False
eval:  dos forwards iguales? True

Lo demás que fija el apéndice C

• Contexto de 2.000 tokens
• Prompts de más de 1.000 tokens, afuera
• Respuestas de hasta 1.000 tokens
• Pesos en fp16, con copia maestra en fp32

Tres detalles más del apéndice C, que valen para todos los modelos del paper. El contexto es de dos mil tokens. Los prompts de más de mil tokens se descartan. Y las respuestas se cortan en mil. Así, prompt más respuesta siempre entran en el contexto. Y los pesos y activaciones van en fp dieciséis, punto flotante de dieciséis bits, con una copia maestra de los pesos en treinta y dos bits: el paso del optimizador se aplica sobre la copia precisa.

Parte 3 · Overfitting a propósito

Ahora la decisión más interesante del paso uno: seguir entrenando cuando la validación dice que pares.

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

1 epoch → 16 · «our SFT models overfit on validation loss after 1 epoch; however, we find that training for more epochs helps both the RM score and human preference ratings, despite this overfitting.»

La frase textual. Los modelos SFT hacen overfitting en la pérdida de validación después de una sola epoch. Sin embargo, entrenar más epochs mejora el puntaje del reward model y la preferencia humana, a pesar del overfitting. Y en el apéndice agregan por qué eligen con el RM: predice mejor la preferencia humana que la pérdida de validación.

SFT de juguete, 16 epochs

El mismo experimento, en chico. Doscientas demostraciones, un GPT dos de cuatro capas, residual dropout cero coma dos y coseno. Dieciséis epochs. Cada epoch es una pasada por las doscientas, en batches de ocho, en orden aleatorio. Al final de cada epoch medimos la pérdida en mil ejemplos de validación. Y una métrica que hace de reward model: qué fracción de la validación ordena bien, mirando solo la parte ordenada. La pérdida de train baja siempre. La de validación baja hasta cerca del epoch diez y después sube. Y la fracción que ordena bien sigue subiendo: sesenta y dos, sesenta y ocho, setenta y cuatro por ciento.

print("epoch   train    val  ordena")
for epoch in range(1, EPOCHS + 1):
    model.train()
    for i in torch.randperm(200).split(B):        # una pasada = un epoch
        loss = model(input_ids=tr[i], labels=tr_lab[i]).loss
        loss.backward(); opt.step(); opt.zero_grad(); sched.step()
    model.eval()
    with torch.no_grad():
        l_tr = model(input_ids=tr, labels=tr_lab).loss.item()
        l_va = model(input_ids=va, labels=va_lab).loss.item()
    acc = metrica(model)                          # proxy del RM score
    hist.append((epoch, l_tr, l_va, acc))
    if epoch in (1, 4, 7, 10, 13, 16):
        print(f"{epoch:5d} {l_tr:7.2f} {l_va:6.2f} {acc:7.0%}")

epoch   train    val  ordena
    1    1.51   1.56      0%
    4    0.81   0.97     25%
    7    0.60   0.84     41%
   10    0.48   0.79     62%
   13    0.39   0.81     68%
   16    0.34   0.81     74%

Las tres curvas del juguete

Las tres curvas, con los números del snippet. La pérdida de train, punteada, baja todo el tiempo: el modelo se aprende sus doscientas demostraciones. La de validación, en ámbar, toca fondo en el epoch once y después sube: overfitting. En el paper pasa después del primero; en el juguete, más tarde. Y la métrica, en verde, sigue subiendo hasta el epoch dieciséis. Una explicación posible, que el paper no da: la pérdida castiga no adivinar el comentario libre de cada demostración, y la métrica solo mira si hizo la tarea.

Dos reglas para elegir el checkpoint

$k$: epoch · argmin: el $k$ que da el mínimo

acá el RM score lo reemplaza «ordena bien»

Las dos reglas, en símbolos. k es el número de epoch. Argmin quiere decir: el k donde la función da su mínimo. La regla de siempre elige el epoch con menor pérdida de validación. En el dibujo, una barra por epoch, del nueve al dieciséis: la más baja está en el once. La regla del paper es argmax, el k donde la función da su máximo, del puntaje que el reward model les pone a las respuestas de validación. En el juguete, ese puntaje lo hace la fracción que ordena bien. En el dibujo, la barra más alta está en el dieciséis.

Cuál elige cada regla

Y en código son dos líneas sobre el historial. La regla de siempre: el mínimo de la pérdida de validación. La del paper: el máximo de la métrica. Por pérdida, el epoch once: ordena bien el sesenta y ocho por ciento. Por métrica, el dieciséis: setenta y cuatro, con una pérdida de validación peor. Si hubiéramos parado donde dice la pérdida, nos quedábamos con un modelo que hace peor la tarea.

# hist: (epoch, loss de train, loss de validación, métrica) por epoch
por_loss = min(hist, key=lambda h: h[2])       # la regla de siempre
por_metrica = max(hist, key=lambda h: h[3])    # la regla del paper
for nombre, (ep, _, l_va, acc) in [("por val loss", por_loss),
                                   ("por métrica", por_metrica)]:
    print(f"{nombre:13} → epoch {ep:2d}  val {l_va:.3f}  ordena {acc:.0%}")

por val loss  → epoch 11  val 0.778  ordena 68%
por métrica   → epoch 16  val 0.810  ordena 74%

Parte 4 · Un segundo SFT

Cuidado al leer el paper: el modelo del que arranca PPO, proximal policy optimization, no es exactamente este SFT.

El SFT que inicializa PPO

El apéndice C punto tres describe otro SFT, el que sirve de punto de partida para PPO. Dos epochs en lugar de dieciséis, y con un diez por ciento de datos de preentrenamiento mezclados. El paper dice diez por ciento, no dice cómo; acá lo hacemos por ejemplos dentro de cada batch. Un batch de treinta y dos. El diez por ciento, redondeado, son tres. Y esos tres salen del texto de preentrenamiento. Veintinueve demostraciones y tres de preentrenamiento. Según el apéndice E, de las variantes que probaron, solo el diez por ciento de mezcla se destacó. Esa misma idea, mezclar preentrenamiento, vuelve en grande en el episodio cinco.

def batch_mezclado(B=32, frac=0.10):           # init de PPO: 10 % de preentreno
    n_pre = round(B * frac)
    i = torch.randint(len(demos), (B - n_pre,))
    j = torch.randint(len(preentreno), (n_pre,))
    return [demos[k] for k in i] + [preentreno[k] for k in j]

b = batch_mezclado()
n_demo = sum(x.startswith("demo") for x in b)
print(f"{len(b)} ejemplos: {n_demo} demos + {len(b) - n_demo} preentreno")

32 ejemplos: 29 demos + 3 preentreno

Resumiendo

• Pérdida solo sobre la respuesta: sin máscara, 95 % sería prompt
• lr 9,65e-6 con coseno al 10 %, sin warmup; Adam β₂ = 0,95
• Residual dropout 0,2: en lo que cada subcapa suma
• 16 epochs aunque la val loss sube desde el primero
• Checkpoint por RM score; en el juguete, epoch 16 contra 11
• Próximo: paso 2, el reward model y sus rankings de K

Papers del episodio

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