RLHF con PPO

Alineación a bajo nivel. Episodio cuatro: RLHF con PPO. Es decir, reinforcement learning from human feedback, con proximal policy optimization. Ya tenemos un modelo de SFT, supervised fine-tuning, que contesta y un modelo de recompensa que puntúa respuestas. Hoy los juntamos: vamos a ajustar la política para que saque más puntaje, con el algoritmo que usó InstructGPT. Y vamos a ver por qué el ingrediente más importante no es la recompensa, sino el freno.

Por qué hace falta aprendizaje por refuerzo

• La recompensa llega por respuesta entera, no por token
• No hay respuesta correcta para imitar: hay que probar
• Generar es muestrear: no se puede derivar a través del muestreo
• El RM es un proxy: si lo exprimís, se rompe

¿Por qué no seguir con entropía cruzada? Porque el modelo de recompensa puntúa la respuesta entera, no cada token. No hay una respuesta correcta para copiar: la política tiene que generar, recibir un puntaje y aprender de eso. Y generar es muestrear tokens, una operación discreta: el gradiente no pasa a través de un sorteo. Eso es exactamente el problema que resuelve el aprendizaje por refuerzo. Pero hay una trampa que vimos al final del episodio tres: el modelo de recompensa es un proxy. Si lo optimizás sin límite, encontrás sus agujeros.

Parte 1 · El objetivo

Antes del algoritmo, la cuenta que queremos maximizar. Tiene dos términos, y el segundo es el importante.

Recompensa menos distancia

$\pi_{\text{ref}}$: el modelo de SFT, congelado

cuánto se alejó la distribución

la de la respuesta es la suma de los tokens

El objetivo de RLHF es este. Maximizar la recompensa esperada de las respuestas que genera la política, pi, menos beta por la divergencia KL, de Kullback-Leibler, entre la política y el modelo de referencia, pi sub ref, que es el de SFT congelado. Beta es la perilla: cuánto te importa no alejarte. La KL es el valor esperado del logaritmo del cociente de probabilidades, sobre respuestas que genera la política. Y como la probabilidad de una respuesta es el producto de las de sus tokens, su logaritmo es una suma. Eso nos va a permitir calcular la KL token por token. En el dibujo, una respuesta real: el log cociente de cada token, uno es negativo, y la suma da casi diez.

Las piezas de RLHF

Este es el mapa del episodio. Tomamos prompts y la política genera respuestas. El modelo de recompensa las puntúa, el de referencia calcula sus log probabilidades, y combinamos las dos cosas en una recompensa penalizada. Un cuarto modelo, el crítico, estima cuánto puntaje se espera desde cada token, y con eso calculamos ventajas con GAE, generalized advantage estimation: qué tokens fueron mejores de lo esperado. Y PPO da un paso, con cuidado de no moverse demasiado. Cuatro modelos en memoria: política, referencia, recompensa y crítico. En InstructGPT, la recompensa y el crítico eran de seis mil millones de parámetros cada uno.

prompts → política π genera → RM puntúa → π_ref: log-probs → recompensa − β·KL → crítico V → ventajas (GAE) → paso de PPO

Parte 2 · La KL, token por token

El freno se calcula con las mismas log-probabilidades que ya sabés sacar de los logits.

Log-probs y KL por token

Para este episodio la política es un modelo de juguete al que le movimos un poco los pesos respecto de la referencia. Calculamos log softmax sobre los logits: la distribución completa en cada posición. Con gather nos quedamos con el log-prob del token que efectivamente salió. La diferencia entre la política y la referencia, en cada token, es una estimación de la KL con una sola muestra. Sumada y promediada sobre muchas respuestas, converge a la KL de verdad. En la primera secuencia hay un token negativo: un token suelto puede dar menos que cero, aunque la KL nunca lo es. Pasa en el trece por ciento de los tokens. Pero el promedio da ocho coma tres nueve, contra ocho coma tres seis de la KL exacta.

Cada línea, su término:

• log_softmax(-1) → $\log \pi(\cdot \mid x, y_{<t})$ — toda la distribución en cada posición
• gather → $\log \pi(y_t \mid x, y_{<t})$ — nos quedamos con el token que salió
• kl_tok = → $k_t = \log \pi(y_t) - \log \pi_{\text{ref}}(y_t)$ — estimador de la KL con una muestra
• kl_tok.sum(-1).mean() → $\mathbb{E}_{y\sim\pi}\big[\textstyle\sum_t k_t\big] = \mathrm{KL}$ — en promedio, da la KL exacta

