El pipeline entero, en chico

InstructGPT. Episodio nueve: el pipeline entero, en chico. Hasta ahora vimos cada paso por separado. Hoy los encadenamos sobre el mismo modelo de juguete del episodio siete: SFT, supervised fine-tuning; después el reward model, entrenado con rankings; después PPO, proximal policy optimization, con mezcla de preentrenamiento. Y al final, la pregunta del episodio seis: ¿el resultado le gana al SFT?

Los tres pasos, con los números de hoy

El mapa. Partimos del modelo que sabe contar para adelante, y le hacemos un SFT corto con demostraciones del pedido: después de INST, contar para atrás. Para el reward model, el SFT genera cuatro respuestas por prompt y una persona simulada las ordena. Con ese reward model como recompensa, PPO-ptx mejora la política, y medimos el win rate contra el SFT con dos jueces: el reward model y el gusto real.

preentrenado → SFT: oro 4,38 → 4 respuestas → ranking + ruido → RM: 94,5 % → PPO-ptx → win rate: 0,72

El paper y el juguete, en números

La escala, para no engañarnos. La política del paper va de mil trescientos millones a ciento setenta y cinco mil millones de parámetros; la nuestra tiene ochenta y tres mil. El SFT del paper usó unos trece mil prompts y dieciséis epochs; el nuestro, ochocientas demostraciones. Su reward model tiene seis mil millones de parámetros y rankings de entre cuatro y nueve respuestas; el nuestro, siempre cuatro. Y PPO: doscientos cincuenta y seis mil episodios contra diecinueve mil. El beta es el mismo; el gamma no, porque nuestras pérdidas tienen otra escala.

paper: secciones 3.2 y 3.5, apéndice C

Parte 1 · Paso 1: SFT

El mismo SFT del episodio siete, pero más corto: queremos dejarle trabajo al aprendizaje por refuerzo.

Un SFT corto

El código del episodio tres, en chico. Solo veinticinco pasos. Cada batch tiene treinta y dos demostraciones: la instrucción y la cuenta correcta para atrás. Entropía cruzada solo sobre la respuesta. Y el diez por ciento de preentrenamiento que el paper usó para inicializar PPO. Medimos muestreando. Oro cuatro coma tres ocho de seis: cumple bastante, no siempre. Y la capacidad de contar para adelante cayó a cero coma tres tres: el SFT corto ya cobra impuesto. En la respuesta de ejemplo, solo dos de los seis tokens siguen la cuenta.

sft = copy.deepcopy(base)                       # paso 1: SFT, ahora más corto
opt = torch.optim.AdamW(sft.parameters(), lr=1e-3)
for paso in range(25):
    x = j.demostraciones(32)                    # INST a a−s → la cuenta para atrás
    lg = sft(input_ids=x).logits[:, j.P - 1:-1]
    loss = cross_entropy(lg.reshape(-1, j.V), x[:, j.P:].reshape(-1))
    loss = loss + 0.1 * j.nll_pre(sft, j.batch_pre(32))
    loss.backward(); opt.step(); opt.zero_grad()

y = j.muestrear(sft, j.prompts(512))
print(f"SFT: oro {j.oro(y).mean():.2f} de 6, capacidad {j.capacidad(sft):.3f}")
print("una respuesta:", y[0].tolist())

SFT: oro 4.38 de 6, capacidad 0.333
una respuesta: [10, 4, 3, 6, 3, 1, 0, 9, 2]

Parte 2 · Paso 2: comparaciones y reward model

Nadie le dice al reward model cuál es la regla. Solo ve rankings de respuestas.

Rankings de cuatro, y la pérdida sobre los seis pares

El paso dos, como en el episodio cuatro. El SFT genera K igual a cuatro respuestas por prompt. Quien etiqueta las juzga por el oro más un ruido, para que no sea perfecto. Y las ordena de mejor a peor. El reward model hace un forward por respuesta, no uno por par. La pérdida recorre los seis pares de cada ranking. Y al final, el sesgo: lo que vale en promedio una demostración. Entrenamos dos: uno fuerte, trescientos pasos, que ordena bien el noventa y cuatro coma cinco por ciento de los pares nuevos; y uno débil, sesenta pasos, ochenta y seis coma nueve. El débil lo vamos a usar al final.

