El impuesto de alineación

InstructGPT. Episodio siete: el impuesto de alineación. En el episodio cinco agregamos un término raro al objetivo de PPO, proximal policy optimization: un poco de preentrenamiento mezclado con el aprendizaje por refuerzo. Hoy vemos por qué hizo falta: qué se rompe sin esa mezcla, cuánto arregla, y lo reproducimos en un modelo de juguete.

Parte 1 · Qué se pierde

Alinear con preferencias mejora lo que la gente pide. Pero el modelo también sabía hacer otras cosas.

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

alignment tax · «This is an example of an “alignment tax” since our alignment procedure comes at the cost of lower performance on certain tasks that we may care about.»

El nombre está en el paper. Durante el RLHF, reinforcement learning from human feedback, el modelo empeora en datasets públicos que antes resolvía. Lo llaman impuesto de alineación: alinear cuesta rendimiento en tareas que también nos importan. Y explican por qué preocupa: si alinear cuesta caro, hay un incentivo para usar el modelo sin alinear.

Cuatro benchmarks, cuatro modelos de 175B

Los números, de la Tabla catorce, todos en ciento setenta y cinco mil millones y sin ejemplos en el prompt. DROP: preguntas que piden contar o restar sobre un párrafo. F1 mide cuántas palabras de la respuesta coinciden con la correcta. SQuAD: la respuesta está escrita en el texto. HellaSwag: elegir el final más sensato de una frase. Y traducir del francés al inglés, con BLEU, bilingual evaluation understudy, que cuenta grupos de palabras en común con una traducción de referencia. En las cuatro filas, PPO queda por debajo de GPT tres. Y fijate en DROP: el SFT, el supervised fine-tuning, ya había perdido mucho.

Tabla 14 del paper, zero-shot

Cuánto cobra el impuesto, y cuánto devuelve la mezcla

Dos cuentas con esos números. La tabla, copiada del paper. El impuesto: cuánto pierde PPO respecto de GPT tres, en proporción. Y qué parte de esa pérdida devuelve PPO-ptx. DROP pierde la mitad, y la mezcla devuelve solo el quince por ciento. SQuAD pierde un tercio y recupera casi ocho décimos. HellaSwag casi no perdía, y PPO-ptx termina por encima de GPT tres: recupera más del cien por ciento. Traducción: recupera dos tercios. La mezcla ayuda mucho, pero no en todo.

tabla14 = {"DROP (F1)":      [27.53, 15.79, 13.08, 15.23],
           "SQuADv2 (F1)":   [64.30, 57.67, 43.68, 59.85],
           "HellaSwag (acc)": [0.781, 0.753, 0.743, 0.807],
           "FR→EN (BLEU)":   [38.92, 36.90, 24.16, 34.28]}
for nombre, (gpt, sft, ppo, ptx) in tabla14.items():
    impuesto = (gpt - ppo) / gpt             # cuánto pierde PPO, relativo a GPT-3
    recupera = (ptx - ppo) / (gpt - ppo)     # qué parte devuelve la mezcla
    print(f"{nombre:16} PPO {-impuesto:+.0%}   ptx recupera {recupera:.0%}")

DROP (F1)        PPO -52%   ptx recupera 15%
SQuADv2 (F1)     PPO -32%   ptx recupera 78%
HellaSwag (acc)  PPO -5%   ptx recupera 168%
FR→EN (BLEU)     PPO -38%   ptx recupera 69%

El impuesto y lo recuperado

$m_k$: la métrica del benchmark $k$

1 = vuelve a GPT-3 · más de 1 = lo supera

En fórmulas. m sub k es la métrica del benchmark k: F1, acierto o BLEU. El impuesto es la caída de PPO dividida por lo que tenía GPT tres. En el dibujo, DROP: de veintisiete coma cinco a trece. Lo recuperado compara dos distancias: cuánto subió PPO-ptx respecto de PPO, sobre cuánto había bajado PPO. Si da uno, volvió al nivel de GPT tres; más de uno, lo superó. En el dibujo, los cuatro benchmarks: solo HellaSwag pasa el uno.

Por qué un modelo alineado puede olvidar

• PPO solo ve prompts de la API
• Los pesos son los mismos para todas las tareas
• La recompensa no premia lo que no pregunta
• El freno de KL mira solo donde la política responde

