El problema

InstructGPT. Episodio uno: el problema. En esta serie leemos de punta a punta un solo paper: el de OpenAI de dos mil veintidós que convirtió a GPT tres en un modelo que sigue instrucciones. Cada decisión del paper, con sus números, y cada paso reimplementado en chico. Hoy, el problema que quiso resolver, visto en un modelo de verdad.

El mapa de la serie

El recorrido. Los dos primeros episodios son el planteo: el problema y los datos que juntaron. Después, los tres pasos del método. SFT, supervised fine-tuning: ajustar con ejemplos escritos por personas. Un reward model, un modelo de recompensa que aprende qué respuesta prefiere la gente. Y PPO, proximal policy optimization, con un agregado propio del paper. Luego, qué salió: cómo lo evaluaron, qué capacidades se pierden, y qué no resolvió. Y al final, el pipeline entero en un modelo de juguete, y lo que vino después.

El problema → Los datos → Paso 1: SFT → Paso 2: RM → Paso 3: PPO-ptx → Cómo se evaluó → El impuesto → Los límites → Todo, en chico → Lo que vino

Parte 1 · Un modelo que completa

Antes de leer el paper, veamos el problema con nuestros ojos, en Gemma dos de dos mil millones de parámetros, la versión base: sin ningún ajuste.

Le pedimos algo a un modelo base

Cargamos Gemma dos, dos B, base: el modelo tal como salió del preentrenamiento. Le damos tres pedidos. El primero, escribí un cuento sobre una rana sabia, es el ejemplo que usa el paper. Los otros dos son nuestros. Greedy quiere decir que en cada paso elige el token más probable, sin azar. Y mostramos lo que agregó. Mirá. A la rana le contesta repitiendo el pedido, igual, palabra por palabra. Al de la gravedad le contesta con algo que parece un foro: usuario cero cero cero uno. Y al de traducir, otro foro, donde alguien arranca con hola: en el texto completo, pregunta cómo se traduce esa misma frase. No contesta nada: completa.

nombre = "google/gemma-2-2b"                         # modelo base: sin ajuste
tok = AutoTokenizer.from_pretrained(nombre)
model = AutoModelForCausalLM.from_pretrained(nombre, dtype=torch.float32)

pedidos = ["Write a story about a wise frog.",       # el ejemplo del paper
           "Explain gravity to a 6 year old.",
           "Traducí al inglés: el gato duerme."]
for p in pedidos:
    ids = tok(p, return_tensors="pt")
    out = model.generate(**ids, max_new_tokens=24,
                         do_sample=False)            # greedy: lo más probable
    nuevo = tok.decode(out[0, ids["input_ids"].shape[1]:])
    print(corto(p, 32))
    print("→", corto(nuevo))

Write a story about a wise frog.
→ ⏎⏎Write a story about a 
Explain gravity to a 6 year old.
→ ⏎⏎[User 0001]⏎⏎I'm tryin
Traducí al inglés: el gato duerm
→ ⏎⏎[User 0001]⏎⏎Hola,⏎⏎¿C

¿Qué cree que viene después?

Miremos un solo paso. Los logits del último token: un número por cada uno de los doscientos cincuenta y seis mil tokens del vocabulario. Softmax los convierte en probabilidades que suman uno. Y miramos las cinco más altas. Lo más probable, quince coma siete por ciento, es un doble salto de línea: arrancar un párrafo nuevo. Después The, con quince coma dos. Y después Use, e Include: el comienzo de consignas como usá estas palabras en tu cuento. El modelo leyó esto como el enunciado de una tarea escolar, y sigue la tarea.

x = "Write a story about a wise frog."
ids = tok(x, return_tensors="pt")
with torch.no_grad():
    logits = model(**ids).logits[0, -1]              # [256000]: un logit por token
p = logits.softmax(-1)                               # probabilidades que suman 1
top = p.topk(5)
for prob, i in zip(top.values, top.indices):
    print(f"{tok.decode(i)!r:12} {prob:.3f}")

'\n\n'       0.157
' The'       0.152
' Use'       0.084
' Include'   0.045
' It'        0.041