def comparar(n):                     # K respuestas del SFT, ordenadas
    x = j.prompts(n).repeat_interleave(K, 0)
    y = j.muestrear(sft, x).view(n, K, -1)
    gusto = j.oro(y.flatten(0, 1)).view(n, K)        # quien etiqueta: oro
    gusto = gusto + gumbel.sample((n, K))            # … más ruido
    orden = gusto.argsort(-1, descending=True)        # de la mejor a la peor
    return y.gather(1, orden[..., None].expand_as(y))

def entrenar_rm(pasos):
    rm = j.RM(); opt = torch.optim.AdamW(rm.parameters(), lr=1e-3)
    for paso in range(pasos):
        r = rm(comparar(64).flatten(0, 1)).view(64, K)  # un forward por respuesta
        pares = [F.logsigmoid(r[:, a] - r[:, b]) for a, b in PARES]
        loss = -torch.stack(pares).mean()               # los 6 pares
        loss.backward(); opt.step(); opt.zero_grad()
    with torch.no_grad():
        sesgo = rm(j.demostraciones(512)).mean().item()   # las demos valen 0
    return rm, sesgo

La ecuación 1 del paper

$y_w$: la preferida · $y_l$: la otra · $\sigma$: sigmoide

todos los pares de un ranking, en un solo batch

las demostraciones quedan en 0

En fórmulas. r theta de x e y es el número que el reward model le da a la respuesta y. Para cada par, la preferida y sub w y la otra y sub l. Sigma de la diferencia es la probabilidad de ordenarlas bien, y la pérdida es menos su logaritmo. En el dibujo: bien ordenado, la sigmoide está cerca de uno y la pérdida es chica. Con cuatro respuestas hay seis pares: cuatro por tres sobre dos. El paper los pone todos en el mismo elemento del batch y divide por seis. En el dibujo, el promedio de los seis. La pérdida solo mira diferencias, así que sumarle una constante a r no cambia nada. Por eso el paper resta un sesgo: que las demostraciones valgan cero.

Ouyang et al., 2022 — sección 3.5 · arXiv:2203.02155

media 0 · «since the RM loss is invariant to shifts in reward, we normalize the reward model using a bias so that the labeler demonstrations achieve a mean score of 0 before doing RL.»

Ese detalle está en el paper. Como la pérdida no cambia si se corre la recompensa, normalizan el reward model con un sesgo para que las demostraciones tengan puntaje medio cero antes del RL, el reinforcement learning. Así, un puntaje positivo quiere decir «mejor que una demostración».

Parte 3 · Paso 3: PPO-ptx contra el reward model

El loop del episodio siete, con un cambio: la recompensa ya no es el oro. Es el reward model.

El reward model como recompensa

El loop, función por función. El puntaje es el reward model menos su sesgo. La recompensa le resta beta por el KL, la divergencia de Kullback-Leibler contra el SFT, con el beta del paper, cero coma cero dos. La mezcla de preentrenamiento, con gamma uno. Y cada veinticinco iteraciones anotamos el puntaje, el oro y el KL. El oro no entra en la cuenta: lo miramos para saber la verdad. Con el reward model fuerte: oro cinco coma nueve cinco, KL cero coma siete. Las otras dos filas usan el reward model débil; quedan para la parte cinco.

Cada línea, su término:

• puntaje = rm(y) - sesgo → $r_\theta(x,y) - b$ — el RM juzga; $b$: el sesgo de las demos
• r = puntaje - beta * kl → $r_\theta(x,y) - \beta \log \tfrac{\pi^{RL}_\phi(y\mid x)}{\pi^{SFT}(y\mid x)}$ — $\beta = 0{,}02$, como el paper
• loss = loss + gamma * j.nll_pre(pol, j.batch_pre(64)) → $+\,\gamma\,\mathbb{E}_{x\sim D_{\text{pretrain}}}\big[\log \pi^{RL}_\phi(x)\big]$ — ptx: seguir sabiendo contar
• fila = [puntaje.mean(), j.oro(y).mean(), kl.mean()] → $\big(\, r_\theta,\; \text{oro},\; \mathrm{KL} \,\big)$ — el oro solo se mira; no entrena