Antes del experimento, las hipótesis. Durante PPO, el modelo solo ve prompts que clientes mandaron por la API, application programming interface: treinta y un mil, ninguno de DROP ni de SQuAD. Pero los pesos que cambian son los mismos que resolvían esas tareas. Cada paso de gradiente mueve todo el modelo, no solo la parte que responde a la API. La recompensa no dice nada de lo que no pregunta: si saber restar sobre un párrafo se degrada, ningún término lo nota. Y el freno de KL, la divergencia de Kullback-Leibler, que en teoría impide alejarse, se mide sobre las mismas respuestas. Ahora lo comprobamos.

Parte 2 · Un experimento de juguete

Para entender por qué pasa, lo reproducimos en chico: un modelo que sabe contar, y un pedido que no tiene nada que ver.

El mundo de juguete

El experimento. Un transformer chico, de dos capas, se preentrena con secuencias que cuentan para adelante: tres, cinco, siete, nueve, uno, con un paso fijo y dando la vuelta en diez. Su capacidad, nuestro benchmark, es acertar el próximo número de secuencias nuevas. El pedido del usuario: después de una marca INST y dos números, seguir contando para atrás. El SFT lo aprende con demostraciones y, como en el paper, con un diez por ciento de preentrenamiento mezclado. Después, aprendizaje por refuerzo en tres variantes. La recompensa es el oro: cuántos de los seis tokens de la respuesta cuentan bien para atrás.

preentrenar → 3 5 7 9 1 … → SFT + 10 % mezcla → INST 7 5 → 3 1 9 → RL: tres variantes → oro y capacidad

Preentrenar y hacer el SFT

El modelo preentrenado viene de antes. Su capacidad. El SFT: batches de demostraciones, la instrucción y la cuenta para atrás. La pérdida de siempre, la entropía cruzada, solo sobre los tokens de la respuesta. Y la mezcla: un décimo de la pérdida de preentrenamiento sobre secuencias que cuentan para adelante. Medimos el oro del SFT. El preentrenado acierta todo. El SFT logra cinco coma tres tres de seis en el pedido nuevo, y su capacidad ya bajó a cero coma ocho uno. Como DROP en la tabla: el SFT ya cobra parte del impuesto.

print(f"preentrenado: capacidad {j.capacidad(base):.3f}")    # benchmark: contar
sft = copy.deepcopy(base)
opt = torch.optim.AdamW(sft.parameters(), lr=1e-3)
for paso in range(60):
    x = j.demostraciones(32)                    # [INST, a, a−s] + contar 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))   # 10 % de preentrenamiento
    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}")

preentrenado: capacidad 1.000
SFT: oro 5.33 de 6, capacidad 0.812

Tres variantes del mismo aprendizaje por refuerzo

El loop, en una función. La política genera respuestas a sesenta y cuatro prompts. La recompensa es el oro menos beta por el KL: cuánto se alejó del SFT en esas respuestas. La ventaja A es esa recompensa centrada. Como cada muestra se usa una sola vez, el cociente de PPO vale uno y el recorte no actúa: queda un policy gradient. El PPO completo está en el episodio cinco. Y la línea de ptx: gamma por la pérdida de preentrenamiento. La corremos tres veces. Las tres variantes llegan a casi el mismo oro, cinco coma siete. La capacidad no: PPO, cero coma tres cuatro. Con la mezcla, cero coma nueve nueve seis. Y con beta veinticinco veces más grande, cero coma dos ocho.

Cada línea, su término:

• y = j.muestrear(pol, j.prompts(64)) → $y \sim \pi^{RL}_\phi(\cdot \mid x)$ — la política responde prompts con INST
• r = j.oro(y) - beta * kl → $r(x,y) - \beta \log \frac{\pi^{RL}_\phi(y\mid x)}{\pi^{SFT}(y\mid x)}$ — recompensa menos el freno de KL
• loss = -(A * j.logprobs(pol, y).sum(-1)).mean() → $-A \log \pi^{RL}_\phi(y\mid x)$ — una sola pasada: el ratio vale 1
• gamma * j.nll_pre(pol, j.batch_pre(64)) → $-\gamma\, \log \pi^{RL}_\phi(x_{\text{pre}})$ — ptx: contar para adelante

