Paso 2: el reward model

InstructGPT. Episodio cuatro: paso dos, el reward model, el modelo de recompensa. La pérdida de Bradley-Terry ya la viste en Alineación, episodio tres: la sigmoide de la diferencia de recompensas. Hoy, lo que el paper le agrega: rankings de cuatro a nueve respuestas, cómo se convierten en pares, y por qué todos los pares de un prompt van juntos en el batch.

Las decisiones del paso 2

• Un solo RM de 6B para todas las políticas
• Rankings de K = 4 a 9 respuestas por prompt
• Las C(K,2) comparaciones de un prompt, juntas en el batch
• Una sola epoch; al final, un sesgo que centra en 0

Cuatro decisiones. Un solo RM, reward model, de seis mil millones de parámetros, para entrenar políticas de todos los tamaños, incluso la de ciento setenta y cinco mil millones. En lugar de comparar dos respuestas, el etiquetador ordena entre cuatro y nueve. Todas las comparaciones que salen de un mismo prompt viajan juntas, como un solo elemento del batch. Se entrena una sola epoch. Y antes de usarlo, se le resta una constante para que las demostraciones humanas den, en promedio, cero.

Parte 1 · De un ranking a pares

Ordenar cuatro respuestas lleva poco más que comparar dos, y rinde mucho más. Veamos cuánto.

Un ranking de 4 → pares

Cuatro respuestas, A, B, C y D, con el puesto que les dio una persona: B primera, A y D empatadas en el segundo, C última. Es un ranking inventado. Recorremos todos los pares posibles con combinations. Los empates se descartan: el paper lo dice en el apéndice C punto dos. Y cada par queda como ganadora, perdedora. Seis pares posibles, cinco sin el empate. El símbolo de mayor curvado se lee es preferida a: B a A, A a C, B a C, B a D, D a C.

respuestas = ["A", "B", "C", "D"]
puesto = {"A": 2, "B": 1, "C": 3, "D": 2}     # 1 = la mejor; A y D empatan

pares = []
for i, j in combinations(respuestas, 2):       # los C(K,2) pares posibles
    if puesto[i] == puesto[j]:
        continue                                # los empates se descartan
    w, l = (i, j) if puesto[i] < puesto[j] else (j, i)
    pares.append((w, l))                        # (ganadora, perdedora)

print("K =", len(respuestas), "→ pares posibles:", comb(len(respuestas), 2))
print("sin empates:", len(pares))
print(" ".join(f"{w}≻{l}" for w, l in pares))

K = 4 → pares posibles: 6
sin empates: 5
B≻A A≻C B≻C B≻D D≻C

Cuántos pares da un ranking

«K sobre 2»: cuántos pares distintos hay entre $K$

el ranking del snippet

K sobre dos es la cantidad de pares distintos que se pueden formar con K cosas: cada una de las K con cada una de las otras K menos uno, dividido dos porque A con B es el mismo par que B con A. En el dibujo: con cuatro respuestas, seis pares; con nueve, treinta y seis. Una sola tarea de ordenar nueve respuestas rinde treinta y seis comparaciones. Con nuestro ranking: seis posibles, y como A y D empatan, quedan cinco.

Quién le gana a quién

Los mismos cinco pares como matriz. Cada fila es una respuesta, y un uno en la columna j quiere decir que le ganó a j. B, la primera, le gana a las otras tres. La fila de C está vacía: perdió con todas. Y entre A y D no hay uno en ninguna dirección: es el empate descartado.

Parte 2 · Un número por respuesta

El reward model es el mismo transformer, con una salida distinta: en vez de un logit por token, un solo número.

Del SFT al reward model

La sección tres punto cinco lo describe en una frase: el modelo SFT, el del supervised fine-tuning, sin la última capa de unembedding. En chico: el unembedding de nuestro GPT dos de juguete es una matriz de dieciséis, el vocabulario, por sesenta y cuatro, la dimensión del residual stream. Nos quedamos con todo el cuerpo. Y en lugar del unembedding, una proyección de sesenta y cuatro a uno. Leemos el vector del último token, que ya vio el prompt y la respuesta entera. El SFT da un logit por token del vocabulario en cada posición; el RM, un solo número por secuencia.

