RLHF: mejorar la recompensa sin alejarse

Con el modelo de recompensaBloque 2 · Lección 4Preferencias y el modelo de recompensaUna nota se lee en la escala de quien la pone; una comparación entre dos respuestas, no. El modelo de Bradley–Terry convierte comparaciones en un número por respuesta: el anotador prefiere una con la sigmoide de la diferencia de recompensas, que queda definida salvo una constante por instrucción. Un modelo de recompensa se entrena con la entropía cruzada binaria de esas comparaciones. La lección lo monta sobre el mini-GPT congelado, que ordena bien lo que vio y adivina en lo que no conoce. ya entrenado y congelado, $r_\phi$, el bloque tiene con qué puntuar respuestas y una política que mejorar, la del ajuste supervisado. Lo que cambia ahora es de dónde salen las respuestas con que se entrena. En el ajuste las escribía alguien y el modelo aprendía a imitarlas; aquí las escribe el propio modelo, $r_\phi$ las puntúa, y el modelo aprende de lo que reciben. Eso es el aprendizaje por refuerzo: aprender de lo que pasa después de actuar, no de un ejemplo de cómo actuar.

El método completo se llama aprendizaje por refuerzo con retroalimentación humana (reinforcement learning from human feedback, RLHF), porque la recompensa sale de comparaciones de personas, y tiene un riesgo conocido desde sus primeros experimentos. En uno de ellos, un brazo robótico debía aprender a agarrar un objeto a partir de comparaciones de personas que lo veían por una cámara. Aprendió a poner la mano entre la cámara y el objeto, donde parecía que lo agarraba. Esta lección escribe el objetivo que RLHF optimiza, que lleva dentro el freno contra eso, y el gradiente que lo sube.

Dos dificultades separan la idea del algoritmo. La primera es el sorteo, que no deja pasar ningún gradiente: un cambio pequeño de $\theta$ deja el token sorteado igual o lo cambia por otro, y la recompensa salta con él. Lo que sí cambia poco a poco es la probabilidad de cada respuesta. De ahí sale el método: sortear unas cuantas, puntuarlas, y hacer más probables las que puntúan por encima de la media y menos las de debajo.

La segunda dificultad es la propia recompensa. El modelo de la lección anterior acertaba con rivales nuevos en las preguntas que conocía y echaba una moneda al aire en las que no. Una política que sólo persigue su número hace lo que hizo el brazo: encuentra dónde se equivoca la recompensa a su favor. Por eso RLHF cobra la distancia: cada respuesta paga por lo improbable que la encontraba el modelo de partida, y la política sólo se aleja de él donde lo que gana compensa lo que paga. Es el freno que el ajuste supervisadoBloque 2 · Lección 3Ajuste supervisado: máxima verosimilitud sobre un formatoSeguir entrenando un modelo preentrenado sobre pares de instrucción y respuesta es la entropía cruzada de siempre, cobrada sólo en la respuesta: máxima verosimilitud sobre el conjunto de pares. La lección deriva su óptimo, que copia el conjunto, y su gradiente, que es el del preentrenamiento con las filas del prompt a cero. Después ajusta el mini-GPT con diez preguntas: aprende el formato, lee la pregunta, inventa el contenido y olvida su novela. no tenía, y mide la distancia en las distribuciones, no en los pesos.

El objetivo: recompensa menos distancia

Fijemos los objetos. $\pi_\theta$ es la política que se entrena, y empieza siendo el modelo del ajuste supervisado. $\pi_{\text{ref}}$ es ese mismo modelo congelado, la referencia contra la que se mide la distancia. $r_\phi$ es el modelo de recompensa, congelado también: $\phi$ no se mueve en toda la lección, y por eso los dos conjuntos de parámetros llevan letras distintas.

La distancia es la divergencia de Kullback–Leibler (KL) de la política a la referencia: la media, sobre las respuestas de la política, del logaritmo del cociente de sus probabilidades,

Nunca es negativa, y vale cero sólo si las dos políticas coinciden: es la desigualdad de la lección sobre la perplejidadBloque 1 · Lección 6Perplejidad: medir un modelo de lenguajeUn modelo que sólo cuenta bytes paga, sobre las páginas de Marianela que el mini-GPT no vio, la misma pérdida por token que el mini-GPT. La lección explica por qué eso no dice nada: convierte la pérdida en perplejidad, el número de candidatas entre las que duda el modelo, y en bits por byte, la medida que no depende del tokenizador; le pone al lado dos varas de medir, el modelo uniforme y el de unigramas; y demuestra por qué la pérdida sobre el texto de entrenamiento engaña y la del texto reservado no., con $\pi_\theta$ en el papel de las frecuencias y $\pi_{\text{ref}}$ en el del modelo. No es simétrica, y la asimetría es lo que interesa. La media lleva los pesos de $\pi_\theta$, así que una respuesta cuesta mucho cuando la política la escribe a menudo y la referencia casi nunca. Dejar de escribir algo que la referencia sí escribiría no se cobra por sí mismo: esa respuesta entra en la media con peso cero.