El objetivo con el que se entrenó

$\theta$: todos los pesos · $\ell$: los logits

preentrenamiento: predecir el próximo token de internet

el paper lo llama misaligned

¿Por qué hace eso? Por el objetivo con el que se entrenó. El modelo da, para la posición t, una probabilidad para cada token posible, x sub t, dado todo lo anterior, x menor que t. Theta son todos sus pesos. En el dibujo, los logits reales de Gemma después del pedido, y el softmax. Preentrenar es minimizar esta pérdida: la suma, sobre todos los tokens de internet, de menos el logaritmo de la probabilidad del token que de verdad vino. En el dibujo, ese menos log p para cada token de la continuación que repite el pedido: casi todos chicos, porque es muy probable. Y el problema: el mínimo de esa pérdida es imitar internet, no seguir instrucciones. Repetir le cuesta siete coma siete; el cuento, catorce coma siete.

Cuánto le cuesta cada respuesta

Ese número sale de esta función. La probabilidad de una respuesta entera y, de T tokens, es el producto de las probabilidades de cada token dado lo anterior. Con logaritmos, el producto se vuelve suma. Log softmax da el logaritmo de la probabilidad de cada token del vocabulario, en cada posición. Gather elige, en cada posición, el del token que de verdad sigue. Y se suman. Repetir el pedido: nueve tokens, menos siete coma siete en total. El cuento: doce tokens, menos catorce coma siete. Casi todo ese costo está en la primera palabra, Once, que sola vale menos nueve. Para el modelo base, arrancar un cuento es raro; repetir la consigna, normal.

Cada línea, su término:

• def logp(x, y) → $\log p(y \mid x) = \sum_{t=1}^{T} \log p(y_t \mid x, y_{<t})$ — la respuesta entera: suma de sus tokens
• log_softmax(-1) → $\log \mathrm{softmax}(\ell_t)$ — log-probs de todo el vocabulario
• gather(-1 → $\log p(y_t \mid \cdot)$ — el del token que efectivamente sigue
• lp_y.sum() → $\textstyle\sum_t$ — sumar los T términos

def logp(x, y):                                      # log p(y | x), token a token
    n = len(tok(x)["input_ids"])
    ids = tok(x + y, return_tensors="pt")["input_ids"]
    with torch.no_grad():
        lp = model(input_ids=ids).logits[0, :-1].log_softmax(-1)
    lp_y = lp[n - 1:].gather(-1, ids[0, n:, None])   # el log-prob de cada token de y
    return lp_y.sum().item(), len(lp_y)

x = "Write a story about a wise frog."
repite = "\n\nWrite a story about a wise frog."
cuento = "\n\nOnce upon a time, there was a wise frog."
for nombre_y, y in [("repetir", repite), ("cuento", cuento)]:
    total, T = logp(x, y)
    print(f"{nombre_y:8} T={T:2d}  log p = {total:6.1f}  por token {total/T:5.2f}")

repetir  T= 9  log p =   -7.7  por token -0.85
cuento   T=12  log p =  -14.7  por token -1.23

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

misaligned · «Making language models bigger does not inherently make them better at following a user's intent.»

El paper arranca exactamente acá. Primera oración: hacer más grandes los modelos no los hace, por sí solo, mejores para seguir la intención de quien los usa. Y en la introducción le pone nombre: el objetivo de predecir el próximo token de una página de internet es distinto del objetivo seguí las instrucciones del usuario de forma útil y segura. Por eso dicen que ese objetivo está misaligned: desalineado.

Parte 2 · La pregunta del paper

Si el objetivo de preentrenamiento no alcanza, ¿con qué se reemplaza? El paper responde con datos humanos y tres pasos.

Qué quiere decir «alineado» en el paper

• Intención explícita: seguir la instrucción
• Intenciones implícitas: no inventar, no dañar
• Helpful · honest · harmless (Askell et al., 2021)
• Lo mide gente: el juicio de los etiquetadores

El paper define qué busca. Que el modelo actúe según la intención del usuario: la explícita, seguir la instrucción. Y las implícitas: ser veraz, y no ser sesgado, tóxico ni dañino. Lo resume con tres palabras en inglés, tomadas de Askell y otros: helpful, que ayude a resolver la tarea; honest, que no invente información ni engañe; y harmless, que no cause daño. Y como la intención de un pedido puede ser ambigua, la métrica principal no es automática: es la preferencia de un grupo de personas contratadas para etiquetar.

El truco previo: un prompt few-shot

Antes de entrenar nada, hay un truco conocido. Anteponer un prefijo con ejemplos de instrucción y respuesta. Dos ejemplos inventados por nosotros: traducir y nombrar colores. Y después, el mismo pedido de la rana. Esto se llama prompt few-shot: pocos ejemplos en el contexto. Ahora sí escribe un cuento: la rana era muy sabia. Flojo y repetitivo, pero responde. El paper usa esto como línea de base, GPT tres prompted, con un prefijo que salió de una competencia interna entre dos autores. Ya vamos a ver que no alcanza.

prefijo = ("Instruction: Translate to Spanish: good morning.\n"
           "Answer: Buenos días.\n\n"
           "Instruction: Name three colors.\n"
           "Answer: Red, green and blue.\n\n")
pedido = "Instruction: Write a story about a wise frog.\nAnswer:"

ids = tok(prefijo + pedido, return_tensors="pt")
out = model.generate(**ids, max_new_tokens=40, do_sample=False)
nuevo = tok.decode(out[0, ids["input_ids"].shape[1]:])
for linea in textwrap.wrap(nuevo.split("\n")[0], 34)[:5]:
    print(linea)

 The frog was very wise. He knew
that he had to be careful because
he was a frog. He knew that he had
to be careful because he was a
frog. He knew that he had to

Una distribución, no una intención

Una cosa más del modelo base. En vez de elegir siempre lo más probable, muestreamos: sorteamos cada token según su probabilidad. Temperatura uno: las probabilidades tal cual. Y pedimos cuatro respuestas al mismo pedido. Cuatro continuaciones distintas, y ninguna es un cuento: usá estos términos en tu cuento, usá un adjetivo en tu cuento. Son consignas de una tarea. El modelo no tiene una intención que cumplir: tiene una distribución, y el pedido cae en la zona de los enunciados escolares. Muestrear varias respuestas al mismo prompt es justo lo que el paper hace en el paso dos, para que la gente las compare.

torch.manual_seed(0)
x = "Write a story about a wise frog."
ids = tok(x, return_tensors="pt")
out = model.generate(**ids, max_new_tokens=20, do_sample=True,
                     temperature=1.0,                # muestrear, no elegir
                     num_return_sequences=4)
for k, fila in enumerate(out):
    nuevo = tok.decode(fila[ids["input_ids"].shape[1]:],
                       skip_special_tokens=True)
    print(k + 1, corto(nuevo))

1 The frog can be wise for any re
2 Use the following terms in your
3 This is a good time for you to 
4 Use an adjective in your story.

El método: tres pasos

El método, que viene de trabajos anteriores de OpenAI sobre resúmenes. Es RLHF, reinforcement learning from human feedback: aprendizaje por refuerzo con feedback humano. Paso uno: de una lista de prompts, personas escriben la respuesta deseada, las demostraciones, y con eso se ajusta GPT tres de forma supervisada. Paso dos: el modelo genera entre cuatro y nueve respuestas para un mismo prompt, una persona las ordena de mejor a peor, y con esos rankings se entrena el reward model. Paso tres: se optimiza el modelo para que saque mucha recompensa, con PPO. Lo que sale es InstructGPT.

prompts → demostraciones → SFT → K respuestas → ranking humano → RM → PPO-ptx → InstructGPT

Los tres pasos, en tres fórmulas

paso 1: solo cuentan los tokens de la respuesta

paso 2: $y_w$ ganó, $y_l$ perdió

paso 3: recompensa menos un freno ($\beta$)

Los tres pasos en fórmulas, que ya viste en Alineación, episodios dos a cuatro. Pi es el modelo, al que en esta jerga se le dice política. El paso uno es la misma pérdida de antes, pero solo sobre los tokens de la respuesta: en el dibujo, el prompt queda tachado. Paso dos. r es el reward model: un número para cada prompt x y respuesta y. y sub w es la que ganó; y sub l, la que perdió. Sigma, la sigmoide, lleva la diferencia a una probabilidad. En el dibujo: si ordena bien, pérdida chica. Paso tres: la recompensa menos beta por el logaritmo del cociente entre pi RL, el modelo que se entrena con aprendizaje por refuerzo, y pi SFT, el del paso uno. Ese freno es la KL, la divergencia de Kullback-Leibler, y se paga por token.

Parte 3 · El resultado

Ahora el resultado que hizo famoso al paper, y lo que significa exactamente.

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

1,3B > 175B · «outputs from the 1.3B parameter InstructGPT model are preferred to outputs from the 175B GPT-3, despite having over 100x fewer parameters.»

Si viste Alineación, episodio uno, ya citamos este resultado; ahora lo abrimos. Los etiquetadores prefirieron las respuestas del InstructGPT de mil trescientos millones de parámetros a las de GPT tres de ciento setenta y cinco mil millones: más de cien veces más grande. Y el paper subraya que tienen la misma arquitectura. Lo único que los separa es el ajuste con datos humanos.

Qué quiere decir «preferido»

$\mathbb{1}[\cdot]$ vale 1 si es cierto, 0 si no

«over 100x fewer parameters»

¿Preferido cómo? Para un conjunto de prompts de prueba, de clientes que no aparecen en el entrenamiento, se generan las respuestas de dos modelos, A y B, y una persona elige. El win rate, la tasa de victorias, es la fracción de prompts donde ganó A. El uno con corchetes vale uno si lo de adentro es cierto y cero si no; el símbolo de mayor curvado se lee es preferida a. En el dibujo, ocho prompts inventados: A gana seis, win rate cero coma setenta y cinco. Y los tamaños: ciento setenta y cinco dividido uno coma tres da unas ciento treinta y cinco veces. En el dibujo, los tres tamaños que entrena el paper.

175B contra 175B

Con el mismo tamaño, la diferencia es grande. InstructGPT de ciento setenta y cinco mil millones le gana a GPT tres del mismo tamaño el ochenta y cinco por ciento de las veces, más o menos tres. Y contra GPT tres con el prefijo few-shot, el truco que vimos, gana el setenta y uno por ciento, más o menos cuatro. El truco ayuda, pero no alcanza. Como referencia, cincuenta por ciento sería no poder distinguirlos.

• vs. GPT-3: 85 ± 3 %
• vs. GPT-3 few-shot: 71 ± 4 %
• empate: 50 %

Lo demás que encontró, en una línea

• Veracidad: el doble de respuestas veraces en TruthfulQA
• Inventa menos: 21 % contra 41 % en tareas cerradas
• Toxicidad: ~25 % menos, si se le pide ser respetuoso
• Sesgo: sin mejora significativa; y un «alignment tax»

El paper reporta más cosas, que abrimos en los episodios seis a ocho. En TruthfulQA, un benchmark de preguntas trampa, da respuestas veraces e informativas el doble de veces. En tareas cerradas, como resumir, donde no hay que agregar nada que no esté en la entrada, inventa información en el veintiuno por ciento de los casos, contra cuarenta y uno de GPT tres. Genera alrededor de un veinticinco por ciento menos de texto tóxico, cuando se le pide ser respetuoso. En sesgo no mejora. Y aparece un costo: rinde peor en algunos benchmarks públicos. El paper lo llama alignment tax, el impuesto de alineación.

Resumiendo

• Un modelo base completa documentos: no sigue pedidos
• Gemma base repite la consigna: $\log p$ −7,7 contra −14,7 del cuento
• El objetivo de preentrenamiento está «misaligned»
• Método: SFT → reward model → PPO-ptx (Alineación ep. 2–4)
• 1,3B preferido a 175B; 175B gana 85 % contra GPT-3
• Próximo: los datos, de dónde salen los prompts y quién etiqueta

Papers del episodio

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