Modelos de recompensa

Alineación a bajo nivel. Episodio tres: modelos de recompensa. Terminamos el episodio anterior con un problema: el ajuste supervisado solo imita, y nunca ve una respuesta mala. Hoy vemos la pieza que lo resuelve: un modelo que mira una respuesta y devuelve un número que dice qué tan buena es.

Por qué comparar en vez de escribir

• Elegir entre dos respuestas es más fácil que escribir la ideal
• Las comparaciones muestran qué es peor, no solo qué es bueno
• Las personas son ruidosas: hace falta un modelo probabilístico
• El RM reemplaza a la persona dentro del loop de RL

La idea de fondo es práctica. Para una persona es mucho más fácil decir cuál de dos respuestas es mejor que escribir la respuesta perfecta. Y una comparación trae información que el SFT, supervised fine-tuning, no tiene: qué hace que una respuesta sea peor. Pero las personas no son consistentes. Dos etiquetadores pueden elegir distinto, y el mismo etiquetador puede cambiar de opinión. Necesitamos un modelo que trate la elección como algo probabilístico. Y el objetivo final es que ese modelo, el RM, reward model, reemplace a la persona: en el próximo episodio, el RL, reinforcement learning, le va a pedir millones de puntajes, y nadie puede etiquetar millones de respuestas a mano.

Cómo se ve un dato de preferencias

Así se ve un conjunto de preferencias, con ejemplos inventados. Cada fila tiene un prompt y dos respuestas, casi siempre generadas por el mismo modelo de SFT. Una persona marcó cuál prefiere: la elegida, y sub w, y la rechazada, y sub l. Fijate que no hace falta que la elegida sea perfecta: alcanza con que sea mejor que la otra. Y la persona no tiene que explicar por qué. Toda la información está en el orden. El trabajo del modelo de recompensa es reconstruir, a partir de miles de estos órdenes, una función que puntúe cualquier respuesta.

Parte 1 · De comparaciones a números

El puente entre «prefiero esta» y un número tiene nombre: el modelo de Bradley-Terry.

Bradley-Terry

la probabilidad de preferir depende solo de la diferencia

la sigmoide: de cualquier real a $(0, 1)$

la escala absoluta no está definida

El modelo es de mil novecientos cincuenta y dos, pensado para ordenar cosas a partir de comparaciones de a pares. Cada respuesta tiene un puntaje oculto, r. La probabilidad de que una persona prefiera la respuesta ganadora, y sub w, sobre la perdedora, y sub l, es la sigmoide de la diferencia de puntajes. La sigmoide lleva cualquier número real a una probabilidad. En el dibujo: diferencia cero da un medio, más dos da cero coma ochenta y ocho, y menos dos, cero coma doce. Y una consecuencia que vamos a necesitar: si le sumás la misma constante a todos los puntajes de un prompt, nada cambia. Solo importan las diferencias.

Un etiquetador sintético

Para entrenar sin datos de personas, armamos un etiquetador de mentira que cumple Bradley-Terry al pie de la letra. Tiene un gusto oculto, r estrella: un peso por cada token, y el puntaje de una respuesta es la suma. Generamos dos respuestas de juguete, secuencias de largo variable. Calculamos la probabilidad de que prefiera la primera con la sigmoide de la diferencia. Y tiramos una moneda cargada con esa probabilidad. Mirá el primer par: b tenía más puntaje, la probabilidad de elegir a era cero coma tres ocho, y ganó b. En el segundo, los puntajes empatan: es cara o ceca. El ruido es parte de los datos, como con personas reales.

def pares(n):                          # dos respuestas por prompt, y quién ganó
    a, ma = respuestas(n)
    b, mb = respuestas(n)
    p_a = torch.sigmoid(r_oculta(a) - r_oculta(b))   # Bradley-Terry
    gana_a = torch.rand(n) < p_a                     # la persona elige
    return a, b, p_a, gana_a

a, b, p_a, gana_a = pares(4)
for i in range(4):
    print(f"r*(a)={r_oculta(a[i]):+.2f} r*(b)={r_oculta(b[i]):+.2f}"
          f"  P(a≻b)={p_a[i]:.2f}  ganó {'a' if gana_a[i] else 'b'}")