Una respuesta es una secuencia de tokens, y la regla de la cadena escribe su probabilidad como un producto con un factor por token, así que el logaritmo del cociente es una suma. Escribamos $\pi_\theta(y_j \mid x, y_{<j})$ para el factor $j$ de $\pi_\theta(y \mid x)$, y lo mismo para la referencia:

Cada respuesta sorteada de la política da así una estimación de la KL, sumada token a token, cuya media es la KL exacta. Una estimación sola puede salir negativa, si la respuesta es de las que la referencia prefiere, aunque su media nunca lo sea. Y hay un caso que la celda va a encontrar. Si la política escribe siempre la misma respuesta $y$, la media tiene un único término y la KL vale $-\log \pi_{\text{ref}}(y \mid x)$: lo improbable que era esa respuesta, contado en nats.

El objetivo de RLHF suma la recompensa media de la política y le resta la KL multiplicada por $\beta$,

para una instrucción $x$. Lo que se maximiza es su media sobre las instrucciones del entrenamiento, y todo lo que sigue vale instrucción a instrucción. $\beta$ es el precio de un nat en unidades de recompensa. Con $\beta = 0$, la política persigue la recompensa sin mirar dónde; con $\beta$ muy grande, no se mueve de la referencia. En medio hay un trato: la política se aleja sólo donde gana más recompensa de la que paga en KL.

El gradiente de una media sobre un sorteo

Empecemos por la recompensa. Como suma sobre todas las respuestas posibles, su media es $\sum_y \pi_\theta(y \mid x)\, r_\phi(x, y)$, y $\theta$ sólo aparece en las probabilidades. La derivada del logaritmo da $\nabla_\theta \pi_\theta = \pi_\theta \nabla_\theta \log \pi_\theta$, y con ella la suma vuelve a ser una media sobre las respuestas de la política:

Es el truco de la derivada logarítmica (log-derivative trick), y el estimador que sale de él se llama REINFORCE desde que Williams lo publicó en 1992. La derivada de una media es ahora la media de algo que se calcula sobre una respuesta, y una media se estima sorteando. Con $B$ respuestas $y^{(1)}, \dots, y^{(B)}$ de la política,

Nada en esa expresión atraviesa el sorteo. La recompensa de cada respuesta es un número, y $\nabla_\theta \log \pi_\theta(y^{(i)} \mid x)$ es un gradiente que ya conoces: con el signo cambiado, el de la pérdida del ajuste supervisado para el par $(x, y^{(i)})$. Un paso de REINFORCE es un paso de ajuste supervisado sobre las respuestas que la política acaba de escribir, con la pérdida de cada una multiplicada por su recompensa. El ajuste imitaba respuestas ajenas, todas por igual; éste imita las propias, más cuanto mejor puntúan.

Lo que se puede restar sin cambiar nada

Hay una media que vale cero, y es la que hace práctico el estimador:

Restar a todas las recompensas un número que dependa de $x$ y no de $y$ deja, por tanto, el gradiente igual en media. Es el $\kappa(x)$ de la lección anterior en otro papel: allí no cambiaba ninguna comparación, aquí no cambia el gradiente esperado. Lo que sí cambia es cuánto se equivoca un batch. Con una recompensa que nunca es negativa, como la de la celda, cada respuesta sorteada se empuja hacia arriba, también las peores. Restando la recompensa media, las de encima suben y las de debajo bajan. Esa resta es la línea base (baseline).

La celda usa la media de su propio batch, que incluye a la respuesta a la que se le resta, y eso no rompe la cuenta. La diferencia entre una recompensa y la media de las $B$ es $(B - 1)/B$ veces su diferencia con la media de las otras $B - 1$, que no depende de ella, y un factor fijo no cambia la dirección.

La KL, descontada de cada recompensa

Falta el término de la KL, que también es una media sobre las respuestas de la política, con una diferencia: lo que se promedia depende de $\theta$. La regla del producto da dos sumas, y la segunda es la que acaba de valer cero:

con $\mathbb{D}_{\text{KL}}$ abreviando la de la definición. Juntando las tres piezas,

La KL entra como un descuento en la recompensa de cada respuesta: $\beta$ veces su logaritmo del cociente, la suma token a token de antes, tratada como un número más. Una respuesta que la referencia encontraba mucho menos probable que la política llega al estimador rebajada. InstructGPT se entrenó con esta recompensa descontada, pero no con este estimador tal cual. Usaba PPO (proximal policy optimization), que aprovecha cada batch de respuestas en varios pasos y recorta en ellos cuánto puede cambiar el cociente entre la política nueva y la que las escribió; está en las lecturas.

Una recompensa que cuenta letras, en el mini-GPT