def rl(beta, gamma, iters=150):
    pol, ref, cap = 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))             # la política responde
        with torch.no_grad():
            kl = (j.logprobs(pol, y) - j.logprobs(ref, y)).sum(-1)
            r = j.oro(y) - beta * kl                     # recompensa − β·KL
            A = (r - r.mean()) / (r.std() + 1e-8)        # ventaja
        loss = -(A * j.logprobs(pol, y).sum(-1)).mean() # una pasada: ratio 1
        loss = loss + gamma * j.nll_pre(pol, j.batch_pre(64))  # ptx
        loss.backward(); opt.step(); opt.zero_grad()
        if it % 10 == 0 or it == iters - 1:
            cap.append(j.capacidad(pol))                 # el benchmark
    return pol, cap, j.oro(y).mean()

La capacidad durante el aprendizaje por refuerzo

Las curvas, medidas cada diez iteraciones. PPO: la capacidad cae de cero coma ocho a cero coma tres, mientras la recompensa sube. Ese es el impuesto. Con beta veinticinco veces más grande cae igual. Más freno de KL no protege lo que se pierde. PPO-ptx sube: termina contando mejor que el SFT, igual que HellaSwag en el paper. La línea punteada es el SFT, el punto de partida.

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

$\mathbb{E}$: promedio · $D_{\pi^{RL}}$: prompts de la API y lo que responde la política

$\gamma$: cuánto pesa · el paper: 27,8

dos gradientes que se suman en los mismos pesos

La ecuación dos del paper. E es un promedio. Adentro, para un prompt x de la API y una respuesta y que genera la política pi RL, con pesos phi: la recompensa del reward model menos beta por el logaritmo del cociente entre la política y el SFT. Ese logaritmo, sumado sobre los tokens de la respuesta, es el KL de esa muestra. En el dibujo, uno por token. El término nuevo: gamma por el log de la probabilidad de un texto x del preentrenamiento. Maximizarlo es lo mismo que minimizar la pérdida de preentrenamiento de siempre. En el dibujo, un log por token: cuanto más cerca de cero, mejor lo predice. Y el gradiente del total es la suma de los dos gradientes. En el dibujo, dos vectores que se suman.

Parte 3 · Por qué el KL no alcanza

El freno de KL existe para no alejarse del SFT. ¿Por qué no protege la capacidad?

Dónde mira el KL

Medimos la divergencia de cada política contra el SFT en dos lugares. El KL del próximo token, promediado sobre posiciones. Primero, en respuestas a prompts con INST: es lo que ve el término de KL durante el entrenamiento. Después, en secuencias que cuentan para adelante: donde vive el benchmark. PPO: cero coma uno donde mira, uno coma dos donde no. Beta alto baja el primero a cero coma cero siete cinco, pero el segundo sube a uno coma cinco. Y PPO-ptx también se aleja del SFT al contar, cero coma cuatro dos, pero hacia el preentrenamiento: cuenta mejor.

@torch.no_grad()
def kl_por_token(pol, x):                  # KL(π‖π_SFT) del próximo token, promediada
    lp, lq = pol(input_ids=x).logits.log_softmax(-1), sft(input_ids=x).logits.log_softmax(-1)
    return (lp.exp() * (lp - lq)).sum(-1).mean().item()

for nombre, pol in pols.items():
    x_rl = j.muestrear(pol, j.prompts(256))          # donde PPO mide el KL
    x_pre = j.batch_pre(256)                         # donde vive el benchmark
    print(f"{nombre:8} KL en prompts RL {kl_por_token(pol, x_rl):.3f}"
          f"   en contar {kl_por_token(pol, x_pre):.3f}")

PPO      KL en prompts RL 0.108   en contar 1.212
PPO-ptx  KL en prompts RL 0.088   en contar 0.420
β alto   KL en prompts RL 0.075   en contar 1.510

El KL es un promedio sobre los prompts de RL

solo cuentan los prompts de RL

$\beta = 0{,}5$: 25 veces el default

la mezcla entrena donde vive la capacidad

La explicación está en la fórmula. El KL es un promedio sobre prompts x de la distribución de RL, y respuestas y de la política. Las entradas que nunca aparecen ahí no cuentan. En el dibujo, las tres de contar están tachadas: el promedio sale cero coma uno, aunque al contar el modelo cambió doce veces más. Subir beta castiga más fuerte, pero en el mismo lugar. En el dibujo, el KL de contar sigue alto. La mezcla pone un término justo sobre la distribución que queremos conservar. Mientras el modelo comparta pesos entre tareas, lo que no se entrena se puede perder.