def logprobs(model, ids):              # log π(y_t | x, y_<t) por token
    lp = model(input_ids=ids).logits[:, :-1].log_softmax(-1)
    lp = lp.gather(-1, ids[:, 1:, None]).squeeze(-1)  # el que salió
    return lp[:, P - 1:]                   # solo la respuesta

y = muestrear(pol, prompts(2000))                    # respuestas de la política
with torch.no_grad():
    kl_tok = logprobs(pol, y) - logprobs(ref, y)     # log π − log π_ref
print("seq 0, por token:", kl_tok[0].numpy().round(2))
print(f"tokens negativos: {(kl_tok < 0).float().mean():.0%}")
print(f"estimador (promedio): {kl_tok.sum(-1).mean():.3f}")

seq 0, por token: [ 1.14  2.18  1.89  2.66 -0.15  2.27]
tokens negativos: 13%
estimador (promedio): 8.389

La recompensa con el freno adentro

Ahora metemos el freno en la recompensa. Cada token recibe menos beta por su estimación de KL: si la política le dio a ese token mucha más probabilidad que la referencia, paga. Y el puntaje del modelo de recompensa se suma solo en el último token, porque recién ahí la respuesta está completa. La suma de todo es exactamente el objetivo de la escena anterior, para una respuesta. Con beta cero coma uno, los tres primeros tokens reciben recompensas chiquitas, positivas o negativas, y el último uno coma diecinueve. Total: uno coma ciento cincuenta y siete, el uno coma dos del RM menos cero coma cero cuarenta y tres de freno.

Cada línea, su término:

• r = -beta * kl_tok → $r_t = -\beta\, k_t$ — cada token paga por alejarse
• r[:, -1] += score → $r_T \mathrel{+}= r_{\text{RM}}(x, y)$ — el RM puntúa al final
• print(f"total → $\textstyle\sum_t r_t = r_{\text{RM}} - \beta \sum_t k_t$ — la cuenta del objetivo, por respuesta

def recompensa(score, kl_tok, beta=0.1):
    r = -beta * kl_tok                  # penalización en cada token
    r[:, -1] += score                   # el RM puntúa al final
    return r

r = recompensa(score, kl_tok)
print("r por token:", r.numpy().round(3))
print(f"total: {r.sum():.3f} = {score.item():.1f} − 0.1 · {kl_tok.sum():.2f}")

r por token: [[-0.005  0.002 -0.03   1.19 ]]
total: 1.157 = 1.2 − 0.1 · 0.43

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

β = 0,02 · InstructGPT agrega una penalización KL por token respecto del modelo de SFT para mitigar la sobreoptimización del RM, con β = 0,02; ratio de recorte de PPO 0,2.

Esto es lo que hizo InstructGPT. Una penalización de KL por token respecto del modelo de SFT, con la razón explícita de mitigar la sobreoptimización del modelo de recompensa. Usaron beta igual a cero coma cero dos, y un recorte de PPO de cero coma dos, que vamos a ver en un momento.

Parte 3 · Repartir el crédito

La recompensa llega al final: ¿qué token tuvo la culpa? Para eso está el crítico.

Ventajas con GAE

sorpresa de un paso: lo que pasó menos lo esperado

suma de sorpresas futuras, descontadas

el blanco para entrenar al crítico

El crítico, V, es un value head: un número, como el reward head del episodio tres, pero en cada posición: estima cuánto puntaje total se espera desde ese punto de la respuesta. La sorpresa de un paso, delta, es la recompensa que llegó más lo que se espera desde el siguiente estado, menos lo que se esperaba antes. La ventaja estimada, con GAE, acumula las sorpresas futuras con un descuento gama por lambda. En el dibujo, el primer token: cuatro sorpresas, cada una más descontada, suman cero coma siete. Lambda negocia entre ruido y sesgo. InstructGPT usa gama igual a uno, sin descuento. Y la ventaja más el valor da el retorno, que es el blanco con el que se entrena al crítico.

GAE en ocho líneas

En código. El cálculo va de atrás para adelante, porque cada ventaja depende de la siguiente. La sorpresa de un paso. La recurrencia: la ventaja de este token es su sorpresa más gama lambda por la ventaja del siguiente. Y devolvemos ventajas y retornos. Usamos las recompensas de la escena anterior y un crítico que esperaba poco al principio. El primer token tiene ventaja cero coma siete: el crítico esperaba cero coma cuatro y la respuesta terminó valiendo bastante más. El último, cero coma diecinueve: ahí el crítico ya esperaba casi todo. La recompensa del final se repartió hacia atrás.

def gae(r, v, gamma=1.0, lam=0.95):
    A, ultimo, T = torch.zeros_like(r), 0.0, r.shape[1]
    for t in reversed(range(T)):                     # de atrás para adelante
        v_sig = v[:, t + 1] if t + 1 < T else 0.0
        delta = r[:, t] + gamma * v_sig - v[:, t]    # δ_t: sorpresa de un paso
        ultimo = delta + gamma * lam * ultimo        # A_t = δ_t + γλ A_{t+1}
        A[:, t] = ultimo
    return A, A + v                                  # ventajas y retornos

A, R = gae(r, v)
print("ventajas A:", A.numpy().round(3))
print("retornos R:", R.numpy().round(3))

ventajas A: [[0.703 0.64  0.251 0.19 ]]
retornos R: [[1.103 1.14  1.15  1.19 ]]

Parte 4 · El paso recortado de PPO

Con las ventajas ya sabemos qué tokens reforzar. PPO decide cuánto, sin dar pasos demasiado largos.

Gradiente de política y recorte

subir la log-prob de lo que salió mejor de lo esperado

cuánto cambió la política desde que generó

$\epsilon = 0{,}2$: el ratio vive en $[0{,}8;\,1{,}2]$

La base es el gradiente de política. Multiplicás el gradiente de la log probabilidad de cada token por su ventaja: si salió mejor de lo esperado, lo hacés más probable. Así se esquiva el muestreo: no derivás a través del sorteo, derivás la probabilidad de lo que salió. PPO reutiliza cada lote para varios pasos, y para eso mide cuánto cambió la política desde que generó: el ratio rho, nueva sobre vieja. Y la pérdida toma el mínimo entre el término normal y uno con el ratio recortado entre cero coma ocho y uno coma dos. Una vez que un token ya cambió lo suficiente en la dirección buena, deja de recibir empuje.

El recorte, caso por caso

Veámoslo con seis tokens. El ratio es la exponencial de la diferencia de log-probs. La versión recortada. El mínimo de las dos, con el signo cambiado. Y miramos qué tokens reciben gradiente. Con ventaja positiva y ratio uno coma cinco, no empuja: ese token ya subió de más. Con ventaja negativa y ratio cero coma cinco, tampoco: ya bajó bastante. Pero con ventaja negativa y ratio uno coma cinco, sí empuja: la política subió un token que era malo, y el mínimo elige el término sin recortar para corregirlo. El recorte frena las mejoras, nunca los errores.

def perdida_ppo(lp, lp_old, A, eps=0.2):
    ratio = (lp - lp_old).exp()                       # π / π_old
    s1 = ratio * A
    s2 = ratio.clamp(1 - eps, 1 + eps) * A            # recortado a [0.8, 1.2]
    return -torch.min(s1, s2).mean(), ratio

loss, ratio = perdida_ppo(lp, lp_old, A)
loss.backward()
print("ratio:   ", ratio.detach().numpy().round(2))
print("ventaja: ", A.numpy())
print("¿empuja?:", (lp.grad != 0).numpy())

ratio:    [0.5 1.  1.5 0.5 1.  1.5]
ventaja:  [ 1.  1.  1. -1. -1. -1.]
¿empuja?: [ True  True False False  True  True]

Schulman et al., 2017 — Proximal Policy Optimization Algorithms · arXiv:1707.06347

0,82 vs. −0,39 · En su benchmark de control continuo (7 entornos, 21 corridas), recortar con $\epsilon = 0{,}2$ dio un puntaje normalizado promedio de 0,82; sin recorte ni penalización, −0,39.

El paper de PPO comparó variantes en tareas de control continuo, con puntajes normalizados donde cero es la política al azar. Recortando con épsilon cero coma dos, el promedio fue cero coma ochenta y dos, el mejor de su tabla. Sin recorte ni penalización, menos cero coma treinta y nueve: peor que no entrenar. Reusar datos sin freno rompe la política.

Parte 5 · El loop completo

Ya tenemos todas las piezas. Juntémoslas y miremos qué hace beta.

RLHF con PPO, de punta a punta

Este es el loop entero, con beta cero coma uno. La política genera respuestas de seis tokens. El modelo de recompensa las puntúa; acá, un gusto oculto como el del episodio tres. Restamos beta por la KL con la referencia. El crítico da las ventajas, que además normalizamos, una práctica habitual. Cuatro pasadas sobre el mismo lote. Y la pérdida suma la de PPO con el error cuadrático del crítico. El ratio de la primera pasada da exactamente uno: política vieja y nueva son la misma. Ese es el chequeo de cordura. La recompensa sube de menos cero coma tres cuatro a cuatro coma nueve uno, cerca del máximo de cinco coma dos. Y la KL, de cero a once.

for it in range(61):
    y = muestrear(pol, prompts(64))                  # 1. la política genera
    with torch.no_grad():
        lp_old, lp_ref = logprobs(pol, y), logprobs(ref, y)
        score = rm(y[:, P:])                          # 2. el RM puntúa
        r = recompensa(score, lp_old - lp_ref, beta)  # 3. menos β·KL
        A, R = gae(r, critico(y))                     # 4. ventajas
        A = (A - A.mean()) / (A.std() + 1e-8)
    for epoca in range(4):                            # 5. varias pasadas
        lp = logprobs(pol, y)
        if it == 0 and epoca == 0:
            print("ratio inicial:", (lp - lp_old).exp().mean().item())
        loss = perdida_ppo(lp, lp_old, A) + 0.5 * F.mse_loss(critico(y), R)
        loss.backward(); opt.step(); opt.zero_grad()
    if it % 20 == 0:
        kl = (lp_old - lp_ref).sum(-1).mean()
        print(f"it {it:2d}  RM {score.mean():+.2f}  KL {kl:.2f}")

ratio inicial: 1.0
it  0  RM -0.34  KL 0.00
it 20  RM +4.88  KL 10.99
it 40  RM +4.94  KL 11.76
it 60  RM +4.91  KL 11.28

La perilla β

Ahora corremos el mismo loop con cuatro valores de beta. De cero, sin freno, a uno. Para cada uno, medimos la recompensa media, la KL, y mostramos una respuesta de ejemplo. Con beta cero la recompensa es la máxima, cinco coma dos, y la respuesta es el token trece repetido seis veces: el preferido del modelo de recompensa. Es un colapso. Con cero coma uno, lo mismo, con algo menos de KL. Con cero coma tres, la respuesta vuelve a variar y la KL cae a cinco. Con beta uno, casi no se movió.

for beta in [0.0, 0.1, 0.3, 1.0]:
    score, kl, ejemplo = rlhf(beta)
    print(f"β={beta:<4} RM {score:+.2f}  KL {kl:5.2f}  {ejemplo}")

β=0.0  RM +5.20  KL 17.22  [13, 13, 13, 13, 13, 13]
β=0.1  RM +4.93  KL 11.60  [13, 13, 13, 13, 13, 13]
β=0.3  RM +3.74  KL  5.23  [2, 13, 5, 4, 13, 4]
β=1.0  RM +1.63  KL  1.09  [2, 13, 8, 2, 0, 4]

Cuánta recompensa compra cada unidad de KL

Puesto en un gráfico. Cada punto es un beta: a menos freno, más KL y más recompensa, pero con rendimientos decrecientes. Y la recompensa se pega contra el máximo del modelo de recompensa justo cuando la política colapsa. En un modelo real, ese máximo no es una respuesta perfecta: es un agujero del proxy. El paper de Anthropic sobre RLHF reporta una relación aproximadamente lineal entre la recompensa y la raíz cuadrada de la KL; en el episodio seis vemos qué pasa cuando se sigue empujando.

Un punto por β (1; 0,3; 0,1; 0). Punteada: máximo del RM, 5,20. Salida real de ep04_beta.

Lo que cuesta RLHF con PPO

• Cuatro modelos en memoria: π, π_ref, RM y crítico
• Muchos hiperparámetros: β, ε, λ, épocas, tasa
• Sin β, colapsa sobre el agujero del RM
• ¿Hace falta el RM explícito? (episodio 5)

Antes de cerrar, el costo. Cuatro modelos en memoria a la vez. Una lista larga de hiperparámetros que interactúan: beta, épsilon, lambda, épocas, tasa de aprendizaje. Y cada iteración necesita generar, que es lo más lento. Sin el freno de la KL, la política encuentra el agujero del modelo de recompensa y se queda ahí. La pregunta natural es si todo este aparato hace falta. Resulta que el objetivo con KL tiene una solución en forma cerrada, y de ahí sale DPO, direct preference optimization, que entrena directo sobre las preferencias, sin modelo de recompensa ni crítico.

Lo que te llevás

• Objetivo: $\mathbb{E}[r] - \beta\,\mathrm{KL}(\pi\|\pi_{\text{ref}})$
• KL por token: $\log\pi - \log\pi_{\text{ref}}$ del token muestreado
• Recompensa: $-\beta k_t$ en cada token, más el RM al final
• GAE: sorpresas $\delta_t$ acumuladas hacia atrás
• PPO: ratio recortado a $[1-\epsilon, 1+\epsilon]$
• Próximo: DPO y la constitución

Papers del episodio

• Ouyang et al., 2022 — Training language models to follow instructions with human feedback
• Schulman et al., 2017 — Proximal Policy Optimization Algorithms