Paso 3: PPO y PPO-ptx

InstructGPT. Episodio cinco: paso tres, PPO, proximal policy optimization, y PPO-ptx. El loop de aprendizaje por refuerzo con freno de KL, la divergencia de Kullback-Leibler, ya lo escribimos entero en Alineación, episodio cuatro. Hoy, lo que eligió el paper: un entorno de un solo paso, un freno por token, una función de valor que nace del reward model, y un tercer término, el de preentrenamiento, que le da nombre al modelo final.

Las decisiones del paso 3

• Entorno bandit: un prompt, una respuesta, una recompensa
• Freno de KL contra el SFT, en cada token (β = 0,02)
• Función de valor inicializada desde el RM
• PPO-ptx: sumar gradientes de preentrenamiento (γ = 27,8)

Cuatro decisiones. El entorno es un bandit: llega un prompt de un cliente, el modelo responde, el reward model, el RM, puntúa, y el episodio termina. A esa recompensa se le resta una penalidad por token: la KL contra el modelo SFT, el del supervised fine-tuning; en concreto, el SFT con mezcla de preentrenamiento del episodio tres, del que también arranca la política. Beta, el peso del freno, vale cero coma cero dos. La función de valor de PPO arranca como una copia del RM. Y PPO-ptx: en cada paso se suman también gradientes de preentrenamiento, con peso gamma igual a veintisiete coma ocho. El modelo sin esta mezcla se llama PPO, a secas.

Parte 1 · Un bandit con freno

Primero, qué recompensa recibe cada token de una respuesta.

Una iteración de PPO-ptx

Una iteración completa. La parte PPO: un prompt de cliente, pi RL, la política que se entrena con aprendizaje por refuerzo, responde; el RM le pone un número; se le resta el freno de KL por token, y sale el gradiente de PPO. La parte ptx: un pedazo de texto de preentrenamiento, el mismo tipo de texto con que se entrenó GPT tres, y el gradiente de predecirlo, multiplicado por gamma. Los dos gradientes se acumulan en el mismo lugar, y recién ahí el optimizador, Adam, adaptive moment estimation, da un paso.

prompt de cliente → π_RL responde → RM: r(x, y) → − β·KL por token → gradiente PPO → texto de preentreno → + γ · gradiente → paso de Adam

La recompensa de cada token

Una respuesta de cinco tokens, con log-probs sintéticos. Por cada token, el log de la probabilidad según la política menos el log según el SFT. El freno: menos beta por eso, en cada token. El RM puntúa una sola vez, al final: eso es lo que hace que sea un bandit. Y el retorno de cada posición, la suma de las recompensas desde ahí hasta el final, sin descuento. Los frenos son de milésimas, porque beta es chico. El último token carga el uno coma siete del RM. Y el retorno desde el primer token, uno coma sesenta y ocho, es exactamente r menos beta por la suma de los log ratios.

kl_tok = lp_rl - lp_sft                    # log π_RL − log π_SFT, por token
recompensa = -beta * kl_tok                # el freno, en cada token
recompensa[-1] += r_rm                     # el RM puntúa al final: un bandit
retorno = recompensa.flip(0).cumsum(0).flip(0)   # G_t = Σ_{k≥t} r_k, sin descuento

print("kl:", *[f"{v:+.2f}" for v in kl_tok])
print("r: ", *[f"{v:+.3f}" for v in recompensa])
print("G: ", *[f"{v:+.3f}" for v in retorno])
print(f"G_0 = {retorno[0]:.3f}")
print(f"r − β·Σkl = {r_rm - beta * kl_tok.sum():.3f}")

kl: +0.30 +0.20 -0.20 +0.50 +0.20
r:  -0.006 -0.004 +0.004 -0.010 +1.696
G:  +1.680 +1.686 +1.690 +1.686 +1.696
G_0 = 1.680
r − β·Σkl = 1.680

Recompensa por token y retorno

$\beta = 0{,}02$; el RM solo en el último token

retorno: lo que falta cobrar desde t, sin descuento

ojo: este $\gamma$ de GAE no es el $\gamma$ del paper

La recompensa del token t es menos beta por el logaritmo del cociente entre las dos políticas en ese token, más la del RM si t es el último, T mayúscula. En el dibujo, las cinco barras: cuatro casi invisibles y la última, grande. G sub t, el retorno, es la suma desde t hasta el final. Es lo que la función de valor aprende a predecir. Y un aviso de notación. En Alineación cuatro, GAE, generalized advantage estimation, descontaba el futuro con una letra gamma. El paper dice que no aplica descuento, o sea, ese factor vale uno: en el dibujo, cada recompensa futura pesa lo mismo. La gamma del paper es otra cosa: el peso del preentrenamiento, que vemos enseguida.