r*(a)=+1.31 r*(b)=+1.78  P(a≻b)=0.38  ganó b
r*(a)=-0.83 r*(b)=-0.83  P(a≻b)=0.50  ganó a
r*(a)=-0.45 r*(b)=-1.43  P(a≻b)=0.73  ganó a
r*(a)=-1.62 r*(b)=+1.00  P(a≻b)=0.07  ganó b

La sigmoide de la diferencia

Esta es la curva. Con diferencia cero, cincuenta por ciento. Con diferencia dos, ochenta y ocho por ciento. Y nunca llega a uno: siempre queda una chance de que la persona elija la otra. Otra forma de leerla: la diferencia de recompensas es el logaritmo de las chances. Diferencia uno significa que la ganadora es e veces más probable que la perdedora, unas dos coma siete veces.

La diferencia de recompensas es el logaritmo de las chances de preferir.

Parte 2 · El reward head

Ahora necesitamos una red que, dada una respuesta, produzca ese puntaje r.

Un transformer que devuelve un número

La arquitectura es casi la misma del episodio uno. Entra el prompt con la respuesta, pasa por un transformer, y cada token termina con su vector en el residual stream. En InstructGPT, ese transformer arranca desde el modelo de SFT, sin la capa final de logits. De todos los vectores nos quedamos con uno: el del último token real. Por la máscara causal, es el único que vio la secuencia entera. Y en vez de proyectarlo a doscientos cincuenta y seis mil logits, lo proyectamos a un solo número con una capa lineal. Ese número es la recompensa.

prompt + respuesta → transformer (SFT) → residual stream → último token real → Linear d → 1 → r(x, y)

El reward head, en código

En código. El cuerpo es un LlamaModel: el transformer sin la capa de logits. El reward head es una lineal de sesenta y cuatro a uno. El detalle que más bugs trae: en un lote, las respuestas tienen largos distintos y se rellenan a la derecha. El último token real está en la suma de la máscara menos uno, no en la última columna. Indexamos ese vector en cada fila. Y probamos que el relleno no cambia nada. Tres respuestas, de ocho, siete y cinco tokens. Les agregamos cuatro posiciones más de relleno y las recompensas dan idénticas. Si hubieras tomado la última columna, estarías puntuando el relleno.