print("unembedding del SFT:", tuple(sft.lm_head.weight.shape))  # [V, d]

class RewardModel(nn.Module):
    def __init__(self, lm):
        super().__init__()
        self.cuerpo = lm.transformer                # todo menos el unembedding
        self.cabeza = nn.Linear(lm.config.n_embd, 1)  # proyección a un número
    def forward(self, ids):
        h = self.cuerpo(input_ids=ids).last_hidden_state   # [B, T, d]
        return self.cabeza(h[:, -1]).squeeze(-1)    # último token → r(x, y)

rm = RewardModel(sft)
ids = torch.tensor([[5, 9, 3, 3, 12, 7, 4]])        # prompt + respuesta
print("logits del SFT:", tuple(sft(input_ids=ids).logits.shape))
print("salida del RM: ", tuple(rm(ids).shape))
print(f"r(x, y) = {rm(ids).item():+.3f}")

unembedding del SFT: (16, 64)
logits del SFT: (1, 7, 16)
salida del RM:  (1,)
r(x, y) = -0.295

La arquitectura del RM

El camino completo. Entra el prompt seguido de la respuesta. Pasa por el transformer de seis mil millones. Ojo con un detalle del paper: la sección tres punto cinco dice que se parte del SFT, pero el apéndice C punto dos dice que el RM final se inicializó de un GPT tres de seis mil millones ajustado en datasets públicos, por razones históricas, y que partir del SFT da resultados parecidos. Se toma el vector del último token. La proyección lo lleva a un número. Ese número es r de x, y: la recompensa.

prompt + respuesta → transformer de 6B → último token: h → proyección d → 1 → r(x, y)

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

6B · «In this paper we only use 6B RMs, as this saves a lot of compute, and we found that 175B RM training could be unstable»

¿Por qué seis mil millones y no ciento setenta y cinco? El paper da dos razones. Ahorra mucho cómputo, porque durante PPO, proximal policy optimization, el RM se evalúa en cada respuesta. Y el de ciento setenta y cinco mil millones era inestable de entrenar, lo que lo hacía peor como punto de partida para la función de valor de PPO, que vemos en el episodio cinco.

Parte 3 · La pérdida con K respuestas

Con K recompensas para un mismo prompt, la pérdida se arma sobre todos sus pares a la vez.

Todos los pares de un prompt, de una vez

Supongamos que el RM ya dio cuatro recompensas para las cuatro respuestas, con el ranking de antes. Restamos todas contra todas con broadcasting: D sub i j es r i menos r j, una matriz de cuatro por cuatro. G sub i j marca los pares donde i quedó mejor que j. Para esos pares, menos el logaritmo de la sigmoide de la diferencia: la pérdida de Bradley-Terry de Alineación tres. Y el promedio. Cinco pares, todas las diferencias positivas: el RM ya ordena bien. Las pérdidas van de cero coma catorce, el par más claro, a cero coma cuarenta y cuatro. Promedio: cero coma tres.

Cada línea, su término:

• dif = r[:, None] - r[None, :] → $D_{ij} = r_i - r_j$ — una matriz de K × K
• gana = puesto[:, None] < puesto[None, :] → $G_{ij} = \mathbb{1}[\,y_i \succ y_j\,]$ — los pares que de verdad ganó i
• perdidas = -F.logsigmoid(dif[gana]) → $-\log \sigma(r_w - r_l)$ — Bradley-Terry, par por par
• loss = perdidas.mean() → $\tfrac{1}{\#\text{pares}} \textstyle\sum$ — el $1/\binom{K}{2}$ de la ecuación 1

r = torch.tensor([0.9, 1.5, -0.4, 0.2])      # r(x, y) de las K = 4 respuestas
puesto = torch.tensor([2, 1, 3, 2])          # el ranking humano (A y D empatan)

dif = r[:, None] - r[None, :]                # dif[i, j] = r_i − r_j
gana = puesto[:, None] < puesto[None, :]     # i quedó mejor que j
perdidas = -F.logsigmoid(dif[gana])          # −log σ(r_w − r_l), por par
loss = perdidas.mean()                       # promedio: el 1 / C(K,2)

print("pares:", int(gana.sum()))
print("r_w − r_l:", [f"{v:+.1f}" for v in dif[gana]])
print("pérdida:  ", [f"{v:.2f}" for v in perdidas])
print(f"loss del prompt = {loss:.3f}")

pares: 5
r_w − r_l: ['+1.3', '+0.6', '+1.9', '+1.3', '+0.6']
pérdida:   ['0.24', '0.44', '0.14', '0.24', '0.44']
loss del prompt = 0.299

La ecuación 1 del paper

$\theta$: pesos del RM · $D$: las comparaciones humanas

los números del snippet

La ecuación uno, tal cual. Theta son los pesos del RM. D es el dataset de comparaciones humanas, y la E grande es el promedio sobre ese dataset de tripletes: el prompt x, la respuesta ganadora y sub w y la perdedora y sub l. Adentro, Bradley-Terry. Y adelante, uno sobre K sobre dos: cada prompt pesa lo mismo, tenga cuatro respuestas o nueve. En el dibujo, la sigmoide con dos pares del snippet, y uno inventado al revés: cuanto más grande la diferencia, más cerca de uno, y menos pérdida. Con los números del snippet: cinco pérdidas, y su promedio, cero coma treinta.

La matriz de diferencias

La matriz de diferencias del snippet. Es antisimétrica: la celda i j es menos la j i, y la diagonal es cero. De las dieciséis celdas, la pérdida usa solo las cinco donde la fila ganó: B contra A, C y D, A contra C, y D contra C. El costo de calcularla es tener las cuatro recompensas: con eso sale todo lo demás.

Por qué agrupar por prompt

• Los pares de un mismo prompt están muy correlacionados
• Mezclados: cada respuesta entra en K−1 updates distintos
• Resultado: overfitting en una sola pasada
• Agrupados: un forward por respuesta, y no overfittea

La decisión que el paper explica con más cuidado. Las comparaciones que salen de un mismo ranking están muy correlacionadas: comparten el prompt y comparten respuestas. Si las mezclás en un solo dataset, cada respuesta aparece en K menos uno pares distintos, y el modelo la ve K menos uno veces en la misma epoch. El paper encontró que así el RM hace overfitting en una sola pasada. En cambio, si todas las comparaciones de un prompt son un solo elemento del batch, cada respuesta pasa una vez por el modelo, y dicen que mejora mucho la accuracy y la log loss de validación.

Forwards y comparaciones

un forward por respuesta, no uno por comparación

batch de 64 prompts: hasta 2.304 comparaciones

La cuenta de cómputo. Con nueve respuestas, nueve forwards dan treinta y seis comparaciones: las diferencias salen de restar números ya calculados. El paper lo pone así: un forward por respuesta, en lugar de K sobre dos. Y el apéndice C punto dos da el tamaño: el batch es de sesenta y cuatro prompts. Si todos tuvieran nueve respuestas, serían sesenta y cuatro por treinta y seis: hasta dos mil trescientas cuatro comparaciones por paso. Con learning rate nueve por diez a la menos seis, coseno hasta el diez por ciento, y una sola epoch: con más, dicen, hace overfitting enseguida.

Parte 4 · Entrenarlo, en chico

Un reward model de juguete, entrenado con rankings ruidosos, una sola epoch y un batch por prompt.

Una epoch, agrupada por prompt

El gusto real es sintético: cada token suma un peso oculto. Cada etiquetador ordena entre cuatro y nueve respuestas según ese gusto, más ruido. Mil quinientos prompts. Una sola epoch, en grupos de dieciséis prompts. Juntamos todas las respuestas del grupo en un tensor. Un forward, y separamos las recompensas por prompt. La pérdida de cada prompt es la que vimos; promediamos entre prompts. Y contamos. Nueve mil setecientos cuarenta y cinco forwards para veintiocho mil novecientas veintitrés comparaciones: casi tres por forward.

forwards = comparaciones = 0
for b in range(0, len(train), B):                 # una sola epoch
    grupo = train[b:b + B]                        # B prompts, cada uno con K
    ids = torch.cat([s for s, _, _ in grupo])     # todas las respuestas juntas
    r = rm(ids).split([len(s) for s, _, _ in grupo])   # un forward por respuesta
    loss = torch.stack([perdida_prompt(ri, p)
                        for ri, (_, p, _) in zip(r, grupo)]).mean()
    loss.backward(); opt.step(); opt.zero_grad(); sched.step()
    forwards += len(ids)
    comparaciones += sum(len(s) * (len(s) - 1) // 2 for s, _, _ in grupo)
print(f"forwards {forwards:,} | pares {comparaciones:,}")

forwards 9,745 | pares 28,923

¿Qué tan bien ordena?

Medido en trescientos prompts de validación. El RM ordena como el etiquetador el ochenta y cuatro por ciento de los pares. Pero el etiquetador mismo, con su ruido, coincide con el gusto real solo el ochenta y cinco por ciento: ese es el techo de imitarlo. Y contra el gusto real, el RM acierta el noventa y tres: promediar miles de comparaciones ruidosas cancela el ruido. Cuidado con generalizar: acá el ruido es independiente y sin sesgo. Si los etiquetadores comparten un sesgo, el RM lo aprende igual de bien.

• RM vs. etiquetas: 84 %
• etiquetador vs. gusto real: 85 %
• RM vs. gusto real: 93 %

Centrar en las demostraciones

El último paso antes de usarlo. La pérdida solo mira diferencias, así que a cualquier RM se le puede sumar una constante sin cambiar nada. El paper la fija: resta un sesgo para que las demostraciones de los etiquetadores den cero en promedio. Puntuamos doscientas demostraciones sintéticas. Su promedio es el sesgo b. Y comparamos la pérdida con r y con r menos b. El promedio pasa de tres coma seis a cero. La pérdida, cero coma cinco cero nueve en los dos casos: idéntica.

with torch.no_grad():
    r_demos = rm(demos)                       # el RM puntúa las demostraciones
    b = r_demos.mean()                        # el sesgo a restar
    print(f"antes:   promedio en demos {r_demos.mean():+.3f}")
    print(f"después: promedio en demos {(r_demos - b).mean().abs():.3f}")

    r = rm(par)                               # un prompt con 4 respuestas
    puesto = torch.tensor([0, 1, 2, 3])
    gana = puesto[:, None] < puesto[None, :]
    for nombre, rr in [("r", r), ("r − b", r - b)]:
        loss = -F.logsigmoid((rr[:, None] - rr[None, :])[gana]).mean()
        print(f"loss con {nombre:5}: {loss:.4f}")

antes:   promedio en demos +3.595
después: promedio en demos 0.000
loss con r    : 0.5090
loss con r − b: 0.5090

Por qué se puede restar b

$b$: un solo número

la b se cancela: la pérdida no cambia

r prima es r menos b, y b es un solo número: el promedio del RM sobre las demostraciones. En el dibujo, el promedio pasa de tres coma seis a cero. ¿Por qué no cambia nada? Porque la pérdida solo usa diferencias, y al restar, la b de la ganadora se cancela con la b de la perdedora. Entonces, ¿para qué centrar? El paper no lo explica. Lo que sí cambia es el nivel: en PPO, una respuesta tan buena como una demostración recibe recompensa cero, y la recompensa se suma con la penalidad por token del episodio cinco.

Resumiendo

• Ranking de K = 4 a 9 → $\binom{K}{2}$ pares (de 6 a 36); empates afuera
• RM = transformer sin unembedding + proyección a 1 número
• Ecuación 1: Bradley-Terry promediado con $1/\binom{K}{2}$
• Pares de un prompt juntos: K forwards, sin overfitting
• 6B, 1 epoch, batch de 64 prompts; al final, $r - b$
• Próximo: paso 3, PPO con el freno de KL y PPO-ptx

Papers del episodio

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