La función de valor nace del RM

PPO necesita una función de valor: un modelo que, a mitad de una respuesta, estime el retorno que falta. El paper la inicializa desde el RM. Misma arquitectura que el RM: el transformer y una proyección a un número. La diferencia: el RM lee solo el último token, y la función de valor da un número en cada posición. Y copiamos los pesos del RM. Ocho valores, uno por posición. El último coincide con el RM: menos cero coma treinta y uno. Arranca sabiendo puntuar respuestas terminadas; durante PPO aprende a estimar las que están a medio camino. Para todas las políticas, el RM y la función de valor son de seis mil millones.

class FuncionDeValor(nn.Module):              # V(prefijo): un número por posición
    def __init__(self):
        super().__init__()
        self.cuerpo, self.cabeza = GPT2Model(cfg), nn.Linear(64, 1)
    def forward(self, ids):
        h = self.cuerpo(input_ids=ids).last_hidden_state
        return self.cabeza(h).squeeze(-1)     # [B, T]: un valor por token

valor = FuncionDeValor().eval()
valor.load_state_dict(rm.state_dict())        # se inicializa desde el RM
ids = torch.tensor([[3, 7, 1, 1, 9, 4, 2, 5]])
with torch.no_grad():
    v = valor(ids)[0]
    print("V:", *[f"{x:+.2f}" for x in v[:4]])
    print("  ", *[f"{x:+.2f}" for x in v[4:]])
    print(f"RM de la respuesta: {rm(ids).item():+.2f}")

V: +0.55 +0.58 +0.08 -0.11
   +0.40 -0.82 -0.28 -0.31
RM de la respuesta: -0.31

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

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

Por qué el freno. El paper lo dice en una línea: la penalidad de KL por token contra el SFT está para mitigar la sobreoptimización del reward model. El RM es una aproximación de lo que prefieren cuarenta personas, y PPO es muy bueno encontrando respuestas que lo engañan. El freno obliga a quedarse cerca de un modelo que todavía habla bien.

Parte 2 · La ecuación 2

El objetivo completo que maximiza PPO-ptx tiene tres términos. Los calculamos con números.

El objetivo de PPO-ptx, término por término

$\phi$: pesos de la política · respuestas de la propia política

el freno: cuánto se alejó del SFT

$\gamma$: peso del preentrenamiento · sin prompt

La ecuación dos, de a un término. Phi son los pesos de la política. El primer término es la recompensa promedio, con x e y sacados de la propia política: prompts de clientes y respuestas que genera pi RL. En el dibujo, el RM de cuatro respuestas del snippet. Dentro del mismo promedio, el freno: beta por el log del cociente, que para una respuesta entera es la suma de los log ratios de sus tokens. En el dibujo, los de la primera respuesta: suman cero coma veinticuatro. Y el término nuevo: gamma por el log de la probabilidad que la política le da a texto de preentrenamiento. x acá es un documento entero, sin prompt. Es exactamente menos la pérdida de preentrenar, y suma menos diez coma sesenta y uno.

Los tres términos, en código

Con la política que vamos a entrenar en un rato. Beta acá es cero coma cero cinco y gamma, dos: valores del juguete. r, una por respuesta. El log ratio de cada respuesta, sumando sobre sus tokens. El log de la probabilidad del texto de preentrenamiento. Y el objetivo: el promedio de r menos beta por el log ratio, más gamma por ese log. Las cuatro recompensas, de uno coma treinta y cinco a cuatro coma dieciocho. Los log ratios, positivos: la política ya se alejó un poco del SFT. Y el objetivo: dos coma veintisiete, más dos por menos diez coma sesenta y uno. El término de preentrenamiento domina la suma, pero lo que importa no es su valor sino su gradiente.

Cada línea, su término:

• r = rm(y[:, P:]) → $r_\theta(x, y)$ — una recompensa por respuesta
• log_ratio = (logp_tokens(pol, y, P) → $\log \pi^{\text{RL}}(y|x) - \log \pi^{\text{SFT}}(y|x)$ — suma sobre los tokens de y
• log_pre = logp_tokens(pol, x_pre, 1).sum() → $\log \pi^{\text{RL}}(x)$ — x: texto de preentrenamiento
• objetivo = (r - beta * log_ratio).mean() → $\overline{r - \beta \log\text{ratio}} + \gamma \log \pi(x)$ — lo que PPO-ptx maximiza

with torch.no_grad():
    r = rm(y[:, P:])                                   # r_θ(x, y), una por respuesta
    log_ratio = (logp_tokens(pol, y, P)
                 - logp_tokens(ref, y, P)).sum(-1)     # log π_RL(y|x) − log π_SFT(y|x)
    log_pre = logp_tokens(pol, x_pre, 1).sum()         # log π_RL(x), texto de preentreno
    objetivo = (r - beta * log_ratio).mean() + gamma * log_pre
print("r:         ", [f"{v:+.2f}" for v in r])
print("log ratio: ", [f"{v:+.2f}" for v in log_ratio])
print(f"log π(x_pre) = {log_pre:.2f}")
print(f"objetivo = {(r - beta * log_ratio).mean():.2f} + {gamma}·({log_pre:.2f})"
      f" = {objetivo:.2f}")

r:          ['+2.05', '+1.35', '+4.18', '+1.63']
log ratio:  ['+0.24', '+0.62', '+1.06', '+0.68']
log π(x_pre) = -10.61
objetivo = 2.27 + 2.0·(-10.61) = -18.95

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

PPO-ptx · «For "PPO" models, γ is set to 0. Unless otherwise specified, in this paper InstructGPT refers to the PPO-ptx models.»

El nombre. Si gamma vale cero, el modelo se llama PPO. Si no, PPO-ptx. Y el paper aclara: cuando dice InstructGPT, sin más, se refiere a los PPO-ptx. El tercer término no es un detalle: es el modelo que se publicó.

Parte 3 · Dos gradientes, un buffer

Cómo se implementa el tercer término: no se suman pérdidas en una sola línea, se acumulan gradientes.

Gradientes que se acumulan

$\nabla_\phi$: el gradiente respecto de los pesos

la pérdida de preentrenar, sobre texto nuevo

Nabla sub phi es el gradiente respecto de los pesos: un número por peso. El apéndice C punto cuatro lo describe así: para cada minibatch, calculan el gradiente de PPO y el de preentrenamiento en pasos consecutivos, y los acumulan en el mismo buffer, multiplicando el segundo por gamma. En el dibujo, dos columnas de gradientes, la segunda pesada, sumadas en una. Y L ptx es la pérdida de siempre del preentrenamiento: menos el log de cada token dado los anteriores, promediado. Maximizar el tercer término es minimizar esto. En el dibujo, los mismos tokens de antes, ahora con signo cambiado.

Los números del apéndice C.4

Doscientos cincuenta y seis mil episodios: respuestas generadas, sobre unos treinta y un mil prompts. Batches de quinientas doce, partidos en minibatches de sesenta y cuatro. Beta cero coma cero dos, gamma veintisiete coma ocho. Y ocho veces más ejemplos de preentrenamiento que episodios. Hacemos las cuentas. Quinientas iteraciones, de ocho minibatches. Dos millones cuarenta y ocho mil ejemplos de preentrenamiento. Si se reparten parejo, quinientos doce por cada minibatch de sesenta y cuatro respuestas: eso es suposición nuestra. Una sola pasada de PPO por batch, recorte de cero coma dos, y un exponential moving average de los pesos, un promedio móvil exponencial, con decaimiento cero coma novecientos noventa y dos.

episodios = 256_000          # respuestas generadas en todo el entrenamiento
batch, minibatch = 512, 64
beta, gamma = 0.02, 27.8     # freno de KL y peso del preentrenamiento
ptx_por_episodio = 8         # 8 veces más ejemplos de preentrenamiento

iteraciones = episodios // batch
print("iteraciones:", iteraciones)
print("minibatches por iteración:", batch // minibatch)
print("ejemplos de ptx:", f"{ptx_por_episodio * episodios:,}")
print("parejo: 64 respuestas +", ptx_por_episodio * minibatch, "ptx")
print("1 pasada | clip 0,2 | EMA 0,992")

iteraciones: 500
minibatches por iteración: 8
ejemplos de ptx: 2,048,000
parejo: 64 respuestas + 512 ptx
1 pasada | clip 0,2 | EMA 0,992

Parte 4 · El efecto de γ, en chico

Un PPO-ptx de juguete donde el reward model y el lenguaje tiran para lados distintos.

PPO-ptx de juguete

El juguete. El preentrenamiento es una cadena de Markov: cada token tiene dos sucesores probables. El RM premia tokens sueltos, sin respetar ese lenguaje. La política responde a prompts del mismo lenguaje. Recompensa menos beta por la KL. Simplificamos: sin función de valor, la ventaja es la recompensa normalizada en el batch. El gradiente de PPO, minibatch por minibatch. Y enseguida, gamma por el gradiente de preentrenamiento, al mismo buffer. CE es la entropía cruzada en texto de preentrenamiento nuevo. Con gamma cero, el RM llega a ocho coma dos, casi su máximo, pero la CE pasa de uno coma trece a siete coma dieciocho. Con gamma dos, se queda en uno coma dieciséis, y el RM en uno coma cinco.

def entrenar(gamma, beta=0.05, iters=60):
    pol, opt = nueva_politica()                       # parte del modelo base
    for it in range(iters):
        y = muestrear(pol, texto(64, P))              # 1. la política responde
        with torch.no_grad():
            lp_old, lp_ref = logprobs(pol, y), logprobs(ref, y)
            R = rm(y[:, P:]) - beta * (lp_old - lp_ref).sum(-1)  # 2. r − β·KL
            A = (R - R.mean()) / (R.std() + 1e-8)     # 3. ventaja, sin crítico
        for mb in torch.arange(64).split(32):         # 4. minibatches, 1 pasada
            perdida_ppo(logprobs(pol, y[mb]), lp_old[mb], A[mb]).backward()
            z = texto(64)                             # texto de preentrenamiento
            (gamma * pol(input_ids=z, labels=z).loss).backward()  # + γ·∇ptx
            opt.step(); opt.zero_grad()
        anotar(gamma, it, y, pol)
    return pol

print(f"antes de RL:  CE {ce_preentreno(base):.2f}")
for gamma in [0, 0.5, 2]:
    pol = entrenar(gamma)
    _, r, ce = historia[gamma][-1]
    print(f"γ = {gamma:<3}  RM {r:+.1f}  CE {ce:.2f}")

antes de RL:  CE 1.13
γ = 0    RM +8.2  CE 7.18
γ = 0.5  RM +5.2  CE 1.39
γ = 2    RM +1.5  CE 1.16

Lo que se olvida del preentrenamiento

La pérdida en texto de preentrenamiento a lo largo de las sesenta iteraciones. Con gamma cero, en rojo, se dispara: la política se olvida del lenguaje que sabía. Con gamma cero coma cinco, sube un poco. Con gamma dos, casi no se mueve. Esa subida es nuestra versión de juguete de lo que el paper llama alignment tax, el impuesto de alineación, que medimos de verdad en el episodio siete.

Y lo que cuesta en recompensa

La otra cara. Sin gamma, la recompensa sube hasta el techo del RM. Con gamma cero coma cinco, se queda en el medio. Con gamma dos, sube poco. En el juguete el conflicto es total, porque el RM premia justamente romper el lenguaje. El paper reporta algo más suave: subir gamma baja la recompensa de validación, pero la preferencia humana casi no cambió.

Por qué no alcanza con subir β

• Subir β no arregla las regresiones (ni ×100)
• β óptimo: alrededor de 0,01 y 0,02
• Un solo γ = 27,8 para 1,3B, 6B y 175B
• La preferencia humana, casi insensible a γ

Una pregunta natural: si el problema es alejarse, ¿por qué no frenar más fuerte con beta? El apéndice E lo probó. Con gamma en cero, subieron beta hasta dos, cien veces el valor por defecto, y las regresiones en benchmarks públicos no se arreglaron, mientras la recompensa caía. El mejor beta para la preferencia humana anda entre cero coma cero uno y cero coma cero dos. Y un único gamma, veintisiete coma ocho, funcionó para los tres tamaños. La preferencia humana casi no se movió con gamma. Una forma de verlo: la KL te mantiene cerca solo en las respuestas que generás; el preentrenamiento cuida todo lo demás.

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

γ = 27,8 · «We multiply the pretraining gradients by a coefficient, γ = 27.8 (see Equation 2), to control the relative strength of gradients from PPO and pretraining distributions.»

El número final. Gamma, veintisiete coma ocho, multiplica los gradientes de preentrenamiento para balancear su fuerza contra los de PPO. Ojo, no compares con nuestros cero coma cinco y dos: depende de la escala de cada pérdida, y nuestro juguete no tiene nada que ver con GPT tres.

Resumiendo

• Bandit: el RM puntúa al final; freno de KL por token, β = 0,02
• Sin descuento; función de valor copiada del RM (6B)
• Ecuación 2: $r - \beta \log(\pi^{RL}/\pi^{SFT}) + \gamma \log \pi^{RL}(x)$
• Dos gradientes al mismo buffer: PPO y γ·preentreno
• Juguete: γ = 0 → CE 7,18; γ = 2 → 1,16
• Próximo: cómo se evaluó (episodio 6)

Papers del episodio

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