Ouyang et al., 2022 — apéndice E.6 · arXiv:2203.02155

β = 2,0 (100×) · «even by increasing the KL reward coefficient to 2.0, which is 100 times of the default value, the regressions still cannot be fixed.»

El paper lo probó en grande, con el modelo de mil trescientos millones. Subieron beta hasta dos, cien veces el valor por defecto, y las regresiones en DROP y SQuAD no se arreglaron; la recompensa de validación cayó. En cambio, la mezcla las revierte con una pérdida chica de recompensa. Nuestro juguete da lo mismo.

Más freno contra más mezcla

Las dos herramientas, lado a lado. Subir beta es un freno general: castiga cualquier cambio, también el que la recompensa quiere. Y como solo mira las respuestas a prompts de RL, deja sin vigilar justo lo que se rompe. El resultado: menos recompensa y la regresión sigue ahí. La mezcla es un ancla puesta en otro lugar: entrena directamente sobre la distribución del preentrenamiento. Cuesta cómputo, ocho veces más ejemplos, y trae lo que haya en ese texto. Pero ataca la causa.

Parte 4 · Cómo lo implementa el paper

Dos gradientes por minibatch, uno detrás del otro, acumulados en el mismo buffer.

Acumular es lo mismo que sumar

El apéndice C cuenta el detalle. Gamma vale veintisiete coma ocho, y usan ocho veces más ejemplos de preentrenamiento que episodios de RL. Para cada minibatch calculan el gradiente de PPO. Después el de preentrenamiento, multiplicado por gamma, y los dos se acumulan en el mismo lugar: en PyTorch, backward suma en punto grad en vez de reemplazar. Lo comparamos con una sola pérdida sumada. Los tres números coinciden. Acumular dos backward es exactamente el gradiente de la suma: la ecuación dos, implementada en dos pasos.

gamma = 27.8                                        # el γ del paper (apéndice C.4)
def perdidas():                                     # dos pérdidas cualquiera
    return capa(x_rl).pow(2).mean(), capa(x_pre).abs().mean()

l_ppo, l_pre = perdidas()                           # como el paper: en pasos seguidos
l_ppo.backward()                                    # 1) gradiente de PPO
(gamma * l_pre).backward()                          # 2) se ACUMULA en .grad
acumulado = capa.weight.grad.clone()

capa.zero_grad()
l_ppo, l_pre = perdidas()
(l_ppo + gamma * l_pre).backward()                  # una sola pérdida sumada
print("acumulado:", acumulado.numpy().round(3))
print("sumado:   ", capa.weight.grad.numpy().round(3))

acumulado: [[  2.246 -17.996 -22.897]]
sumado:    [[  2.246 -17.996 -22.897]]

Lo que la mezcla no arregla

• DROP, SQuADv2 y traducción siguen debajo de GPT-3
• Entrenar el doble: DROP y SQuAD vuelven a caer
• Puede traer conductas del preentrenamiento
• Sube la capacidad, no el win rate

El propio paper marca los límites. Con la mezcla, PPO-ptx sigue por debajo de GPT tres en DROP, SQuAD y traducción; dicen que hace falta más trabajo. Con quinientos doce mil episodios en vez de doscientos cincuenta y seis mil, las tareas vuelven a caer un poco debajo de GPT tres. Mezclar texto de internet puede reintroducir conductas no deseadas que están en ese texto; proponen filtrarlo. Y la mezcla no mejora la preferencia humana: el win rate de PPO y PPO-ptx es parecido. Su trabajo es no perder lo que ya estaba.

Resumiendo

• Alignment tax: PPO empeora DROP, SQuAD, HellaSwag y traducción
• PPO-ptx: $+\gamma\,\mathbb{E}[\log \pi(x_{\text{pre}})]$, con $\gamma = 27{,}8$
• Recupera mucho, no todo: DROP 15 %, HellaSwag lo supera
• Juguete: capacidad 0,34 con PPO, 0,996 con la mezcla
• El KL solo mira los prompts de RL; subir β no alcanza
• Próximo: los límites del paper

Papers del episodio

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