Las tres celdas corren sobre el mini-GPT, tu mini-GPT del bloque 1 cargado desde el checkpoint, sin la plantilla: la instrucción es el principio de una frase de la novela, La Nela, y la respuesta, los 16 tokens que la continúan. La referencia es ese mismo modelo, congelado. En RLHF es el modelo ajustado; aquí no puede serlo, por la razón de la lección anterior: el ajustado olvidó la novela y sólo sabe escribir frases de colores.

La política tampoco es la red entera. He congelado el mini-GPT, y la política es él más un vector $\theta$ de 512 números, uno por entrada, que se suma a sus logits en cada posición: $\pi_\theta(v \mid x, y_{<j})$ es proporcional a $\pi_{\text{ref}}(v \mid x, y_{<j})\, e^{\theta_v}$. Con la red entera, un paso cuesta en el navegador unas tres veces más, y en mis pruebas ocho respuestas por paso no bastaban para mover sus 136 000 pesos sin que la política se alejara por puro ruido, con el freno puesto y sin ganar recompensa. Con $\theta$, el estimador es el mismo y $\nabla_\theta \log \pi_\theta(y \mid x)$ cabe en una línea: en cada token, el one-hot del token escrito menos el softmax.

La recompensa de la celda no sale de ninguna comparación. Cuenta las letras a de la respuesta, y está ahí porque se equivoca a la vista. Un modelo de recompensa de verdad tiene sesgos que nadie ha elegido, como premiar las respuestas largas, y la política los encuentra antes que nadie. Éste tiene uno solo, exagerado, para que se vea qué hace la política con él. La primera celda prepara el muestreo, ocho respuestas a la vez con la caché de claves y valores, y mide la referencia: con $\theta = 0$ la política es $\pi_{\text{ref}}$.

El mini-GPT escribe de media $3.47$ letras a en 16 tokens, y la KL sale exactamente cero, porque con $\theta = 0$ cada cociente vale 1. La segunda celda es el estimador de la lección: treinta pasos de ascenso por gradiente sobre $\theta$, ocho respuestas por paso, la línea base y el descuento de la KL. La ejecuta con $\beta = 0$, sin freno.

La recompensa y la KL suben juntas: $9.12$ letras con $18.91$ nats en el paso 10, y $16$ con $73.32$ en el 15, donde se quedan. Las 32 respuestas de la evaluación son la misma, a a a a a a a a a a a a a a a a, el token a dieciséis veces. Ni siquiera es la de más recompensa posible, porque aba repetido daría el doble. Es la primera que encontró el ascenso, y de ahí ya no sale: ocho respuestas iguales tienen la misma recompensa, la línea base las deja en cero y el gradiente con ellas.

Su KL, $73.34$ nats, es la del caso de una sola respuesta. La referencia la escribiría con probabilidad $e^{-73.3}$, del orden de $10^{-32}$. Esto es reward hacking: la recompensa ha pasado de $3.47$ a $16$, y lo que debía medir, que la respuesta fuera buena, se ha perdido por el camino. La tercera celda repite el entrenamiento con $\beta = 0.5$, de modo que cada nat que la política se aleja le cuesta media letra.

Con el freno, la recompensa sube a $7.44$, más del doble de la de partida, por $5.48$ nats de KL. Las respuestas siguen siendo novela, ahora con más a que antes: a la triste bajar a Florentina, a. La política las ha buscado donde la referencia también las pone, en la, en a la, en las palabras que las llevan.

La cuenta del objetivo dice por qué. Con $\beta = 0.5$, la respuesta de la segunda celda vale $16 - 0.5 \cdot 73.3 = -20.7$, mucho menos que quedarse en la referencia, que vale $3.47$, y que lo que logra la política entrenada, $7.44 - 0.5 \cdot 5.48 = 4.70$. El freno no prohíbe aquella respuesta: la hace salir cara, y con un $\beta$ lo bastante pequeño volvería a compensar.

Comprueba tu intuición

Cinco preguntas, la última con los números de las celdas.

RLHF funciona, y es caro. Necesita un modelo de recompensa entrenado aparte, sortear respuestas en cada paso y tener a la vez en memoria la política, la referencia y el modelo de recompensa, y estimar un gradiente con unas pocas respuestas es tan ruidoso que esta lección ha tenido que congelar la red para que ocho por paso bastaran. Y lo que se optimiza no es la recompensa sino un trato entre ella y la KL, cuyo óptimo esta lección no ha buscado: el ascenso sólo se acerca a él paso a paso.

La lección siguiente, sobre la optimización directa de preferencias (direct preference optimization, DPO), lo busca. El objetivo de esta lección tiene una solución exacta, escrita con la referencia y la recompensa, y esa solución, leída al revés, convierte la política en su propio modelo de recompensa: con ella, las comparaciones entrenan la política directamente, sin modelo de recompensa y sin sorteos.