def ppo_ptx(rm, sesgo, beta, gamma=1.0, iters=300):
    pol, ref, hist = copy.deepcopy(sft), copy.deepcopy(sft).eval(), []
    opt = torch.optim.AdamW(pol.parameters(), lr=3e-4)
    for it in range(iters):
        y = j.muestrear(pol, j.prompts(64))
        with torch.no_grad():
            puntaje = rm(y) - sesgo                     # paso 2: el RM juzga
            kl = (j.logprobs(pol, y) - j.logprobs(ref, y)).sum(-1)
            r = puntaje - beta * kl
            A = (r - r.mean()) / (r.std() + 1e-8)
        loss = -(A * j.logprobs(pol, y).sum(-1)).mean()
        loss = loss + gamma * j.nll_pre(pol, j.batch_pre(64))
        loss.backward(); opt.step(); opt.zero_grad()
        if it % 25 == 0 or it == iters - 1:            # RM y oro, lado a lado
            fila = [puntaje.mean(), j.oro(y).mean(), kl.mean()]
            hist.append([float(v) for v in fila])
    return pol, hist

El oro durante PPO-ptx

El oro, la métrica que la política nunca ve. Con el reward model fuerte y el beta del paper, sube de cuatro coma tres a casi seis. Los tres pasos funcionan. Con el reward model débil y un freno de KL cinco veces más fuerte, también sube, más despacio. Con el débil y sin freno, el oro se desploma. La línea punteada es el SFT. A esa tercera curva volvemos enseguida.

Parte 4 · Medir con dos jueces

El paper mide con personas. Acá tenemos dos jueces baratos: el reward model, y el oro que generó los datos.

Win rate contra el SFT

Mil prompts de test que ningún paso vio. El win rate del episodio seis, con un detalle: el oro es un entero y hay empates, que cuentan medio. El SFT responde los mismos prompts. Primer juez: el reward model con el que se entrenó esa política. Segundo juez: el oro. Con el reward model fuerte, los dos jueces coinciden: cero coma setenta y cinco y cero coma setenta y dos. Con el débil sin freno, el reward model dice cero coma ochenta y nueve, y el oro, cero coma dos ocho: la política es peor que el SFT. Con freno, cero coma setenta y tres contra cero coma sesenta y cuatro.

def win_rate(a, b):                         # 1 si gana A, ½ si empatan, 0 si pierde
    return ((a > b).float() + 0.5 * (a == b).float()).mean().item()

with torch.no_grad():
    y_sft = j.muestrear(sft, x)             # el baseline responde los mismos prompts
    for nombre, rm_nombre, beta in casos:
        pol = j.cargar(f"ep09_pol_{rm_nombre}_{beta}.pt").eval()
        rm, _ = j.cargar_rm(f"ep09_rm_{rm_nombre}.pt")
        y_pol = j.muestrear(pol, x)
        juez_rm = win_rate(rm(y_pol), rm(y_sft))       # el RM como juez
        juez_oro = win_rate(j.oro(y_pol), j.oro(y_sft))  # el gusto real
        print(f"{nombre:18} juez RM {juez_rm:.2f}   juez oro {juez_oro:.2f}")

RM fuerte, β=0,02  juez RM 0.75   juez oro 0.72
RM débil, β=0      juez RM 0.89   juez oro 0.28
RM débil, β=0,1    juez RM 0.73   juez oro 0.64

Win rate con empates

$s_i$: el puntaje que da el juez en el prompt $i$

RM débil, sin freno de KL

En fórmula. s sub i es el puntaje que el juez le da a cada respuesta en el prompt i. Si la política gana, suma uno; si empata, un medio; si pierde, cero. Y promediamos sobre los n prompts. En el dibujo, seis prompts: tres victorias, dos empates y una derrota dan cero coma seis siete. Lo importante es quién juzga. La misma política, con dos jueces, da cero coma ochenta y nueve y cero coma veintiocho. Si el juez es el mismo reward model que se optimizó, el win rate no sirve como evaluación.

Tres políticas, dos jueces

Las seis barras. Reward model fuerte, juzgado por sí mismo. Y por el oro: casi igual. Reward model débil sin freno: según él mismo, una gran mejora. Según el oro, un desastre. Débil con freno, según él mismo. Y según el oro: mejora real, menor. El juez entrenado se equivoca justo donde la política lo empujó.

• fuerte · RM: 0,75
• fuerte · oro: 0,72
• débil · RM: 0,89
• débil · oro: 0,28
• β 0,1 · RM: 0,73
• β 0,1 · oro: 0,64

Parte 5 · Reward hacking

La política no busca el gusto real. Busca lo que el reward model premia, y si hay un atajo, lo encuentra.

Qué aprendió a responder