class ModeloRecompensa(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cuerpo = LlamaModel(cfg)           # transformer sin logits
        self.head = nn.Linear(cfg.hidden_size, 1)    # reward head
    def forward(self, ids, mask):
        h = self.cuerpo(input_ids=ids, attention_mask=mask).last_hidden_state
        ultimo = mask.sum(1) - 1                # último token real
        h_ult = h[torch.arange(len(ids)), ultimo]   # [batch, d]
        return self.head(h_ult).squeeze(-1)     # [batch]

rm = ModeloRecompensa(cfg)
ids, mask = respuestas(3)
print(mask)
print("r:", rm(ids, mask).detach().numpy().round(3))
mas_relleno = torch.cat([ids, torch.zeros(3, 4, dtype=torch.long)], 1)
mask2 = torch.cat([mask, torch.zeros(3, 4, dtype=torch.long)], 1)
print("con más relleno:", rm(mas_relleno, mask2).detach().numpy().round(3))

tensor([[1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 0, 0, 0]])
r: [ 0.539  0.727 -0.218]
con más relleno: [ 0.539  0.727 -0.218]

Parte 3 · La pérdida

Tenemos el modelo de preferencias y la red. Falta la pérdida que los une.

Máxima verosimilitud de las preferencias

$-\log$ de la probabilidad de lo que eligió la persona

empuja más cuanto peor ordenado está el par

InstructGPT: $K$ respuestas rankeadas, $\binom{K}{2}$ pares por prompt

La pérdida sale directo de Bradley-Terry. Queremos que el modelo le dé alta probabilidad a lo que la persona eligió de verdad. Es menos el logaritmo de la sigmoide de la diferencia, promediado sobre los pares. En el dibujo, tres pares: el bien ordenado cuesta cero coma trece, el empatado cero coma sesenta y nueve, y el par al revés, dos coma trece. Es una regresión logística sobre la diferencia. El gradiente es simétrico: sube la ganadora y baja la perdedora en la misma cantidad, y esa cantidad es grande cuando el par está mal ordenado y casi cero cuando ya está bien. En InstructGPT, cada persona rankeaba entre cuatro y nueve respuestas a la vez, lo que da K tomados de a dos pares por prompt, y los entrenaban todos juntos en el mismo lote.

La pérdida en tres líneas

En PyTorch. Tres pares con recompensas inventadas. La diferencia. Menos log sigmoide: usamos logsigmoid y no el logaritmo de la sigmoide, porque es estable con diferencias grandes. Y comparamos con la entropía cruzada binaria con etiqueta uno: es exactamente lo mismo. Por último, sumamos diez a todo. Diferencia más dos: probabilidad cero coma ocho ocho, pérdida cero coma uno dos siete. Diferencia cero: cero coma seis nueve tres, el logaritmo de dos. Diferencia menos dos, el par al revés: dos coma uno dos siete. Coincide con la binaria, y sumar diez no cambia nada.

Cada línea, su término:

• delta = r_w - r_l → $\Delta r = r_w - r_l$ — solo importa la diferencia
• loss = -F.logsigmoid → $-\log\sigma(\Delta r)$ — la pérdida por par
• binary_cross_entropy_with_logits → $\mathrm{BCE}(\Delta r,\, 1)$ — es una clasificación binaria
• sumo 10 → $\Delta r \text{ invariante a } r + c$ — la escala no está fijada

delta = r_w - r_l                                   # Δr = r(x,y_w) − r(x,y_l)
p = torch.sigmoid(delta)                            # P(y_w ≻ y_l)
loss = -F.logsigmoid(delta)                         # −log σ(Δr), por par
print("Δr:  ", [f"{v:+.1f}" for v in delta])
print("P:   ", [f"{v:.3f}" for v in p])
print("loss:", [f"{v:.3f}" for v in loss])

bce = F.binary_cross_entropy_with_logits(delta, torch.ones(3), reduction="none")
print("igual a BCE con etiqueta 1:", torch.allclose(loss, bce))
corrida = -F.logsigmoid((r_w + 10) - (r_l + 10))
print("¿y si sumo 10 a todo?", torch.allclose(loss, corrida))

Δr:   ['+2.0', '+0.0', '-2.0']
P:    ['0.881', '0.500', '0.119']
loss: ['0.127', '0.693', '2.127']
igual a BCE con etiqueta 1: True
¿y si sumo 10 a todo? True

Entrenar el modelo de recompensa

Y el loop completo, igual al del episodio dos salvo por la pérdida. Cada paso pide sesenta y cuatro pares nuevos al etiquetador sintético. Dos forwards con el mismo modelo: uno para la elegida y otro para la rechazada. La pérdida de Bradley-Terry. Y cada doscientos pasos medimos la accuracy de preferencias en pares de validación: qué fracción de los pares ordena igual que la persona. La pérdida arranca cerca del logaritmo de dos, cero coma siete siete, y se estanca alrededor de cero coma cuatro ocho. La accuracy sube a cero coma siete siete y ahí se queda. En un modelo de verdad, este loop es igual, con dos diferencias: el cuerpo tiene miles de millones de parámetros, y en InstructGPT se entrenaba una sola época, porque con más se sobreajustaba enseguida. ¿Por qué no llega a uno?

rm = ModeloRecompensa(cfg)
opt = torch.optim.AdamW(rm.parameters(), lr=1e-3)
for paso in range(601):
    w, mw, l, ml, _ = pares(64)                     # elegida, rechazada
    r_w, r_l = rm(w, mw), rm(l, ml)                 # dos forwards, mismo modelo
    loss = -F.logsigmoid(r_w - r_l).mean()          # L_RM = −E log σ(Δr)
    loss.backward(); opt.step(); opt.zero_grad()
    if paso % 200 == 0:
        with torch.no_grad():                       # ¿ordena bien pares nuevos?
            acc = (rm(vw, vmw) > rm(vl, vml)).float().mean()
        print(f"paso {paso:3d}  loss {loss.item():.3f}  acc {acc:.2f}")

paso   0  loss 0.768  acc 0.65
paso 200  loss 0.475  acc 0.76
paso 400  loss 0.495  acc 0.77
paso 600  loss 0.481  acc 0.77

Parte 4 · ¿Qué tan bueno es?

Para juzgar un modelo de recompensa hay que saber cuál es el máximo posible.

Accuracy contra el techo

Evaluamos con cuatro mil pares que el modelo nunca vio. Medimos la accuracy. Y calculamos el techo: como el etiquetador es ruidoso, ni siquiera alguien que conociera r estrella exacta acertaría siempre. Lo mejor posible es apostar al más probable en cada par. Además, como acá sí conocemos el gusto oculto, medimos la correlación con él. Y miramos la media. Accuracy cero coma siete siete cinco, techo cero coma siete siete seis. El modelo está en el techo. La correlación con r estrella es cero coma nueve seis: aprendió el gusto casi entero. Y la media es arbitraria, porque Bradley-Terry no fija la escala. Ojo: en un caso real no conocés r estrella, así que la correlación no se puede medir. Lo único que tenés son pares reservados y cuánto coinciden las personas entre sí.

w, mw, l, ml, p_a = pares(4000)                     # pares nuevos, nunca vistos
with torch.no_grad():
    acc = (rm(w, mw) > rm(l, ml)).float().mean()
techo = torch.maximum(p_a, 1 - p_a).mean()   # apostar al más probable
print(f"accuracy del RM:   {acc:.3f}")
print(f"techo alcanzable:  {techo:.3f}")

y, m = respuestas(4000)
with torch.no_grad():
    r = rm(y, m)
c = torch.corrcoef(torch.stack([r, r_oculta(y)]))[0, 1]
print(f"correlación con r*: {c:.3f}")
print(f"media de r: {r.mean():+.2f}  (la escala absoluta no importa)")

accuracy del RM:   0.775
techo alcanzable:  0.776
correlación con r*: 0.960
media de r: +0.10  (la escala absoluta no importa)

El ruido de quien etiqueta pone el techo

Poné los números en perspectiva. Adivinar al azar da cero coma cinco. Nuestro modelo, cero coma setenta y siete. Y el techo, casi lo mismo. Una accuracy de setenta y siete por ciento parece mediocre, pero acá es perfecta. La lección vale para modelos reales: la accuracy de un RM se lee contra cuánto coinciden las personas entre sí, no contra el cien por ciento. Y al revés: si tu RM da noventa y cinco por ciento con personas que coinciden setenta, desconfiá. Probablemente encontró un atajo en los datos, como el largo de la respuesta.

• azar: 0,50
• RM entrenado: 0,775
• *techo (conoce r)**: 0,776

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

72,6 ± 1,5 % · Los etiquetadores de entrenamiento coinciden entre sí el 72,6 ± 1,5 % de las veces; los etiquetadores reservados, el 77,3 ± 1,3 %.

¿Y cuánto coinciden las personas reales? En InstructGPT, los etiquetadores que producían los datos de entrenamiento coincidían entre sí el setenta y dos coma seis por ciento de las veces. Ese es, en la práctica, el orden del techo de cualquier modelo de recompensa entrenado con esos datos.

Parte 5 · Lo que el RM no es

Antes de optimizar contra este número, conviene tener claro qué es y qué no es.

Cuidados antes de usarlo

• Es un proxy aprendido de preferencias, no «lo bueno»
• Su escala es arbitraria: se normaliza antes del RL
• Fuera de su distribución puede dar puntajes absurdos
• Optimizarlo demasiado lo rompe (episodio 6)

Cuatro cuidados. Es un proxy: aprendió a predecir qué eligen ciertas personas en ciertos datos, no qué es bueno en general. Si a quienes etiquetan les gustan las respuestas largas, al modelo también. Su escala es arbitraria. InstructGPT le suma un sesgo para que las demostraciones humanas tengan puntaje medio cero antes del RL. Solo vio respuestas parecidas a las de su entrenamiento. Una política que explora puede encontrar textos raros donde el puntaje es alto sin razón. Y por eso, optimizarlo demasiado lo rompe. Lo vamos a medir en el episodio seis.

Lo que te llevás

• Comparar es más fácil y más informativo que escribir
• Bradley-Terry: $P(y_w \succ y_l) = \sigma(r_w - r_l)$
• RM = transformer + Linear(d, 1) sobre el último token real
• Pérdida: $-\log\sigma(\Delta r)$, una logística sobre la diferencia
• Accuracy se compara con el acuerdo entre personas
• Próximo: optimizar la política contra el RM con PPO

Papers del episodio

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