Miremos las respuestas. El SFT y la política entrenada con el reward model débil sin freno. Tres prompts, sin azar. Y el puntaje del reward model. El SFT, sin azar, responde perfecto: oro seis, con puntajes entre cero coma tres y cero coma ocho. La política hackeada responde cosas como ocho, cinco, seis, tres, seis, siete: oro cero. Y el reward model les da más puntaje, hasta uno coma tres cuatro. Encontró secuencias que el reward model nunca vio durante su entrenamiento, donde se equivoca a favor.

for nombre in ["ep09_sft.pt", "ep09_pol_débil_0.0.pt"]:      # antes y después
    pol = j.cargar(nombre).eval()
    with torch.no_grad():
        y = j.muestrear(pol, x, greedy=True)
        puntaje = rm(y) - sesgo
    print(nombre.removesuffix(".pt"))
    for fila, p, o in zip(y.tolist(), puntaje, j.oro(y)):
        print(f"  {fila[1:3]} → {fila[3:]}  RM {p:+.2f}  oro {o:.0f}")

ep09_sft
  [3, 1] → [9, 7, 5, 3, 1, 9]  RM +0.31  oro 6
  [3, 2] → [1, 0, 9, 8, 7, 6]  RM +0.74  oro 6
  [6, 3] → [0, 7, 4, 1, 8, 5]  RM +0.84  oro 6
ep09_pol_débil_0.0
  [3, 1] → [8, 5, 6, 3, 6, 7]  RM +0.81  oro 0
  [3, 2] → [5, 6, 7, 4, 5, 4]  RM +1.34  oro 1
  [6, 3] → [4, 5, 4, 5, 4, 7]  RM +1.32  oro 0

Ley de Goodhart, en una corrida

La corrida sin freno, en una figura. El puntaje del reward model sube durante toda la corrida: termina uno coma cinco por encima del SFT. El oro baja casi dos puntos. Es la ley de Goodhart que vimos en Alineación, episodio seis: cuando una medida se vuelve objetivo, deja de medir. El reward model es una aproximación aprendida, y optimizarlo fuerte encuentra sus errores.

Lo que hace el freno de KL

sin freno ($\beta = 0$), al final de la corrida

con β = 0,1, el atajo casi no paga

La cuenta del freno. El KL mide cuánto se alejó la política del SFT, en nats: la unidad del logaritmo natural. Sin freno, la política hackeada termina a trece coma seis nats; las otras dos, a menos de dos. En el dibujo, las tres. Ese atajo le sube el puntaje del reward model en uno coma cinco cuatro. Con beta cero coma uno, alejarse trece coma seis nats le costaría uno coma tres seis. Casi todo lo que gana. El freno no impide aprender: hace que las respuestas raras, donde el reward model no sabe, salgan caras.

Ouyang et al., 2022 — sección 3.5 · arXiv:2203.02155

«over-optimization» · «we add a per-token KL penalty from the SFT model at each token to mitigate over-optimization of the reward model.»

Es la razón que da el paper para ese término. Agregan una penalización de KL por token contra el SFT para mitigar la sobreoptimización del reward model. En Alineación, episodio seis, vimos cómo se mide esa sobreoptimización en grande.

Qué se parece al paper y qué no

• Igual: SFT → RM con K respuestas → PPO con β y γ
• Distinto: el gusto real acá es una regla conocida
• Distinto: sin value function; una pasada por muestra
• Igual: el RM como juez infla el win rate

Antes de cerrar, la honestidad del juguete. La estructura es la del paper: demostraciones, rankings de varias respuestas con todos sus pares, y PPO con freno de KL y mezcla. Pero acá el gusto real es una regla que conocemos, y podemos medirla gratis. En el paper, el único juez es una persona, y es cara. Nuestro PPO no tiene value function, la red que estima cuánto va a valer cada estado, y usa cada muestra una vez. Lo que sí se traslada: evaluar con el mismo reward model que optimizaste te dice lo que querés escuchar.

Resumiendo

• SFT corto: oro 4,38 de 6; la capacidad ya cae
• RM con rankings de 4 y sus 6 pares: 94,5 % de acierto
• Sesgo restado: las demostraciones valen 0
• PPO-ptx con β = 0,02: win rate 0,72 según el oro
• RM débil sin freno: 0,89 según el RM, 0,28 según el oro
• Próximo: lo que vino después, DPO y CAI

Papers del episodio

• Ouyang et al., 2022 — sección 3.5